fix(openai): WS 会话被抢占时先发关闭帧再取消

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

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() {

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) {

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 {

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))