refactor(session): streamline tenant context handling for async operations

- Updated session title generation and message completion to utilize the session's tenant ID directly, ensuring correct context is used in background operations.
- Removed unnecessary context value extraction and background context creation, simplifying the code and improving clarity.
- Enhanced the handling of stop events and message updates to consistently apply the session's tenant ID, improving the accuracy of operations in shared agent scenarios.

This update improves the reliability of async operations by ensuring the correct tenant context is consistently applied.
This commit is contained in:
wizardchen
2026-02-05 17:10:20 +08:00
committed by lyingbug
parent c806eaf7f9
commit 03005d71fa
3 changed files with 15 additions and 11 deletions
+4 -7
View File
@@ -342,15 +342,12 @@ func (s *sessionService) GenerateTitleAsync(
modelID string,
eventBus *event.EventBus,
) {
// Extract values from context before cloning
tenantID := ctx.Value(types.TenantIDContextKey)
// Use session's tenant for title generation (session belongs to session.TenantID; ctx may have effectiveTenantID when using shared agent)
sessionTenantID := session.TenantID
requestID := ctx.Value(types.RequestIDContextKey)
go func() {
// Create new background context and copy values
bgCtx := context.Background()
if tenantID != nil {
bgCtx = context.WithValue(bgCtx, types.TenantIDContextKey, tenantID)
}
// Create new background context with session owner's tenant so ListModels/GetChatModel/Update use correct scope
bgCtx := context.WithValue(context.Background(), types.TenantIDContextKey, sessionTenantID)
if requestID != nil {
bgCtx = context.WithValue(bgCtx, types.RequestIDContextKey, requestID)
}
+4 -1
View File
@@ -160,6 +160,7 @@ func (h *Handler) setupStreamHandler(
func (h *Handler) setupStopEventHandler(
eventBus *event.EventBus,
sessionID string,
sessionTenantID uint64,
assistantMessage *types.Message,
cancel context.CancelFunc,
) {
@@ -167,7 +168,9 @@ func (h *Handler) setupStopEventHandler(
logger.Infof(ctx, "Received stop event, cancelling async operations for session: %s", sessionID)
cancel()
assistantMessage.Content = "用户停止了本次对话"
h.completeAssistantMessage(ctx, assistantMessage)
// Use session's tenant for message update (ctx may have effectiveTenantID when using shared agent)
updateCtx := context.WithValue(ctx, types.TenantIDContextKey, sessionTenantID)
h.completeAssistantMessage(updateCtx, assistantMessage)
return nil
})
}
+7 -3
View File
@@ -210,7 +210,7 @@ func (h *Handler) setupSSEStream(reqCtx *qaRequestContext, generateTitle bool) *
}
// Setup stop event handler
h.setupStopEventHandler(eventBus, reqCtx.sessionID, reqCtx.assistantMessage, cancel)
h.setupStopEventHandler(eventBus, reqCtx.sessionID, reqCtx.session.TenantID, reqCtx.assistantMessage, cancel)
// Setup stream handler
h.setupStreamHandler(asyncCtx, reqCtx.sessionID, reqCtx.assistantMessage.ID,
@@ -413,7 +413,9 @@ func (h *Handler) executeNormalModeQA(reqCtx *qaRequestContext, generateTitle bo
logger.Infof(streamCtx.asyncCtx, "Knowledge QA service completed for session: %s", sessionID)
// Content already contains <think>...</think> tags from chat_completion_stream.go
h.completeAssistantMessage(streamCtx.asyncCtx, streamCtx.assistantMessage)
// Use session's tenant for message update (asyncCtx may have effectiveTenantID when using shared agent)
updateCtx := context.WithValue(streamCtx.asyncCtx, types.TenantIDContextKey, reqCtx.session.TenantID)
h.completeAssistantMessage(updateCtx, streamCtx.assistantMessage)
// Emit EventAgentComplete - this will trigger handleComplete which sends the SSE complete event
// Note: Don't cancel context here, let the SSE handler close naturally after receiving the complete event
streamCtx.eventBus.Emit(streamCtx.asyncCtx, event.Event{
@@ -517,7 +519,9 @@ func (h *Handler) executeAgentModeQA(reqCtx *qaRequestContext) {
errors.NewInternalServerError(fmt.Sprintf("Agent QA service panicked: %v\n%s", r, string(buf))),
map[string]interface{}{"session_id": sessionID})
}
h.completeAssistantMessage(streamCtx.asyncCtx, streamCtx.assistantMessage)
// Use session's tenant for message update (session belongs to session.TenantID; asyncCtx may have effectiveTenantID when using shared agent)
updateCtx := context.WithValue(streamCtx.asyncCtx, types.TenantIDContextKey, reqCtx.session.TenantID)
h.completeAssistantMessage(updateCtx, streamCtx.assistantMessage)
logger.Infof(streamCtx.asyncCtx, "Agent QA service completed for session: %s", sessionID)
}()