From cd1ee1d1ad735888ed63ae72b21ecbe9105ec8a8 Mon Sep 17 00:00:00 2001 From: xlplbo Date: Thu, 10 Sep 2026 17:50:24 +0800 Subject: [PATCH] =?UTF-8?q?fix(openai):=20WS=20=E6=8E=A5=E5=85=A5=E4=BC=9A?= =?UTF-8?q?=E8=AF=9D=E7=8A=B6=E6=80=81=E6=94=B9=E7=94=A8=E6=89=A7=E8=A1=8C?= =?UTF-8?q?=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, "应恰好有一条会话被抢占") +}