1663 lines · 48.3 KB · Go
RawDownload
// Package executor provides runtime execution capabilities for various AI service providers.// This file implements a Codex executor that uses the Responses API WebSocket transport.package executor import (	"bytes"	"context"	"fmt"	"io"	"net"	"net/http"	"net/url"	"strconv"	"strings"	"sync"	"time" 	"github.com/gin-gonic/gin"	"github.com/google/uuid"	"github.com/gorilla/websocket"	"github.com/router-for-me/CLIProxyAPI/v6/internal/config"	"github.com/router-for-me/CLIProxyAPI/v6/internal/misc"	"github.com/router-for-me/CLIProxyAPI/v6/internal/runtime/executor/helps"	"github.com/router-for-me/CLIProxyAPI/v6/internal/thinking"	"github.com/router-for-me/CLIProxyAPI/v6/internal/util"	cliproxyauth "github.com/router-for-me/CLIProxyAPI/v6/sdk/cliproxy/auth"	cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v6/sdk/cliproxy/executor"	"github.com/router-for-me/CLIProxyAPI/v6/sdk/proxyutil"	sdktranslator "github.com/router-for-me/CLIProxyAPI/v6/sdk/translator"	log "github.com/sirupsen/logrus"	"github.com/tidwall/gjson"	"github.com/tidwall/sjson"	"golang.org/x/net/proxy") const (	codexResponsesWebsocketBetaHeaderValue = "responses_websockets=2026-02-06"	codexResponsesWebsocketIdleTimeout     = 5 * time.Minute	codexResponsesWebsocketHandshakeTO     = 30 * time.Second) // CodexWebsocketsExecutor executes Codex Responses requests using a WebSocket transport.//// It preserves the existing CodexExecutor HTTP implementation as a fallback for endpoints// not available over WebSocket (e.g. /responses/compact) and for websocket upgrade failures.type CodexWebsocketsExecutor struct {	*CodexExecutor 	store *codexWebsocketSessionStore} type codexWebsocketSessionStore struct {	mu       sync.Mutex	sessions map[string]*codexWebsocketSession} var globalCodexWebsocketSessionStore = &codexWebsocketSessionStore{	sessions: make(map[string]*codexWebsocketSession),} type codexWebsocketSession struct {	sessionID string 	reqMu sync.Mutex 	connMu sync.Mutex	conn   *websocket.Conn	wsURL  string	authID string 	writeMu sync.Mutex 	activeMu     sync.Mutex	activeCh     chan codexWebsocketRead	activeDone   <-chan struct{}	activeCancel context.CancelFunc 	readerConn *websocket.Conn 	upstreamDisconnectOnce sync.Once	upstreamDisconnectCh   chan error} func NewCodexWebsocketsExecutor(cfg *config.Config) *CodexWebsocketsExecutor {	return &CodexWebsocketsExecutor{		CodexExecutor: NewCodexExecutor(cfg),		store:         globalCodexWebsocketSessionStore,	}} type codexWebsocketRead struct {	conn    *websocket.Conn	msgType int	payload []byte	err     error} func (s *codexWebsocketSession) setActive(ch chan codexWebsocketRead) {	if s == nil {		return	}	s.activeMu.Lock()	if s.activeCancel != nil {		s.activeCancel()		s.activeCancel = nil		s.activeDone = nil	}	s.activeCh = ch	if ch != nil {		activeCtx, activeCancel := context.WithCancel(context.Background())		s.activeDone = activeCtx.Done()		s.activeCancel = activeCancel	}	s.activeMu.Unlock()} func (s *codexWebsocketSession) clearActive(ch chan codexWebsocketRead) {	if s == nil {		return	}	s.activeMu.Lock()	if s.activeCh == ch {		s.activeCh = nil		if s.activeCancel != nil {			s.activeCancel()		}		s.activeCancel = nil		s.activeDone = nil	}	s.activeMu.Unlock()} func (s *codexWebsocketSession) writeMessage(conn *websocket.Conn, msgType int, payload []byte) error {	if s == nil {		return fmt.Errorf("codex websockets executor: session is nil")	}	if conn == nil {		return fmt.Errorf("codex websockets executor: websocket conn is nil")	}	s.writeMu.Lock()	defer s.writeMu.Unlock()	return conn.WriteMessage(msgType, payload)} func (s *codexWebsocketSession) configureConn(conn *websocket.Conn) {	if s == nil || conn == nil {		return	}	conn.SetPingHandler(func(appData string) error {		s.writeMu.Lock()		defer s.writeMu.Unlock()		// Reply pongs from the same write lock to avoid concurrent writes.		return conn.WriteControl(websocket.PongMessage, []byte(appData), time.Now().Add(10*time.Second))	})} func (s *codexWebsocketSession) notifyUpstreamDisconnect(err error) {	if s == nil {		return	}	s.upstreamDisconnectOnce.Do(func() {		if s.upstreamDisconnectCh == nil {			return		}		select {		case s.upstreamDisconnectCh <- err:		default:		}		close(s.upstreamDisconnectCh)	})} func (e *CodexWebsocketsExecutor) Execute(ctx context.Context, auth *cliproxyauth.Auth, req cliproxyexecutor.Request, opts cliproxyexecutor.Options) (resp cliproxyexecutor.Response, err error) {	if ctx == nil {		ctx = context.Background()	}	if opts.Alt == "responses/compact" {		return e.CodexExecutor.executeCompact(ctx, auth, req, opts)	} 	baseModel := thinking.ParseSuffix(req.Model).ModelName	apiKey, baseURL := codexCreds(auth)	if baseURL == "" {		baseURL = "https://chatgpt.com/backend-api/codex"	} 	reporter := helps.NewUsageReporter(ctx, e.Identifier(), baseModel, auth)	defer reporter.TrackFailure(ctx, &err) 	from := opts.SourceFormat	to := sdktranslator.FromString("codex")	originalPayloadSource := req.Payload	if len(opts.OriginalRequest) > 0 {		originalPayloadSource = opts.OriginalRequest	}	originalPayload := originalPayloadSource	originalTranslated := sdktranslator.TranslateRequest(from, to, baseModel, originalPayload, false)	body := sdktranslator.TranslateRequest(from, to, baseModel, req.Payload, false) 	body, err = thinking.ApplyThinking(body, req.Model, from.String(), to.String(), e.Identifier())	if err != nil {		return resp, err	} 	requestedModel := helps.PayloadRequestedModel(opts, req.Model)	requestPath := helps.PayloadRequestPath(opts)	body = helps.ApplyPayloadConfigWithRoot(e.cfg, baseModel, to.String(), "", body, originalTranslated, requestedModel, requestPath)	body, _ = sjson.SetBytes(body, "model", baseModel)	body, _ = sjson.SetBytes(body, "stream", true)	body, _ = sjson.DeleteBytes(body, "prompt_cache_retention")	body, _ = sjson.DeleteBytes(body, "safety_identifier")	body = normalizeCodexInstructions(body)	if e.cfg == nil || e.cfg.DisableImageGeneration == config.DisableImageGenerationOff {		body = ensureImageGenerationTool(body, baseModel, auth)	} 	httpURL := strings.TrimSuffix(baseURL, "/") + "/responses"	wsURL, err := buildCodexResponsesWebsocketURL(httpURL)	if err != nil {		return resp, err	} 	body, wsHeaders := applyCodexPromptCacheHeaders(from, req, body)	wsHeaders = applyCodexWebsocketHeaders(ctx, wsHeaders, auth, apiKey, e.cfg) 	var authID, authLabel, authType, authValue string	if auth != nil {		authID = auth.ID		authLabel = auth.Label		authType, authValue = auth.AccountInfo()	} 	executionSessionID := executionSessionIDFromOptions(opts)	var sess *codexWebsocketSession	if executionSessionID != "" {		sess = e.getOrCreateSession(executionSessionID)		sess.reqMu.Lock()		defer sess.reqMu.Unlock()	} 	wsReqBody := buildCodexWebsocketRequestBody(body)	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) 	conn, respHS, errDial := e.ensureUpstreamConn(ctx, auth, sess, authID, wsURL, wsHeaders)	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 {			return e.CodexExecutor.Execute(ctx, auth, req, opts)		}		if respHS != nil && respHS.StatusCode > 0 {			return resp, statusErr{code: respHS.StatusCode, msg: string(bodyErr)}		}		helps.RecordAPIWebsocketError(ctx, e.cfg, "dial", errDial)		return resp, errDial	}	recordAPIWebsocketHandshake(ctx, e.cfg, respHS)	if sess == nil {		logCodexWebsocketConnected(executionSessionID, authID, wsURL)		defer func() {			reason := "completed"			if err != nil {				reason = "error"			}			logCodexWebsocketDisconnected(executionSessionID, authID, wsURL, reason, err)			if errClose := conn.Close(); errClose != nil {				log.Errorf("codex websockets executor: close websocket error: %v", errClose)			}		}()	} 	var readCh chan codexWebsocketRead	if sess != nil {		readCh = make(chan codexWebsocketRead, 4096)		sess.setActive(readCh)		defer sess.clearActive(readCh)	} 	if errSend := writeCodexWebsocketMessage(sess, conn, wsReqBody); errSend != nil {		if sess != nil {			e.invalidateUpstreamConn(sess, conn, "send_error", errSend) 			// Retry once with a fresh websocket connection. This is mainly to handle			// upstream closing the socket between sequential requests within the same			// execution session.			connRetry, respHSRetry, errDialRetry := e.ensureUpstreamConn(ctx, auth, sess, authID, wsURL, wsHeaders)			if errDialRetry == nil && connRetry != nil {				wsReqBodyRetry := buildCodexWebsocketRequestBody(body)				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)				if errSendRetry := writeCodexWebsocketMessage(sess, connRetry, wsReqBodyRetry); errSendRetry == nil {					conn = connRetry					wsReqBody = wsReqBodyRetry				} else {					e.invalidateUpstreamConn(sess, connRetry, "send_error", errSendRetry)					helps.RecordAPIWebsocketError(ctx, e.cfg, "send_retry", errSendRetry)					return resp, errSendRetry				}			} else {				closeHTTPResponseBody(respHSRetry, "codex websockets executor: close handshake response body error")				helps.RecordAPIWebsocketError(ctx, e.cfg, "dial_retry", errDialRetry)				return resp, errDialRetry			}		} else {			helps.RecordAPIWebsocketError(ctx, e.cfg, "send", errSend)			return resp, errSend		}	} 	for {		if ctx != nil && ctx.Err() != nil {			return resp, ctx.Err()		}		msgType, payload, errRead := readCodexWebsocketMessage(ctx, sess, conn, readCh)		if errRead != nil {			helps.RecordAPIWebsocketError(ctx, e.cfg, "read", errRead)			return resp, errRead		}		if msgType != websocket.TextMessage {			if msgType == websocket.BinaryMessage {				err = fmt.Errorf("codex websockets executor: unexpected binary message")				if sess != nil {					e.invalidateUpstreamConn(sess, conn, "unexpected_binary", err)				}				helps.RecordAPIWebsocketError(ctx, e.cfg, "unexpected_binary", err)				return resp, err			}			continue		} 		payload = bytes.TrimSpace(payload)		if len(payload) == 0 {			continue		}		helps.AppendAPIWebsocketResponse(ctx, e.cfg, payload) 		if wsErr, ok := parseCodexWebsocketError(payload); ok {			if sess != nil {				e.invalidateUpstreamConn(sess, conn, "upstream_error", wsErr)			}			helps.RecordAPIWebsocketError(ctx, e.cfg, "upstream_error", wsErr)			return resp, wsErr		} 		payload = normalizeCodexWebsocketCompletion(payload)		eventType := gjson.GetBytes(payload, "type").String()		if eventType == "response.completed" {			if detail, ok := helps.ParseCodexUsage(payload); ok {				reporter.Publish(ctx, detail)			}			var param any			out := sdktranslator.TranslateNonStream(ctx, to, from, req.Model, originalPayload, body, payload, &param)			resp = cliproxyexecutor.Response{Payload: out}			return resp, nil		}	}} 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.NewUsageReporter(ctx, e.Identifier(), baseModel, auth)	defer reporter.TrackFailure(ctx, &err) 	from := opts.SourceFormat	to := sdktranslator.FromString("codex")	body := req.Payload 	body, err = thinking.ApplyThinking(body, req.Model, from.String(), to.String(), e.Identifier())	if err != nil {		return nil, err	} 	requestedModel := helps.PayloadRequestedModel(opts, req.Model)	requestPath := helps.PayloadRequestPath(opts)	body = helps.ApplyPayloadConfigWithRoot(e.cfg, baseModel, to.String(), "", body, body, requestedModel, requestPath)	body = normalizeCodexInstructions(body)	if e.cfg == nil || e.cfg.DisableImageGeneration == config.DisableImageGenerationOff {		body = ensureImageGenerationTool(body, baseModel, auth)	} 	httpURL := strings.TrimSuffix(baseURL, "/") + "/responses"	wsURL, err := buildCodexResponsesWebsocketURL(httpURL)	if err != nil {		return nil, err	} 	body, wsHeaders := applyCodexPromptCacheHeaders(from, req, body)	wsHeaders = applyCodexWebsocketHeaders(ctx, wsHeaders, auth, apiKey, e.cfg) 	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()		}	} 	wsReqBody := buildCodexWebsocketRequestBody(body)	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) 	conn, 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 {			return e.CodexExecutor.ExecuteStream(ctx, auth, req, opts)		}		if respHS != nil && respHS.StatusCode > 0 {			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	}	recordAPIWebsocketHandshake(ctx, e.cfg, respHS) 	if sess == nil {		logCodexWebsocketConnected(executionSessionID, authID, wsURL)	} 	var readCh chan codexWebsocketRead	if sess != nil {		readCh = make(chan codexWebsocketRead, 4096)		sess.setActive(readCh)	} 	if errSend := writeCodexWebsocketMessage(sess, conn, wsReqBody); errSend != nil {		helps.RecordAPIWebsocketError(ctx, e.cfg, "send", errSend)		if sess != nil {			e.invalidateUpstreamConn(sess, conn, "send_error", errSend) 			// Retry once with a new websocket connection for the same execution session.			connRetry, 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(readCh)				sess.reqMu.Unlock()				return nil, errDialRetry

Lines 1–500 of 1663.