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:
@@ -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)
|
||||
}
|
||||
|
||||
@@ -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
|
||||
})
|
||||
}
|
||||
|
||||
@@ -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)
|
||||
}()
|
||||
|
||||
|
||||
Reference in New Issue
Block a user