From b8ed24508555406a34b0bbc9c980ff158362d678 Mon Sep 17 00:00:00 2001 From: r266-tech Date: Sun, 6 Sep 2026 18:42:06 +0800 Subject: [PATCH] fix(openai-ws): recover later turns from quota exhaustion --- .../service/openai_ws_v2/passthrough_relay.go | 25 +++++++- .../openai_ws_v2/passthrough_relay_test.go | 58 +++++++++++++++++++ .../openai_ws_v2_passthrough_adapter.go | 7 +++ ...openai_ws_v2_passthrough_lifecycle_test.go | 56 ++++++++++++++++++ 4 files changed, 143 insertions(+), 3 deletions(-) diff --git a/backend/internal/service/openai_ws_v2/passthrough_relay.go b/backend/internal/service/openai_ws_v2/passthrough_relay.go index fdb539a073..c37847e561 100644 --- a/backend/internal/service/openai_ws_v2/passthrough_relay.go +++ b/backend/internal/service/openai_ws_v2/passthrough_relay.go @@ -102,6 +102,7 @@ type RelayTraceEvent struct { type relayState struct { usage Usage turnUsage Usage + turnWroteDownstream atomic.Bool requestModelMu sync.RWMutex requestModel string pendingTurnStart atomic.Pointer[time.Time] @@ -204,7 +205,8 @@ func Relay( return upstreamConn.WriteFrame(writeCtx, msgType, payload) } writeClientFrameUpstream := func(msgType coderws.MessageType, payload []byte) error { - if isClientResponseCreateFrame(msgType, payload) { + isResponseCreate := isClientResponseCreateFrame(msgType, payload) + if isResponseCreate { state.setRequestModel(strings.TrimSpace(gjson.GetBytes(payload, "model").String())) turnStartedAt := time.Time{} if options.TakeNextTurnStartedAt != nil { @@ -214,8 +216,18 @@ func Relay( turnStartedAt = nowFn() } state.setPendingTurnStartedAt(turnStartedAt) + // The policy-enforcing client connection has accepted this turn. + // Reset before the write so an immediate upstream response cannot race + // with the transport returning from WriteFrame. + state.turnWroteDownstream.Store(false) } - return writeUpstream(msgType, payload) + err := writeUpstream(msgType, payload) + if err != nil && isResponseCreate { + // The relay exits on this error, but retain the previous turn's state + // for accurate diagnostics while the two relay goroutines settle. + state.turnWroteDownstream.Store(true) + } + return err } writeClient := func(msgType coderws.MessageType, payload []byte) error { // 下行写超时故意不挂在 relayCtx 上：coder/websocket 在已武装的 write @@ -559,7 +571,11 @@ func runUpstreamToClient( } markActivity() if beforeWriteClient != nil { - if err := beforeWriteClient(msgType, payload, wroteDownstream); err != nil { + wroteDownstreamInTurn := wroteDownstream + if state != nil { + wroteDownstreamInTurn = state.turnWroteDownstream.Load() + } + if err := beforeWriteClient(msgType, payload, wroteDownstreamInTurn); err != nil { emitRelayTrace(onTrace, RelayTraceEvent{ Stage: "upstream_message_rejected", Direction: "upstream_to_client", @@ -630,6 +646,9 @@ func runUpstreamToClient( return } wroteDownstream = true + if state != nil { + state.turnWroteDownstream.Store(true) + } if afterWriteClient != nil { afterWriteClient(msgType, payload) } diff --git a/backend/internal/service/openai_ws_v2/passthrough_relay_test.go b/backend/internal/service/openai_ws_v2/passthrough_relay_test.go index a5505b4981..dced90de33 100644 --- a/backend/internal/service/openai_ws_v2/passthrough_relay_test.go +++ b/backend/internal/service/openai_ws_v2/passthrough_relay_test.go @@ -4,6 +4,7 @@ import ( "context" "errors" "io" + "strings" "sync" "sync/atomic" "testing" @@ -903,6 +904,63 @@ func TestRelay_OnTurnComplete_UsesSubsequentResponseCreateTimeAcrossPricingBound } } +func TestRelay_BeforeWriteClientTracksDownstreamPerTurn(t *testing.T) { + t.Parallel() + + clientConn := newPassthroughTestFrameConn(nil, false) + upstreamConn := newPassthroughTestFrameConn(nil, false) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + wroteStates := make(chan bool, 3) + done := make(chan *RelayExit, 1) + stopErr := errors.New("stop after second-turn error") + go func() { + _, relayExit := Relay( + ctx, + clientConn, + upstreamConn, + []byte(`{"type":"response.create","model":"gpt-5.3-codex","input":[]}`), + RelayOptions{BeforeWriteClient: func(_ coderws.MessageType, payload []byte, wroteDownstream bool) error { + wroteStates <- wroteDownstream + if strings.Contains(string(payload), `"type":"error"`) { + return stopErr + } + return nil + }}, + ) + done <- relayExit + }() + + require.Eventually(t, func() bool { return len(upstreamConn.Writes()) == 1 }, time.Second, time.Millisecond) + upstreamConn.readCh <- passthroughTestFrame{ + msgType: coderws.MessageText, + payload: []byte(`{"type":"response.completed","response":{"id":"resp_first","usage":{"input_tokens":1,"output_tokens":1}}}`), + } + require.False(t, <-wroteStates) + require.Eventually(t, func() bool { return len(clientConn.Writes()) == 1 }, time.Second, time.Millisecond) + + clientConn.readCh <- passthroughTestFrame{ + msgType: coderws.MessageText, + payload: []byte(`{"type":"response.create","model":"gpt-5.3-codex","input":[]}`), + } + require.Eventually(t, func() bool { return len(upstreamConn.Writes()) == 2 }, time.Second, time.Millisecond) + upstreamConn.readCh <- passthroughTestFrame{ + msgType: coderws.MessageText, + payload: []byte(`{"type":"error","error":{"type":"usage_limit_reached"}}`), + } + require.False(t, <-wroteStates, "the next turn must not inherit the first turn's downstream write") + + select { + case relayExit := <-done: + require.NotNil(t, relayExit) + require.ErrorIs(t, relayExit.Err, stopErr) + require.True(t, relayExit.WroteDownstream, "connection-wide diagnostics must retain prior output") + case <-time.After(time.Second): + t.Fatal("relay did not stop after the rejected second-turn event") + } +} + func TestRelay_BinaryFramePassthrough(t *testing.T) { t.Parallel() diff --git a/backend/internal/service/openai_ws_v2_passthrough_adapter.go b/backend/internal/service/openai_ws_v2_passthrough_adapter.go index b7c2940965..5b25f31e74 100644 --- a/backend/internal/service/openai_ws_v2_passthrough_adapter.go +++ b/backend/internal/service/openai_ws_v2_passthrough_adapter.go @@ -1299,6 +1299,13 @@ func (s *OpenAIGatewayService) proxyResponsesWebSocketV2Passthrough( truncateOpenAIWSLogValue(errTypeRaw, openAIWSLogValueMaxLen), truncateOpenAIWSLogValue(errMsgRaw, openAIWSLogValueMaxLen), ) + if completedTurns.Load() > 0 { + return NewOpenAIWSClientCloseError( + coderws.StatusTryAgainLater, + "upstream rate limit exceeded; please reconnect", + errors.New("later passthrough turn was rate limited before output"), + ) + } return s.newOpenAIWSRateLimitFailoverError(account, handshakeHeaders, payload, errMsgRaw) }, OnTrace: func(event openaiwsv2.RelayTraceEvent) { diff --git a/backend/internal/service/openai_ws_v2_passthrough_lifecycle_test.go b/backend/internal/service/openai_ws_v2_passthrough_lifecycle_test.go index 5fa6a7def9..24afbeda09 100644 --- a/backend/internal/service/openai_ws_v2_passthrough_lifecycle_test.go +++ b/backend/internal/service/openai_ws_v2_passthrough_lifecycle_test.go @@ -4,6 +4,7 @@ import ( "context" "encoding/json" "errors" + "fmt" "net/http" "net/http/httptest" "strings" @@ -210,6 +211,61 @@ func startPassthroughLifecycleServerWithHooks( return server, serverErr } +func TestPassthroughLifecycle_LaterTurnPreOutputRateLimitRequestsReconnect(t *testing.T) { + gin.SetMode(gin.TestMode) + controlCtx, cancelControl := context.WithCancelCause(context.Background()) + defer cancelControl(context.Canceled) + upstream := newStagedPassthroughConn() + account := passthroughLifecycleAccount() + repo := &openAIWSRateLimitSignalRepo{stubOpenAIAccountRepo: stubOpenAIAccountRepo{accounts: []Account{*account}}} + svc := newPassthroughLifecycleService(passthroughLifecycleConfig(), upstream) + svc.accountRepo = repo + svc.rateLimitService = &RateLimitService{accountRepo: repo} + + server, serverErr := startPassthroughLifecycleServer(t, controlCtx, svc, account) + defer server.Close() + clientConn := dialPassthroughLifecycleClient(t, server) + defer func() { _ = clientConn.CloseNow() }() + firstRequest := requirePassthroughUpstreamWrite(t, upstream, time.Second) + require.Equal(t, "response.create", gjson.GetBytes(firstRequest, "type").String()) + + upstream.Send(`{"type":"response.completed","response":{"id":"resp_first","model":"gpt-5.1","usage":{"input_tokens":1,"output_tokens":1}}}`) + completed, err := readPassthroughLifecycleFrame(t, clientConn, time.Second) + require.NoError(t, err) + require.Equal(t, "response.completed", gjson.GetBytes(completed, "type").String()) + + writeCtx, cancelWrite := context.WithTimeout(context.Background(), time.Second) + err = clientConn.Write(writeCtx, coderws.MessageText, []byte(`{"type":"response.create","model":"gpt-5.1","stream":false}`)) + cancelWrite() + require.NoError(t, err) + secondRequest := requirePassthroughUpstreamWrite(t, upstream, time.Second) + require.Equal(t, "response.create", gjson.GetBytes(secondRequest, "type").String()) + + resetAt := time.Now().Add(90 * time.Minute).Unix() + upstream.Send(fmt.Sprintf(`{"type":"error","error":{"code":"rate_limit_exceeded","type":"usage_limit_reached","message":"The usage limit has been reached","resets_at":%d}}`, resetAt)) + _, err = readPassthroughLifecycleFrame(t, clientConn, time.Second) + var websocketCloseErr coderws.CloseError + require.ErrorAs(t, err, &websocketCloseErr) + require.Equal(t, coderws.StatusTryAgainLater, websocketCloseErr.Code) + require.Equal(t, "upstream rate limit exceeded; please reconnect", websocketCloseErr.Reason) + require.Len(t, repo.rateLimitCalls, 1) + require.WithinDuration(t, time.Unix(resetAt, 0), repo.rateLimitCalls[0], 2*time.Second) + + select { + case err := <-serverErr: + var closeErr *OpenAIWSClientCloseError + require.ErrorAs(t, err, &closeErr) + require.Equal(t, coderws.StatusTryAgainLater, closeErr.StatusCode()) + case <-time.After(time.Second): + t.Fatal("later-turn rate limit did not terminate passthrough") + } + select { + case replay := <-upstream.writes: + t.Fatalf("later-turn reconnect must not replay the retained first request: %s", replay) + default: + } +} + func TestPassthroughLifecycle_CyberTerminalEventsMarkBeforeAfterTurn(t *testing.T) { gin.SetMode(gin.TestMode)