mirror of
https://github.com/router-for-me/CLIProxyAPI.git
synced 2026-09-03 06:35:00 +08:00
* feat(codex): add opt-in stream bootstrap buffering
The upstream smuggles capacity rejections into an HTTP 200 stream. The
handshake events arrive normally and only a later event carries
{"error":{"type":"service_unavailable_error","code":
"server_is_overloaded"}}. By then the executor has already handed the
first chunk downstream, the response is committed, and the conductor can
no longer retry on another credential, so the request fails even though
other credentials were available.
When codex.stream-bootstrap-buffering is enabled the executor holds back
the handshake events until it can tell whether the stream carries real
output or a rejection. An overload rejection then fails the attempt
before any chunk is delivered, letting the conductor retry on another
credential; every other terminal failure is flushed in order and
delivered in-stream exactly as before.
Detection uses an event-type allow-list rather than a fixed count. On the
websocket transport codex.rate_limits and codex.response.metadata arrive
before response.created, making the first generated event the fifth
frame, so a small counter would release the stream before the rejection
is visible. Buffering is bounded and hitting the bound degrades to the
original unbuffered behaviour.
Two details are load-bearing. The error must be returned synchronously:
delivering it as the first stream chunk makes ExecuteStream downgrade it
into a committed 200 and the status is lost. And the websocket path must
not signal an upstream disconnect for a rejection it intends to retry,
because the downstream handler closes the client connection on that
signal and the retry would have nowhere to deliver.
The 503 status is produced only on this path rather than in the shared
codexTerminalFailureStatus mapping, so disabling the feature restores the
previous behaviour exactly, including cooldown classification and
retry-after parsing.
Defaults to false: response headers are withheld until generation
starts, which can trip client or reverse-proxy read timeouts.
* test(codex): pin bootstrap overload failover through the conductor
Executor-level tests cannot show what the client finally receives. These
exercise ExecuteStream end to end to pin three properties that are easy
to regress:
- consecutive overloaded credentials are skipped until one serves the
request, and retries are capped by max-retry-credentials rather than
multiplying with request-retry
- exhausting the pool surfaces the upstream status instead of a
committed 200 stream
- with buffering disabled the rejection stays an in-stream error on a
committed stream, which is the behaviour the feature must preserve
The third case also documents why the executor returns its error
synchronously: an error arriving as the first stream chunk is wrapped and
downgraded into a committed 200, silently losing the status.
634 lines
24 KiB
Go
634 lines
24 KiB
Go
package executor
|
|
|
|
import (
|
|
"bytes"
|
|
"context"
|
|
"fmt"
|
|
"net/http"
|
|
"strings"
|
|
|
|
"github.com/gorilla/websocket"
|
|
"github.com/router-for-me/CLIProxyAPI/v7/internal/config"
|
|
"github.com/router-for-me/CLIProxyAPI/v7/internal/runtime/executor/helps"
|
|
"github.com/router-for-me/CLIProxyAPI/v7/internal/thinking"
|
|
cliproxyauth "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/auth"
|
|
cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executor"
|
|
sdktranslator "github.com/router-for-me/CLIProxyAPI/v7/sdk/translator"
|
|
log "github.com/sirupsen/logrus"
|
|
"github.com/tidwall/gjson"
|
|
)
|
|
|
|
func (e *CodexWebsocketsExecutor) ExecuteStream(ctx context.Context, auth *cliproxyauth.Auth, req cliproxyexecutor.Request, opts cliproxyexecutor.Options) (_ *cliproxyexecutor.StreamResult, err error) {
|
|
log.Debugf("Executing Codex Websockets stream request with auth ID: %s, model: %s", auth.ID, req.Model)
|
|
if ctx == nil {
|
|
ctx = context.Background()
|
|
}
|
|
if opts.Alt == "responses/compact" {
|
|
return nil, statusErr{code: http.StatusBadRequest, msg: "streaming not supported for /responses/compact"}
|
|
}
|
|
|
|
baseModel := thinking.ParseSuffix(req.Model).ModelName
|
|
apiKey, baseURL := codexCreds(auth)
|
|
if baseURL == "" {
|
|
baseURL = "https://chatgpt.com/backend-api/codex"
|
|
}
|
|
|
|
reporter := helps.NewExecutorUsageReporter(ctx, e, baseModel, auth)
|
|
defer reporter.TrackFailure(ctx, &err)
|
|
|
|
from := opts.SourceFormat
|
|
responseFormat := cliproxyexecutor.ResponseFormatOrSource(opts)
|
|
to := sdktranslator.FromString("codex")
|
|
originalPayloadSource := req.Payload
|
|
if len(opts.OriginalRequest) > 0 {
|
|
originalPayloadSource = opts.OriginalRequest
|
|
}
|
|
originalPayload := originalPayloadSource
|
|
originalTranslated, body := translateCodexRequestPair(from, to, baseModel, originalPayload, req.Payload, true)
|
|
|
|
body, err = helps.ApplyRequestThinking(body, req, opts, from.String(), to.String(), e.Identifier())
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
requestedModel := helps.PayloadRequestedModel(opts, req.Model)
|
|
requestPath := helps.PayloadRequestPath(opts)
|
|
body = helps.ApplyPayloadConfigWithRequest(e.cfg, baseModel, to.String(), from.String(), "", body, originalTranslated, requestedModel, requestPath, opts.Headers)
|
|
body = helps.SetStringIfDifferent(body, "model", baseModel)
|
|
body = normalizeCodexInstructions(body)
|
|
if e.cfg == nil || e.cfg.DisableImageGeneration == config.DisableImageGenerationOff {
|
|
body = ensureImageGenerationTool(body, baseModel, auth, opts.Headers)
|
|
}
|
|
body = sanitizeOpenAIResponsesReasoningEncryptedContent(ctx, "codex websockets executor", body)
|
|
body = normalizeCodexWebsocketParallelToolCalls(body, opts.Headers)
|
|
multiAgentV2Conflict := helps.HasCodexMultiAgentV2NamespaceConflict(body)
|
|
body, optimizeMultiAgentV2 := helps.OptimizeCodexMultiAgentV2RequestForAuth(ctx, opts.Headers, body, e.cfg, auth, baseModel)
|
|
body, replayScope, errReplay := applyCodexReasoningReplayCacheRequired(ctx, from, req, opts, body)
|
|
if errReplay != nil {
|
|
return nil, errReplay
|
|
}
|
|
|
|
httpURL := strings.TrimSuffix(baseURL, "/") + "/responses"
|
|
wsURL, err := buildCodexResponsesWebsocketURL(httpURL)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
body, wsHeaders, errPromptCache := applyCodexPromptCacheHeadersWithContext(ctx, from, req, body, opts.Headers)
|
|
if errPromptCache != nil {
|
|
return nil, errPromptCache
|
|
}
|
|
clientBody := body
|
|
var identityState codexIdentityConfuseState
|
|
upstreamBody, identityState := applyCodexIdentityConfuseBody(e.cfg, auth, originalPayloadSource, body)
|
|
reporter.SetTranslatedReasoningEffort(clientBody, to.String())
|
|
wsHeaders = applyCodexWebsocketHeaders(ctx, wsHeaders, auth, apiKey, e.cfg, opts.Headers)
|
|
applyModelHeaderOverrides(wsHeaders, baseModel)
|
|
applyCodexIdentityConfuseHeaders(wsHeaders, &identityState)
|
|
|
|
var authID, authLabel, authType, authValue string
|
|
authID = auth.ID
|
|
authLabel = auth.Label
|
|
authType, authValue = auth.AccountInfo()
|
|
|
|
executionSessionID := executionSessionIDFromOptions(opts)
|
|
var sess *codexWebsocketSession
|
|
if executionSessionID != "" {
|
|
sess = e.getOrCreateSession(executionSessionID)
|
|
if sess != nil {
|
|
sess.reqMu.Lock()
|
|
}
|
|
}
|
|
streamSessionLocked := sess != nil
|
|
unlockStreamSession := func() {
|
|
if sess != nil && streamSessionLocked {
|
|
sess.reqMu.Unlock()
|
|
streamSessionLocked = false
|
|
}
|
|
}
|
|
|
|
wsReqBody := buildCodexWebsocketRequestBody(upstreamBody)
|
|
wsReqLog := helps.UpstreamRequestLog{
|
|
URL: wsURL,
|
|
Method: "WEBSOCKET",
|
|
Headers: wsHeaders.Clone(),
|
|
Body: wsReqBody,
|
|
Provider: e.Identifier(),
|
|
AuthID: authID,
|
|
AuthLabel: authLabel,
|
|
AuthType: authType,
|
|
AuthValue: authValue,
|
|
}
|
|
helps.RecordAPIWebsocketRequest(ctx, e.cfg, wsReqLog)
|
|
|
|
var conn *websocket.Conn
|
|
var closer *websocketConnectionCloser
|
|
var respHS *http.Response
|
|
var errDial error
|
|
if cliproxyexecutor.RequiredUpstreamWebsocket(ctx) {
|
|
conn, closer = existingWebsocketSessionConn(sess, authID, wsURL)
|
|
if conn == nil {
|
|
if sess != nil {
|
|
sess.reqMu.Unlock()
|
|
}
|
|
return nil, cliproxyexecutor.NewUpstreamWebsocketReplayRequiredError()
|
|
}
|
|
} else {
|
|
conn, closer, respHS, errDial = e.ensureUpstreamConn(ctx, auth, sess, authID, wsURL, wsHeaders)
|
|
}
|
|
var upstreamHeaders http.Header
|
|
if respHS != nil {
|
|
upstreamHeaders = respHS.Header.Clone()
|
|
}
|
|
if errDial != nil {
|
|
bodyErr := websocketHandshakeBody(respHS)
|
|
if respHS != nil {
|
|
helps.RecordAPIWebsocketUpgradeRejection(ctx, e.cfg, websocketUpgradeRequestLog(wsReqLog), respHS.StatusCode, respHS.Header.Clone(), bodyErr)
|
|
}
|
|
if respHS != nil && respHS.StatusCode == http.StatusUpgradeRequired {
|
|
if sess != nil {
|
|
sess.reqMu.Unlock()
|
|
}
|
|
if opts.ExecutionLifecycle != nil || cliproxyexecutor.DownstreamWebsocket(ctx) {
|
|
return nil, statusErr{code: respHS.StatusCode, msg: string(bodyErr)}
|
|
}
|
|
return e.CodexExecutor.ExecuteStream(ctx, auth, req, opts)
|
|
}
|
|
if respHS != nil && respHS.StatusCode > 0 {
|
|
if sess != nil {
|
|
sess.reqMu.Unlock()
|
|
}
|
|
return nil, statusErr{code: respHS.StatusCode, msg: string(bodyErr)}
|
|
}
|
|
helps.RecordAPIWebsocketError(ctx, e.cfg, "dial", errDial)
|
|
if sess != nil {
|
|
sess.reqMu.Unlock()
|
|
}
|
|
return nil, errDial
|
|
}
|
|
if errBind := sess.bindExecutionLifecycle(opts, conn, closer, req.Model); errBind != nil {
|
|
if sess != nil {
|
|
sess.reqMu.Unlock()
|
|
}
|
|
closeWebsocketAfterBindFailure(sess, conn, closer)
|
|
return nil, errBind
|
|
}
|
|
recordAPIWebsocketHandshake(ctx, e.cfg, respHS)
|
|
reporter.StartResponseTTFT()
|
|
|
|
if sess == nil {
|
|
logCodexWebsocketConnected(executionSessionID, authID, wsURL)
|
|
}
|
|
|
|
var readCh chan codexWebsocketRead
|
|
if sess != nil {
|
|
readCh = sess.activate(conn)
|
|
}
|
|
restoreMultiAgentV2 := !multiAgentV2Conflict && (optimizeMultiAgentV2 || sess.isMultiAgentV2Optimized(conn))
|
|
|
|
if errSend := writeCodexWebsocketMessage(sess, conn, wsReqBody); errSend != nil {
|
|
errSend = mapCodexWebsocketWriteError(sess, conn, errSend)
|
|
helps.RecordAPIWebsocketError(ctx, e.cfg, "send", errSend)
|
|
if sess != nil {
|
|
if cliproxyexecutor.RequiredUpstreamWebsocket(ctx) {
|
|
e.invalidateUpstreamConnWithoutDisconnectNotify(sess, conn, "send_error", errSend)
|
|
sess.clearActive(conn, readCh)
|
|
sess.reqMu.Unlock()
|
|
if !shouldRetryCodexWebsocketSend(errSend) {
|
|
return nil, errSend
|
|
}
|
|
return nil, cliproxyexecutor.NewUpstreamWebsocketReplayRequiredError()
|
|
}
|
|
e.invalidateUpstreamConn(sess, conn, "send_error", errSend)
|
|
if !shouldRetryCodexWebsocketSend(errSend) {
|
|
sess.clearActive(conn, readCh)
|
|
sess.reqMu.Unlock()
|
|
return nil, errSend
|
|
}
|
|
|
|
// Retry once with a new websocket connection for the same execution session.
|
|
connRetry, closerRetry, respHSRetry, errDialRetry := e.ensureUpstreamConn(ctx, auth, sess, authID, wsURL, wsHeaders)
|
|
if errDialRetry != nil || connRetry == nil {
|
|
closeHTTPResponseBody(respHSRetry, "codex websockets executor: close handshake response body error")
|
|
helps.RecordAPIWebsocketError(ctx, e.cfg, "dial_retry", errDialRetry)
|
|
sess.clearActive(conn, readCh)
|
|
sess.reqMu.Unlock()
|
|
return nil, errDialRetry
|
|
}
|
|
previousConn, previousReadCh := conn, readCh
|
|
conn = connRetry
|
|
closer = closerRetry
|
|
if errBind := sess.bindExecutionLifecycle(opts, conn, closer, req.Model); errBind != nil {
|
|
clearRetryActiveState(sess, previousConn, previousReadCh)
|
|
sess.reqMu.Unlock()
|
|
closeWebsocketAfterBindFailure(sess, conn, closer)
|
|
return nil, errBind
|
|
}
|
|
readCh = sess.activate(conn)
|
|
restoreMultiAgentV2 = !multiAgentV2Conflict && (optimizeMultiAgentV2 || sess.isMultiAgentV2Optimized(conn))
|
|
wsReqBodyRetry := buildCodexWebsocketRequestBody(upstreamBody)
|
|
helps.RecordAPIWebsocketRequest(ctx, e.cfg, helps.UpstreamRequestLog{
|
|
URL: wsURL,
|
|
Method: "WEBSOCKET",
|
|
Headers: wsHeaders.Clone(),
|
|
Body: wsReqBodyRetry,
|
|
Provider: e.Identifier(),
|
|
AuthID: authID,
|
|
AuthLabel: authLabel,
|
|
AuthType: authType,
|
|
AuthValue: authValue,
|
|
})
|
|
recordAPIWebsocketHandshake(ctx, e.cfg, respHSRetry)
|
|
reporter.StartResponseTTFT()
|
|
if errSendRetry := writeCodexWebsocketMessage(sess, conn, wsReqBodyRetry); errSendRetry != nil {
|
|
errSendRetry = mapCodexWebsocketWriteError(sess, conn, errSendRetry)
|
|
helps.RecordAPIWebsocketError(ctx, e.cfg, "send_retry", errSendRetry)
|
|
e.invalidateUpstreamConn(sess, conn, "send_error", errSendRetry)
|
|
sess.clearActive(conn, readCh)
|
|
sess.reqMu.Unlock()
|
|
return nil, errSendRetry
|
|
}
|
|
wsReqBody = wsReqBodyRetry
|
|
} else {
|
|
logCodexWebsocketDisconnected(executionSessionID, authID, wsURL, "send_error", errSend)
|
|
if errClose := closer.Close(); errClose != nil {
|
|
log.Errorf("codex websockets executor: close websocket error: %v", errClose)
|
|
}
|
|
return nil, errSend
|
|
}
|
|
}
|
|
|
|
if optimizeMultiAgentV2 || multiAgentV2Conflict {
|
|
sess.setMultiAgentV2Optimized(conn, optimizeMultiAgentV2 && !multiAgentV2Conflict)
|
|
}
|
|
|
|
buffering := e.cfg != nil && e.cfg.Codex.StreamBootstrapBuffering
|
|
|
|
claudeInputTokens := helps.NewClaudeInputTokenState(from, to, responseFormat, originalPayload)
|
|
var param any
|
|
outputItemsByIndex := make(map[int64][]byte)
|
|
var outputItemsFallback [][]byte
|
|
|
|
var bufferedChunks [][]byte
|
|
var initialChunks [][]byte
|
|
immediateTerminal := false
|
|
// bootstrapTerminalErr holds a non-overload terminal failure seen while buffering. It is
|
|
// delivered as an in-stream chunk after the buffered handshake so downstream behaviour stays
|
|
// identical to the unbuffered path instead of silently turning into a credential failover.
|
|
var bootstrapTerminalErr error
|
|
|
|
if buffering {
|
|
for {
|
|
if ctx != nil && ctx.Err() != nil {
|
|
if sess != nil {
|
|
sess.clearActive(conn, readCh)
|
|
unlockStreamSession()
|
|
} else {
|
|
_ = closer.Close()
|
|
}
|
|
return nil, ctx.Err()
|
|
}
|
|
msgType, payload, errRead := readCodexWebsocketMessage(ctx, sess, conn, readCh)
|
|
if errRead != nil {
|
|
mappedErr := mapCodexWebsocketReadError(errRead)
|
|
if sess != nil {
|
|
e.invalidateUpstreamConn(sess, conn, "read_error", mappedErr)
|
|
sess.clearActive(conn, readCh)
|
|
unlockStreamSession()
|
|
} else {
|
|
logCodexWebsocketDisconnected(executionSessionID, authID, wsURL, "read_error", mappedErr)
|
|
_ = closer.Close()
|
|
}
|
|
helps.RecordAPIWebsocketError(ctx, e.cfg, "read", mappedErr)
|
|
reporter.PublishFailure(ctx, mappedErr)
|
|
return nil, mappedErr
|
|
}
|
|
if msgType != websocket.TextMessage {
|
|
if msgType == websocket.BinaryMessage {
|
|
errBinary := fmt.Errorf("codex websockets executor: unexpected binary message")
|
|
if sess != nil {
|
|
e.invalidateUpstreamConn(sess, conn, "unexpected_binary", errBinary)
|
|
sess.clearActive(conn, readCh)
|
|
unlockStreamSession()
|
|
} else {
|
|
logCodexWebsocketDisconnected(executionSessionID, authID, wsURL, "unexpected_binary", errBinary)
|
|
_ = closer.Close()
|
|
}
|
|
helps.RecordAPIWebsocketError(ctx, e.cfg, "unexpected_binary", errBinary)
|
|
reporter.PublishFailure(ctx, errBinary)
|
|
return nil, errBinary
|
|
}
|
|
continue
|
|
}
|
|
|
|
payload = bytes.TrimSpace(payload)
|
|
if len(payload) == 0 {
|
|
continue
|
|
}
|
|
reporter.MarkFirstResponseByte()
|
|
payload = applyCodexIdentityConfuseResponsePayload(payload, identityState)
|
|
helps.AppendAPIWebsocketResponse(ctx, e.cfg, payload)
|
|
payload = helps.RestoreCodexMultiAgentV2Response(payload, restoreMultiAgentV2)
|
|
|
|
if wsErr, ok := parseCodexWebsocketError(payload); ok {
|
|
if sess != nil {
|
|
e.invalidateUpstreamConn(sess, conn, "upstream_error", wsErr)
|
|
sess.clearActive(conn, readCh)
|
|
unlockStreamSession()
|
|
} else {
|
|
logCodexWebsocketDisconnected(executionSessionID, authID, wsURL, "upstream_error", wsErr)
|
|
_ = closer.Close()
|
|
}
|
|
if errClearReplay := clearCodexReasoningReplayOnWebsocketError(ctx, replayScope, payload); errClearReplay != nil {
|
|
helps.RecordAPIWebsocketError(ctx, e.cfg, "replay_clear_error", errClearReplay)
|
|
reporter.PublishFailure(ctx, errClearReplay)
|
|
return nil, errClearReplay
|
|
}
|
|
helps.RecordAPIWebsocketError(ctx, e.cfg, "upstream_error", wsErr)
|
|
reporter.PublishFailure(ctx, wsErr)
|
|
return nil, wsErr
|
|
}
|
|
if streamErr, terminalBody, ok := codexTerminalFailureErr(payload); ok {
|
|
// A transient capacity rejection is retried on another credential, so the
|
|
// downstream websocket session must survive this upstream teardown. Notifying
|
|
// the disconnect here would close the client connection before the retry can
|
|
// deliver anything. Every other terminal failure is forwarded in-stream and
|
|
// legitimately terminates the session, so it keeps the notifying variant.
|
|
failoverPending := isCodexOverloadBootstrapFailure(terminalBody)
|
|
if sess != nil {
|
|
unlockStreamSession()
|
|
if failoverPending {
|
|
e.invalidateUpstreamConnWithoutDisconnectNotify(sess, conn, "terminal_failure", streamErr)
|
|
} else {
|
|
e.invalidateUpstreamConn(sess, conn, "terminal_failure", streamErr)
|
|
}
|
|
sess.clearActive(conn, readCh)
|
|
} else {
|
|
logCodexWebsocketDisconnected(executionSessionID, authID, wsURL, "terminal_failure", streamErr)
|
|
_ = closer.Close()
|
|
}
|
|
if errClearReplay := clearCodexReasoningReplayOnInvalidSignature(ctx, replayScope, streamErr.StatusCode(), terminalBody); errClearReplay != nil {
|
|
helps.RecordAPIWebsocketError(ctx, e.cfg, "replay_clear_error", errClearReplay)
|
|
reporter.PublishFailure(ctx, errClearReplay)
|
|
return nil, errClearReplay
|
|
}
|
|
helps.RecordAPIWebsocketError(ctx, e.cfg, "upstream_error", streamErr)
|
|
reporter.PublishFailure(ctx, streamErr)
|
|
if failoverPending {
|
|
// Fail the attempt before the downstream headers are committed so the
|
|
// conductor can transparently retry on another credential, and report the
|
|
// status the upstream refused to put on the wire.
|
|
helps.LogWithRequestID(ctx).Debugf("codex websockets executor: bootstrap overload rejection after %d buffered handshake events, failing over", len(bufferedChunks))
|
|
return nil, newCodexBootstrapOverloadErr(terminalBody)
|
|
}
|
|
bootstrapTerminalErr = streamErr
|
|
break
|
|
}
|
|
|
|
eventType := gjson.GetBytes(payload, "type").String()
|
|
isTerminalEvent := eventType == "response.completed" || eventType == "response.done" || eventType == "error"
|
|
if eventType == "response.output_item.done" {
|
|
collectCodexOutputItemDone(payload, outputItemsByIndex, &outputItemsFallback)
|
|
}
|
|
completedPayload := payload
|
|
if eventType == "response.completed" || eventType == "response.done" {
|
|
completedPayload = normalizeCodexWebsocketCompletion(completedPayload)
|
|
completedPayload = patchCodexCompletedOutput(completedPayload, outputItemsByIndex, outputItemsFallback)
|
|
cacheCodexReasoningReplayFromCompleted(replayScope, completedPayload)
|
|
if detail, ok := helps.ParseCodexUsage(completedPayload); ok {
|
|
reporter.Publish(ctx, detail)
|
|
}
|
|
}
|
|
|
|
var currentChunks [][]byte
|
|
if cliproxyexecutor.DownstreamWebsocket(ctx) {
|
|
clientPayload := applyCodexIdentityExposeResponsePayload(payload, identityState)
|
|
downstreamPayload := helps.EnsureResponsesUsageDetails(clientPayload)
|
|
currentChunks = [][]byte{downstreamPayload}
|
|
} else {
|
|
payload = normalizeCodexWebsocketCompletion(payload)
|
|
if eventType == "response.completed" || eventType == "response.done" {
|
|
payload = completedPayload
|
|
}
|
|
clientPayload := applyCodexIdentityExposeResponsePayload(payload, identityState)
|
|
line := encodeCodexWebsocketAsSSE(clientPayload)
|
|
currentChunks = helps.TranslateStreamWithClaudeInputTokens(ctx, to, responseFormat, req.Model, originalPayload, clientBody, line, ¶m, claudeInputTokens)
|
|
}
|
|
|
|
if isCodexHandshakeMetadataEvent(eventType) && !isTerminalEvent {
|
|
if len(bufferedChunks) < codexBootstrapMaxBufferedEvents {
|
|
bufferedChunks = append(bufferedChunks, currentChunks...)
|
|
continue
|
|
}
|
|
helps.LogWithRequestID(ctx).Debugf("codex websockets executor: bootstrap buffer limit %d reached, releasing stream without overload probing", codexBootstrapMaxBufferedEvents)
|
|
}
|
|
|
|
initialChunks = currentChunks
|
|
if isTerminalEvent {
|
|
immediateTerminal = true
|
|
}
|
|
break
|
|
}
|
|
}
|
|
|
|
chanCapacity := len(bufferedChunks) + len(initialChunks)
|
|
if bootstrapTerminalErr != nil {
|
|
chanCapacity++
|
|
}
|
|
out := make(chan cliproxyexecutor.StreamChunk, chanCapacity)
|
|
for _, chunk := range bufferedChunks {
|
|
out <- cliproxyexecutor.StreamChunk{Payload: chunk}
|
|
}
|
|
for _, chunk := range initialChunks {
|
|
out <- cliproxyexecutor.StreamChunk{Payload: chunk}
|
|
}
|
|
if bootstrapTerminalErr != nil {
|
|
// The upstream connection was already invalidated and released in the terminal-failure
|
|
// branch above, so only the buffered payloads plus the in-stream error remain to emit.
|
|
out <- cliproxyexecutor.StreamChunk{Err: bootstrapTerminalErr}
|
|
close(out)
|
|
return &cliproxyexecutor.StreamResult{Headers: upstreamHeaders, Chunks: out}, nil
|
|
}
|
|
if immediateTerminal {
|
|
if sess != nil {
|
|
sess.clearActive(conn, readCh)
|
|
unlockStreamSession()
|
|
} else {
|
|
logCodexWebsocketDisconnected(executionSessionID, authID, wsURL, "completed", nil)
|
|
if errClose := closer.Close(); errClose != nil {
|
|
log.Errorf("codex websockets executor: close websocket error: %v", errClose)
|
|
}
|
|
}
|
|
close(out)
|
|
return &cliproxyexecutor.StreamResult{Headers: upstreamHeaders, Chunks: out}, nil
|
|
}
|
|
|
|
go func() {
|
|
terminateReason := "completed"
|
|
var terminateErr error
|
|
|
|
defer close(out)
|
|
defer func() {
|
|
if sess != nil {
|
|
sess.clearActive(conn, readCh)
|
|
unlockStreamSession()
|
|
return
|
|
}
|
|
logCodexWebsocketDisconnected(executionSessionID, authID, wsURL, terminateReason, terminateErr)
|
|
if errClose := closer.Close(); errClose != nil {
|
|
log.Errorf("codex websockets executor: close websocket error: %v", errClose)
|
|
}
|
|
}()
|
|
|
|
send := func(chunk cliproxyexecutor.StreamChunk) bool {
|
|
if ctx == nil {
|
|
out <- chunk
|
|
return true
|
|
}
|
|
select {
|
|
case out <- chunk:
|
|
return true
|
|
case <-ctx.Done():
|
|
return false
|
|
}
|
|
}
|
|
|
|
for {
|
|
if ctx != nil && ctx.Err() != nil {
|
|
terminateReason = "context_done"
|
|
terminateErr = ctx.Err()
|
|
_ = send(cliproxyexecutor.StreamChunk{Err: ctx.Err()})
|
|
return
|
|
}
|
|
msgType, payload, errRead := readCodexWebsocketMessage(ctx, sess, conn, readCh)
|
|
if errRead != nil {
|
|
if sess != nil && ctx != nil && ctx.Err() != nil {
|
|
terminateReason = "context_done"
|
|
terminateErr = ctx.Err()
|
|
_ = send(cliproxyexecutor.StreamChunk{Err: ctx.Err()})
|
|
return
|
|
}
|
|
mappedErr := mapCodexWebsocketReadError(errRead)
|
|
terminateReason = "read_error"
|
|
terminateErr = mappedErr
|
|
helps.RecordAPIWebsocketError(ctx, e.cfg, "read", mappedErr)
|
|
reporter.PublishFailure(ctx, mappedErr)
|
|
_ = send(cliproxyexecutor.StreamChunk{Err: mappedErr})
|
|
return
|
|
}
|
|
if msgType != websocket.TextMessage {
|
|
if msgType == websocket.BinaryMessage {
|
|
err = fmt.Errorf("codex websockets executor: unexpected binary message")
|
|
terminateReason = "unexpected_binary"
|
|
terminateErr = err
|
|
helps.RecordAPIWebsocketError(ctx, e.cfg, "unexpected_binary", err)
|
|
reporter.PublishFailure(ctx, err)
|
|
if sess != nil {
|
|
e.invalidateUpstreamConn(sess, conn, "unexpected_binary", err)
|
|
}
|
|
_ = send(cliproxyexecutor.StreamChunk{Err: err})
|
|
return
|
|
}
|
|
continue
|
|
}
|
|
|
|
payload = bytes.TrimSpace(payload)
|
|
if len(payload) == 0 {
|
|
continue
|
|
}
|
|
reporter.MarkFirstResponseByte()
|
|
payload = applyCodexIdentityConfuseResponsePayload(payload, identityState)
|
|
helps.AppendAPIWebsocketResponse(ctx, e.cfg, payload)
|
|
payload = helps.RestoreCodexMultiAgentV2Response(payload, restoreMultiAgentV2)
|
|
|
|
if wsErr, ok := parseCodexWebsocketError(payload); ok {
|
|
terminateReason = "upstream_error"
|
|
terminateErr = wsErr
|
|
if sess != nil {
|
|
e.invalidateUpstreamConn(sess, conn, "upstream_error", wsErr)
|
|
}
|
|
if errClearReplay := clearCodexReasoningReplayOnWebsocketError(ctx, replayScope, payload); errClearReplay != nil {
|
|
terminateErr = errClearReplay
|
|
helps.RecordAPIWebsocketError(ctx, e.cfg, "replay_clear_error", errClearReplay)
|
|
reporter.PublishFailure(ctx, errClearReplay)
|
|
_ = send(cliproxyexecutor.StreamChunk{Err: errClearReplay})
|
|
return
|
|
}
|
|
helps.RecordAPIWebsocketError(ctx, e.cfg, "upstream_error", wsErr)
|
|
reporter.PublishFailure(ctx, wsErr)
|
|
_ = send(cliproxyexecutor.StreamChunk{Err: wsErr})
|
|
return
|
|
}
|
|
if streamErr, terminalBody, ok := codexTerminalFailureErr(payload); ok {
|
|
terminateReason = "upstream_error"
|
|
terminateErr = streamErr
|
|
if sess != nil {
|
|
unlockStreamSession()
|
|
e.invalidateUpstreamConn(sess, conn, "terminal_failure", streamErr)
|
|
}
|
|
if errClearReplay := clearCodexReasoningReplayOnInvalidSignature(ctx, replayScope, streamErr.StatusCode(), terminalBody); errClearReplay != nil {
|
|
terminateErr = errClearReplay
|
|
helps.RecordAPIWebsocketError(ctx, e.cfg, "replay_clear_error", errClearReplay)
|
|
reporter.PublishFailure(ctx, errClearReplay)
|
|
_ = send(cliproxyexecutor.StreamChunk{Err: errClearReplay})
|
|
return
|
|
}
|
|
helps.RecordAPIWebsocketError(ctx, e.cfg, "upstream_error", streamErr)
|
|
reporter.PublishFailure(ctx, streamErr)
|
|
_ = send(cliproxyexecutor.StreamChunk{Err: streamErr})
|
|
return
|
|
}
|
|
|
|
eventType := gjson.GetBytes(payload, "type").String()
|
|
isTerminalEvent := eventType == "response.completed" || eventType == "response.done" || eventType == "error"
|
|
if eventType == "response.output_item.done" {
|
|
collectCodexOutputItemDone(payload, outputItemsByIndex, &outputItemsFallback)
|
|
}
|
|
completedPayload := payload
|
|
if eventType == "response.completed" || eventType == "response.done" {
|
|
completedPayload = normalizeCodexWebsocketCompletion(completedPayload)
|
|
completedPayload = patchCodexCompletedOutput(completedPayload, outputItemsByIndex, outputItemsFallback)
|
|
cacheCodexReasoningReplayFromCompleted(replayScope, completedPayload)
|
|
if detail, ok := helps.ParseCodexUsage(completedPayload); ok {
|
|
reporter.Publish(ctx, detail)
|
|
}
|
|
}
|
|
|
|
clientPayload := applyCodexIdentityExposeResponsePayload(payload, identityState)
|
|
if cliproxyexecutor.DownstreamWebsocket(ctx) {
|
|
downstreamPayload := helps.EnsureResponsesUsageDetails(clientPayload)
|
|
if !send(cliproxyexecutor.StreamChunk{Payload: downstreamPayload}) {
|
|
terminateReason = "context_done"
|
|
terminateErr = ctx.Err()
|
|
return
|
|
}
|
|
if isTerminalEvent {
|
|
return
|
|
}
|
|
continue
|
|
}
|
|
|
|
payload = normalizeCodexWebsocketCompletion(payload)
|
|
if eventType == "response.completed" || eventType == "response.done" {
|
|
payload = completedPayload
|
|
}
|
|
eventType = gjson.GetBytes(payload, "type").String()
|
|
clientPayload = applyCodexIdentityExposeResponsePayload(payload, identityState)
|
|
line := encodeCodexWebsocketAsSSE(clientPayload)
|
|
chunks := helps.TranslateStreamWithClaudeInputTokens(ctx, to, responseFormat, req.Model, originalPayload, clientBody, line, ¶m, claudeInputTokens)
|
|
for i := range chunks {
|
|
if !send(cliproxyexecutor.StreamChunk{Payload: chunks[i]}) {
|
|
terminateReason = "context_done"
|
|
terminateErr = ctx.Err()
|
|
return
|
|
}
|
|
}
|
|
if eventType == "response.completed" || eventType == "response.done" {
|
|
return
|
|
}
|
|
}
|
|
}()
|
|
|
|
return &cliproxyexecutor.StreamResult{Headers: upstreamHeaders, Chunks: out}, nil
|
|
}
|