From 1dc0a09009ecfa412763e4ab38673a03692a68f3 Mon Sep 17 00:00:00 2001 From: li Date: Mon, 31 Aug 2026 16:53:57 +0800 Subject: [PATCH] =?UTF-8?q?fix(openai):=20ctx=5Fpool=20WS=20ingress=20?= =?UTF-8?q?=E4=B8=8B=E5=8F=91=E5=89=8D=E6=94=B9=E5=86=99=E4=B8=8A=E6=B8=B8?= =?UTF-8?q?=E5=AE=B9=E9=87=8F=E9=99=8D=E8=BD=BD=E9=94=99=E8=AF=AF=E7=A0=81?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 上游容量降载（server_is_overloaded / slow_down）被 ctx_pool 的 ingress 直写路径 原样转发给客户端。Codex CLI 按闭集判定，这两个码属致命集，客户端打印 "Selected model is at capacity" 并直接终止会话，不退避重试。 仓库已有 sanitizeOpenAICapacityShedErrorCodeForClient 处理这件事，且注释写明了 Codex 的闭集判定与"账号状态判定使用改写前原始 payload"的前提。它此前只在 HTTP/SSE（openai_gateway_response_handling.go）与 http_bridge （openai_ws_http_bridge.go）两条路径被调用，ctx_pool ingress 直写路径是唯一漏掉 的一条——同一账号同一时刻切到 http_bridge 就能正常退避重试。 改写结果写入独立的 clientMessage，不动 upstreamMessage：写出点之后的 markOpenAIWSClientVisibleFailure 与 handleOpenAIWSTerminalTransientFailure 仍按原始 payload 判定账号状态，与 http_bridge 的写法一致。 新增用例覆盖改写生效与改写范围两侧（非容量类错误码必须原样下发）。 --- .../service/openai_ws_forwarder_ingress.go | 19 +- .../openai_ws_ingress_capacity_shed_test.go | 184 ++++++++++++++++++ 2 files changed, 202 insertions(+), 1 deletion(-) create mode 100644 backend/internal/service/openai_ws_ingress_capacity_shed_test.go diff --git a/backend/internal/service/openai_ws_forwarder_ingress.go b/backend/internal/service/openai_ws_forwarder_ingress.go index 5f5976c7d0..b6af8fc06c 100644 --- a/backend/internal/service/openai_ws_forwarder_ingress.go +++ b/backend/internal/service/openai_ws_forwarder_ingress.go @@ -1125,7 +1125,24 @@ func (s *OpenAIGatewayService) ProxyResponsesWebSocketFromClient( } } replayCollector.AddEvent(eventType, upstreamMessage) - if err := writeClientMessage(upstreamMessage); err != nil { + // 客户端写出副本改写容量降载码：Codex 对 error/response.failed 中的 + // server_is_overloaded / slow_down 判致命并终止会话，改写后走客户端 + // 内置退避重试。HTTP/SSE（openai_gateway_response_handling.go）与 + // http_bridge（openai_ws_http_bridge.go）两条路径早已这么做， + // ctx_pool 的 ingress 直写路径是唯一漏掉的一条 —— 同一个上游降载 + // 事件在这里会让会话就地终止，切到 http_bridge 却能正常退避重试。 + // + // 必须写进独立变量而不是原地改 upstreamMessage：下面的 + // markOpenAIWSClientVisibleFailure 与 handleOpenAIWSTerminalTransientFailure + // 仍要按未改写的原始 payload 判定账号状态，这正是 + // sanitizeOpenAICapacityShedErrorCodeForClient 注释里写明的前提。 + clientMessage := upstreamMessage + if eventType == "error" || eventType == "response.failed" { + if rewritten, changed := sanitizeOpenAICapacityShedErrorCodeForClient(clientMessage); changed { + clientMessage = rewritten + } + } + if err := writeClientMessage(clientMessage); err != nil { if isOpenAIWSClientDisconnectError(err) { clientDisconnected = true closeStatus, closeReason := summarizeOpenAIWSReadCloseError(err) diff --git a/backend/internal/service/openai_ws_ingress_capacity_shed_test.go b/backend/internal/service/openai_ws_ingress_capacity_shed_test.go new file mode 100644 index 0000000000..d37b7f2cda --- /dev/null +++ b/backend/internal/service/openai_ws_ingress_capacity_shed_test.go @@ -0,0 +1,184 @@ +package service + +import ( + "context" + "net/http" + "net/http/httptest" + "strings" + "testing" + "time" + + coderws "github.com/coder/websocket" + "github.com/gin-gonic/gin" + "github.com/stretchr/testify/require" +) + +// openAIWSIngressCapacityShedRepo 补齐 SetError，避免非容量类错误（如 +// workspace_suspended）走到账号状态副作用时打空指针。 +type openAIWSIngressCapacityShedRepo struct { + stubOpenAIAccountRepo +} + +func (r *openAIWSIngressCapacityShedRepo) SetError(context.Context, int64, string) error { return nil } + +func (r *openAIWSIngressCapacityShedRepo) SetRateLimited(context.Context, int64, time.Time) error { + return nil +} + +func (r *openAIWSIngressCapacityShedRepo) UpdateExtra(context.Context, int64, map[string]any) error { + return nil +} + +// ctx_pool 的 ingress 直写路径把 error / response.failed 交给 WS 客户端前，必须和 +// HTTP/SSE（openai_gateway_response_handling.go）与 http_bridge +// （openai_ws_http_bridge.go）两条路径一样，把容量降载码改写为可重试的 +// server_error：Codex 按闭集判定，server_is_overloaded / slow_down 属致命集， +// 客户端会打印 "Selected model is at capacity" 并直接终止会话而不是退避重试。 +// +// 第二个用例锁住改写范围：非容量类错误码必须原样下发，客户端依赖原码各自处理。 +func TestProxyResponsesWebSocketFromClient_RewritesCapacityShedCodeForClient(t *testing.T) { + gin.SetMode(gin.TestMode) + + tests := []struct { + name string + upstreamEvents [][]byte + wantContains []string + wantAbsent []string + }{ + { + name: "capacity_shed_error_and_failed_are_rewritten", + upstreamEvents: [][]byte{ + []byte(`{"type":"error","error":{"type":"service_unavailable_error","code":"server_is_overloaded","message":"Our servers are currently overloaded. Please try again later."}}`), + []byte(`{"type":"response.failed","response":{"id":"resp_shed","status":"failed","error":{"code":"server_is_overloaded","message":"Our servers are currently overloaded. Please try again later."}}}`), + }, + wantContains: []string{ + `"code":"server_error"`, + "Our servers are currently overloaded", + }, + wantAbsent: []string{"server_is_overloaded"}, + }, + { + name: "non_capacity_error_code_is_passed_through", + upstreamEvents: [][]byte{ + []byte(`{"type":"error","error":{"type":"invalid_request_error","code":"workspace_suspended","message":"workspace is suspended"}}`), + []byte(`{"type":"response.failed","response":{"id":"resp_suspended","status":"failed","error":{"code":"workspace_suspended","message":"workspace is suspended"}}}`), + }, + wantContains: []string{ + `"code":"workspace_suspended"`, + "workspace is suspended", + }, + wantAbsent: []string{"server_error"}, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + cfg := newOpenAIWSV2TestConfig() + cfg.Security.URLAllowlist.Enabled = false + cfg.Security.URLAllowlist.AllowInsecureHTTP = true + cfg.Gateway.OpenAIWS.MaxConnsPerAccount = 1 + cfg.Gateway.OpenAIWS.MinIdlePerAccount = 0 + cfg.Gateway.OpenAIWS.MaxIdlePerAccount = 1 + cfg.Gateway.OpenAIWS.QueueLimitPerConn = 8 + cfg.Gateway.OpenAIWS.DialTimeoutSeconds = 3 + cfg.Gateway.OpenAIWS.ReadTimeoutSeconds = 3 + cfg.Gateway.OpenAIWS.WriteTimeoutSeconds = 3 + + events := make([][]byte, 0, len(tt.upstreamEvents)) + for _, event := range tt.upstreamEvents { + events = append(events, append([]byte(nil), event...)) + } + captureConn := &openAIWSCaptureConn{events: events} + pool := newOpenAIWSConnPool(cfg) + pool.setClientDialerForTest(&openAIWSCaptureDialer{conn: captureConn}) + + account := Account{ + ID: 5401, + Name: "openai-ingress-capacity-shed", + 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}, + } + repo := &openAIWSIngressCapacityShedRepo{stubOpenAIAccountRepo: stubOpenAIAccountRepo{accounts: []Account{account}}} + svc := &OpenAIGatewayService{ + accountRepo: repo, + rateLimitService: &RateLimitService{accountRepo: repo}, + httpUpstream: &httpUpstreamRecorder{}, + cache: &stubGatewayCache{}, + cfg: cfg, + openaiWSResolver: NewOpenAIWSProtocolResolver(cfg), + toolCorrector: NewCodexToolCorrector(), + openaiWSPool: pool, + } + + serverDone := make(chan struct{}) + wsServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + defer close(serverDone) + conn, err := coderws.Accept(w, r, &coderws.AcceptOptions{CompressionMode: coderws.CompressionContextTakeover}) + if err != nil { + 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 + + readCtx, cancel := context.WithTimeout(r.Context(), 3*time.Second) + msgType, firstMessage, readErr := conn.Read(readCtx) + cancel() + if readErr != nil || (msgType != coderws.MessageText && msgType != coderws.MessageBinary) { + return + } + _ = svc.ProxyResponsesWebSocketFromClient(r.Context(), ginCtx, conn, &account, "sk-test", firstMessage, nil) + })) + defer wsServer.Close() + + dialCtx, cancelDial := context.WithTimeout(context.Background(), 3*time.Second) + clientConn, _, err := coderws.Dial(dialCtx, "ws"+strings.TrimPrefix(wsServer.URL, "http"), nil) + cancelDial() + require.NoError(t, err) + defer func() { _ = clientConn.CloseNow() }() + + writeCtx, cancelWrite := context.WithTimeout(context.Background(), 3*time.Second) + err = clientConn.Write(writeCtx, coderws.MessageText, []byte(`{"type":"response.create","model":"gpt-5.1","stream":false}`)) + cancelWrite() + require.NoError(t, err) + + var frames []string + for len(frames) < len(tt.upstreamEvents) { + readCtx, cancel := context.WithTimeout(context.Background(), 3*time.Second) + _, message, readErr := clientConn.Read(readCtx) + cancel() + if readErr != nil { + break + } + frames = append(frames, string(message)) + } + // 本轮已终止，主动断开客户端让 ingress 退出 turn 循环。 + _ = clientConn.CloseNow() + + require.NotEmpty(t, frames, "客户端应至少收到一个下发事件") + joined := strings.Join(frames, "\n") + for _, want := range tt.wantContains { + require.Contains(t, joined, want, "客户端收到的事件:\n%s", joined) + } + for _, absent := range tt.wantAbsent { + require.NotContains(t, joined, absent, "客户端收到的事件:\n%s", joined) + } + + select { + case <-serverDone: + case <-time.After(5 * time.Second): + t.Fatal("等待 ingress websocket 结束超时") + } + }) + } +}