diff --git a/backend/internal/handler/openai_gateway_handler.go b/backend/internal/handler/openai_gateway_handler.go index 15a7cca8a2..c168bf6de8 100644 --- a/backend/internal/handler/openai_gateway_handler.go +++ b/backend/internal/handler/openai_gateway_handler.go @@ -2992,7 +2992,7 @@ func (h *OpenAIGatewayHandler) ResponsesWebSocket(c *gin.Context) { // WebSocket 首包可能很大，hash 必须在 hooks 外算成字符串，避免 AfterTurn 闭包保活请求体。 requestPayloadHash = service.HashUsageRequestPayload(wsFirstMessage) - if preemptCtx, cleanupPreempt, armed := h.gatewayService.BeginOpenAIWSIngressSessionPreemption(ctx, c, account, wsFirstMessage); armed { + if preemptCtx, cleanupPreempt, armed := h.gatewayService.BeginOpenAIWSIngressSessionPreemptionWithClient(ctx, c, account, wsFirstMessage, wsConn); armed { ctx = preemptCtx defer cleanupPreempt() } @@ -3004,6 +3004,8 @@ func (h *OpenAIGatewayHandler) ResponsesWebSocket(c *gin.Context) { return } if service.IsOpenAIWSSessionPreemptedError(err) { + // 关闭帧已由抢占登记在取消前发给本连接，这里只记录并释放。 + reqLog.Info("openai.websocket_ingress_preempted", zap.Int64("account_id", account.ID)) return } var failoverErr *service.UpstreamFailoverError diff --git a/backend/internal/service/openai_ws_forwarder_ingress.go b/backend/internal/service/openai_ws_forwarder_ingress.go index 0def1bd0cb..4658708911 100644 --- a/backend/internal/service/openai_ws_forwarder_ingress.go +++ b/backend/internal/service/openai_ws_forwarder_ingress.go @@ -104,7 +104,7 @@ func (s *OpenAIGatewayService) ProxyResponsesWebSocketFromClient( // The handler normally owns this registration across retry attempts. Direct // callers still get the same session-scoped preemption behavior here. - if preemptCtx, cleanupPreempt, armed := s.BeginOpenAIWSIngressSessionPreemption(ctx, c, account, firstClientMessage); armed { + if preemptCtx, cleanupPreempt, armed := s.BeginOpenAIWSIngressSessionPreemptionWithClient(ctx, c, account, firstClientMessage, clientConn); armed { ctx = preemptCtx defer cleanupPreempt() defer func() { diff --git a/backend/internal/service/openai_ws_forwarder_ingress_execution_scope_test.go b/backend/internal/service/openai_ws_forwarder_ingress_execution_scope_test.go index ef2594535d..c9a18fc748 100644 --- a/backend/internal/service/openai_ws_forwarder_ingress_execution_scope_test.go +++ b/backend/internal/service/openai_ws_forwarder_ingress_execution_scope_test.go @@ -2,6 +2,7 @@ package service import ( "context" + "errors" "io" "net/http" "net/http/httptest" @@ -339,6 +340,10 @@ func TestOpenAIGatewayService_ProxyResponsesWebSocketFromClient_CodexThreadsDoNo func TestOpenAIGatewayService_ProxyResponsesWebSocketFromClient_SameCodexThreadStillPreempts(t *testing.T) { serverErrs, aReadErr := runOpenAIWSCodexThreadPair(t, "thread-a", "thread-a") require.Error(t, aReadErr, "同线程重连必须取代旧连接") + var closeErr coderws.CloseError + require.True(t, errors.As(aReadErr, &closeErr), "被取代的连接应收到关闭帧而不是裸断开: %v", aReadErr) + require.Equal(t, coderws.StatusTryAgainLater, closeErr.Code) + require.Equal(t, openAIWSSessionPreemptedCloseReason, closeErr.Reason) preempted := 0 for _, err := range serverErrs { if IsOpenAIWSSessionPreemptedError(err) { diff --git a/backend/internal/service/openai_ws_session_preemption.go b/backend/internal/service/openai_ws_session_preemption.go index 22ed8fd4f4..f1f55c34f9 100644 --- a/backend/internal/service/openai_ws_session_preemption.go +++ b/backend/internal/service/openai_ws_session_preemption.go @@ -6,8 +6,10 @@ import ( "fmt" "strings" "sync" + "sync/atomic" "time" + coderws "github.com/coder/websocket" "github.com/gin-gonic/gin" "github.com/google/uuid" ) @@ -18,8 +20,23 @@ const ( openAIWSSessionPreemptOwnerTTL = 2 * time.Hour openAIWSSessionPreemptWatchInterval = 2 * time.Second openAIWSSessionPreemptCachePrefix = "wspreempt:" + // openAIWSSessionPreemptCloseGrace 是被抢占连接从收到关闭帧到被取消的最长等待： + // 关闭帧在 Close 一开始就写出，其余时间只是等对端回应，到点直接取消。 + openAIWSSessionPreemptCloseGrace = time.Second + openAIWSSessionPreemptedCloseReason = "session preempted by a newer connection" ) +// openAIWSPreemptClientCloser 是被抢占时用来给旧客户端发关闭帧的连接，*coderws.Conn 满足该接口。 +type openAIWSPreemptClientCloser interface { + Close(code coderws.StatusCode, reason string) error +} + +// openAIWSSessionPreemptState 随抢占上下文传递：preempted 在关闭帧发出前就置位， +// 让旧连接在被取消之前遇到的读写错误也能归类为“被抢占”。 +type openAIWSSessionPreemptState struct { + preempted atomic.Bool +} + // OpenAIWSSessionPreemptionCache is an optional GatewayCache capability. The // production Redis cache implements all operations atomically; cache stubs do // not need to implement it for ordinary gateway tests. @@ -49,13 +66,32 @@ func (s *OpenAIGatewayService) BeginOpenAIWSIngressSessionPreemption( c *gin.Context, account *Account, firstClientMessage []byte, +) (context.Context, func(), bool) { + return s.BeginOpenAIWSIngressSessionPreemptionWithClient(ctx, c, account, firstClientMessage, nil) +} + +// BeginOpenAIWSIngressSessionPreemptionWithClient 与 BeginOpenAIWSIngressSessionPreemption 相同， +// 另外登记客户端连接：本连接被更新的连接取代时，先给它发带原因的关闭帧，再取消上下文， +// 客户端因此能立即重试，而不是在裸断开后静默等到空闲超时。 +func (s *OpenAIGatewayService) BeginOpenAIWSIngressSessionPreemptionWithClient( + ctx context.Context, + c *gin.Context, + account *Account, + firstClientMessage []byte, + clientConn openAIWSPreemptClientCloser, ) (context.Context, func(), bool) { if ctx == nil { ctx = context.Background() } - if armed, _ := ctx.Value(openAIWSSessionPreemptContextKey{}).(bool); armed { + if state, _ := ctx.Value(openAIWSSessionPreemptContextKey{}).(*openAIWSSessionPreemptState); state != nil { return ctx, func() {}, true } + var notifyPreempted func() + if clientConn != nil { + notifyPreempted = func() { + _ = clientConn.Close(coderws.StatusTryAgainLater, openAIWSSessionPreemptedCloseReason) + } + } if s != nil && s.cfg != nil && s.cfg.Gateway.OpenAIWS.ModeRouterV2Enabled && account != nil { switch account.ResolveOpenAIResponsesWebSocketV2Mode(s.cfg.Gateway.OpenAIWS.IngressModeDefault) { case OpenAIWSIngressModePassthrough, OpenAIWSIngressModeHTTPBridge: @@ -85,6 +121,7 @@ func (s *OpenAIGatewayService) BeginOpenAIWSIngressSessionPreemption( preemptAPIKeyID, preemptScope, false, + notifyPreempted, ) if !armed { return ctx, func() {}, false @@ -103,7 +140,7 @@ func (s *OpenAIGatewayService) BeginOpenAIWSIngressSessionPreemption( truncateOpenAIWSLogValue(preemptThreadID, openAIWSIDValueMaxLen), ) } - return context.WithValue(preemptCtx, openAIWSSessionPreemptContextKey{}, true), cleanup, true + return preemptCtx, cleanup, true } func newOpenAIWSSessionPreemptKey(groupID, apiKeyID int64, sessionHash string) (openAIWSSessionPreemptKey, bool) { @@ -161,6 +198,7 @@ func (s *OpenAIGatewayService) beginOpenAIWSSessionPreemptContext( groupID, apiKeyID int64, sessionHash string, httpIngressWSOneShot bool, + notifyPreempted func(), ) (context.Context, func(), bool, bool) { if ctx == nil { ctx = context.Background() @@ -173,16 +211,35 @@ func (s *OpenAIGatewayService) beginOpenAIWSSessionPreemptContext( return ctx, func() {}, false, false } - preemptCtx, cancel := context.WithCancelCause(ctx) + state := &openAIWSSessionPreemptState{} + preemptCtx, cancel := context.WithCancelCause(context.WithValue(ctx, openAIWSSessionPreemptContextKey{}, state)) ownerToken := uuid.NewString() var preemptOnce sync.Once preempt := func() { preemptOnce.Do(func() { + state.preempted.Store(true) if stateStore := s.getOpenAIWSStateStore(); stateStore != nil { stateStore.DeleteSessionTurnState(key.groupID, key.sessionHash) stateStore.DeleteSessionConn(key.groupID, key.sessionHash) } - cancel(errOpenAIWSSessionPreempted) + if notifyPreempted == nil { + cancel(errOpenAIWSSessionPreempted) + return + } + // 取消会让 coder 立即关掉正在读写的客户端连接，关闭帧必须先于取消发出； + // 通知在独立协程里做，不阻塞取代它的新连接。 + notified := make(chan struct{}) + go func() { + defer close(notified) + notifyPreempted() + }() + go func() { + select { + case <-notified: + case <-time.After(openAIWSSessionPreemptCloseGrace): + } + cancel(errOpenAIWSSessionPreempted) + }() }) } previousRemoteOwner, remoteClaimed := s.claimOpenAIWSSessionPreemptOwner(ctx, key, ownerToken) @@ -284,7 +341,13 @@ func (s *OpenAIGatewayService) watchOpenAIWSSessionPreemptOwner(ctx context.Cont } func isOpenAIWSSessionPreempted(ctx context.Context) bool { - return ctx != nil && errors.Is(context.Cause(ctx), errOpenAIWSSessionPreempted) + if ctx == nil { + return false + } + if state, _ := ctx.Value(openAIWSSessionPreemptContextKey{}).(*openAIWSSessionPreemptState); state != nil && state.preempted.Load() { + return true + } + return errors.Is(context.Cause(ctx), errOpenAIWSSessionPreempted) } func IsOpenAIWSSessionPreemptedError(err error) bool { diff --git a/backend/internal/service/openai_ws_session_preemption_test.go b/backend/internal/service/openai_ws_session_preemption_test.go index f424e0f7b2..243cbcab68 100644 --- a/backend/internal/service/openai_ws_session_preemption_test.go +++ b/backend/internal/service/openai_ws_session_preemption_test.go @@ -12,10 +12,105 @@ import ( "time" "github.com/Wei-Shaw/sub2api/internal/config" + coderws "github.com/coder/websocket" "github.com/gin-gonic/gin" "github.com/stretchr/testify/require" ) +type openAIWSPreemptCloseRecorder struct { + mu sync.Mutex + watchCtx context.Context + ctxErrAtSend error + codes []coderws.StatusCode + reason string + release chan struct{} + closed chan struct{} +} + +func (r *openAIWSPreemptCloseRecorder) Close(code coderws.StatusCode, reason string) error { + r.mu.Lock() + r.codes = append(r.codes, code) + r.reason = reason + if r.watchCtx != nil { + r.ctxErrAtSend = r.watchCtx.Err() + } + r.mu.Unlock() + if r.release != nil { + <-r.release + } + close(r.closed) + return nil +} + +func TestOpenAIWSIngressSessionPreemptionSendsCloseFrameBeforeCancel(t *testing.T) { + gin.SetMode(gin.TestMode) + svc := &OpenAIGatewayService{} + account := &Account{ID: 1, Platform: PlatformOpenAI, Type: AccountTypeOAuth} + firstMessage := []byte(`{"type":"response.create","input":"hello"}`) + closer := &openAIWSPreemptCloseRecorder{closed: make(chan struct{})} + + firstCtx, firstCleanup, armed := svc.BeginOpenAIWSIngressSessionPreemptionWithClient( + context.Background(), newOpenAIWSPreemptCodexContext(11, "thread-a"), account, firstMessage, closer, + ) + require.True(t, armed) + defer firstCleanup() + closer.watchCtx = firstCtx + + _, secondCleanup, armed := svc.BeginOpenAIWSIngressSessionPreemptionWithClient( + context.Background(), newOpenAIWSPreemptCodexContext(11, "thread-a"), account, firstMessage, nil, + ) + require.True(t, armed) + defer secondCleanup() + + require.True(t, isOpenAIWSSessionPreempted(firstCtx), "关闭帧发出前旧连接就应被判定为已抢占") + select { + case <-closer.closed: + case <-time.After(time.Second): + t.Fatal("旧客户端应收到关闭帧") + } + closer.mu.Lock() + require.Equal(t, []coderws.StatusCode{coderws.StatusTryAgainLater}, closer.codes) + require.Equal(t, openAIWSSessionPreemptedCloseReason, closer.reason) + require.NoError(t, closer.ctxErrAtSend, "取消必须等关闭帧发出之后") + closer.mu.Unlock() + select { + case <-firstCtx.Done(): + case <-time.After(2 * time.Second): + t.Fatal("旧连接应在关闭帧之后被取消") + } + require.ErrorIs(t, context.Cause(firstCtx), errOpenAIWSSessionPreempted) +} + +func TestOpenAIWSIngressSessionPreemptionCancelsAfterCloseGrace(t *testing.T) { + gin.SetMode(gin.TestMode) + svc := &OpenAIGatewayService{} + account := &Account{ID: 1, Platform: PlatformOpenAI, Type: AccountTypeOAuth} + firstMessage := []byte(`{"type":"response.create","input":"hello"}`) + closer := &openAIWSPreemptCloseRecorder{closed: make(chan struct{}), release: make(chan struct{})} + defer close(closer.release) + + firstCtx, firstCleanup, armed := svc.BeginOpenAIWSIngressSessionPreemptionWithClient( + context.Background(), newOpenAIWSPreemptCodexContext(11, "thread-a"), account, firstMessage, closer, + ) + require.True(t, armed) + defer firstCleanup() + + started := time.Now() + _, secondCleanup, armed := svc.BeginOpenAIWSIngressSessionPreemptionWithClient( + context.Background(), newOpenAIWSPreemptCodexContext(11, "thread-a"), account, firstMessage, nil, + ) + require.True(t, armed) + defer secondCleanup() + require.Less(t, time.Since(started), openAIWSSessionPreemptCloseGrace, "新连接不得等待旧客户端的关闭握手") + + select { + case <-firstCtx.Done(): + case <-time.After(openAIWSSessionPreemptCloseGrace + time.Second): + t.Fatal("对端不回应关闭帧时也必须在宽限期内取消旧连接") + } + require.ErrorIs(t, context.Cause(firstCtx), errOpenAIWSSessionPreempted) +} + func TestOpenAIWSSessionPreemptRegistryCancelsSameScopedSessionOnly(t *testing.T) { var registry openAIWSSessionPreemptRegistry key := openAIWSSessionPreemptKey{groupID: 7, apiKeyID: 11, sessionHash: "sess"} @@ -86,22 +181,22 @@ func TestOpenAIWSSessionPreemptContextEligibilityAndLocalCancellation(t *testing apiKey := &Account{ID: 2, Platform: PlatformOpenAI, Type: AccountTypeAPIKey} grok := &Account{ID: 3, Platform: PlatformGrok, Type: AccountTypeOAuth} - _, cleanup, armed, _ := svc.beginOpenAIWSSessionPreemptContext(context.Background(), apiKey, 7, 11, "sess", false) + _, cleanup, armed, _ := svc.beginOpenAIWSSessionPreemptContext(context.Background(), apiKey, 7, 11, "sess", false, nil) cleanup() require.False(t, armed) - _, cleanup, armed, _ = svc.beginOpenAIWSSessionPreemptContext(context.Background(), grok, 7, 11, "sess", false) + _, cleanup, armed, _ = svc.beginOpenAIWSSessionPreemptContext(context.Background(), grok, 7, 11, "sess", false, nil) cleanup() require.False(t, armed) - _, cleanup, armed, _ = svc.beginOpenAIWSSessionPreemptContext(context.Background(), oauth, 7, 11, "sess", true) + _, cleanup, armed, _ = svc.beginOpenAIWSSessionPreemptContext(context.Background(), oauth, 7, 11, "sess", true, nil) cleanup() require.False(t, armed, "HTTP-ingress one-shot must not participate") - firstCtx, firstCleanup, armed, replaced := svc.beginOpenAIWSSessionPreemptContext(context.Background(), oauth, 7, 11, "sess", false) + firstCtx, firstCleanup, armed, replaced := svc.beginOpenAIWSSessionPreemptContext(context.Background(), oauth, 7, 11, "sess", false, nil) require.True(t, armed) require.False(t, replaced) stateStore.BindSessionTurnState(7, "sess", "turn-state", time.Hour) stateStore.BindSessionConn(7, "sess", "conn-1", time.Hour) - _, secondCleanup, armed, replaced := svc.beginOpenAIWSSessionPreemptContext(context.Background(), oauth, 7, 11, "sess", false) + _, secondCleanup, armed, replaced := svc.beginOpenAIWSSessionPreemptContext(context.Background(), oauth, 7, 11, "sess", false, nil) require.True(t, armed) require.True(t, replaced) require.True(t, isOpenAIWSSessionPreempted(firstCtx))