diff --git a/backend/internal/service/openai_ws_forwarder_ingress.go b/backend/internal/service/openai_ws_forwarder_ingress.go index 09e16dc704..5f5976c7d0 100644 --- a/backend/internal/service/openai_ws_forwarder_ingress.go +++ b/backend/internal/service/openai_ws_forwarder_ingress.go @@ -133,6 +133,10 @@ func (s *OpenAIGatewayService) ProxyResponsesWebSocketFromClient( if wsDecision.Transport != OpenAIUpstreamTransportResponsesWebsocketV2 { return fmt.Errorf("websocket ingress requires ws_v2 transport, got=%s", wsDecision.Transport) } + if s.shouldBridgeOpenAIWSPassthroughFirstMessage(account, firstClientMessage) { + forceHTTPBridge = true + break + } // 透传 relay 通过 TurnStarted 记录每个 turn 的开始时刻，但不触发 // BeforeTurn；因此仍只有建连时的利润准入门，没有 turn 级复核。 // handler 计费在 turn 定价未冻结时回退到对应的 turn 开始时刻。 diff --git a/backend/internal/service/openai_ws_forwarder_ingress_session_test.go b/backend/internal/service/openai_ws_forwarder_ingress_session_test.go index b05a2efd10..150a87b079 100644 --- a/backend/internal/service/openai_ws_forwarder_ingress_session_test.go +++ b/backend/internal/service/openai_ws_forwarder_ingress_session_test.go @@ -1118,6 +1118,141 @@ func TestOpenAIGatewayService_ProxyResponsesWebSocketFromClient_PassthroughModeR require.Len(t, upstreamConn.writes, 1, "passthrough 模式应透传首条 response.create") } +func TestOpenAIGatewayService_ProxyResponsesWebSocketFromClient_PassthroughBridgeBoundaries(t *testing.T) { + tests := []struct { + name string + payload string + threshold int64 + wantRelayReject bool + }{ + { + name: "small response create", + payload: `{"type":"response.create","model":"gpt-5.1"}`, + threshold: 1024, + }, + { + name: "continuation response create", + payload: `{"type":"response.create","previous_response_id":"resp_previous","model":"gpt-5.1"}`, + threshold: 1, + }, + { + name: "other event type", + payload: `{"type":"session.update","padding":"xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx"}`, + threshold: 1, + }, + { + name: "malformed data", + payload: `{"type":"response.create","padding":"xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx"`, + threshold: 1, + wantRelayReject: true, + }, + { + name: "duplicate type", + payload: `{"type":"response.create","type":"response.create","model":"gpt-5.1"}`, + threshold: 1, + }, + { + name: "duplicate previous response id", + payload: `{"type":"response.create","previous_response_id":null,"previous_response_id":null,"model":"gpt-5.1"}`, + threshold: 1, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + gin.SetMode(gin.TestMode) + cfg := &config.Config{} + cfg.Security.URLAllowlist.Enabled = false + cfg.Security.URLAllowlist.AllowInsecureHTTP = true + cfg.Gateway.OpenAIWS.Enabled = true + cfg.Gateway.OpenAIWS.APIKeyEnabled = true + cfg.Gateway.OpenAIWS.ResponsesWebsocketsV2 = true + cfg.Gateway.OpenAIWS.ModeRouterV2Enabled = true + cfg.Gateway.OpenAIWS.IngressModeDefault = OpenAIWSIngressModeCtxPool + cfg.Gateway.OpenAIWS.HTTPBridgeEnabled = true + cfg.Gateway.OpenAIWS.HTTPBridgeThresholdBytes = tt.threshold + + upstreamConn := &openAIWSCaptureConn{events: [][]byte{ + []byte(`{"type":"response.completed","response":{"id":"resp_duplicate_keys","model":"gpt-5.1","usage":{"input_tokens":1,"output_tokens":1}}}`), + }} + dialer := &openAIWSCaptureDialer{conn: upstreamConn} + httpUpstream := &httpUpstreamRecorder{} + svc := &OpenAIGatewayService{ + cfg: cfg, + httpUpstream: httpUpstream, + cache: &stubGatewayCache{}, + openaiWSResolver: NewOpenAIWSProtocolResolver(cfg), + toolCorrector: NewCodexToolCorrector(), + openaiWSPassthroughDialer: dialer, + } + account := &Account{ + ID: 453, + Platform: PlatformOpenAI, + Type: AccountTypeAPIKey, + Status: StatusActive, + Schedulable: true, + Concurrency: 1, + Credentials: map[string]any{"api_key": "sk-test"}, + Extra: map[string]any{ + "openai_apikey_responses_websockets_v2_mode": OpenAIWSIngressModePassthrough, + }, + } + + errCh := make(chan error, 1) + wsServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + conn, err := coderws.Accept(w, r, nil) + if err != nil { + errCh <- err + return + } + defer func() { _ = conn.CloseNow() }() + _, firstMessage, err := conn.Read(r.Context()) + if err != nil { + errCh <- err + return + } + rec := httptest.NewRecorder() + ginCtx, _ := gin.CreateTestContext(rec) + ginCtx.Request = r + errCh <- svc.ProxyResponsesWebSocketFromClient(r.Context(), ginCtx, conn, account, "sk-test", firstMessage, nil) + })) + defer wsServer.Close() + + clientConn, _, err := coderws.Dial(context.Background(), "ws"+strings.TrimPrefix(wsServer.URL, "http"), nil) + require.NoError(t, err) + defer func() { _ = clientConn.CloseNow() }() + require.NoError(t, clientConn.Write(context.Background(), coderws.MessageText, []byte(tt.payload))) + _, event, err := clientConn.Read(context.Background()) + if tt.wantRelayReject { + require.Error(t, err) + } else { + require.NoError(t, err) + require.Equal(t, "resp_duplicate_keys", gjson.GetBytes(event, "response.id").String()) + require.NoError(t, clientConn.Close(coderws.StatusNormalClosure, "done")) + } + + select { + case proxyErr := <-errCh: + if tt.wantRelayReject { + require.Error(t, proxyErr) + } else if proxyErr != nil { + require.Contains(t, proxyErr.Error(), "StatusNormalClosure") + } + case <-time.After(5 * time.Second): + t.Fatal("等待 boundary passthrough websocket 结束超时") + } + require.Empty(t, httpUpstream.requests, "boundary frame must not use the HTTP bridge") + if tt.wantRelayReject { + require.Zero(t, dialer.DialCount()) + require.Empty(t, upstreamConn.writes) + } else { + require.Equal(t, 1, dialer.DialCount()) + require.Len(t, upstreamConn.writes, 1) + } + }) + } +} + func TestOpenAIGatewayService_ProxyResponsesWebSocketFromClient_PassthroughHeadersUsePromptCacheAndTurnState(t *testing.T) { gin.SetMode(gin.TestMode) diff --git a/backend/internal/service/openai_ws_http_bridge.go b/backend/internal/service/openai_ws_http_bridge.go index ef5c73e129..8e5cbf6159 100644 --- a/backend/internal/service/openai_ws_http_bridge.go +++ b/backend/internal/service/openai_ws_http_bridge.go @@ -115,6 +115,131 @@ func (s *OpenAIGatewayService) shouldBridgeOpenAIWSHTTP(account *Account, payloa return threshold > 0 && int64(payloadBytes) >= threshold } +func (s *OpenAIGatewayService) shouldBridgeOpenAIWSPassthroughFirstMessage(account *Account, payload []byte) bool { + if account != nil && account.Platform == PlatformGrok { + return true + } + if !s.openAIWSHTTPBridgeEnabled() || int64(len(payload)) < s.openAIWSHTTPBridgeThresholdBytes() { + return false + } + if !json.Valid(payload) { + return false + } + + i := skipOpenAIWSJSONSpace(payload, 0) + if i >= len(payload) || payload[i] != '{' { + return false + } + i++ + eventType := "response.create" + previousResponseID := "" + typeSeen, previousResponseIDSeen := false, false + for { + i = skipOpenAIWSJSONSpace(payload, i) + if payload[i] == '}' { + break + } + keyStart := i + keyEnd := scanOpenAIWSJSONString(payload, keyStart) + i = skipOpenAIWSJSONSpace(payload, keyEnd) + i++ // json.Valid guarantees the colon. + i = skipOpenAIWSJSONSpace(payload, i) + valueStart := i + i = skipOpenAIWSJSONValue(payload, i) + + key := "" + // A critical key is at most 20 decoded bytes. The generous encoded bound + // covers escaped spellings without allocating attacker-sized key strings. + if keyEnd-keyStart <= 128 { + _ = json.Unmarshal(payload[keyStart:keyEnd], &key) + } + switch key { + case "type": + if typeSeen { + return false + } + typeSeen = true + var value *string + if err := json.Unmarshal(payload[valueStart:i], &value); err != nil { + return false + } + if value == nil || strings.TrimSpace(*value) == "" { + eventType = "response.create" + } else { + eventType = strings.TrimSpace(*value) + } + case "previous_response_id": + if previousResponseIDSeen { + return false + } + previousResponseIDSeen = true + var value *string + if err := json.Unmarshal(payload[valueStart:i], &value); err != nil { + return false + } + if value != nil { + previousResponseID = strings.TrimSpace(*value) + } + } + i = skipOpenAIWSJSONSpace(payload, i) + if payload[i] == ',' { + i++ + } + } + return eventType == "response.create" && previousResponseID == "" +} + +func skipOpenAIWSJSONSpace(payload []byte, i int) int { + for i < len(payload) { + switch payload[i] { + case ' ', '\t', '\r', '\n': + i++ + default: + return i + } + } + return i +} + +func scanOpenAIWSJSONString(payload []byte, i int) int { + for i++; i < len(payload); i++ { + switch payload[i] { + case '\\': + i++ + case '"': + return i + 1 + } + } + return len(payload) +} + +func skipOpenAIWSJSONValue(payload []byte, i int) int { + if payload[i] == '"' { + return scanOpenAIWSJSONString(payload, i) + } + if payload[i] != '{' && payload[i] != '[' { + for i < len(payload) && payload[i] != ',' && payload[i] != '}' { + i++ + } + return i + } + depth := 0 + for ; i < len(payload); i++ { + switch payload[i] { + case '"': + i = scanOpenAIWSJSONString(payload, i) - 1 + case '{', '[': + depth++ + case '}', ']': + depth-- + if depth == 0 { + return i + 1 + } + } + } + return len(payload) +} + func prepareOpenAIWSHTTPBridgeBody(account *Account, payload []byte) ([]byte, error) { var body map[string]any if err := decodeOpenAIJSONUseNumber(payload, &body); err != nil { diff --git a/backend/internal/service/openai_ws_http_bridge_test.go b/backend/internal/service/openai_ws_http_bridge_test.go index bbff773d76..2a52b2ac39 100644 --- a/backend/internal/service/openai_ws_http_bridge_test.go +++ b/backend/internal/service/openai_ws_http_bridge_test.go @@ -720,6 +720,43 @@ func TestOpenAIWSHTTPBridgeDecisionKeepsSmallFramesOnWS(t *testing.T) { require.True(t, svc.shouldBridgeOpenAIWSHTTP(&Account{Platform: PlatformGrok}, 1, "resp_existing")) } +func TestOpenAIWSPassthroughFirstMessageBridgeDecision(t *testing.T) { + const threshold = 100 + svc := &OpenAIGatewayService{cfg: &config.Config{Gateway: config.GatewayConfig{ + OpenAIWS: config.GatewayOpenAIWSConfig{ + HTTPBridgeEnabled: true, + HTTPBridgeThresholdBytes: threshold, + }, + }}} + exactThresholdPayload := `{"type":"response.create","input":"` + + strings.Repeat("x", threshold-len(`{"type":"response.create","input":""}`)) + `"}` + + tests := []struct { + name string + payload string + want bool + }{ + {name: "oversized without previous response id bridges", payload: `{"type":"response.create","input":"` + strings.Repeat("x", 100) + `"}`, want: true}, + {name: "omitted type bridges as response create", payload: `{"input":"` + strings.Repeat("x", 100) + `"}`, want: true}, + {name: "blank type bridges as response create", payload: `{"type":" ","padding":"` + strings.Repeat("x", 100) + `"}`, want: true}, + {name: "exact threshold bridges", payload: exactThresholdPayload, want: true}, + {name: "small stays passthrough", payload: `{"type":"response.create","input":"x"}`}, + {name: "previous response id stays passthrough", payload: `{"type":"response.create","previous_response_id":"resp_previous","input":"` + strings.Repeat("x", 100) + `"}`}, + {name: "duplicate type stays passthrough", payload: `{"type":"response.create","type":"response.create","input":"` + strings.Repeat("x", 100) + `"}`}, + {name: "duplicate previous response id stays passthrough", payload: `{"type":"response.create","previous_response_id":null,"previous_response_id":null,"input":"` + strings.Repeat("x", 100) + `"}`}, + {name: "null type bridges as response create", payload: `{"type":null,"padding":"` + strings.Repeat("x", 100) + `"}`, want: true}, + {name: "non-string type stays passthrough", payload: `{"type":123,"padding":"` + strings.Repeat("x", 100) + `"}`}, + {name: "oversized response cancel stays passthrough", payload: `{"type":"response.cancel","padding":"` + strings.Repeat("x", 100) + `"}`}, + {name: "oversized other event stays passthrough", payload: `{"type":"session.update","padding":"` + strings.Repeat("x", 100) + `"}`}, + {name: "oversized malformed JSON stays passthrough", payload: `{"type":"response.create","padding":"` + strings.Repeat("x", 100)}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + require.Equal(t, tt.want, svc.shouldBridgeOpenAIWSPassthroughFirstMessage(nil, []byte(tt.payload))) + }) + } +} + func TestProxyOpenAIWSHTTPBridgeTurnTransportErrorFailoverSafety(t *testing.T) { gin.SetMode(gin.TestMode) @@ -1665,9 +1702,11 @@ func TestOpenAIWSHTTPBridgeAcceptsFirstFrameAboveLegacy16MiB(t *testing.T) { Enabled: true, APIKeyEnabled: true, ResponsesWebsocketsV2: true, + ModeRouterV2Enabled: true, + IngressModeDefault: OpenAIWSIngressModeCtxPool, ClientReadLimitBytes: 64 * 1024 * 1024, HTTPBridgeEnabled: true, - HTTPBridgeThresholdBytes: 15 * 1024 * 1024, + HTTPBridgeThresholdBytes: 17*1024*1024 + 512, }, }, } @@ -1677,22 +1716,40 @@ func TestOpenAIWSHTTPBridgeAcceptsFirstFrameAboveLegacy16MiB(t *testing.T) { toolCorrector: NewCodexToolCorrector(), } account := &Account{ - ID: 9, - Name: "api-key", - Platform: PlatformOpenAI, - Type: AccountTypeAPIKey, - Credentials: map[string]any{"api_key": "sk-upstream"}, + ID: 9, + Name: "api-key", + Platform: PlatformOpenAI, + Type: AccountTypeAPIKey, + Credentials: map[string]any{ + "api_key": "sk-upstream", + "base_url": "https://env-openai.example/v1", + "model_provider": "env-openai", + }, Extra: map[string]any{ "openai_apikey_responses_websockets_v2_enabled": true, + "openai_apikey_responses_websockets_v2_mode": OpenAIWSIngressModePassthrough, }, Concurrency: 1, Status: StatusActive, } - payload := []byte(`{"type":"response.create","generate":true,"model":"gpt-5","stream":true,"input":"` + strings.Repeat("x", 17*1024*1024) + `"}`) + payload := []byte(strings.Repeat(" ", 1024) + `{"type":"response.create","generate":true,"model":"gpt-5","stream":true,"input":"` + strings.Repeat("x", 17*1024*1024) + `"}`) require.Greater(t, len(payload), 16*1024*1024) + require.GreaterOrEqual(t, int64(len(payload)), cfg.Gateway.OpenAIWS.HTTPBridgeThresholdBytes) require.Less(t, int64(len(payload)), ResolveOpenAIWSClientReadLimitBytes(cfg)) + type turnOutcome struct { + turn int + result *OpenAIForwardResult + err error + } + turnOutcomeCh := make(chan turnOutcome, 1) + hooks := &OpenAIWSIngressHooks{ + AfterTurn: func(turn int, result *OpenAIForwardResult, turnErr error) { + turnOutcomeCh <- turnOutcome{turn: turn, result: result, err: turnErr} + }, + } + errCh := make(chan error, 1) wsServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { conn, err := coderws.Accept(w, r, &coderws.AcceptOptions{CompressionMode: coderws.CompressionContextTakeover}) @@ -1724,7 +1781,7 @@ func TestOpenAIWSHTTPBridgeAcceptsFirstFrameAboveLegacy16MiB(t *testing.T) { proxyCtx, cancelProxy := context.WithTimeout(r.Context(), 20*time.Second) defer cancelProxy() - errCh <- svc.ProxyResponsesWebSocketFromClient(proxyCtx, ginCtx, conn, account, "sk-test", firstMessage, nil) + errCh <- svc.ProxyResponsesWebSocketFromClient(proxyCtx, ginCtx, conn, account, "sk-test", firstMessage, hooks) })) defer wsServer.Close() @@ -1764,8 +1821,23 @@ func TestOpenAIWSHTTPBridgeAcceptsFirstFrameAboveLegacy16MiB(t *testing.T) { t.Fatal("timed out waiting for websocket bridge proxy to finish") } + select { + case outcome := <-turnOutcomeCh: + require.NoError(t, outcome.err) + require.Equal(t, 1, outcome.turn) + require.NotNil(t, outcome.result) + require.Equal(t, 9, outcome.result.Usage.InputTokens) + require.Equal(t, 1, outcome.result.Usage.OutputTokens) + require.Equal(t, "response.completed", outcome.result.UpstreamTerminalEvent) + default: + t.Fatal("AfterTurn was not called for websocket HTTP bridge turn") + } + require.NotNil(t, upstream.lastReq) require.Equal(t, http.MethodPost, upstream.lastReq.Method) + require.Equal(t, "https://env-openai.example/v1/responses", upstream.lastReq.URL.String()) + require.Equal(t, "env-openai", account.GetCredential("model_provider")) + require.Less(t, int64(len(upstream.lastBody)), cfg.Gateway.OpenAIWS.HTTPBridgeThresholdBytes) require.Greater(t, len(upstream.lastBody), 16*1024*1024) require.False(t, gjson.GetBytes(upstream.lastBody, "type").Exists()) require.False(t, gjson.GetBytes(upstream.lastBody, "generate").Exists())