From 03005d71faa2b49cced989c4b82ee1f3131b76cf Mon Sep 17 00:00:00 2001 From: wizardchen Date: Thu, 5 Feb 2026 17:10:20 +0800 Subject: [PATCH] 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. --- internal/application/service/session.go | 11 ++++------- internal/handler/session/helpers.go | 5 ++++- internal/handler/session/qa.go | 10 +++++++--- 3 files changed, 15 insertions(+), 11 deletions(-) diff --git a/internal/application/service/session.go b/internal/application/service/session.go index 26bf9ff4..f8c8643a 100644 --- a/internal/application/service/session.go +++ b/internal/application/service/session.go @@ -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) } diff --git a/internal/handler/session/helpers.go b/internal/handler/session/helpers.go index d0cd6614..81752abf 100644 --- a/internal/handler/session/helpers.go +++ b/internal/handler/session/helpers.go @@ -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 }) } diff --git a/internal/handler/session/qa.go b/internal/handler/session/qa.go index 71385cdd..651101b6 100644 --- a/internal/handler/session/qa.go +++ b/internal/handler/session/qa.go @@ -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 ... 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) }()