diff --git a/backend/internal/handler/openai_gateway_handler.go b/backend/internal/handler/openai_gateway_handler.go index d327710423..8936204bf2 100644 --- a/backend/internal/handler/openai_gateway_handler.go +++ b/backend/internal/handler/openai_gateway_handler.go @@ -3016,7 +3016,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() } @@ -3028,6 +3028,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_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 new file mode 100644 index 0000000000..e15ef0a8ba --- /dev/null +++ b/backend/internal/service/openai_ws_execution_scope.go @@ -0,0 +1,156 @@ +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" + + // 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 多智能体会话里 +// 父线程与全部子智能体共用同一个 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()) +} + +// 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 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.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(openAIWSExecutionScopeSeed(apiKeyID, "thread", threadID, lane)) + return scope, threadID + } + if sessionID := strings.TrimSpace(explicitOpenAIRequestSessionID(c, body)); 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 new file mode 100644 index 0000000000..10193a9ae3 --- /dev/null +++ b/backend/internal/service/openai_ws_execution_scope_test.go @@ -0,0 +1,244 @@ +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 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{ + "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 必须参与隔离") + + 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) + 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)) + 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) + + require.Equal(t, "", scopeOf(nil, bodyFirst, 11), "只有请求内容时不得产生执行作用域") + nilScope, nilThread := resolveOpenAIWSExecutionScope(nil, bodyFirst, 11) + require.Equal(t, "", nilScope) + require.Equal(t, "", nilThread) +} diff --git a/backend/internal/service/openai_ws_forwarder_ingress.go b/backend/internal/service/openai_ws_forwarder_ingress.go index 889f921595..ca6135352b 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() { @@ -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..c9a18fc748 --- /dev/null +++ b/backend/internal/service/openai_ws_forwarder_ingress_execution_scope_test.go @@ -0,0 +1,356 @@ +package service + +import ( + "context" + "errors" + "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, "同线程重连必须取代旧连接") + 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) { + preempted++ + } else { + require.NoError(t, err) + } + } + require.Equal(t, 1, preempted, "应恰好有一条会话被抢占") +} 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 d37f41bbd4..facff98ff9 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, @@ -124,6 +125,12 @@ func (s *OpenAIGatewayService) forwardOpenAIWSV2( sessionHash, legacySessionHash = openAIWSSessionHashesFromID(promptCacheKey) attachOpenAILegacySessionHashToGin(c, legacySessionHash) } + // 与 WS 接入路径共用执行作用域：codex 多智能体共用 session-id，turn state 与 + // store=false 的连接绑定必须按线程隔离，同一线程在两条路径之间也才能共享状态。 + // 作用域由 Forward 从改写前的原始请求算出后传入，reqBody 此时已带账号 namespace。 + if executionScope = strings.TrimSpace(executionScope); executionScope != "" { + sessionHash = executionScope + } 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..da13cd1fcb --- /dev/null +++ b/backend/internal/service/openai_ws_forwarder_v2_execution_scope_test.go @@ -0,0 +1,201 @@ +package service + +import ( + "context" + "fmt" + "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) +} + +// 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 改写影响") +} diff --git a/backend/internal/service/openai_ws_session_preemption.go b/backend/internal/service/openai_ws_session_preemption.go index 95fa49e7f7..32b49ccbbc 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: @@ -66,29 +102,50 @@ 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, + notifyPreempted, ) if !armed { return ctx, func() {}, false } if preemptedPrevious { if stateStore := s.getOpenAIWSStateStore(); stateStore != nil { - stateStore.DeleteSessionTurnState(preemptGroupID, preemptSessionHash) - stateStore.DeleteSessionConn(preemptGroupID, preemptSessionHash) + 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 lane=%s", + account.ID, + preemptGroupID, + preemptAPIKeyID, + truncateOpenAIWSLogValue(preemptScope, 12), + truncateOpenAIWSLogValue(preemptThreadID, openAIWSIDValueMaxLen), + truncateOpenAIWSLogValue(lane, openAIWSIDValueMaxLen), + ) } - return context.WithValue(preemptCtx, openAIWSSessionPreemptContextKey{}, true), cleanup, true + return preemptCtx, cleanup, true } func newOpenAIWSSessionPreemptKey(groupID, apiKeyID int64, sessionHash string) (openAIWSSessionPreemptKey, bool) { @@ -146,6 +203,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() @@ -158,16 +216,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) @@ -269,7 +346,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 d00adf63b4..f378758478 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)) @@ -288,3 +383,154 @@ 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 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) + 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") +}