mirror of
https://github.com/Wei-Shaw/sub2api.git
synced 2026-10-14 12:02:27 +08:00
fix(openai): record WS cyber policy failures across paths
Cyber-policy detection only inspected response.failed, and the ingress check ran after error-specific early returns. As a result, error-shaped failures and the connection-pool, HTTP-bridge, and WSv2 paths could miss cyber-policy hits. Centralize event marking and invoke it before retry and failover returns across the WebSocket implementations. Preserve the recorded guard during account failover and defer connection blocking until the logical failover chain completes, avoiding duplicate records while ensuring subsequent turns are blocked.
This commit is contained in:
@@ -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.
|
||||
|
||||
@@ -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 {
|
||||
@@ -2310,6 +2325,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) {
|
||||
@@ -2753,10 +2769,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 := ""
|
||||
@@ -2778,10 +2796,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
|
||||
@@ -4088,15 +4110,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) {
|
||||
|
||||
@@ -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
|
||||
}
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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" {
|
||||
|
||||
@@ -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)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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) {
|
||||
|
||||
@@ -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" +
|
||||
|
||||
@@ -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")
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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(
|
||||
|
||||
Reference in New Issue
Block a user