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)