Merge pull request #6636 from specialpointcentral/cyber-ws-recording-fix

fix(openai): record WS cyber policy failures across paths

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.

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) {