feat(openai): 解析 codex 线程标识派生 WS 执行作用域

2026-10-20 08:19:58 +00:00 · 2026-09-10 17:50:23 +08:00
parent 98d86915be
commit 06b61ba96f
2 changed files with 198 additions and 0 deletions
@@ -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 "", ""
+}
@@ -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)
+}