diff --git a/backend/internal/service/openai_gateway_response_handling.go b/backend/internal/service/openai_gateway_response_handling.go index 2563dbc3e3..efb9f2e0d8 100644 --- a/backend/internal/service/openai_gateway_response_handling.go +++ b/backend/internal/service/openai_gateway_response_handling.go @@ -1464,12 +1464,23 @@ func (s *OpenAIGatewayService) bindHTTPResponseAccount(ctx context.Context, c *g if store == nil { return } + // The client may close the stream immediately after receiving its terminal + // event, canceling the request context before these durable affinity writes + // run. Preserve request values, but give all Redis writes one bounded budget + // independent of the downstream connection lifecycle. + bindBaseCtx := context.Background() + if ctx != nil { + bindBaseCtx = context.WithoutCancel(ctx) + } + bindCtx, cancel := context.WithTimeout(bindBaseCtx, openAIWSStateStoreRedisTimeout) + defer cancel() + groupID := getOpenAIGroupIDFromContext(c) ttl := s.openAIWSResponseStickyTTL() - logOpenAIWSBindResponseAccountWarn(groupID, account.ID, responseID, store.BindResponseAccount(ctx, groupID, responseID, account.ID, ttl)) + logOpenAIWSBindResponseAccountWarn(groupID, account.ID, responseID, store.BindResponseAccount(bindCtx, groupID, responseID, account.ID, ttl)) if rawOwner, ok := c.Get(openAIHTTPResponseOwnerContextKey); ok { if owner, ok := rawOwner.(openAIHTTPResponseOwner); ok && owner.userID > 0 && owner.apiKeyID > 0 { - if err := s.BindOpenAIHTTPResponseOwner(ctx, groupID, responseID, owner.userID, owner.apiKeyID); err != nil { + if err := s.BindOpenAIHTTPResponseOwner(bindCtx, groupID, responseID, owner.userID, owner.apiKeyID); err != nil { logger.L().Warn( "openai.http_bind_response_owner_failed", zap.Int64("group_id", groupID), diff --git a/backend/internal/service/openai_gateway_service_test.go b/backend/internal/service/openai_gateway_service_test.go index 69113a59d9..cf4611ccce 100644 --- a/backend/internal/service/openai_gateway_service_test.go +++ b/backend/internal/service/openai_gateway_service_test.go @@ -589,6 +589,37 @@ func TestOpenAIGatewayService_BindHTTPResponseAccount(t *testing.T) { require.False(t, owned) } +func TestOpenAIGatewayService_BindHTTPResponseAccount_DetachesCanceledRequestContext(t *testing.T) { + gin.SetMode(gin.TestMode) + rec := httptest.NewRecorder() + c, _ := gin.CreateTestContext(rec) + c.Request = httptest.NewRequest(http.MethodPost, "/openai/v1/responses", nil) + groupID := int64(4201) + c.Set("api_key", &APIKey{ID: 501, GroupID: &groupID}) + SetOpenAIHTTPResponseOwner(c, 601, 501) + + cache := &responseBindContextProbeCache{} + svc := &OpenAIGatewayService{cache: cache} + account := &Account{ID: 37001, Platform: PlatformOpenAI, Type: AccountTypeAPIKey} + requestCtx, cancelRequest := context.WithCancel(context.Background()) + cancelRequest() + startedAt := time.Now() + + svc.bindHTTPResponseAccount(requestCtx, c, account, "resp_http_canceled_001") + + require.Len(t, cache.setContextErrors, 3) + for _, contextErr := range cache.setContextErrors { + require.NoError(t, contextErr, "response affinity writes must survive downstream cancellation") + } + require.Len(t, cache.setDeadlines, 3) + for _, deadline := range cache.setDeadlines { + require.True(t, deadline.After(startedAt)) + require.LessOrEqual(t, deadline.Sub(startedAt), openAIWSStateStoreRedisTimeout+100*time.Millisecond) + } + require.Len(t, cache.sessionBindings, 3) + require.Contains(t, cache.sessionBindings, openAIWSResponseAccountCacheKey("resp_http_canceled_001")) +} + func TestOpenAIGatewayService_GenerateExplicitSessionHash_SkipsContentFallback(t *testing.T) { gin.SetMode(gin.TestMode) svc := &OpenAIGatewayService{} @@ -710,6 +741,23 @@ type stubGatewayCache struct { deletedSessions map[string]int } +type responseBindContextProbeCache struct { + stubGatewayCache + setContextErrors []error + setDeadlines []time.Time +} + +func (c *responseBindContextProbeCache) SetSessionAccountID(ctx context.Context, groupID int64, sessionHash string, accountID int64, ttl time.Duration) error { + c.setContextErrors = append(c.setContextErrors, ctx.Err()) + if deadline, ok := ctx.Deadline(); ok { + c.setDeadlines = append(c.setDeadlines, deadline) + } + if err := ctx.Err(); err != nil { + return err + } + return c.stubGatewayCache.SetSessionAccountID(ctx, groupID, sessionHash, accountID, ttl) +} + func (c *stubGatewayCache) GetSessionAccountID(ctx context.Context, groupID int64, sessionHash string) (int64, error) { if id, ok := c.sessionBindings[sessionHash]; ok { return id, nil