diff --git a/backend/internal/handler/openai_gateway_handler.go b/backend/internal/handler/openai_gateway_handler.go index e57cc759e6..08625facf8 100644 --- a/backend/internal/handler/openai_gateway_handler.go +++ b/backend/internal/handler/openai_gateway_handler.go @@ -1616,10 +1616,10 @@ const ( // 由 BeforeTurn 在每个 turn 开始时冻结，AfterTurn 的用量提交读取它；turn 在 // 连接内串行推进，互斥锁只为跨用量提交 goroutine 的读取安全。 // -// ws_v2 passthrough ingress 没有 BeforeTurn，因此本值会保持零；AfterTurn 必须 -// 以 TurnStarted 已记录的所属 turn 开始时刻为回退，而不是用建连或记录时刻。 -// 这样每个 passthrough turn 都按自己的开始时刻计价，但不改变其仅在建连时执行 -// 准入门、没有 turn 级利润复核的既有行为。 +// 零值语义（重要）：首轮准入由握手路径完成，不调用 BeforeTurn，因此首轮保持 +// 零值并回退到 TurnStarted 记录的首轮开始时刻。后续 turn 在 response.create +// 写入上游前调用 BeforeTurn，按当时的利润门复核并冻结定价。绝不能用建连时刻 +// 初始化，否则会把长连接的所有 turn 钉死在建连时的峰谷因子。 type openAIWSTurnPricing struct { mu sync.Mutex at time.Time @@ -2321,8 +2321,8 @@ func (h *OpenAIGatewayHandler) ResponsesWebSocket(c *gin.Context) { // turn-tagged slot preserves the exact mapping used for the in-flight request. var turnChannelMapping atomic.Pointer[openAIWSTurnChannelMappingSnapshot] turnChannelMapping.Store(&openAIWSTurnChannelMappingSnapshot{turn: 1, mapping: channelMappingWS}) - // turn 级定价：BeforeTurn 重新冻结 pricingAt 并按最新门复核当前账号； - // passthrough 没有 BeforeTurn 时，AfterTurn 回退到 TurnStarted 的所属 turn 时刻。 + // turn 级定价：首轮回退到 TurnStarted 的所属 turn 时刻；后续 turn 由 + // BeforeTurn 重新冻结 pricingAt 并按最新门复核当前账号。 var turnPricing openAIWSTurnPricing hooks := &service.OpenAIWSIngressHooks{ ClientLifecycleContext: clientLifecycleCtx, @@ -2335,10 +2335,9 @@ func (h *OpenAIGatewayHandler) ResponsesWebSocket(c *gin.Context) { c.Set(securityAuditWSTurnContextKey, turn) service.BeginOpsStreamTurn(c, turn) setCyberTurnBody(turn, payload) - // Passthrough ingress intentionally skips BeforeTurn, so enforce only - // the connection-level cyber session gate here as well. Native ingress - // visits this hook first and gets the same side-effect-free close error; - // its original BeforeTurn guard remains as defense in depth. + // 连接级 cyber session gate 也在 BeforeRequest 先执行，使 native 与 + // passthrough ingress 都能在 BeforeTurn 及上游写入前无副作用地拒绝。 + // BeforeTurn 中保留同一检查作为防御式兜底。 if cyberBlockedThisConn { return service.NewOpenAIWSClientCloseError(coderws.StatusPolicyViolation, cyberSessionBlockedClientMsg, nil) } diff --git a/backend/internal/service/openai_ws_forwarder_ingress.go b/backend/internal/service/openai_ws_forwarder_ingress.go index 5f5976c7d0..4e2a1f13a2 100644 --- a/backend/internal/service/openai_ws_forwarder_ingress.go +++ b/backend/internal/service/openai_ws_forwarder_ingress.go @@ -137,9 +137,9 @@ func (s *OpenAIGatewayService) ProxyResponsesWebSocketFromClient( forceHTTPBridge = true break } - // 透传 relay 通过 TurnStarted 记录每个 turn 的开始时刻，但不触发 - // BeforeTurn；因此仍只有建连时的利润准入门，没有 turn 级复核。 - // handler 计费在 turn 定价未冻结时回退到对应的 turn 开始时刻。 + // 首轮准入由握手路径完成；后续 response.create 会在写入上游前 + // 依次回调 BeforeRequest 和 BeforeTurn，并在终止或失败时回调 + // AfterTurn，从而覆盖 turn 级利润复核、定价冻结和并发槽位释放。 return s.proxyResponsesWebSocketV2Passthrough( ctx, c, diff --git a/backend/internal/service/openai_ws_passthrough_turn_pricing_test.go b/backend/internal/service/openai_ws_passthrough_turn_pricing_test.go index 16be2bd4c6..a405fbb29a 100644 --- a/backend/internal/service/openai_ws_passthrough_turn_pricing_test.go +++ b/backend/internal/service/openai_ws_passthrough_turn_pricing_test.go @@ -60,50 +60,38 @@ func startPassthroughHookRecordingServer( return server, serverErr } -// TestPassthroughIngressReportsTurnStartedBeforeAfterTurnWithoutBeforeTurn 钉死 -// ws_v2 透传 ingress 与 handler 侧 turn 定价的耦合：透传 relay 不触发 -// BeforeTurn，但会在每个 AfterTurn 前通过 TurnStarted 报告同一 turn 的开始时刻。 -// -// handler 的 recordTurnStart 保存该时刻，AfterTurn 再用 currentOr(turnStart) -// 作为计费 PricingAt；不触发 BeforeTurn 也意味着透传仍没有 turn 级利润复核。 -func TestPassthroughIngressReportsTurnStartedBeforeAfterTurnWithoutBeforeTurn(t *testing.T) { +func TestPassthroughIngressFollowUpCallsBeforeTurnAfterBeforeRequest(t *testing.T) { gin.SetMode(gin.TestMode) controlCtx, cancelControl := context.WithCancelCause(context.Background()) defer cancelControl(context.Canceled) upstream := newStagedPassthroughConn() - upstream.Send(`{"type":"response.completed","response":{"id":"resp_pricing","model":"gpt-5.1","usage":{"input_tokens":1,"output_tokens":1}}}`) + upstream.Send(`{"type":"response.completed","response":{"id":"resp_pricing_1","model":"gpt-5.1","usage":{"input_tokens":1,"output_tokens":1}}}`) var hooksMu sync.Mutex - beforeTurnCalls := 0 - expectedTurnStartedAt := time.Date(2026, time.August, 17, 9, 59, 59, 0, time.UTC) - type hookEvent struct { - name string - turn int - startedAt time.Time - } - var hookEvents []hookEvent + var callbacks []string + afterTurnCalls := 0 hooks := &OpenAIWSIngressHooks{ - InitialTurnStartedAt: expectedTurnStartedAt, - TurnStarted: func(turn int, startedAt time.Time) { + BeforeRequest: func(int, []byte, string) error { hooksMu.Lock() - hookEvents = append(hookEvents, hookEvent{name: "TurnStarted", turn: turn, startedAt: startedAt}) - hooksMu.Unlock() - }, - BeforeTurn: func(int) error { - hooksMu.Lock() - beforeTurnCalls++ + callbacks = append(callbacks, "before_request") hooksMu.Unlock() return nil }, - AfterTurn: func(turn int, _ *OpenAIForwardResult, _ error) { + BeforeTurn: func(int) error { hooksMu.Lock() - hookEvents = append(hookEvents, hookEvent{name: "AfterTurn", turn: turn}) + callbacks = append(callbacks, "before_turn") + hooksMu.Unlock() + return nil + }, + AfterTurn: func(int, *OpenAIForwardResult, error) { + hooksMu.Lock() + afterTurnCalls++ hooksMu.Unlock() }, } - server, serverErr := startPassthroughHookRecordingServer( + server, _ := startPassthroughHookRecordingServer( t, controlCtx, newPassthroughLifecycleService(passthroughLifecycleConfig(), upstream), @@ -113,30 +101,90 @@ func TestPassthroughIngressReportsTurnStartedBeforeAfterTurnWithoutBeforeTurn(t defer server.Close() clientConn := dialPassthroughLifecycleClient(t, server) defer func() { _ = clientConn.CloseNow() }() + require.Equal(t, "response.create", gjson.GetBytes(requirePassthroughUpstreamWrite(t, upstream, time.Second), "type").String()) event, err := readPassthroughLifecycleFrame(t, clientConn, 3*time.Second) require.NoError(t, err) require.Equal(t, "response.completed", gjson.GetBytes(event, "type").String()) - // 等待连接自然结束（inter-turn idle 超时），确保 AfterTurn 已提交。 - _, _ = readPassthroughLifecycleFrame(t, clientConn, 3*time.Second) + writeCtx, cancelWrite := context.WithTimeout(context.Background(), time.Second) + err = clientConn.Write(writeCtx, coderws.MessageText, []byte(`{"type":"response.create","model":"gpt-5.1"}`)) + cancelWrite() + require.NoError(t, err) + require.Equal(t, "response.create", gjson.GetBytes(requirePassthroughUpstreamWrite(t, upstream, time.Second), "type").String()) + upstream.Send(`{"type":"response.completed","response":{"id":"resp_pricing_2","model":"gpt-5.1","usage":{"input_tokens":1,"output_tokens":1}}}`) + event, err = readPassthroughLifecycleFrame(t, clientConn, 3*time.Second) + require.NoError(t, err) + require.Equal(t, "response.completed", gjson.GetBytes(event, "type").String()) + + hooksMu.Lock() + gotCallbacks := append([]string(nil), callbacks...) + gotAfter := afterTurnCalls + hooksMu.Unlock() + + require.Equal(t, []string{"before_request", "before_turn"}, gotCallbacks) + require.Equal(t, 2, gotAfter) +} + +func TestPassthroughIngressBeforeTurnRejectionDoesNotForwardFollowUp(t *testing.T) { + gin.SetMode(gin.TestMode) + controlCtx, cancelControl := context.WithCancelCause(context.Background()) + defer cancelControl(context.Canceled) + + upstream := newStagedPassthroughConn() + upstream.Send(`{"type":"response.completed","response":{"id":"resp_reject_1","model":"gpt-5.1","usage":{"input_tokens":1,"output_tokens":1}}}`) + rejection := errors.New("turn rejected") + var hooksMu sync.Mutex + afterTurnCalls := 0 + var finalErr error + beforeTurnTurn := 0 + hooks := &OpenAIWSIngressHooks{ + BeforeTurn: func(turn int) error { + hooksMu.Lock() + beforeTurnTurn = turn + hooksMu.Unlock() + return rejection + }, + AfterTurn: func(_ int, _ *OpenAIForwardResult, turnErr error) { + hooksMu.Lock() + afterTurnCalls++ + if turnErr != nil { + finalErr = turnErr + } + hooksMu.Unlock() + }, + } + + server, serverErr := startPassthroughHookRecordingServer(t, controlCtx, newPassthroughLifecycleService(passthroughLifecycleConfig(), upstream), passthroughLifecycleAccount(), hooks) + defer server.Close() + clientConn := dialPassthroughLifecycleClient(t, server) + defer func() { _ = clientConn.CloseNow() }() + requirePassthroughUpstreamWrite(t, upstream, time.Second) + _, err := readPassthroughLifecycleFrame(t, clientConn, 3*time.Second) + require.NoError(t, err) + + writeCtx, cancelWrite := context.WithTimeout(context.Background(), time.Second) + err = clientConn.Write(writeCtx, coderws.MessageText, []byte(`{"type":"response.create","model":"gpt-5.1"}`)) + cancelWrite() + require.NoError(t, err) select { - case <-serverErr: + case payload := <-upstream.writes: + t.Fatalf("rejected response.create was forwarded upstream: %s", payload) + case <-time.After(100 * time.Millisecond): + } + select { + case err = <-serverErr: + require.ErrorIs(t, err, rejection) case <-time.After(3 * time.Second): - t.Fatal("passthrough ingress did not exit") + t.Fatal("passthrough ingress did not exit after BeforeTurn rejection") } hooksMu.Lock() - gotBefore := beforeTurnCalls - gotEvents := append([]hookEvent(nil), hookEvents...) + gotBeforeTurn, gotAfter, gotFinalErr := beforeTurnTurn, afterTurnCalls, finalErr hooksMu.Unlock() - - require.Zero(t, gotBefore, "透传 ingress 不应调用 BeforeTurn") - require.GreaterOrEqual(t, len(gotEvents), 2, "透传 ingress 应报告 TurnStarted 和 AfterTurn") - require.Equal(t, "TurnStarted", gotEvents[0].name) - require.Equal(t, expectedTurnStartedAt, gotEvents[0].startedAt, "TurnStarted 必须携带入口冻结的首轮开始时刻") - require.Equal(t, "AfterTurn", gotEvents[1].name) - require.Equal(t, gotEvents[0].turn, gotEvents[1].turn, "TurnStarted 后应提交同一 turn 的 AfterTurn") + require.Equal(t, 2, gotBeforeTurn) + require.Equal(t, 2, gotAfter, "each started turn must be finalized exactly once") + require.ErrorIs(t, gotFinalErr, rejection) } func TestPassthroughIngressFreezesSubsequentTurnBeforeRequestPolicy(t *testing.T) { diff --git a/backend/internal/service/openai_ws_v2_passthrough_adapter.go b/backend/internal/service/openai_ws_v2_passthrough_adapter.go index 95de10df71..cf8c0de7cb 100644 --- a/backend/internal/service/openai_ws_v2_passthrough_adapter.go +++ b/backend/internal/service/openai_ws_v2_passthrough_adapter.go @@ -1054,6 +1054,11 @@ func (s *OpenAIGatewayService) proxyResponsesWebSocketV2Passthrough( return payload, nil, err } } + if hooks != nil && hooks.BeforeTurn != nil { + if err := hooks.BeforeTurn(turnNo); err != nil { + return payload, nil, err + } + } if hooks != nil && hooks.MapRequestModel != nil { upstreamModel, err := hooks.MapRequestModel(turnNo, requestModelForThisFrame) if err != nil {