github/sub2api
4
0
Fork 0
mirror of https://github.com/Wei-Shaw/sub2api.git synced 2026-10-19 16:17:56 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix(openai): reacquire passthrough turn slots

Browse Source
This commit is contained in:
wucm667
2026-09-01 08:13:29 +08:00
parent 200602b41b
commit 5fb5dfb34c
4 changed files with 105 additions and 53 deletions
Download Patch File Download Diff File Expand all files Collapse all files
backend/internal/handler/openai_gateway_handler.go
+9 -10
View File
@@ -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)
}
backend/internal/service/openai_ws_forwarder_ingress.go
+3 -3
View File
@@ -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,
backend/internal/service/openai_ws_passthrough_turn_pricing_test.go
+88 -40
View File
@@ -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) {
backend/internal/service/openai_ws_v2_passthrough_adapter.go
+5
View File
@@ -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 {