Merge pull request #6417 from fengshao1227/fix/ws-ingress-capacity-shed-client-code

fix(openai): ctx_pool WS ingress 下发前改写上游容量降载错误码（Codex 会话被硬终止）
2026-10-14 03:52:26 +08:00 · 2026-09-01 08:54:00 +08:00
parent 9d9ed1cc62 1dc0a09009
commit 1cce2b38ee
2 changed files with 202 additions and 1 deletions
@@ -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)
@@ -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 结束超时")
+			}
+		})
+	}
+}