776 lines · 30.6 KB · Go
RawDownload
package executor import (	"bytes"	"context"	"fmt"	"net/http"	"strings"	"time" 	"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)	preserveNativeOutput := helps.IsNativeCodexRequest(req.Payload, 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, preserveNativeOutput)	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)	body = helps.NormalizeCodexToolSchemas(body)	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, preserveNativeOutput, 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	isEphemeralSession := false	if executionSessionID != "" {		sess = e.getOrCreateSession(executionSessionID)		if sess != nil {			sess.reqMu.Lock()		}	} else {		isEphemeralSession = true		sess = newEphemeralCodexWebsocketSession()	}	streamSessionLocked := sess != nil && !isEphemeralSession	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	dialCtx := ctx	if cliproxyexecutor.RequiredUpstreamWebsocket(ctx) {		conn, closer = existingWebsocketSessionConn(sess, authID, wsURL)		if conn == nil {			unlockStreamSession()			return nil, cliproxyexecutor.NewUpstreamWebsocketReplayRequiredError()		}	} else {		dialCtx = cliproxyexecutor.WithUpstreamAttemptTracker(ctx)		conn, closer, respHS, errDial = e.ensureUpstreamConn(dialCtx, 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 {			unlockStreamSession()			if opts.ExecutionLifecycle == nil && !cliproxyexecutor.DownstreamWebsocket(ctx) {				return e.CodexExecutor.ExecuteStream(ctx, auth, req, opts)			}			if cliproxyexecutor.UpstreamAttempted(dialCtx) {				cliproxyexecutor.MarkUpstreamAttempt(ctx)			}			return nil, statusErr{code: respHS.StatusCode, msg: string(bodyErr)}		}		if cliproxyexecutor.UpstreamAttempted(dialCtx) {			cliproxyexecutor.MarkUpstreamAttempt(ctx)		}		if respHS != nil && respHS.StatusCode > 0 {			unlockStreamSession()			return nil, newCodexStatusErrWithCooling(respHS.StatusCode, bodyErr, e.modelLevelCooling())		}		helps.RecordAPIWebsocketError(ctx, e.cfg, "dial", errDial)		unlockStreamSession()		return nil, errDial	}	if errBind := sess.bindExecutionLifecycle(opts, conn, closer, req.Model); errBind != nil {		unlockStreamSession()		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)) 	cliproxyexecutor.MarkUpstreamAttempt(ctx)	if errSend := writeCodexWebsocketMessage(sess, conn, wsReqBody); errSend != nil {		errSend = mapCodexWebsocketWriteError(sess, conn, errSend)		helps.RecordAPIWebsocketError(ctx, e.cfg, "send", errSend)		if sess != nil && !isEphemeralSession {			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()			cliproxyexecutor.MarkUpstreamAttempt(ctx)			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 {			if sess != nil {				e.invalidateUpstreamConn(sess, conn, "send_error", errSend)				sess.clearActive(conn, readCh)				if isEphemeralSession {					closeCodexWebsocketSession(sess, "send_error")				}			} else {				logCodexWebsocketDisconnected(executionSessionID, authID, wsURL, "send_error", errSend)				if closer != nil {					_ = closer.Close()				}			}			return nil, errSend		}	} 	if optimizeMultiAgentV2 || multiAgentV2Conflict {		sess.setMultiAgentV2Optimized(conn, optimizeMultiAgentV2 && !multiAgentV2Conflict)	} 	buffering := e.cfg != nil && e.cfg.Codex.StreamBootstrapBuffering	var bootstrapTimeout time.Duration	var bootstrapStart time.Time	var exhaustionLogged bool	if buffering {		bootstrapTimeout = e.cfg.Codex.StreamBootstrapTimeoutDuration()		bootstrapStart = nowCodexBootstrap()	} 	claudeInputTokens := helps.NewClaudeInputTokenState(from, to, responseFormat, originalPayload)	var param any	outputItemsByIndex := make(map[int64][]byte)	var outputItemsFallback [][]byte 	var bufferedChunks [][]byte	// bufferedFrames counts every websocket message read during bootstrap, including the ones the	// loop skips, so a peer that only sends frames the loop ignores cannot keep the window open.	bufferedFrames := 0	bufferedBytes := 0	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	sawOutputDelta := false 	if buffering {		for {			if ctx != nil && ctx.Err() != nil {				if sess != nil {					sess.clearActive(conn, readCh)					unlockStreamSession()					if isEphemeralSession {						closeCodexWebsocketSession(sess, "context_done")					}				} 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()					if isEphemeralSession {						closeCodexWebsocketSession(sess, "read_error")					}				} else {					logCodexWebsocketDisconnected(executionSessionID, authID, wsURL, "read_error", mappedErr)					_ = closer.Close()				}				helps.RecordAPIWebsocketError(ctx, e.cfg, "read", mappedErr)				reporter.PublishFailure(ctx, mappedErr)				return nil, mappedErr			}			// Count every message ReadMessage returns, including the ones this loop goes on to skip,			// so a peer sending only skippable text frames still closes the window. Control frames			// are not counted: the websocket library answers ping and pong inside ReadMessage and			// never returns them, so only the read deadline bounds a peer that sends nothing else.			// windowOpen is carried into the skip branches below rather than breaking here, because			// this message has not been processed yet and dropping it would lose a token, or a			// terminal event, from the turn.			bufferedFrames++			timeSinceStart := nowCodexBootstrap().Sub(bootstrapStart)			timeoutReached := bootstrapTimeout > 0 && timeSinceStart >= bootstrapTimeout			windowOpen := bufferedFrames <= codexBootstrapMaxBufferedFrames && !timeoutReached			if !windowOpen && !exhaustionLogged {				exhaustionLogged = true				exhausted := "frame budget"				if timeoutReached {					exhausted = "time budget"				}				helps.LogWithRequestID(ctx).Debugf("codex websockets executor: bootstrap %s exhausted after %d messages read / %v; this message will be released", exhausted, bufferedFrames, timeSinceStart)			}			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()						if isEphemeralSession {							closeCodexWebsocketSession(sess, "unexpected_binary")						}					} 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				}				// No window check here: ReadMessage only ever returns text or binary, and binary				// returned just above, so nothing reaches this line. The empty-payload skip below				// is the reachable one and does consult the window.				continue			} 			payload = bytes.TrimSpace(payload)			if len(payload) == 0 {				if !windowOpen {					break				}				continue			}			observeCodexTokenEvent(reporter, payload)			payload = applyCodexIdentityConfuseResponsePayload(payload, identityState)			helps.AppendCodexAPIWebsocketResponse(ctx, e.cfg, payload)			helps.EmitWebSocketResponseEvent(ctx, opts, auth, e.Identifier(), req.Model, payload)			payload = helps.RestoreCodexMultiAgentV2Response(payload, restoreMultiAgentV2) 			if wsErr, ok := parseCodexWebsocketErrorWithCooling(payload, e.modelLevelCooling()); 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)				if timeoutReached {					helps.LogWithRequestID(ctx).Debugf("codex websockets executor: bootstrap error after %d messages read / %v, time budget exhausted; delivering in-stream", bufferedFrames, timeSinceStart)					bootstrapTerminalErr = wsErr					break				}				return nil, wsErr			}			if streamErr, terminalBody, ok := codexTerminalFailureErrWithCooling(payload, e.modelLevelCooling()); 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 failoverPending && timeoutReached {					failoverPending = false					helps.LogWithRequestID(ctx).Debugf("codex websockets executor: bootstrap overload rejection after %d messages read / %v, time budget exhausted; delivering in-stream", bufferedFrames, timeSinceStart)				}				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 {					if isEphemeralSession {						closeCodexWebsocketSession(sess, "bootstrap_overload")					}					// 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 messages read, failing over", bufferedFrames)					return nil, newCodexBootstrapOverloadErr(terminalBody)				}				bootstrapTerminalErr = streamErr				break			} 			eventType := gjson.GetBytes(payload, "type").String()			isTerminalEvent := eventType == "response.completed" || eventType == "response.done" || eventType == "response.incomplete" || eventType == "response.failed" || eventType == "error"			if helps.HasMeaningfulCodexOutputDelta(payload) {				sawOutputDelta = true			}			if helps.IsCodexTerminalEmptyIncomplete(payload, len(outputItemsByIndex)+len(outputItemsFallback), sawOutputDelta) {				streamErr := newCodexEmptyIncompleteStreamError()				helps.RecordAPIWebsocketError(ctx, e.cfg, "upstream_error", streamErr)				reporter.PublishFailure(ctx, streamErr)				if sess != nil {					e.invalidateUpstreamConn(sess, conn, "terminal_empty_incomplete", streamErr)					sess.clearActive(conn, readCh)					unlockStreamSession()				} else if closer != nil {					_ = closer.Close()				}				bootstrapTerminalErr = streamErr				break			}			if eventType == "response.output_item.done" {				collectCodexOutputItemDone(payload, outputItemsByIndex, &outputItemsFallback)			}			completedPayload := payload			if eventType == "response.completed" || eventType == "response.done" || eventType == "response.incomplete" {				completedPayload = normalizeCodexWebsocketCompletion(completedPayload)				if !preserveNativeOutput {					completedPayload = patchCodexCompletedOutput(completedPayload, outputItemsByIndex, outputItemsFallback)				}				if eventType != "response.incomplete" {					cacheCodexReasoningReplayFromCompleted(replayScope, completedPayload)				}				if detail, ok := helps.ParseCodexUsage(completedPayload); ok {					reporter.Publish(ctx, detail)				} else {					reporter.EnsurePublished(ctx)				}			} 			var currentChunks [][]byte

Lines 1–500 of 776.