From 06b61ba96f1dffdad0e2387669f1df6cf901f7be Mon Sep 17 00:00:00 2001 From: xlplbo Date: Thu, 10 Sep 2026 17:50:23 +0800 Subject: [PATCH 1/7] =?UTF-8?q?feat(openai):=20=E8=A7=A3=E6=9E=90=20codex?= =?UTF-8?q?=20=E7=BA=BF=E7=A8=8B=E6=A0=87=E8=AF=86=E6=B4=BE=E7=94=9F=20WS?= =?UTF-8?q?=20=E6=89=A7=E8=A1=8C=E4=BD=9C=E7=94=A8=E5=9F=9F?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- .../service/openai_ws_execution_scope.go | 72 ++++++++++ .../service/openai_ws_execution_scope_test.go | 126 ++++++++++++++++++ 2 files changed, 198 insertions(+) create mode 100644 backend/internal/service/openai_ws_execution_scope.go create mode 100644 backend/internal/service/openai_ws_execution_scope_test.go diff --git a/backend/internal/service/openai_ws_execution_scope.go b/backend/internal/service/openai_ws_execution_scope.go new file mode 100644 index 0000000000..730e86cb01 --- /dev/null +++ b/backend/internal/service/openai_ws_execution_scope.go @@ -0,0 +1,72 @@ +package service + +import ( + "encoding/json" + "fmt" + "strings" + + "github.com/gin-gonic/gin" + "github.com/tidwall/gjson" +) + +const ( + openAIWSThreadIDHeader = "thread-id" + openAIWSWindowIDHeader = "x-codex-window-id" +) + +// resolveOpenAIWSClientThreadID 提取 codex 客户端的线程标识。codex 多智能体会话里 +// 父线程与全部子智能体共用同一个 session-id 头，只有线程标识能把它们区分开。 +// 优先级：thread-id 头 → x-codex-turn-metadata 头 → x-codex-window-id 头 → 请求体 client_metadata。 +func resolveOpenAIWSClientThreadID(c *gin.Context, body []byte) string { + if c != nil && c.Request != nil { + if id := strings.TrimSpace(c.GetHeader(openAIWSThreadIDHeader)); id != "" { + return id + } + if id := codexTurnMetadataThreadID(c.GetHeader(openAIWSTurnMetadataHeader)); id != "" { + return id + } + if window := strings.TrimSpace(c.GetHeader(openAIWSWindowIDHeader)); window != "" { + if id := strings.TrimSpace(strings.SplitN(window, ":", 2)[0]); id != "" { + return id + } + } + } + if len(body) == 0 { + return "" + } + if id := strings.TrimSpace(gjson.GetBytes(body, "client_metadata.thread_id").String()); id != "" { + return id + } + return codexTurnMetadataThreadID(gjson.GetBytes(body, "client_metadata."+openAIWSTurnMetadataHeader).String()) +} + +func codexTurnMetadataThreadID(raw string) string { + raw = strings.TrimSpace(raw) + if raw == "" { + return "" + } + var metadata struct { + ThreadID string `json:"thread_id"` + } + if err := json.Unmarshal([]byte(raw), &metadata); err != nil { + return "" + } + return strings.TrimSpace(metadata.ThreadID) +} + +// resolveOpenAIWSExecutionScope 派生 WS 接入用于抢占与会话级状态的执行作用域键， +// 并一并返回解析到的线程标识供日志使用。身份只取客户端声明的线程标识或显式会话标识， +// 不掺入请求内容，因此同线程重连在输入变化后仍命中同一键；键里带 API key，同组不同 key +// 互不影响。两种身份都没有时返回空作用域：按请求内容推导的粘性种子只服务账号亲和， +// 不是可靠身份，不参与抢占。 +func resolveOpenAIWSExecutionScope(c *gin.Context, body []byte, apiKeyID int64) (scope, threadID string) { + if threadID = resolveOpenAIWSClientThreadID(c, body); threadID != "" { + scope, _ = deriveOpenAISessionHashes(fmt.Sprintf("openai_ws_exec:%d|thread=%s", apiKeyID, threadID)) + return scope, threadID + } + if sessionID := strings.TrimSpace(explicitOpenAIRequestSessionID(c, body)); sessionID != "" { + scope, _ = deriveOpenAISessionHashes(fmt.Sprintf("openai_ws_exec:%d|session=%s", apiKeyID, sessionID)) + return scope, "" + } + return "", "" +} diff --git a/backend/internal/service/openai_ws_execution_scope_test.go b/backend/internal/service/openai_ws_execution_scope_test.go new file mode 100644 index 0000000000..349b4cfb34 --- /dev/null +++ b/backend/internal/service/openai_ws_execution_scope_test.go @@ -0,0 +1,126 @@ +package service + +import ( + "net/http" + "net/http/httptest" + "testing" + + "github.com/gin-gonic/gin" + "github.com/stretchr/testify/require" +) + +func newOpenAIWSExecutionScopeTestContext(headers map[string]string) *gin.Context { + c, _ := gin.CreateTestContext(httptest.NewRecorder()) + c.Request = httptest.NewRequest(http.MethodGet, "/v1/responses", nil) + for k, v := range headers { + c.Request.Header.Set(k, v) + } + return c +} + +func TestResolveOpenAIWSClientThreadID(t *testing.T) { + gin.SetMode(gin.TestMode) + cases := []struct { + name string + headers map[string]string + body string + want string + }{ + { + name: "thread-id header wins", + headers: map[string]string{ + openAIWSThreadIDHeader: " t-plain ", + openAIWSTurnMetadataHeader: `{"session_id":"s","thread_id":"t-header"}`, + openAIWSWindowIDHeader: "t-window:0", + }, + body: `{"client_metadata":{"thread_id":"t-body"}}`, + want: "t-plain", + }, + { + name: "turn metadata header second", + headers: map[string]string{ + openAIWSTurnMetadataHeader: `{"session_id":"s","thread_id":"t-header"}`, + openAIWSWindowIDHeader: "t-window:0", + }, + body: `{"client_metadata":{"thread_id":"t-body"}}`, + want: "t-header", + }, + { + name: "window id fallback", + headers: map[string]string{openAIWSWindowIDHeader: " t-window:0 "}, + body: `{"client_metadata":{"thread_id":"t-body"}}`, + want: "t-window", + }, + { + name: "body client_metadata.thread_id", + body: `{"client_metadata":{"thread_id":"t-body"}}`, + want: "t-body", + }, + { + name: "body embedded turn metadata", + body: `{"client_metadata":{"x-codex-turn-metadata":"{\"thread_id\":\"t-embedded\"}"}}`, + want: "t-embedded", + }, + { + name: "malformed header ignored", + headers: map[string]string{openAIWSTurnMetadataHeader: "{not json"}, + body: `{}`, + want: "", + }, + { + name: "none", + body: `{"input":"hi"}`, + want: "", + }, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + got := resolveOpenAIWSClientThreadID(newOpenAIWSExecutionScopeTestContext(tc.headers), []byte(tc.body)) + require.Equal(t, tc.want, got) + }) + } + require.Equal(t, "t", resolveOpenAIWSClientThreadID(nil, []byte(`{"client_metadata":{"thread_id":"t"}}`))) +} + +func TestResolveOpenAIWSExecutionScope(t *testing.T) { + gin.SetMode(gin.TestMode) + threadHeaders := map[string]string{ + "session-id": "root-session", + openAIWSTurnMetadataHeader: `{"session_id":"root-session","thread_id":"thread-a"}`, + } + bodyFirst := []byte(`{"type":"response.create","model":"gpt-5.1","instructions":"sys","input":[{"role":"user","content":"first"}]}`) + bodyRetry := []byte(`{"type":"response.create","model":"gpt-5.1","instructions":"sys","input":[{"role":"user","content":"first"},{"role":"assistant","content":"x"},{"role":"user","content":"second"}]}`) + + scopeOf := func(headers map[string]string, body []byte, apiKeyID int64) string { + scope, _ := resolveOpenAIWSExecutionScope(newOpenAIWSExecutionScopeTestContext(headers), body, apiKeyID) + return scope + } + + byThread, threadID := resolveOpenAIWSExecutionScope(newOpenAIWSExecutionScopeTestContext(threadHeaders), bodyFirst, 11) + require.Len(t, byThread, 16) + require.Equal(t, "thread-a", threadID) + require.Equal(t, byThread, scopeOf(threadHeaders, bodyRetry, 11), "同线程重连的作用域不得随请求内容变化") + + otherThread := map[string]string{ + "session-id": "root-session", + openAIWSTurnMetadataHeader: `{"session_id":"root-session","thread_id":"thread-b"}`, + } + require.NotEqual(t, byThread, scopeOf(otherThread, bodyFirst, 11)) + require.NotEqual(t, byThread, scopeOf(threadHeaders, bodyFirst, 12), "API key 必须参与隔离") + + sessionOnly := map[string]string{"session-id": "root-session"} + bySession, threadID := resolveOpenAIWSExecutionScope(newOpenAIWSExecutionScopeTestContext(sessionOnly), bodyFirst, 11) + require.Len(t, bySession, 16) + require.Equal(t, "", threadID) + require.NotEqual(t, byThread, bySession, "线程标识优先于会话标识") + require.Equal(t, bySession, scopeOf(sessionOnly, bodyRetry, 11)) + require.NotEqual(t, bySession, scopeOf(sessionOnly, bodyFirst, 12)) + + byCacheKey := scopeOf(nil, []byte(`{"type":"response.create","prompt_cache_key":"cache-1","input":"hi"}`), 11) + require.Len(t, byCacheKey, 16) + + require.Equal(t, "", scopeOf(nil, bodyFirst, 11), "只有请求内容时不得产生执行作用域") + nilScope, nilThread := resolveOpenAIWSExecutionScope(nil, bodyFirst, 11) + require.Equal(t, "", nilScope) + require.Equal(t, "", nilThread) +} From 9e7c2d713c8510fe4a2147a7db424747092a8aed Mon Sep 17 00:00:00 2001 From: xlplbo Date: Thu, 10 Sep 2026 17:50:24 +0800 Subject: [PATCH 2/7] =?UTF-8?q?fix(openai):=20WS=20=E4=BC=9A=E8=AF=9D?= =?UTF-8?q?=E6=8A=A2=E5=8D=A0=E6=94=B9=E7=94=A8=E6=89=A7=E8=A1=8C=E4=BD=9C?= =?UTF-8?q?=E7=94=A8=E5=9F=9F=E9=94=AE?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- .../service/openai_ws_session_preemption.go | 27 +++-- .../openai_ws_session_preemption_test.go | 98 +++++++++++++++++++ 2 files changed, 119 insertions(+), 6 deletions(-) diff --git a/backend/internal/service/openai_ws_session_preemption.go b/backend/internal/service/openai_ws_session_preemption.go index 95fa49e7f7..22ed8fd4f4 100644 --- a/backend/internal/service/openai_ws_session_preemption.go +++ b/backend/internal/service/openai_ws_session_preemption.go @@ -66,17 +66,24 @@ func (s *OpenAIGatewayService) BeginOpenAIWSIngressSessionPreemption( } } - preemptSessionHash := "" + preemptScope := "" + preemptThreadID := "" preemptGroupID := getOpenAIGroupIDFromContext(c) + preemptAPIKeyID := getAPIKeyIDFromContext(c) if account != nil && account.Platform == PlatformOpenAI && account.Type == AccountTypeOAuth { - preemptSessionHash = s.GenerateSessionHash(c, firstClientMessage) + // Codex multi-agent sessions share one session-id across the parent thread + // and every sub-agent. Key preemption by the declared execution identity + // (thread first, explicit session otherwise) so siblings never cancel each + // other while a reconnect of the same thread still replaces it. Requests + // without any declared identity are not registered at all. + preemptScope, preemptThreadID = resolveOpenAIWSExecutionScope(c, firstClientMessage, preemptAPIKeyID) } preemptCtx, cleanup, armed, preemptedPrevious := s.beginOpenAIWSSessionPreemptContext( ctx, account, preemptGroupID, - getAPIKeyIDFromContext(c), - preemptSessionHash, + preemptAPIKeyID, + preemptScope, false, ) if !armed { @@ -84,9 +91,17 @@ func (s *OpenAIGatewayService) BeginOpenAIWSIngressSessionPreemption( } if preemptedPrevious { if stateStore := s.getOpenAIWSStateStore(); stateStore != nil { - stateStore.DeleteSessionTurnState(preemptGroupID, preemptSessionHash) - stateStore.DeleteSessionConn(preemptGroupID, preemptSessionHash) + stateStore.DeleteSessionTurnState(preemptGroupID, preemptScope) + stateStore.DeleteSessionConn(preemptGroupID, preemptScope) } + logOpenAIWSModeInfo( + "ingress_ws_session_preempted account_id=%d group_id=%d api_key_id=%d scope=%s thread_id=%s", + account.ID, + preemptGroupID, + preemptAPIKeyID, + truncateOpenAIWSLogValue(preemptScope, 12), + truncateOpenAIWSLogValue(preemptThreadID, openAIWSIDValueMaxLen), + ) } return context.WithValue(preemptCtx, openAIWSSessionPreemptContextKey{}, true), cleanup, true } diff --git a/backend/internal/service/openai_ws_session_preemption_test.go b/backend/internal/service/openai_ws_session_preemption_test.go index d00adf63b4..f424e0f7b2 100644 --- a/backend/internal/service/openai_ws_session_preemption_test.go +++ b/backend/internal/service/openai_ws_session_preemption_test.go @@ -288,3 +288,101 @@ func TestNewOpenAIWSSessionPreemptKeyRequiresFullIsolationScope(t *testing.T) { require.True(t, ok) require.Equal(t, "sess", key.sessionHash) } + +func newOpenAIWSPreemptCodexContext(apiKeyID int64, threadID string) *gin.Context { + groupID := int64(7) + c, _ := gin.CreateTestContext(httptest.NewRecorder()) + c.Request = httptest.NewRequest(http.MethodGet, "/v1/responses", nil) + c.Request.Header.Set("session-id", "root-session") + if threadID != "" { + c.Request.Header.Set(openAIWSTurnMetadataHeader, `{"session_id":"root-session","thread_id":"`+threadID+`"}`) + } + c.Set("api_key", &APIKey{ID: apiKeyID, GroupID: &groupID}) + return c +} + +func TestOpenAIWSIngressSessionPreemptionIsolatesCodexThreads(t *testing.T) { + gin.SetMode(gin.TestMode) + svc := &OpenAIGatewayService{} + account := &Account{ID: 1, Platform: PlatformOpenAI, Type: AccountTypeOAuth} + firstMessage := []byte(`{"type":"response.create","input":"hello"}`) + retryMessage := []byte(`{"type":"response.create","input":[{"role":"user","content":"hello"},{"role":"user","content":"again"}]}`) + + rootCtx, rootCleanup, armed := svc.BeginOpenAIWSIngressSessionPreemption( + context.Background(), newOpenAIWSPreemptCodexContext(11, "root-session"), account, firstMessage, + ) + require.True(t, armed) + defer rootCleanup() + + childCtx, childCleanup, armed := svc.BeginOpenAIWSIngressSessionPreemption( + context.Background(), newOpenAIWSPreemptCodexContext(11, "child-thread"), account, firstMessage, + ) + require.True(t, armed) + defer childCleanup() + require.NoError(t, rootCtx.Err(), "same session but different codex thread must not preempt") + require.NoError(t, childCtx.Err()) + + otherKeyCtx, otherKeyCleanup, armed := svc.BeginOpenAIWSIngressSessionPreemption( + context.Background(), newOpenAIWSPreemptCodexContext(12, "root-session"), account, firstMessage, + ) + require.True(t, armed) + defer otherKeyCleanup() + require.NoError(t, rootCtx.Err(), "same thread under another api key must not preempt") + require.NoError(t, otherKeyCtx.Err()) + + _, rootRetryCleanup, armed := svc.BeginOpenAIWSIngressSessionPreemption( + context.Background(), newOpenAIWSPreemptCodexContext(11, "root-session"), account, retryMessage, + ) + require.True(t, armed) + defer rootRetryCleanup() + require.True(t, IsOpenAIWSSessionPreemptedError(context.Cause(rootCtx)), "same thread reconnect with a different body must still preempt") + require.NoError(t, childCtx.Err(), "root reconnect must leave the child thread alone") + require.NoError(t, otherKeyCtx.Err()) +} + +func TestOpenAIWSIngressSessionPreemptionSkipsContentOnlyIdentity(t *testing.T) { + gin.SetMode(gin.TestMode) + groupID := int64(7) + newContext := func() *gin.Context { + c, _ := gin.CreateTestContext(httptest.NewRecorder()) + c.Request = httptest.NewRequest(http.MethodGet, "/v1/responses", nil) + c.Set("api_key", &APIKey{ID: 11, GroupID: &groupID}) + return c + } + svc := &OpenAIGatewayService{} + account := &Account{ID: 1, Platform: PlatformOpenAI, Type: AccountTypeOAuth} + contentOnly := []byte(`{"type":"response.create","model":"gpt-5.1","instructions":"sys","input":[{"role":"user","content":"same prompt"}]}`) + + firstCtx, firstCleanup, armed := svc.BeginOpenAIWSIngressSessionPreemption(context.Background(), newContext(), account, contentOnly) + require.False(t, armed, "content-derived seed is not an identity and must not arm preemption") + defer firstCleanup() + _, secondCleanup, armed := svc.BeginOpenAIWSIngressSessionPreemption(context.Background(), newContext(), account, contentOnly) + require.False(t, armed) + defer secondCleanup() + require.NoError(t, firstCtx.Err()) +} + +func TestOpenAIWSIngressSessionPreemptionClaimsRemoteOwnerByExecutionScope(t *testing.T) { + gin.SetMode(gin.TestMode) + stub := &openAIWSSessionPreemptCacheStub{} + svc := &OpenAIGatewayService{cache: stub} + account := &Account{ID: 1, Platform: PlatformOpenAI, Type: AccountTypeOAuth} + firstMessage := []byte(`{"type":"response.create","input":"hello"}`) + c := newOpenAIWSPreemptCodexContext(11, "thread-a") + + _, cleanup, armed := svc.BeginOpenAIWSIngressSessionPreemption(context.Background(), c, account, firstMessage) + require.True(t, armed) + defer cleanup() + + scope, _ := resolveOpenAIWSExecutionScope(c, firstMessage, 11) + legacy := svc.GenerateSessionHash(c, firstMessage) + require.NotEmpty(t, scope) + require.NotEqual(t, legacy, scope) + + stub.mu.Lock() + defer stub.mu.Unlock() + _, scoped := stub.owners[stub.key(7, openAIWSSessionPreemptCacheHash(11, scope))] + _, legacyKeyed := stub.owners[stub.key(7, openAIWSSessionPreemptCacheHash(11, legacy))] + require.True(t, scoped, "remote owner must be claimed under the execution scope") + require.False(t, legacyKeyed, "remote owner must not be claimed under the legacy session hash") +} From cd1ee1d1ad735888ed63ae72b21ecbe9105ec8a8 Mon Sep 17 00:00:00 2001 From: xlplbo Date: Thu, 10 Sep 2026 17:50:24 +0800 Subject: [PATCH 3/7] =?UTF-8?q?fix(openai):=20WS=20=E6=8E=A5=E5=85=A5?= =?UTF-8?q?=E4=BC=9A=E8=AF=9D=E7=8A=B6=E6=80=81=E6=94=B9=E7=94=A8=E6=89=A7?= =?UTF-8?q?=E8=A1=8C=E4=BD=9C=E7=94=A8=E5=9F=9F=E9=94=AE?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- .../service/openai_ws_forwarder_ingress.go | 6 + ..._forwarder_ingress_execution_scope_test.go | 351 ++++++++++++++++++ 2 files changed, 357 insertions(+) create mode 100644 backend/internal/service/openai_ws_forwarder_ingress_execution_scope_test.go diff --git a/backend/internal/service/openai_ws_forwarder_ingress.go b/backend/internal/service/openai_ws_forwarder_ingress.go index f24dc25697..0def1bd0cb 100644 --- a/backend/internal/service/openai_ws_forwarder_ingress.go +++ b/backend/internal/service/openai_ws_forwarder_ingress.go @@ -525,12 +525,18 @@ func (s *OpenAIGatewayService) ProxyResponsesWebSocketFromClient( turnState := strings.TrimSpace(c.GetHeader(openAIWSTurnStateHeader)) stateStore := s.getOpenAIWSStateStore() groupID := getOpenAIGroupIDFromContext(c) + apiKeyID := getAPIKeyIDFromContext(c) storeDisabledConnMode := s.openAIWSStoreDisabledConnMode() sessionHash := "" preferredConnID := "" storeDisabled := false refreshIngressRouteState := func(payload openAIWSClientPayload) { + // 会话级状态按执行作用域隔离：codex 多智能体共用 session-id，只有线程标识能把 + // 父线程与子智能体区分开；没有声明身份时沿用原会话哈希。账号粘性仍由 handler 决定。 sessionHash = s.GenerateSessionHash(c, payload.rawForHash) + if scope, _ := resolveOpenAIWSExecutionScope(c, payload.rawForHash, apiKeyID); scope != "" { + sessionHash = scope + } preferredConnID = "" storeDisabled = s.isOpenAIWSStoreDisabledInRequestRaw(payload.payloadRaw, account) if useHTTPBridge { 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 new file mode 100644 index 0000000000..ef2594535d --- /dev/null +++ b/backend/internal/service/openai_ws_forwarder_ingress_execution_scope_test.go @@ -0,0 +1,351 @@ +package service + +import ( + "context" + "io" + "net/http" + "net/http/httptest" + "strings" + "sync" + "testing" + "time" + + "github.com/Wei-Shaw/sub2api/internal/config" + coderws "github.com/coder/websocket" + "github.com/gin-gonic/gin" + "github.com/stretchr/testify/require" + "github.com/tidwall/gjson" +) + +func newOpenAIWSExecutionScopeTestConfig() *config.Config { + cfg := &config.Config{} + cfg.Security.URLAllowlist.Enabled = false + cfg.Security.URLAllowlist.AllowInsecureHTTP = true + cfg.Gateway.OpenAIWS.Enabled = true + cfg.Gateway.OpenAIWS.OAuthEnabled = true + cfg.Gateway.OpenAIWS.APIKeyEnabled = true + cfg.Gateway.OpenAIWS.ResponsesWebsocketsV2 = true + cfg.Gateway.OpenAIWS.MaxConnsPerAccount = 2 + cfg.Gateway.OpenAIWS.MinIdlePerAccount = 0 + cfg.Gateway.OpenAIWS.MaxIdlePerAccount = 2 + cfg.Gateway.OpenAIWS.QueueLimitPerConn = 8 + cfg.Gateway.OpenAIWS.DialTimeoutSeconds = 3 + cfg.Gateway.OpenAIWS.ReadTimeoutSeconds = 3 + cfg.Gateway.OpenAIWS.WriteTimeoutSeconds = 3 + return cfg +} + +func dialOpenAIWSExecutionScopeClient(t *testing.T, wsURL, threadID string) *coderws.Conn { + t.Helper() + header := http.Header{} + header.Set("session-id", "root-session") + if threadID != "" { + header.Set(openAIWSTurnMetadataHeader, `{"session_id":"root-session","thread_id":"`+threadID+`"}`) + } + dialCtx, cancel := context.WithTimeout(context.Background(), 3*time.Second) + defer cancel() + conn, _, err := coderws.Dial(dialCtx, "ws"+strings.TrimPrefix(wsURL, "http"), &coderws.DialOptions{HTTPHeader: header}) + require.NoError(t, err) + return conn +} + +func writeOpenAIWSExecutionScopeRequest(t *testing.T, conn *coderws.Conn, body string) { + t.Helper() + ctx, cancel := context.WithTimeout(context.Background(), 3*time.Second) + defer cancel() + require.NoError(t, conn.Write(ctx, coderws.MessageText, []byte(body))) +} + +// 场景：codex 子智能体与父线程共用 session-id 头，只有 x-codex-turn-metadata 的 thread_id 不同。 +// ctx_pool 下的 turn state 绑定与 store=false 的上游连接绑定必须落在执行作用域键下， +// 不能落在按 session-id 算出的会话哈希下，否则子智能体会覆盖父线程的绑定。 +func TestOpenAIGatewayService_ProxyResponsesWebSocketFromClient_StateBoundToExecutionScope(t *testing.T) { + gin.SetMode(gin.TestMode) + cfg := newOpenAIWSExecutionScopeTestConfig() + + captureConn := &openAIWSCaptureConn{ + events: [][]byte{ + []byte(`{"type":"response.completed","response":{"id":"resp_exec_scope","model":"gpt-5.1","usage":{"input_tokens":1,"output_tokens":1}}}`), + }, + } + handshake := http.Header{} + handshake.Set(openAIWSTurnStateHeader, "turn-state-from-upstream") + pool := newOpenAIWSConnPool(cfg) + pool.setClientDialerForTest(&openAIWSCaptureDialer{conn: captureConn, handshake: handshake}) + defer pool.Close() + + stateStore := NewOpenAIWSStateStore(nil) + svc := &OpenAIGatewayService{ + cfg: cfg, + httpUpstream: &httpUpstreamRecorder{}, + cache: &stubGatewayCache{}, + openaiWSResolver: NewOpenAIWSProtocolResolver(cfg), + toolCorrector: NewCodexToolCorrector(), + openaiWSPool: pool, + openaiWSStateStore: stateStore, + } + groupID := int64(9) + account := &Account{ + ID: 454, + Name: "openai-ingress-exec-scope", + Platform: PlatformOpenAI, + Type: AccountTypeAPIKey, + Status: StatusActive, + Schedulable: true, + Concurrency: 1, + Credentials: map[string]any{"api_key": "sk-test"}, + Extra: map[string]any{"responses_websockets_v2_enabled": true}, + } + + serverErrCh := make(chan error, 1) + keysCh := make(chan [2]string, 1) + wsServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + conn, err := coderws.Accept(w, r, &coderws.AcceptOptions{CompressionMode: coderws.CompressionContextTakeover}) + if err != nil { + serverErrCh <- err + return + } + defer func() { _ = conn.CloseNow() }() + rec := httptest.NewRecorder() + ginCtx, _ := gin.CreateTestContext(rec) + req := r.Clone(r.Context()) + req.Header = req.Header.Clone() + req.Header.Set("User-Agent", "unit-test-agent/1.0") + ginCtx.Request = req + ginCtx.Set("api_key", &APIKey{ID: 21, GroupID: &groupID}) + readCtx, cancel := context.WithTimeout(r.Context(), 3*time.Second) + _, firstMessage, readErr := conn.Read(readCtx) + cancel() + if readErr != nil { + serverErrCh <- readErr + return + } + scope, _ := resolveOpenAIWSExecutionScope(ginCtx, firstMessage, 21) + keysCh <- [2]string{svc.GenerateSessionHash(ginCtx, firstMessage), scope} + serverErrCh <- svc.ProxyResponsesWebSocketFromClient(r.Context(), ginCtx, conn, account, "sk-test", firstMessage, nil) + })) + defer wsServer.Close() + + clientConn := dialOpenAIWSExecutionScopeClient(t, wsServer.URL, "child-thread") + defer func() { _ = clientConn.CloseNow() }() + writeOpenAIWSExecutionScopeRequest(t, clientConn, `{"type":"response.create","model":"gpt-5.1","stream":false,"store":false,"input":[{"role":"user","content":"first"}]}`) + + readCtx, cancelRead := context.WithTimeout(context.Background(), 3*time.Second) + _, completed, readErr := clientConn.Read(readCtx) + cancelRead() + require.NoError(t, readErr) + require.Equal(t, "resp_exec_scope", gjson.GetBytes(completed, "response.id").String()) + require.NoError(t, clientConn.Close(coderws.StatusNormalClosure, "done")) + + select { + case serverErr := <-serverErrCh: + require.NoError(t, serverErr) + case <-time.After(5 * time.Second): + t.Fatal("等待 ingress websocket 结束超时") + } + keys := <-keysCh + legacyHash, scope := keys[0], keys[1] + require.NotEmpty(t, legacyHash) + require.Len(t, scope, 16) + require.NotEqual(t, legacyHash, scope) + + _, connBoundToLegacy := stateStore.GetSessionConn(groupID, legacyHash) + require.False(t, connBoundToLegacy, "上游连接不得绑定到按 session-id 算出的会话哈希") + _, connBoundToScope := stateStore.GetSessionConn(groupID, scope) + require.True(t, connBoundToScope, "上游连接应绑定到执行作用域") + _, turnStateOnLegacy := stateStore.GetSessionTurnState(groupID, legacyHash) + require.False(t, turnStateOnLegacy, "turn state 不得绑定到会话哈希") + turnState, turnStateOnScope := stateStore.GetSessionTurnState(groupID, scope) + require.True(t, turnStateOnScope, "turn state 应绑定到执行作用域") + require.Equal(t, "turn-state-from-upstream", turnState) +} + +// openAIWSGatedConn 是只回一个事件的假上游连接：收到请求后关闭 sent， +// 事件要等 gate 打开才吐出，让用例精确控制“A 在飞时 B 接入”的先后。 +type openAIWSGatedConn struct { + mu sync.Mutex + event []byte + sent chan struct{} + gate chan struct{} + closed bool + consumed bool +} + +func newOpenAIWSGatedConn(event string) *openAIWSGatedConn { + return &openAIWSGatedConn{event: []byte(event), sent: make(chan struct{}), gate: make(chan struct{})} +} + +func (c *openAIWSGatedConn) WriteJSON(ctx context.Context, value any) error { + _ = ctx + _ = value + c.mu.Lock() + defer c.mu.Unlock() + if c.closed { + return errOpenAIWSConnClosed + } + select { + case <-c.sent: + default: + close(c.sent) + } + return nil +} + +func (c *openAIWSGatedConn) ReadMessage(ctx context.Context) ([]byte, error) { + c.mu.Lock() + if c.closed { + c.mu.Unlock() + return nil, errOpenAIWSConnClosed + } + if c.consumed { + c.mu.Unlock() + return nil, io.EOF + } + c.consumed = true + c.mu.Unlock() + select { + case <-ctx.Done(): + return nil, ctx.Err() + case <-c.gate: + } + return c.event, nil +} + +func (c *openAIWSGatedConn) Ping(ctx context.Context) error { + _ = ctx + return nil +} + +func (c *openAIWSGatedConn) Close() error { + c.mu.Lock() + defer c.mu.Unlock() + c.closed = true + return nil +} + +// runOpenAIWSCodexThreadPair 用 OAuth 账号（抢占只对 OAuth ctx_pool 生效）跑两条并发接入： +// A 的请求发到上游后 B 才接入并立即完成，B 完成后才放行 A 的上游事件。 +// 返回 A 与 B 的服务端返回值、A 客户端读结果的错误。 +func runOpenAIWSCodexThreadPair(t *testing.T, threadA, threadB string) (serverErrs []error, aReadErr error) { + t.Helper() + gin.SetMode(gin.TestMode) + cfg := newOpenAIWSExecutionScopeTestConfig() + + gatedConn := newOpenAIWSGatedConn(`{"type":"response.completed","response":{"id":"resp_thread_a","model":"gpt-5.1","usage":{"input_tokens":1,"output_tokens":1}}}`) + fastConn := &openAIWSCaptureConn{ + events: [][]byte{ + []byte(`{"type":"response.completed","response":{"id":"resp_thread_b","model":"gpt-5.1","usage":{"input_tokens":1,"output_tokens":1}}}`), + }, + } + pool := newOpenAIWSConnPool(cfg) + pool.setClientDialerForTest(&openAIWSQueueDialer{conns: []openAIWSClientConn{gatedConn, fastConn}}) + defer pool.Close() + + svc := &OpenAIGatewayService{ + cfg: cfg, + httpUpstream: &httpUpstreamRecorder{}, + cache: &stubGatewayCache{}, + openaiWSResolver: NewOpenAIWSProtocolResolver(cfg), + toolCorrector: NewCodexToolCorrector(), + openaiWSPool: pool, + openaiWSStateStore: NewOpenAIWSStateStore(nil), + } + groupID := int64(9) + account := &Account{ + ID: 455, + Name: "openai-ingress-thread-pair", + Platform: PlatformOpenAI, + Type: AccountTypeOAuth, + Status: StatusActive, + Schedulable: true, + Concurrency: 2, + Credentials: map[string]any{"access_token": "test-token"}, + Extra: map[string]any{"openai_oauth_responses_websockets_v2_enabled": true}, + } + + serverErrCh := make(chan error, 2) + wsServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + conn, err := coderws.Accept(w, r, &coderws.AcceptOptions{CompressionMode: coderws.CompressionContextTakeover}) + if err != nil { + serverErrCh <- err + return + } + defer func() { _ = conn.CloseNow() }() + rec := httptest.NewRecorder() + ginCtx, _ := gin.CreateTestContext(rec) + req := r.Clone(r.Context()) + req.Header = req.Header.Clone() + req.Header.Set("User-Agent", "unit-test-agent/1.0") + ginCtx.Request = req + ginCtx.Set("api_key", &APIKey{ID: 21, GroupID: &groupID}) + readCtx, cancel := context.WithTimeout(r.Context(), 3*time.Second) + _, firstMessage, readErr := conn.Read(readCtx) + cancel() + if readErr != nil { + serverErrCh <- readErr + return + } + serverErrCh <- svc.ProxyResponsesWebSocketFromClient(r.Context(), ginCtx, conn, account, "test-token", firstMessage, nil) + })) + defer wsServer.Close() + + connA := dialOpenAIWSExecutionScopeClient(t, wsServer.URL, threadA) + defer func() { _ = connA.CloseNow() }() + writeOpenAIWSExecutionScopeRequest(t, connA, `{"type":"response.create","model":"gpt-5.1","stream":false,"input":[{"role":"user","content":"thread a"}]}`) + select { + case <-gatedConn.sent: + case <-time.After(3 * time.Second): + t.Fatal("A 的请求应先到达上游") + } + + connB := dialOpenAIWSExecutionScopeClient(t, wsServer.URL, threadB) + defer func() { _ = connB.CloseNow() }() + writeOpenAIWSExecutionScopeRequest(t, connB, `{"type":"response.create","model":"gpt-5.1","stream":false,"input":[{"role":"user","content":"thread b"}]}`) + readCtxB, cancelB := context.WithTimeout(context.Background(), 3*time.Second) + _, completedB, readErrB := connB.Read(readCtxB) + cancelB() + require.NoError(t, readErrB, "B 必须正常完成") + require.Equal(t, "resp_thread_b", gjson.GetBytes(completedB, "response.id").String()) + close(gatedConn.gate) + + readCtxA, cancelA := context.WithTimeout(context.Background(), 5*time.Second) + _, completedA, aReadErr := connA.Read(readCtxA) + cancelA() + if aReadErr == nil { + require.Equal(t, "resp_thread_a", gjson.GetBytes(completedA, "response.id").String()) + require.NoError(t, connA.Close(coderws.StatusNormalClosure, "done")) + } + require.NoError(t, connB.Close(coderws.StatusNormalClosure, "done")) + + for i := 0; i < 2; i++ { + select { + case err := <-serverErrCh: + serverErrs = append(serverErrs, err) + case <-time.After(5 * time.Second): + t.Fatal("等待 ingress websocket 结束超时") + } + } + return serverErrs, aReadErr +} + +func TestOpenAIGatewayService_ProxyResponsesWebSocketFromClient_CodexThreadsDoNotPreemptEachOther(t *testing.T) { + serverErrs, aReadErr := runOpenAIWSCodexThreadPair(t, "thread-a", "thread-b") + require.NoError(t, aReadErr, "子智能体接入后父线程在飞的请求必须继续完成") + for _, err := range serverErrs { + require.NoError(t, err) + } +} + +func TestOpenAIGatewayService_ProxyResponsesWebSocketFromClient_SameCodexThreadStillPreempts(t *testing.T) { + serverErrs, aReadErr := runOpenAIWSCodexThreadPair(t, "thread-a", "thread-a") + require.Error(t, aReadErr, "同线程重连必须取代旧连接") + preempted := 0 + for _, err := range serverErrs { + if IsOpenAIWSSessionPreemptedError(err) { + preempted++ + } else { + require.NoError(t, err) + } + } + require.Equal(t, 1, preempted, "应恰好有一条会话被抢占") +} From e4c369bd56305c9984c01cf9eb2bfb6f1cb8da1a Mon Sep 17 00:00:00 2001 From: xlplbo Date: Fri, 11 Sep 2026 10:25:30 +0800 Subject: [PATCH 4/7] =?UTF-8?q?fix(openai):=20HTTP=20=E7=BB=8F=20WS=20?= =?UTF-8?q?=E6=B1=A0=E8=BD=AC=E5=8F=91=E6=8C=89=E6=89=A7=E8=A1=8C=E4=BD=9C?= =?UTF-8?q?=E7=94=A8=E5=9F=9F=E5=8F=96=E4=BC=9A=E8=AF=9D=E7=8A=B6=E6=80=81?= =?UTF-8?q?=E9=94=AE?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- .../service/openai_ws_execution_scope.go | 22 ++++ .../openai_ws_forwarder_success_test.go | 4 +- .../service/openai_ws_forwarder_v2.go | 5 + ...ai_ws_forwarder_v2_execution_scope_test.go | 104 ++++++++++++++++++ 4 files changed, 134 insertions(+), 1 deletion(-) create mode 100644 backend/internal/service/openai_ws_forwarder_v2_execution_scope_test.go diff --git a/backend/internal/service/openai_ws_execution_scope.go b/backend/internal/service/openai_ws_execution_scope.go index 730e86cb01..fbd1994058 100644 --- a/backend/internal/service/openai_ws_execution_scope.go +++ b/backend/internal/service/openai_ws_execution_scope.go @@ -40,6 +40,28 @@ func resolveOpenAIWSClientThreadID(c *gin.Context, body []byte) string { return codexTurnMetadataThreadID(gjson.GetBytes(body, "client_metadata."+openAIWSTurnMetadataHeader).String()) } +// openAIWSExecutionScopeBodyFromRequest 只序列化执行作用域会读取的字段， +// 供已解析成 map 的 HTTP 请求体取键，避免为此重新编码整个请求体。 +func openAIWSExecutionScopeBodyFromRequest(reqBody map[string]any) []byte { + if len(reqBody) == 0 { + return nil + } + subset := make(map[string]any, 3) + for _, key := range []string{"client_metadata", "prompt_cache_key", "previous_response_id"} { + if value, ok := reqBody[key]; ok { + subset[key] = value + } + } + if len(subset) == 0 { + return nil + } + raw, err := json.Marshal(subset) + if err != nil { + return nil + } + return raw +} + func codexTurnMetadataThreadID(raw string) string { raw = strings.TrimSpace(raw) if raw == "" { diff --git a/backend/internal/service/openai_ws_forwarder_success_test.go b/backend/internal/service/openai_ws_forwarder_success_test.go index 6923352d7d..0401f59f55 100644 --- a/backend/internal/service/openai_ws_forwarder_success_test.go +++ b/backend/internal/service/openai_ws_forwarder_success_test.go @@ -1403,7 +1403,9 @@ func TestOpenAIGatewayService_Forward_WSv2_TurnStateAndMetadataReplayOnReconnect require.NoError(t, err) require.NotNil(t, result1) - sessionHash := svc.GenerateSessionHash(c1, reqBody) + // 会话级状态按执行作用域取键（显式 session_id 也在其中），不再是原会话哈希。 + sessionHash, _ := resolveOpenAIWSExecutionScope(c1, reqBody, getAPIKeyIDFromContext(c1)) + require.NotEmpty(t, sessionHash) store := svc.getOpenAIWSStateStore() turnState, ok := store.GetSessionTurnState(0, sessionHash) require.True(t, ok) diff --git a/backend/internal/service/openai_ws_forwarder_v2.go b/backend/internal/service/openai_ws_forwarder_v2.go index 9f04d42640..30f897928d 100644 --- a/backend/internal/service/openai_ws_forwarder_v2.go +++ b/backend/internal/service/openai_ws_forwarder_v2.go @@ -124,6 +124,11 @@ func (s *OpenAIGatewayService) forwardOpenAIWSV2( sessionHash, legacySessionHash = openAIWSSessionHashesFromID(promptCacheKey) attachOpenAILegacySessionHashToGin(c, legacySessionHash) } + // 与 WS 接入路径共用执行作用域：codex 多智能体共用 session-id，turn state 与 + // store=false 的连接绑定必须按线程隔离，同一线程在两条路径之间也才能共享状态。 + if scope, _ := resolveOpenAIWSExecutionScope(c, openAIWSExecutionScopeBodyFromRequest(reqBody), getAPIKeyIDFromContext(c)); scope != "" { + sessionHash = scope + } if turnState == "" && stateStore != nil && sessionHash != "" { if savedTurnState, ok := stateStore.GetSessionTurnState(groupID, sessionHash); ok { turnState = savedTurnState diff --git a/backend/internal/service/openai_ws_forwarder_v2_execution_scope_test.go b/backend/internal/service/openai_ws_forwarder_v2_execution_scope_test.go new file mode 100644 index 0000000000..59d22036c9 --- /dev/null +++ b/backend/internal/service/openai_ws_forwarder_v2_execution_scope_test.go @@ -0,0 +1,104 @@ +package service + +import ( + "context" + "net/http" + "net/http/httptest" + "testing" + + "github.com/Wei-Shaw/sub2api/internal/config" + "github.com/gin-gonic/gin" + "github.com/gorilla/websocket" + "github.com/stretchr/testify/require" +) + +// 场景：codex 子智能体经 HTTP 入站、网关内部走 WS 池转发时，与父线程共用 session_id 头。 +// turn state 必须落在执行作用域键下，而不是按 session_id 算出的会话哈希下， +// 否则子智能体会覆盖父线程的绑定；同一线程在 WS 接入与 HTTP 路径之间也才能共享状态。 +func TestOpenAIGatewayService_Forward_WSv2_TurnStateBoundToExecutionScope(t *testing.T) { + gin.SetMode(gin.TestMode) + + upgrader := websocket.Upgrader{CheckOrigin: func(r *http.Request) bool { return true }} + wsServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + respHeader := http.Header{} + respHeader.Set("x-codex-turn-state", "turn-state-from-upstream") + conn, err := upgrader.Upgrade(w, r, respHeader) + if err != nil { + t.Errorf("upgrade websocket failed: %v", err) + return + } + defer func() { _ = conn.Close() }() + var request map[string]any + if err := conn.ReadJSON(&request); err != nil { + t.Errorf("read ws request failed: %v", err) + return + } + if err := conn.WriteJSON(map[string]any{ + "type": "response.completed", + "response": map[string]any{ + "id": "resp_exec_scope_http", + "model": "gpt-5.1", + "usage": map[string]any{"input_tokens": 2, "output_tokens": 1}, + }, + }); err != nil { + t.Errorf("write response.completed failed: %v", err) + } + })) + defer wsServer.Close() + + cfg := &config.Config{} + cfg.Security.URLAllowlist.Enabled = false + cfg.Security.URLAllowlist.AllowInsecureHTTP = true + cfg.Gateway.OpenAIWS.Enabled = true + cfg.Gateway.OpenAIWS.OAuthEnabled = true + cfg.Gateway.OpenAIWS.APIKeyEnabled = true + cfg.Gateway.OpenAIWS.ResponsesWebsocketsV2 = true + cfg.Gateway.OpenAIWS.MaxConnsPerAccount = 1 + cfg.Gateway.OpenAIWS.MinIdlePerAccount = 0 + cfg.Gateway.OpenAIWS.MaxIdlePerAccount = 0 + + svc := &OpenAIGatewayService{ + cfg: cfg, + httpUpstream: &httpUpstreamRecorder{}, + cache: &stubGatewayCache{}, + openaiWSResolver: NewOpenAIWSProtocolResolver(cfg), + toolCorrector: NewCodexToolCorrector(), + } + groupID := int64(9) + account := &Account{ + ID: 456, + Name: "openai-http-exec-scope", + Platform: PlatformOpenAI, + Type: AccountTypeAPIKey, + Status: StatusActive, + Schedulable: true, + Concurrency: 1, + Credentials: map[string]any{"api_key": "sk-test", "base_url": wsServer.URL}, + Extra: map[string]any{"responses_websockets_v2_enabled": true}, + } + + reqBody := []byte(`{"model":"gpt-5.1","stream":false,"input":[{"type":"input_text","text":"hello"}]}`) + rec := httptest.NewRecorder() + c, _ := gin.CreateTestContext(rec) + c.Request = httptest.NewRequest(http.MethodPost, "/openai/v1/responses", nil) + c.Request.Header.Set("session_id", "root-session") + c.Request.Header.Set(openAIWSTurnMetadataHeader, `{"session_id":"root-session","thread_id":"child-thread"}`) + c.Set("api_key", &APIKey{ID: 21, GroupID: &groupID}) + + result, err := svc.Forward(context.Background(), c, account, reqBody) + require.NoError(t, err) + require.NotNil(t, result) + + legacyHash := svc.GenerateSessionHash(c, reqBody) + scope, _ := resolveOpenAIWSExecutionScope(c, reqBody, 21) + require.NotEmpty(t, legacyHash) + require.Len(t, scope, 16) + require.NotEqual(t, legacyHash, scope) + + store := svc.getOpenAIWSStateStore() + _, turnStateOnLegacy := store.GetSessionTurnState(groupID, legacyHash) + require.False(t, turnStateOnLegacy, "turn state 不得绑定到按 session_id 算出的会话哈希") + turnState, turnStateOnScope := store.GetSessionTurnState(groupID, scope) + require.True(t, turnStateOnScope, "turn state 应绑定到执行作用域") + require.Equal(t, "turn-state-from-upstream", turnState) +} From d0ca057ca41add88bc9e5a496493da969816125f Mon Sep 17 00:00:00 2001 From: xlplbo Date: Fri, 11 Sep 2026 10:25:30 +0800 Subject: [PATCH 5/7] =?UTF-8?q?fix(openai):=20WS=20=E4=BC=9A=E8=AF=9D?= =?UTF-8?q?=E8=A2=AB=E6=8A=A2=E5=8D=A0=E6=97=B6=E5=85=88=E5=8F=91=E5=85=B3?= =?UTF-8?q?=E9=97=AD=E5=B8=A7=E5=86=8D=E5=8F=96=E6=B6=88?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- .../handler/openai_gateway_handler.go | 4 +- .../service/openai_ws_forwarder_ingress.go | 2 +- ..._forwarder_ingress_execution_scope_test.go | 5 + .../service/openai_ws_session_preemption.go | 73 +++++++++++- .../openai_ws_session_preemption_test.go | 105 +++++++++++++++++- 5 files changed, 177 insertions(+), 12 deletions(-) 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)) From 613722eee434c0a62e9397487b3713c16362e18d Mon Sep 17 00:00:00 2001 From: xlplbo Date: Fri, 11 Sep 2026 11:36:29 +0800 Subject: [PATCH 6/7] =?UTF-8?q?fix(openai):=20HTTP=20=E8=BD=AC=E5=8F=91?= =?UTF-8?q?=E7=9A=84=E6=89=A7=E8=A1=8C=E4=BD=9C=E7=94=A8=E5=9F=9F=E6=94=B9?= =?UTF-8?q?=E4=BB=8E=E5=8E=9F=E5=A7=8B=E8=AF=B7=E6=B1=82=E8=AE=A1=E7=AE=97?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- .../service/openai_gateway_forward.go | 5 + .../service/openai_ws_execution_scope.go | 22 ----- .../service/openai_ws_forwarder_v2.go | 6 +- ...ai_ws_forwarder_v2_execution_scope_test.go | 97 +++++++++++++++++++ 4 files changed, 106 insertions(+), 24 deletions(-) diff --git a/backend/internal/service/openai_gateway_forward.go b/backend/internal/service/openai_gateway_forward.go index 6837cb8236..ef716e0f99 100644 --- a/backend/internal/service/openai_gateway_forward.go +++ b/backend/internal/service/openai_gateway_forward.go @@ -42,6 +42,10 @@ func (s *OpenAIGatewayService) Forward(ctx context.Context, c *gin.Context, acco restrictionResult := s.detectCodexClientRestriction(c, account, body) apiKeyID := getAPIKeyIDFromContext(c) + // 执行作用域必须取自客户端原始身份：后面的账号 namespace 改写与指纹收敛会改掉 + // 请求体里的 client_metadata / prompt_cache_key，用改写后的值取键会让不同会话 + // 落到同一个键，也会与 WS 接入路径按原始报文算出的键对不上。 + wsExecutionScope, _ := resolveOpenAIWSExecutionScope(c, body, apiKeyID) logCodexCLIOnlyDetection(ctx, c, account, apiKeyID, restrictionResult, body) if restrictionResult.Enabled && !restrictionResult.Matched { MarkOpsClientBusinessLimited(c, OpsClientBusinessLimitedReasonLocalPolicyDenied) @@ -874,6 +878,7 @@ func (s *OpenAIGatewayService) Forward(ctx context.Context, c *gin.Context, acco account, wsReqBody, clientPromptCacheKey, + wsExecutionScope, token, wsDecision, isCodexCLI, diff --git a/backend/internal/service/openai_ws_execution_scope.go b/backend/internal/service/openai_ws_execution_scope.go index fbd1994058..730e86cb01 100644 --- a/backend/internal/service/openai_ws_execution_scope.go +++ b/backend/internal/service/openai_ws_execution_scope.go @@ -40,28 +40,6 @@ func resolveOpenAIWSClientThreadID(c *gin.Context, body []byte) string { return codexTurnMetadataThreadID(gjson.GetBytes(body, "client_metadata."+openAIWSTurnMetadataHeader).String()) } -// openAIWSExecutionScopeBodyFromRequest 只序列化执行作用域会读取的字段， -// 供已解析成 map 的 HTTP 请求体取键，避免为此重新编码整个请求体。 -func openAIWSExecutionScopeBodyFromRequest(reqBody map[string]any) []byte { - if len(reqBody) == 0 { - return nil - } - subset := make(map[string]any, 3) - for _, key := range []string{"client_metadata", "prompt_cache_key", "previous_response_id"} { - if value, ok := reqBody[key]; ok { - subset[key] = value - } - } - if len(subset) == 0 { - return nil - } - raw, err := json.Marshal(subset) - if err != nil { - return nil - } - return raw -} - func codexTurnMetadataThreadID(raw string) string { raw = strings.TrimSpace(raw) if raw == "" { diff --git a/backend/internal/service/openai_ws_forwarder_v2.go b/backend/internal/service/openai_ws_forwarder_v2.go index 30f897928d..d908b0ef1a 100644 --- a/backend/internal/service/openai_ws_forwarder_v2.go +++ b/backend/internal/service/openai_ws_forwarder_v2.go @@ -23,6 +23,7 @@ func (s *OpenAIGatewayService) forwardOpenAIWSV2( account *Account, reqBody map[string]any, clientPromptCacheKey string, + executionScope string, token string, decision OpenAIWSProtocolDecision, isCodexCLI bool, @@ -126,8 +127,9 @@ func (s *OpenAIGatewayService) forwardOpenAIWSV2( } // 与 WS 接入路径共用执行作用域：codex 多智能体共用 session-id，turn state 与 // store=false 的连接绑定必须按线程隔离，同一线程在两条路径之间也才能共享状态。 - if scope, _ := resolveOpenAIWSExecutionScope(c, openAIWSExecutionScopeBodyFromRequest(reqBody), getAPIKeyIDFromContext(c)); scope != "" { - sessionHash = scope + // 作用域由 Forward 从改写前的原始请求算出后传入，reqBody 此时已带账号 namespace。 + if executionScope = strings.TrimSpace(executionScope); executionScope != "" { + sessionHash = executionScope } if turnState == "" && stateStore != nil && sessionHash != "" { if savedTurnState, ok := stateStore.GetSessionTurnState(groupID, sessionHash); ok { diff --git a/backend/internal/service/openai_ws_forwarder_v2_execution_scope_test.go b/backend/internal/service/openai_ws_forwarder_v2_execution_scope_test.go index 59d22036c9..da13cd1fcb 100644 --- a/backend/internal/service/openai_ws_forwarder_v2_execution_scope_test.go +++ b/backend/internal/service/openai_ws_forwarder_v2_execution_scope_test.go @@ -2,6 +2,7 @@ package service import ( "context" + "fmt" "net/http" "net/http/httptest" "testing" @@ -102,3 +103,99 @@ func TestOpenAIGatewayService_Forward_WSv2_TurnStateBoundToExecutionScope(t *tes require.True(t, turnStateOnScope, "turn state 应绑定到执行作用域") require.Equal(t, "turn-state-from-upstream", turnState) } + +// Forward 在转发前会按账号 namespace 改写请求体里的 client_metadata，指纹 full 模式还会给每个 +// 请求注入同一个 thread_id。执行作用域必须取自客户端原始请求：否则同一 API key 下不同 +// session_id 的会话会落到同一个键共用 turn state，客户端自带线程标识时也会与 WS 接入路径 +// 按原始报文算出的键对不上。 +func TestOpenAIGatewayService_Forward_WSv2_ExecutionScopeUsesOriginalIdentity(t *testing.T) { + gin.SetMode(gin.TestMode) + + cfg := &config.Config{} + cfg.Security.URLAllowlist.Enabled = false + cfg.Security.URLAllowlist.AllowInsecureHTTP = true + cfg.Gateway.OpenAIWS.Enabled = true + cfg.Gateway.OpenAIWS.OAuthEnabled = true + cfg.Gateway.OpenAIWS.APIKeyEnabled = true + cfg.Gateway.OpenAIWS.ResponsesWebsocketsV2 = true + cfg.Gateway.OpenAIWS.MaxConnsPerAccount = 1 + cfg.Gateway.OpenAIWS.MinIdlePerAccount = 0 + cfg.Gateway.OpenAIWS.MaxIdlePerAccount = 1 + + completed := func(id string) []byte { + return []byte(`{"type":"response.completed","response":{"id":"` + id + `","model":"gpt-5.1","usage":{"input_tokens":1,"output_tokens":1}}}`) + } + captureConn := &openAIWSCaptureConn{events: [][]byte{completed("resp_a"), completed("resp_b"), completed("resp_c")}} + handshake := http.Header{} + handshake.Set(openAIWSTurnStateHeader, "turn-state-from-upstream") + pool := newOpenAIWSConnPool(cfg) + pool.setClientDialerForTest(&openAIWSCaptureDialer{conn: captureConn, handshake: handshake}) + defer pool.Close() + + stateStore := NewOpenAIWSStateStore(nil) + svc := &OpenAIGatewayService{ + cfg: cfg, + httpUpstream: &httpUpstreamRecorder{}, + cache: &stubGatewayCache{}, + openaiWSResolver: NewOpenAIWSProtocolResolver(cfg), + toolCorrector: NewCodexToolCorrector(), + openaiWSPool: pool, + openaiWSStateStore: stateStore, + } + groupID := int64(9) + const apiKeyID = int64(21) + account := &Account{ + ID: 457, + Name: "openai-oauth-fingerprint-full", + Platform: PlatformOpenAI, + Type: AccountTypeOAuth, + Status: StatusActive, + Schedulable: true, + Concurrency: 1, + Credentials: map[string]any{"access_token": "oauth-token-1"}, + Extra: map[string]any{ + "responses_websockets_v2_enabled": true, + codexFingerprintModeExtraKey: "full", + codexFingerprintSeedExtraKey: "11111111-1111-4111-8111-aaaaaaaaaaaa", + }, + } + + forward := func(sessionID, body string) (*gin.Context, []byte) { + rec := httptest.NewRecorder() + c, _ := gin.CreateTestContext(rec) + c.Request = httptest.NewRequest(http.MethodPost, "/openai/v1/responses", nil) + c.Request.Header.Set("session_id", sessionID) + c.Set("api_key", &APIKey{ID: apiKeyID, GroupID: &groupID}) + raw := []byte(body) + result, err := svc.Forward(context.Background(), c, account, raw) + require.NoError(t, err) + require.NotNil(t, result) + return c, raw + } + + plainBody := `{"model":"gpt-5.1","stream":false,"input":[{"type":"input_text","text":"hello"}]}` + cA, rawA := forward("session-a", plainBody) + cB, rawB := forward("session-b", plainBody) + + scopeA, _ := resolveOpenAIWSExecutionScope(cA, rawA, apiKeyID) + scopeB, _ := resolveOpenAIWSExecutionScope(cB, rawB, apiKeyID) + require.NotEmpty(t, scopeA) + require.NotEqual(t, scopeA, scopeB, "不同 session_id 的会话必须落在不同的作用域") + _, boundA := stateStore.GetSessionTurnState(groupID, scopeA) + require.True(t, boundA, "会话 A 的 turn state 应落在按原始 session_id 算出的作用域") + _, boundB := stateStore.GetSessionTurnState(groupID, scopeB) + require.True(t, boundB, "会话 B 的 turn state 应落在按原始 session_id 算出的作用域") + + injected := resolveCodexFingerprintIDs(account, "", codexFingerprintFull) + require.NotNil(t, injected) + injectedScope, _ := deriveOpenAISessionHashes(fmt.Sprintf("openai_ws_exec:%d|thread=%s", apiKeyID, injected.threadID)) + _, boundToInjected := stateStore.GetSessionTurnState(groupID, injectedScope) + require.False(t, boundToInjected, "指纹收敛注入的固定 thread_id 不得成为状态键") + + threadBody := `{"model":"gpt-5.1","stream":false,"client_metadata":{"thread_id":"child-thread"},"input":[{"type":"input_text","text":"hello"}]}` + cC, rawC := forward("session-c", threadBody) + scopeC, threadC := resolveOpenAIWSExecutionScope(cC, rawC, apiKeyID) + require.Equal(t, "child-thread", threadC) + _, boundC := stateStore.GetSessionTurnState(groupID, scopeC) + require.True(t, boundC, "客户端自带线程标识时，键必须与按原始报文算出的一致，不受账号 namespace 改写影响") +} From 4543ddc5c172b3d5f72d35de1c357a886b55526c Mon Sep 17 00:00:00 2001 From: xlplbo Date: Fri, 11 Sep 2026 13:37:05 +0800 Subject: [PATCH 7/7] =?UTF-8?q?fix(openai):=20WS=20=E6=89=A7=E8=A1=8C?= =?UTF-8?q?=E4=BD=9C=E7=94=A8=E5=9F=9F=E6=8C=89=20request=5Fkind=20?= =?UTF-8?q?=E5=88=86=E9=81=93?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- .../service/openai_ws_execution_scope.go | 102 +++++++++++++-- .../service/openai_ws_execution_scope_test.go | 118 ++++++++++++++++++ .../service/openai_ws_session_preemption.go | 7 +- .../openai_ws_session_preemption_test.go | 53 ++++++++ 4 files changed, 270 insertions(+), 10 deletions(-) diff --git a/backend/internal/service/openai_ws_execution_scope.go b/backend/internal/service/openai_ws_execution_scope.go index 730e86cb01..e15ef0a8ba 100644 --- a/backend/internal/service/openai_ws_execution_scope.go +++ b/backend/internal/service/openai_ws_execution_scope.go @@ -12,6 +12,23 @@ import ( const ( openAIWSThreadIDHeader = "thread-id" openAIWSWindowIDHeader = "x-codex-window-id" + + // request_kind 是 codex 放在 x-codex-turn-metadata 里的请求类型， + // 对应源码枚举 CodexResponsesRequestKind，目前四个取值： + // turn 用户轮次的采样请求，含工具调用后的每次续采样 + // prewarm 预热请求：会话启动预开的连接会被首个 turn 复用， + // 新连接上也是先 warmup 再在同一连接发 turn + // compaction 上下文压缩：在两次采样之间顺序执行，复用当前 turn 的 + // client session + // memory 记忆整理：独立 root turn，后台并发发起，却沿用会话的 + // thread_id，与用户 turn 同键时互相切断在飞轮次 + // 前三种与用户 turn 是同一条连接或同一段顺序动作，必须同道； + // 未声明 request_kind 的请求（旧版 codex、只带 thread-id 头）视同 turn， + // 键与此前完全一致；明确声明但认不出的取值各自成道，宁可只在同类间 + // 互相替换，也不让未知副请求打断用户轮次。 + openAIWSRequestKindTurn = "turn" + openAIWSRequestKindPrewarm = "prewarm" + openAIWSRequestKindCompaction = "compaction" ) // resolveOpenAIWSClientThreadID 提取 codex 客户端的线程标识。codex 多智能体会话里 @@ -40,32 +57,99 @@ func resolveOpenAIWSClientThreadID(c *gin.Context, body []byte) string { return codexTurnMetadataThreadID(gjson.GetBytes(body, "client_metadata."+openAIWSTurnMetadataHeader).String()) } -func codexTurnMetadataThreadID(raw string) string { +// codexTurnMetadata 是 x-codex-turn-metadata JSON 里执行作用域用到的字段。 +type codexTurnMetadata struct { + ThreadID string `json:"thread_id"` + RequestKind string `json:"request_kind"` +} + +func parseCodexTurnMetadata(raw string) (codexTurnMetadata, bool) { raw = strings.TrimSpace(raw) if raw == "" { - return "" - } - var metadata struct { - ThreadID string `json:"thread_id"` + return codexTurnMetadata{}, false } + var metadata codexTurnMetadata if err := json.Unmarshal([]byte(raw), &metadata); err != nil { + return codexTurnMetadata{}, false + } + metadata.ThreadID = strings.TrimSpace(metadata.ThreadID) + metadata.RequestKind = strings.ToLower(strings.TrimSpace(metadata.RequestKind)) + return metadata, true +} + +func codexTurnMetadataThreadID(raw string) string { + metadata, _ := parseCodexTurnMetadata(raw) + return metadata.ThreadID +} + +// openAIWSExecutionTurnMetadata 取本次连接的 turn 元数据：x-codex-turn-metadata 头优先， +// 其次请求体 client_metadata 内嵌的同名 JSON（Desktop 走 HTTP 时只有后者）。 +func openAIWSExecutionTurnMetadata(c *gin.Context, body []byte) codexTurnMetadata { + if c != nil && c.Request != nil { + if metadata, ok := parseCodexTurnMetadata(c.GetHeader(openAIWSTurnMetadataHeader)); ok { + return metadata + } + } + if len(body) == 0 { + return codexTurnMetadata{} + } + metadata, _ := parseCodexTurnMetadata(gjson.GetBytes(body, "client_metadata."+openAIWSTurnMetadataHeader).String()) + return metadata +} + +func openAIWSExecutionSubagent(c *gin.Context, body []byte) string { + if c != nil && c.Request != nil { + if subagent := strings.TrimSpace(c.GetHeader(openAISubagentHeader)); subagent != "" { + return strings.ToLower(subagent) + } + } + if len(body) == 0 { return "" } - return strings.TrimSpace(metadata.ThreadID) + return strings.ToLower(strings.TrimSpace(gjson.GetBytes(body, "client_metadata."+openAISubagentHeader).String())) +} + +// resolveOpenAIWSExecutionLane 派生执行作用域里的“道”，同道才互相抢占、共享会话级状态。 +// 主道（turn / prewarm / compaction / 未声明）返回空串，其余 request_kind 各自成道，取值 +// 语义见上面的常量注释；没有线程元数据但声明了 x-openai-subagent 的（guardian 评分器） +// 按子智能体值成道，它们与用户 turn 同键时同样会互相切断在飞轮次。 +func resolveOpenAIWSExecutionLane(c *gin.Context, body []byte) string { + metadata := openAIWSExecutionTurnMetadata(c, body) + switch metadata.RequestKind { + case "", openAIWSRequestKindTurn, openAIWSRequestKindPrewarm, openAIWSRequestKindCompaction: + default: + return "kind=" + metadata.RequestKind + } + if metadata.ThreadID == "" { + if subagent := openAIWSExecutionSubagent(c, body); subagent != "" { + return "subagent=" + subagent + } + } + return "" +} + +func openAIWSExecutionScopeSeed(apiKeyID int64, identity, value, lane string) string { + seed := fmt.Sprintf("openai_ws_exec:%d|%s=%s", apiKeyID, identity, value) + if lane != "" { + seed += "|" + lane + } + return seed } // resolveOpenAIWSExecutionScope 派生 WS 接入用于抢占与会话级状态的执行作用域键， // 并一并返回解析到的线程标识供日志使用。身份只取客户端声明的线程标识或显式会话标识， // 不掺入请求内容，因此同线程重连在输入变化后仍命中同一键；键里带 API key，同组不同 key -// 互不影响。两种身份都没有时返回空作用域：按请求内容推导的粘性种子只服务账号亲和， +// 互不影响；同线程上与用户 turn 并发的独立请求按 resolveOpenAIWSExecutionLane 另成一道。 +// 两种身份都没有时返回空作用域：按请求内容推导的粘性种子只服务账号亲和， // 不是可靠身份，不参与抢占。 func resolveOpenAIWSExecutionScope(c *gin.Context, body []byte, apiKeyID int64) (scope, threadID string) { + lane := resolveOpenAIWSExecutionLane(c, body) if threadID = resolveOpenAIWSClientThreadID(c, body); threadID != "" { - scope, _ = deriveOpenAISessionHashes(fmt.Sprintf("openai_ws_exec:%d|thread=%s", apiKeyID, threadID)) + scope, _ = deriveOpenAISessionHashes(openAIWSExecutionScopeSeed(apiKeyID, "thread", threadID, lane)) return scope, threadID } if sessionID := strings.TrimSpace(explicitOpenAIRequestSessionID(c, body)); sessionID != "" { - scope, _ = deriveOpenAISessionHashes(fmt.Sprintf("openai_ws_exec:%d|session=%s", apiKeyID, sessionID)) + scope, _ = deriveOpenAISessionHashes(openAIWSExecutionScopeSeed(apiKeyID, "session", sessionID, lane)) return scope, "" } return "", "" diff --git a/backend/internal/service/openai_ws_execution_scope_test.go b/backend/internal/service/openai_ws_execution_scope_test.go index 349b4cfb34..10193a9ae3 100644 --- a/backend/internal/service/openai_ws_execution_scope_test.go +++ b/backend/internal/service/openai_ws_execution_scope_test.go @@ -82,6 +82,100 @@ func TestResolveOpenAIWSClientThreadID(t *testing.T) { require.Equal(t, "t", resolveOpenAIWSClientThreadID(nil, []byte(`{"client_metadata":{"thread_id":"t"}}`))) } +func TestResolveOpenAIWSExecutionLane(t *testing.T) { + gin.SetMode(gin.TestMode) + cases := []struct { + name string + headers map[string]string + body string + want string + }{ + { + name: "turn is the main lane", + headers: map[string]string{openAIWSTurnMetadataHeader: `{"thread_id":"t","request_kind":"turn"}`}, + want: "", + }, + { + name: "prewarm shares the turn lane", + headers: map[string]string{openAIWSTurnMetadataHeader: `{"thread_id":"t","request_kind":"prewarm","thread_source":"user"}`}, + want: "", + }, + { + name: "compaction shares the turn lane", + headers: map[string]string{openAIWSTurnMetadataHeader: `{"thread_id":"t","request_kind":"compaction"}`}, + want: "", + }, + { + name: "missing request_kind is the main lane", + headers: map[string]string{openAIWSTurnMetadataHeader: `{"thread_id":"t"}`}, + want: "", + }, + { + name: "memory consolidation gets its own lane", + headers: map[string]string{openAIWSTurnMetadataHeader: `{"thread_id":"t","request_kind":"memory","thread_source":"memory_consolidation"}`}, + want: "kind=memory", + }, + { + name: "request_kind is normalized", + headers: map[string]string{openAIWSTurnMetadataHeader: `{"thread_id":"t","request_kind":" Memory "}`}, + want: "kind=memory", + }, + { + name: "unknown kind never joins the turn lane", + headers: map[string]string{openAIWSTurnMetadataHeader: `{"thread_id":"t","request_kind":"future"}`}, + want: "kind=future", + }, + { + name: "header metadata wins over body metadata", + headers: map[string]string{openAIWSTurnMetadataHeader: `{"thread_id":"t","request_kind":"turn"}`}, + body: `{"client_metadata":{"x-codex-turn-metadata":"{\"thread_id\":\"t\",\"request_kind\":\"memory\"}"}}`, + want: "", + }, + { + name: "body embedded metadata used when header absent", + body: `{"client_metadata":{"x-codex-turn-metadata":"{\"thread_id\":\"t\",\"request_kind\":\"memory\"}"}}`, + want: "kind=memory", + }, + { + name: "malformed header falls back to body", + headers: map[string]string{openAIWSTurnMetadataHeader: "{not json"}, + body: `{"client_metadata":{"x-codex-turn-metadata":"{\"request_kind\":\"memory\"}"}}`, + want: "kind=memory", + }, + { + name: "subagent without thread metadata gets its own lane", + headers: map[string]string{openAIWSThreadIDHeader: "t", openAISubagentHeader: "Guardian"}, + want: "subagent=guardian", + }, + { + name: "subagent thread with its own metadata stays in the turn lane", + headers: map[string]string{ + openAIWSThreadIDHeader: "child", + openAISubagentHeader: "guardian", + openAIWSTurnMetadataHeader: `{"thread_id":"child","request_kind":"turn","subagent_kind":"guardian"}`, + }, + want: "", + }, + { + name: "subagent declared in body client_metadata", + body: `{"client_metadata":{"x-openai-subagent":"guardian"}}`, + want: "subagent=guardian", + }, + { + name: "none", + body: `{"input":"hi"}`, + want: "", + }, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + got := resolveOpenAIWSExecutionLane(newOpenAIWSExecutionScopeTestContext(tc.headers), []byte(tc.body)) + require.Equal(t, tc.want, got) + }) + } + require.Equal(t, "kind=memory", resolveOpenAIWSExecutionLane(nil, []byte(`{"client_metadata":{"x-codex-turn-metadata":"{\"request_kind\":\"memory\"}"}}`))) +} + func TestResolveOpenAIWSExecutionScope(t *testing.T) { gin.SetMode(gin.TestMode) threadHeaders := map[string]string{ @@ -108,6 +202,28 @@ func TestResolveOpenAIWSExecutionScope(t *testing.T) { require.NotEqual(t, byThread, scopeOf(otherThread, bodyFirst, 11)) require.NotEqual(t, byThread, scopeOf(threadHeaders, bodyFirst, 12), "API key 必须参与隔离") + memoryHeaders := map[string]string{ + "session-id": "root-session", + openAIWSTurnMetadataHeader: `{"session_id":"root-session","thread_id":"thread-a","request_kind":"memory","thread_source":"memory_consolidation"}`, + } + byMemory, memoryThread := resolveOpenAIWSExecutionScope(newOpenAIWSExecutionScopeTestContext(memoryHeaders), bodyFirst, 11) + require.Len(t, byMemory, 16) + require.Equal(t, "thread-a", memoryThread) + require.NotEqual(t, byThread, byMemory, "同线程的记忆整理请求不得与用户 turn 同键") + require.Equal(t, byMemory, scopeOf(memoryHeaders, bodyRetry, 11), "记忆整理重连仍命中同一键") + for _, kind := range []string{"turn", "prewarm", "compaction"} { + sameLane := map[string]string{ + "session-id": "root-session", + openAIWSTurnMetadataHeader: `{"session_id":"root-session","thread_id":"thread-a","request_kind":"` + kind + `"}`, + } + require.Equal(t, byThread, scopeOf(sameLane, bodyFirst, 11), kind+" 必须与用户 turn 同键") + } + require.Equal(t, byThread, scopeOf(map[string]string{openAIWSThreadIDHeader: "thread-a"}, bodyFirst, 11), "只有 thread-id 头时仍是主道") + byScorer := scopeOf(map[string]string{openAIWSThreadIDHeader: "thread-a", openAISubagentHeader: "guardian"}, bodyFirst, 11) + require.Len(t, byScorer, 16) + require.NotEqual(t, byThread, byScorer, "无线程元数据的子智能体请求不得与用户 turn 同键") + require.NotEqual(t, byMemory, byScorer) + sessionOnly := map[string]string{"session-id": "root-session"} bySession, threadID := resolveOpenAIWSExecutionScope(newOpenAIWSExecutionScopeTestContext(sessionOnly), bodyFirst, 11) require.Len(t, bySession, 16) @@ -115,6 +231,8 @@ func TestResolveOpenAIWSExecutionScope(t *testing.T) { require.NotEqual(t, byThread, bySession, "线程标识优先于会话标识") require.Equal(t, bySession, scopeOf(sessionOnly, bodyRetry, 11)) require.NotEqual(t, bySession, scopeOf(sessionOnly, bodyFirst, 12)) + sessionMemory := []byte(`{"type":"response.create","input":"hi","client_metadata":{"x-codex-turn-metadata":"{\"request_kind\":\"memory\"}"}}`) + require.NotEqual(t, bySession, scopeOf(sessionOnly, sessionMemory, 11), "只有会话标识时记忆整理也单独成道") byCacheKey := scopeOf(nil, []byte(`{"type":"response.create","prompt_cache_key":"cache-1","input":"hi"}`), 11) require.Len(t, byCacheKey, 16) diff --git a/backend/internal/service/openai_ws_session_preemption.go b/backend/internal/service/openai_ws_session_preemption.go index f1f55c34f9..32b49ccbbc 100644 --- a/backend/internal/service/openai_ws_session_preemption.go +++ b/backend/internal/service/openai_ws_session_preemption.go @@ -131,13 +131,18 @@ func (s *OpenAIGatewayService) BeginOpenAIWSIngressSessionPreemptionWithClient( stateStore.DeleteSessionTurnState(preemptGroupID, preemptScope) stateStore.DeleteSessionConn(preemptGroupID, preemptScope) } + lane := resolveOpenAIWSExecutionLane(c, firstClientMessage) + if lane == "" { + lane = "main" + } logOpenAIWSModeInfo( - "ingress_ws_session_preempted account_id=%d group_id=%d api_key_id=%d scope=%s thread_id=%s", + "ingress_ws_session_preempted account_id=%d group_id=%d api_key_id=%d scope=%s thread_id=%s lane=%s", account.ID, preemptGroupID, preemptAPIKeyID, truncateOpenAIWSLogValue(preemptScope, 12), truncateOpenAIWSLogValue(preemptThreadID, openAIWSIDValueMaxLen), + truncateOpenAIWSLogValue(lane, openAIWSIDValueMaxLen), ) } return preemptCtx, cleanup, true diff --git a/backend/internal/service/openai_ws_session_preemption_test.go b/backend/internal/service/openai_ws_session_preemption_test.go index 243cbcab68..f378758478 100644 --- a/backend/internal/service/openai_ws_session_preemption_test.go +++ b/backend/internal/service/openai_ws_session_preemption_test.go @@ -435,6 +435,59 @@ func TestOpenAIWSIngressSessionPreemptionIsolatesCodexThreads(t *testing.T) { require.NoError(t, otherKeyCtx.Err()) } +func newOpenAIWSPreemptCodexKindContext(apiKeyID int64, threadID, requestKind string) *gin.Context { + c := newOpenAIWSPreemptCodexContext(apiKeyID, threadID) + c.Request.Header.Set(openAIWSTurnMetadataHeader, `{"session_id":"root-session","thread_id":"`+threadID+`","request_kind":"`+requestKind+`"}`) + return c +} + +func TestOpenAIWSIngressSessionPreemptionKeepsDetachedRequestsApart(t *testing.T) { + gin.SetMode(gin.TestMode) + svc := &OpenAIGatewayService{} + account := &Account{ID: 1, Platform: PlatformOpenAI, Type: AccountTypeOAuth} + turnMessage := []byte(`{"type":"response.create","input":"hello"}`) + memoryMessage := []byte(`{"type":"response.create","input":"consolidate"}`) + + turnCtx, turnCleanup, armed := svc.BeginOpenAIWSIngressSessionPreemption( + context.Background(), newOpenAIWSPreemptCodexKindContext(11, "thread-a", "turn"), account, turnMessage, + ) + require.True(t, armed) + defer turnCleanup() + + memoryCtx, memoryCleanup, armed := svc.BeginOpenAIWSIngressSessionPreemption( + context.Background(), newOpenAIWSPreemptCodexKindContext(11, "thread-a", "memory"), account, memoryMessage, + ) + require.True(t, armed) + defer memoryCleanup() + require.NoError(t, turnCtx.Err(), "memory consolidation on the same thread must not preempt the user turn") + require.NoError(t, memoryCtx.Err()) + + prewarmCtx, prewarmCleanup, armed := svc.BeginOpenAIWSIngressSessionPreemption( + context.Background(), newOpenAIWSPreemptCodexKindContext(11, "thread-a", "prewarm"), account, turnMessage, + ) + require.True(t, armed) + defer prewarmCleanup() + require.True(t, IsOpenAIWSSessionPreemptedError(context.Cause(turnCtx)), "prewarm shares the turn lane and replaces the stale turn connection") + require.NoError(t, memoryCtx.Err(), "turn lane reconnect must leave memory consolidation alone") + + _, memoryRetryCleanup, armed := svc.BeginOpenAIWSIngressSessionPreemption( + context.Background(), newOpenAIWSPreemptCodexKindContext(11, "thread-a", "memory"), account, memoryMessage, + ) + require.True(t, armed) + defer memoryRetryCleanup() + require.True(t, IsOpenAIWSSessionPreemptedError(context.Cause(memoryCtx)), "memory reconnect replaces the stale memory connection") + require.NoError(t, prewarmCtx.Err(), "memory reconnect must leave the turn lane alone") + + scorer := newOpenAIWSPreemptCodexContext(11, "") + scorer.Request.Header.Set(openAIWSThreadIDHeader, "thread-a") + scorer.Request.Header.Set(openAISubagentHeader, "guardian") + scorerCtx, scorerCleanup, armed := svc.BeginOpenAIWSIngressSessionPreemption(context.Background(), scorer, account, turnMessage) + require.True(t, armed) + defer scorerCleanup() + require.NoError(t, prewarmCtx.Err(), "guardian scorer without thread metadata must not preempt the user turn") + require.NoError(t, scorerCtx.Err()) +} + func TestOpenAIWSIngressSessionPreemptionSkipsContentOnlyIdentity(t *testing.T) { gin.SetMode(gin.TestMode) groupID := int64(7)