diff --git a/backend/internal/handler/openai_gateway_cyber_test.go b/backend/internal/handler/openai_gateway_cyber_test.go index 1b9fc53a9d..4be214cdc9 100644 --- a/backend/internal/handler/openai_gateway_cyber_test.go +++ b/backend/internal/handler/openai_gateway_cyber_test.go @@ -1,6 +1,7 @@ package handler import ( + "net/http" "net/http/httptest" "strings" "testing" @@ -102,6 +103,44 @@ func TestClearCyberPolicyTurnState(t *testing.T) { require.Equal(t, "turn2", service.GetOpsCyberPolicy(c).Message) } +func TestAdvanceOpenAIWSCyberBlockStateDefersBlockAcrossFailover(t *testing.T) { + failoverErr := &service.UpstreamFailoverError{StatusCode: http.StatusTooManyRequests} + + blocked, pending := advanceOpenAIWSCyberBlockState(false, false, true, failoverErr) + require.False(t, blocked, "the replacement account must receive the current turn") + require.True(t, pending, "the cyber hit must still block later turns") + + blocked, pending = advanceOpenAIWSCyberBlockState(blocked, pending, false, failoverErr) + require.False(t, blocked, "additional failover attempts must remain eligible") + require.True(t, pending) + + blocked, pending = advanceOpenAIWSCyberBlockState(blocked, pending, false, nil) + require.True(t, blocked, "the next client turn must be blocked after failover finishes") + require.False(t, pending) +} + +func TestClearCyberPolicyAttemptStatePreservesRecordedGuardDuringFailover(t *testing.T) { + c := newTestGinContext() + h := &OpenAIGatewayHandler{} + + service.MarkOpsCyberPolicy(c, service.CyberPolicyMark{Message: "account-a", UpstreamStatus: http.StatusOK}) + h.recordCyberPolicyIfMarked(c, nil, nil, nil, "gpt-5", true, nil, service.ChannelUsageFields{}, "") + require.True(t, c.GetBool(cyberPolicyRecordedKey)) + + clearCyberPolicyAttemptState(c, false) + require.Nil(t, service.GetOpsCyberPolicy(c)) + require.True(t, c.GetBool(cyberPolicyRecordedKey), "the same logical turn must not record again after failover") + + service.MarkOpsCyberPolicy(c, service.CyberPolicyMark{Message: "account-b", UpstreamStatus: http.StatusOK}) + h.recordCyberPolicyIfMarked(c, nil, nil, nil, "gpt-5", true, nil, service.ChannelUsageFields{}, "") + require.Equal(t, "account-b", service.GetOpsCyberPolicy(c).Message) + require.True(t, c.GetBool(cyberPolicyRecordedKey)) + + clearCyberPolicyAttemptState(c, true) + require.Nil(t, service.GetOpsCyberPolicy(c)) + require.False(t, c.GetBool(cyberPolicyRecordedKey), "a completed logical turn must reset the guard") +} + // TestBuildCyberSessionBlockedOpsEntry verifies the locally-rejected request is // auditable: 403 / phase=request / type=cyber_policy_session_blocked — distinct // from upstream cyber_policy hits, and it must NOT touch moderation/violation. diff --git a/backend/internal/handler/openai_gateway_handler.go b/backend/internal/handler/openai_gateway_handler.go index 7ad0f72c68..d4e34e3b7b 100644 --- a/backend/internal/handler/openai_gateway_handler.go +++ b/backend/internal/handler/openai_gateway_handler.go @@ -53,6 +53,21 @@ type openAIWSTurnChannelMappingSnapshot struct { mapping service.ChannelMappingResult } +func advanceOpenAIWSCyberBlockState(blocked, pending, marked bool, turnErr error) (bool, bool) { + var failoverErr *service.UpstreamFailoverError + isFailover := errors.As(turnErr, &failoverErr) + if marked { + if isFailover { + return false, true + } + return true, false + } + if pending && !isFailover { + return true, false + } + return blocked, pending +} + var errOpenAIWSUnsupportedModelSwitch = errors.New("selected account does not support websocket model switch") func newOpenAIWSUnsupportedModelSwitchError(model string) error { @@ -2380,6 +2395,7 @@ func (h *OpenAIGatewayHandler) ResponsesWebSocket(c *gin.Context) { return } cyberBlockedThisConn := false + cyberBlockPendingAfterFailover := false var cyberTurnBodiesMu sync.Mutex cyberTurnBodies := map[int][]byte{1: append([]byte(nil), firstMessage...)} setCyberTurnBody := func(turn int, payload []byte) { @@ -2820,10 +2836,12 @@ func (h *OpenAIGatewayHandler) ResponsesWebSocket(c *gin.Context) { AfterTurn: func(turn int, result *service.OpenAIForwardResult, turnErr error) { turnStart := getTurnStart(turn) cyberBlockBody := takeCyberTurnBody(turn) - // F1: cyber 标记按 turn 生命周期清理——defer 保证任意早返回路径都执行； - // CyberBlocked 必须在 submit 前同步预捕获（task 闭包由 worker 池异步执行， - // 届时 defer 已清除标记）。 - defer clearCyberPolicyTurnState(c) + // 每次 attempt 都清 cyber mark；failover 链结束前保留 recorded guard， + // 避免同一逻辑 turn 换号后重复落风控。CyberBlocked 必须在 submit 前 + // 同步预捕获（task 闭包由 worker 池异步执行，届时 mark 已清除）。 + defer func() { + clearCyberPolicyAttemptState(c, !cyberBlockPendingAfterFailover) + }() releaseTurnSlots() turnRequestedModel := reqModel turnUpstreamModel := "" @@ -2845,10 +2863,14 @@ func (h *OpenAIGatewayHandler) ResponsesWebSocket(c *gin.Context) { turnUpstreamModel = turnRequestedModel } turnUsageFields := turnMapping.ToUsageFields(turnRequestedModel, turnUpstreamModel) + cyberMarked := service.GetOpsCyberPolicy(c) != nil h.recordCyberPolicyIfMarked(c, apiKey, account, subscription, turnRequestedModel, turnErr != nil, cyberBlockBody, turnUsageFields, requestPayloadHash) - if service.GetOpsCyberPolicy(c) != nil { - cyberBlockedThisConn = true - } + cyberBlockedThisConn, cyberBlockPendingAfterFailover = advanceOpenAIWSCyberBlockState( + cyberBlockedThisConn, + cyberBlockPendingAfterFailover, + cyberMarked, + turnErr, + ) if turnErr != nil { if result == nil || result.ImageCount <= 0 { return @@ -4155,15 +4177,20 @@ func (h *OpenAIGatewayHandler) recordCyberPolicyIfMarked(c *gin.Context, apiKey }() } -// clearCyberPolicyTurnState resets the cyber mark and the per-request recorded -// guard. WS-only: called at the END of AfterTurn, after recordCyberPolicyIfMarked -// and RecordUsage (which reads CyberBlocked) have both consumed the mark. +// clearCyberPolicyTurnState resets the cyber mark and recorded guard after a +// logical WS turn has finished. func clearCyberPolicyTurnState(c *gin.Context) { + clearCyberPolicyAttemptState(c, true) +} + +func clearCyberPolicyAttemptState(c *gin.Context, resetRecorded bool) { if c == nil { return } service.ClearOpsCyberPolicy(c) - c.Set(cyberPolicyRecordedKey, false) + if resetRecorded { + c.Set(cyberPolicyRecordedKey, false) + } } func summarizeWSCloseErrorForLog(err error) (string, string) { diff --git a/backend/internal/service/openai_cyber_policy.go b/backend/internal/service/openai_cyber_policy.go index 524bade30c..9e51e72483 100644 --- a/backend/internal/service/openai_cyber_policy.go +++ b/backend/internal/service/openai_cyber_policy.go @@ -86,3 +86,22 @@ func detectOpenAICyberPolicy(payload []byte) (bool, string, string) { } return true, "cyber_policy", strings.TrimSpace(msg) } + +func markOpenAICyberPolicyEvent(c *gin.Context, payload []byte, upstreamStatus int, usage *OpenAIUsage) bool { + hit, code, message := detectOpenAICyberPolicy(payload) + if !hit { + return false + } + mark := CyberPolicyMark{ + Code: code, + Message: message, + Body: truncateString(string(payload), 4096), + UpstreamStatus: upstreamStatus, + } + if usage != nil { + mark.UpstreamInTok = usage.InputTokens + mark.UpstreamOutTok = usage.OutputTokens + } + MarkOpsCyberPolicy(c, mark) + return true +} diff --git a/backend/internal/service/openai_ws_forwarder_ingress.go b/backend/internal/service/openai_ws_forwarder_ingress.go index 74eeac702b..aedd078ec6 100644 --- a/backend/internal/service/openai_ws_forwarder_ingress.go +++ b/backend/internal/service/openai_ws_forwarder_ingress.go @@ -1021,6 +1021,12 @@ func (s *OpenAIGatewayService) ProxyResponsesWebSocketFromClient( } lastEventType = eventType } + if openAIWSMessageShouldParseUsage(eventType, upstreamMessage) { + parseOpenAIWSResponseUsageFromCompletedEvent(upstreamMessage, &usage) + } + if eventType == "error" || eventType == "response.failed" { + markOpenAICyberPolicyEvent(c, upstreamMessage, http.StatusOK, &usage) + } if eventType == "error" { s.handleOpenAIWSErrorEventTransientFailure(ctx, account, mappedModel, lease.HandshakeHeaders(), upstreamMessage) errCodeRaw, errTypeRaw, errMsgRaw := parseOpenAIWSErrorEventFields(upstreamMessage) @@ -1135,23 +1141,8 @@ func (s *OpenAIGatewayService) ProxyResponsesWebSocketFromClient( ms := int(time.Since(turnStart).Milliseconds()) firstTokenMs = &ms } - if openAIWSMessageShouldParseUsage(eventType, upstreamMessage) { - parseOpenAIWSResponseUsageFromCompletedEvent(upstreamMessage, &usage) - } imageCounter.AddSSEData(upstreamMessage) - if eventType == "response.failed" { - if hit, code, msg := detectOpenAICyberPolicy(upstreamMessage); hit { - MarkOpsCyberPolicy(c, CyberPolicyMark{ - Code: code, - Message: msg, - Body: truncateString(string(upstreamMessage), 4096), - UpstreamStatus: http.StatusOK, - UpstreamInTok: usage.InputTokens, - UpstreamOutTok: usage.OutputTokens, - }) - } - } if !clientDisconnected { if needModelReplace && len(mappedModelBytes) > 0 && openAIWSEventMayContainModel(eventType) && bytes.Contains(upstreamMessage, mappedModelBytes) { upstreamMessage = replaceOpenAIWSMessageModel(upstreamMessage, mappedModel, originalModel) diff --git a/backend/internal/service/openai_ws_forwarder_v2.go b/backend/internal/service/openai_ws_forwarder_v2.go index 81f8631e96..260ec3eedf 100644 --- a/backend/internal/service/openai_ws_forwarder_v2.go +++ b/backend/internal/service/openai_ws_forwarder_v2.go @@ -585,17 +585,8 @@ func (s *OpenAIGatewayService) forwardOpenAIWSV2( } imageCounter.AddSSEData(message) - if eventType == "response.failed" { - if hit, code, msg := detectOpenAICyberPolicy(message); hit { - MarkOpsCyberPolicy(c, CyberPolicyMark{ - Code: code, - Message: msg, - Body: truncateString(string(message), 4096), - UpstreamStatus: http.StatusOK, - UpstreamInTok: usage.InputTokens, - UpstreamOutTok: usage.OutputTokens, - }) - } + if eventType == "error" || eventType == "response.failed" { + markOpenAICyberPolicyEvent(c, message, http.StatusOK, usage) } if eventType == "error" { diff --git a/backend/internal/service/openai_ws_forwarder_v2_test.go b/backend/internal/service/openai_ws_forwarder_v2_test.go index ab3301cc8d..636535a7eb 100644 --- a/backend/internal/service/openai_ws_forwarder_v2_test.go +++ b/backend/internal/service/openai_ws_forwarder_v2_test.go @@ -2,6 +2,7 @@ package service import ( "context" + "errors" "fmt" "net/http" "net/http/httptest" @@ -94,3 +95,76 @@ func TestForwardOpenAIWSV2_KeepsOutboundAndObservedServiceTiersSeparate(t *testi }) } } + +func TestForwardOpenAIWSV2_MarksCyberPolicyForFailureEventShapes(t *testing.T) { + gin.SetMode(gin.TestMode) + + tests := []struct { + name string + upstreamEvent []byte + wantError bool + wantInput int + wantOutput int + }{ + { + name: "error_before_return", + upstreamEvent: []byte(`{"type":"error","error":{"type":"rate_limit_error","code":"cyber_policy","message":"rate limit exceeded by cyber policy"},"usage":{"input_tokens":5,"output_tokens":1}}`), + wantError: true, + wantInput: 5, + wantOutput: 1, + }, + { + name: "response_failed_terminal", + upstreamEvent: []byte(`{"type":"response.failed","response":{"id":"resp_cyber","status":"failed","error":{"code":"cyber_policy","message":"blocked by cyber policy"},"usage":{"input_tokens":9,"output_tokens":2}}}`), + wantInput: 9, + wantOutput: 2, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + recorder := httptest.NewRecorder() + c, _ := gin.CreateTestContext(recorder) + c.Request = httptest.NewRequest(http.MethodPost, "/openai/v1/responses", nil) + + cfg := newOpenAIWSV2TestConfig() + cfg.Security.URLAllowlist.Enabled = false + cfg.Security.URLAllowlist.AllowInsecureHTTP = true + cfg.Gateway.OpenAIWS.MinIdlePerAccount = 0 + captureConn := &openAIWSCaptureConn{events: [][]byte{append([]byte(nil), tt.upstreamEvent...)}} + pool := newOpenAIWSConnPool(cfg) + pool.setClientDialerForTest(&openAIWSCaptureDialer{conn: captureConn}) + svc := &OpenAIGatewayService{ + cfg: cfg, + httpUpstream: &httpUpstreamRecorder{}, + cache: &stubGatewayCache{}, + openaiWSResolver: NewOpenAIWSProtocolResolver(cfg), + toolCorrector: NewCodexToolCorrector(), + openaiWSPool: pool, + } + account := &Account{ + ID: 5883, Name: "openai-ws-v2-cyber", 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}, + } + + result, err := svc.Forward(context.Background(), c, account, []byte(`{"model":"gpt-5.5","stream":false,"input":"hello"}`)) + if tt.wantError { + require.Error(t, err) + require.Nil(t, result) + var failoverErr *UpstreamFailoverError + require.False(t, errors.As(err, &failoverErr)) + } else { + require.NoError(t, err) + require.NotNil(t, result) + } + + mark := GetOpsCyberPolicy(c) + require.NotNil(t, mark) + require.Equal(t, "cyber_policy", mark.Code) + require.Equal(t, tt.wantInput, mark.UpstreamInTok) + require.Equal(t, tt.wantOutput, mark.UpstreamOutTok) + }) + } +} diff --git a/backend/internal/service/openai_ws_http_bridge.go b/backend/internal/service/openai_ws_http_bridge.go index 4b1c487bef..03339deffc 100644 --- a/backend/internal/service/openai_ws_http_bridge.go +++ b/backend/internal/service/openai_ws_http_bridge.go @@ -563,6 +563,7 @@ func (s *OpenAIGatewayService) proxyOpenAIWSHTTPBridgeTurn( respBody, _ := io.ReadAll(io.LimitReader(resp.Body, openAIWSHTTPBridgeErrorBodyLimitBytes)) _ = resp.Body.Close() + markOpenAICyberPolicyEvent(c, respBody, resp.StatusCode, nil) if resp.StatusCode == http.StatusBadRequest && extractUpstreamErrorCode(respBody) == openAIWSFallbackReasonInvalidEncryptedContent { s.markOpenAIWSInvalidEncryptedContentLineageFromPayload( @@ -776,6 +777,9 @@ func (s *OpenAIGatewayService) proxyOpenAIWSHTTPBridgeTurn( if openAIWSMessageShouldParseUsage(eventType, upstreamMessage) { parseOpenAIWSResponseUsageFromCompletedEvent(upstreamMessage, &usage) } + if eventType == "error" || eventType == "response.failed" { + markOpenAICyberPolicyEvent(c, upstreamMessage, http.StatusOK, &usage) + } imageCounter.AddSSEData(upstreamMessage) if needModelReplace && len(mappedModelBytes) > 0 && openAIWSEventMayContainModel(eventType) && strings.Contains(trimmedData, mappedModel) { diff --git a/backend/internal/service/openai_ws_http_bridge_test.go b/backend/internal/service/openai_ws_http_bridge_test.go index 0cdf64b974..5e23a51985 100644 --- a/backend/internal/service/openai_ws_http_bridge_test.go +++ b/backend/internal/service/openai_ws_http_bridge_test.go @@ -1052,6 +1052,79 @@ func TestProxyOpenAIWSHTTPBridgeTurnBareErrorUsesAuthoritativeFailed(t *testing. require.Equal(t, "response.failed", gjson.GetBytes(writes[1], "type").String()) } +func TestProxyOpenAIWSHTTPBridgeTurnMarksCyberPolicyForFailureShapes(t *testing.T) { + gin.SetMode(gin.TestMode) + + tests := []struct { + name string + statusCode int + body string + wantStatus int + wantInput int + wantOutput int + wantResult bool + wantError bool + }{ + { + name: "http_error_body", + statusCode: http.StatusBadRequest, + body: `{"error":{"code":"cyber_policy","message":"blocked by HTTP cyber policy"}}`, + wantStatus: http.StatusBadRequest, + wantError: true, + }, + { + name: "sse_error_event", + statusCode: http.StatusOK, + body: "data: {\"type\":\"error\",\"error\":{\"code\":\"cyber_policy\",\"message\":\"blocked by error event\"},\"usage\":{\"input_tokens\":5,\"output_tokens\":1}}\n\n", + wantStatus: http.StatusOK, + wantInput: 5, + wantOutput: 1, + wantResult: true, + wantError: true, + }, + { + name: "sse_response_failed_event", + statusCode: http.StatusOK, + body: "data: {\"type\":\"response.failed\",\"response\":{\"id\":\"resp_cyber\",\"error\":{\"code\":\"cyber_policy\",\"message\":\"blocked by failed event\"},\"usage\":{\"input_tokens\":9,\"output_tokens\":2}}}\n\n", + wantStatus: http.StatusOK, + wantInput: 9, + wantOutput: 2, + wantResult: true, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + upstream := &httpUpstreamRecorder{resp: &http.Response{ + StatusCode: tt.statusCode, + Header: make(http.Header), + Body: io.NopCloser(strings.NewReader(tt.body)), + }} + svc := &OpenAIGatewayService{cfg: &config.Config{}, httpUpstream: upstream} + account := &Account{ID: 112, Platform: PlatformOpenAI, Type: AccountTypeAPIKey, Concurrency: 1} + recorder := httptest.NewRecorder() + c, _ := gin.CreateTestContext(recorder) + c.Request = httptest.NewRequest(http.MethodGet, "/v1/responses", nil) + payload := []byte(`{"type":"response.create","model":"gpt-5","input":"hi"}`) + + result, err := svc.proxyOpenAIWSHTTPBridgeTurn( + context.Background(), c, account, "sk-test", payload, len(payload), + "gpt-5", "", "", "", "", 1, func([]byte) error { return nil }, + ) + + require.Equal(t, tt.wantResult, result != nil) + require.Equal(t, tt.wantError, err != nil) + mark := GetOpsCyberPolicy(c) + require.NotNil(t, mark) + require.Equal(t, "cyber_policy", mark.Code) + require.Equal(t, tt.wantStatus, mark.UpstreamStatus) + require.Equal(t, tt.wantInput, mark.UpstreamInTok) + require.Equal(t, tt.wantOutput, mark.UpstreamOutTok) + require.Contains(t, mark.Body, `"code":"cyber_policy"`) + }) + } +} + func TestProxyOpenAIWSHTTPBridgeTurnBareErrorEOFSynthesizesFailed(t *testing.T) { gin.SetMode(gin.TestMode) body := "data: {\"type\":\"response.created\",\"response\":{\"id\":\"resp_eof\",\"status\":\"in_progress\"}}\n\n" + diff --git a/backend/internal/service/openai_ws_ingress_capacity_shed_test.go b/backend/internal/service/openai_ws_ingress_capacity_shed_test.go index d37b7f2cda..c9d041afa6 100644 --- a/backend/internal/service/openai_ws_ingress_capacity_shed_test.go +++ b/backend/internal/service/openai_ws_ingress_capacity_shed_test.go @@ -2,6 +2,7 @@ package service import ( "context" + "errors" "net/http" "net/http/httptest" "strings" @@ -11,6 +12,7 @@ import ( coderws "github.com/coder/websocket" "github.com/gin-gonic/gin" "github.com/stretchr/testify/require" + "github.com/tidwall/gjson" ) // openAIWSIngressCapacityShedRepo 补齐 SetError，避免非容量类错误（如 @@ -182,3 +184,128 @@ func TestProxyResponsesWebSocketFromClient_RewritesCapacityShedCodeForClient(t * }) } } + +func TestProxyResponsesWebSocketFromClient_MarksCyberPolicyBeforeEarlyReturn(t *testing.T) { + gin.SetMode(gin.TestMode) + + tests := []struct { + name string + upstreamEvent []byte + wantFailover bool + wantClientMsg bool + wantInput int + wantOutput int + }{ + { + name: "error_before_rate_limit_failover", + upstreamEvent: []byte(`{"type":"error","error":{"type":"rate_limit_error","code":"cyber_policy","message":"rate limit exceeded by cyber policy"},"usage":{"input_tokens":5,"output_tokens":1}}`), + wantFailover: true, + wantInput: 5, + wantOutput: 1, + }, + { + name: "response_failed_terminal", + upstreamEvent: []byte(`{"type":"response.failed","response":{"id":"resp_cyber","status":"failed","error":{"code":"cyber_policy","message":"blocked by cyber policy"},"usage":{"input_tokens":9,"output_tokens":2}}}`), + wantClientMsg: true, + wantInput: 9, + wantOutput: 2, + }, + } + + 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.ModeRouterV2Enabled = true + cfg.Gateway.OpenAIWS.IngressModeDefault = OpenAIWSIngressModeCtxPool + cfg.Gateway.OpenAIWS.MinIdlePerAccount = 0 + + captureConn := &openAIWSCaptureConn{events: [][]byte{append([]byte(nil), tt.upstreamEvent...)}} + pool := newOpenAIWSConnPool(cfg) + pool.setClientDialerForTest(&openAIWSCaptureDialer{conn: captureConn}) + svc := &OpenAIGatewayService{ + cfg: cfg, + httpUpstream: &httpUpstreamRecorder{}, + cache: &stubGatewayCache{}, + openaiWSResolver: NewOpenAIWSProtocolResolver(cfg), + toolCorrector: NewCodexToolCorrector(), + openaiWSPool: pool, + } + account := &Account{ + ID: 5402, Name: "openai-ingress-cyber", 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": OpenAIWSIngressModeCtxPool, + }, + } + + markCh := make(chan *CyberPolicyMark, 1) + serverErrCh := 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}) + if err != nil { + serverErrCh <- err + return + } + defer func() { _ = conn.CloseNow() }() + + readCtx, cancelRead := context.WithTimeout(r.Context(), 3*time.Second) + _, firstMessage, err := conn.Read(readCtx) + cancelRead() + if err != nil { + serverErrCh <- err + return + } + + recorder := httptest.NewRecorder() + ginCtx, _ := gin.CreateTestContext(recorder) + ginCtx.Request = r.Clone(r.Context()) + hooks := &OpenAIWSIngressHooks{AfterTurn: func(_ int, _ *OpenAIForwardResult, _ error) { + markCh <- GetOpsCyberPolicy(ginCtx) + }} + serverErrCh <- svc.ProxyResponsesWebSocketFromClient(r.Context(), ginCtx, conn, account, "sk-test", firstMessage, hooks) + })) + 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) + + if tt.wantClientMsg { + readCtx, cancelRead := context.WithTimeout(context.Background(), 3*time.Second) + _, message, readErr := clientConn.Read(readCtx) + cancelRead() + require.NoError(t, readErr) + require.Equal(t, "response.failed", gjson.GetBytes(message, "type").String()) + } + + select { + case mark := <-markCh: + require.NotNil(t, mark) + require.Equal(t, "cyber_policy", mark.Code) + require.Equal(t, tt.wantInput, mark.UpstreamInTok) + require.Equal(t, tt.wantOutput, mark.UpstreamOutTok) + case <-time.After(3 * time.Second): + t.Fatal("AfterTurn did not observe the cyber mark") + } + + _ = clientConn.CloseNow() + select { + case serverErr := <-serverErrCh: + var failoverErr *UpstreamFailoverError + require.Equal(t, tt.wantFailover, errors.As(serverErr, &failoverErr)) + case <-time.After(5 * time.Second): + t.Fatal("waiting for ingress websocket shutdown timed out") + } + }) + } +} diff --git a/backend/internal/service/openai_ws_v2_passthrough_adapter.go b/backend/internal/service/openai_ws_v2_passthrough_adapter.go index 8d29839b3b..b7c2940965 100644 --- a/backend/internal/service/openai_ws_v2_passthrough_adapter.go +++ b/backend/internal/service/openai_ws_v2_passthrough_adapter.go @@ -1477,21 +1477,9 @@ func openAIWSPassthroughRelayClientClose(exit openaiwsv2.RelayExit, completedTur } func markOpenAIWSV2PassthroughCyberPolicy(c *gin.Context, payload []byte) bool { - hit, code, message := detectOpenAICyberPolicy(payload) - if !hit { - return false - } usage := OpenAIUsage{} parseOpenAIWSResponseUsageFromCompletedEvent(payload, &usage) - MarkOpsCyberPolicy(c, CyberPolicyMark{ - Code: code, - Message: message, - Body: truncateString(string(payload), 4096), - UpstreamStatus: http.StatusOK, - UpstreamInTok: usage.InputTokens, - UpstreamOutTok: usage.OutputTokens, - }) - return true + return markOpenAICyberPolicyEvent(c, payload, http.StatusOK, &usage) } func (s *OpenAIGatewayService) mapOpenAIWSPassthroughDialError(