Something went wrong. Try again.
sloproxy
Something went wrong. Try again.
52 kB · 1746 lines
Go
at commit 44ea9abc
1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071727374757677787980818283848586878889909192939495969798991001011021031041051061071081091101111121131141151161171181191201211221231241251261271281291301311321331341351361371381391401411421431441451461471481491501511521531541551561571581591601611621631641651661671681691701711721731741751761771781791801811821831841851861871881891901911921931941951961971981992002012022032042052062072082092102112122132142152162172182192202212222232242252262272282292302312322332342352362372382392402412422432442452462472482492502512522532542552562572582592602612622632642652662672682692702712722732742752762772782792802812822832842852862872882892902912922932942952962972982993003013023033043053063073083093103113123133143153163173183193203213223233243253263273283293303313323333343353363373383393403413423433443453463473483493503513523533543553563573583593603613623633643653663673683693703713723733743753763773783793803813823833843853863873883893903913923933943953963973983994004014024034044054064074084094104114124134144154164174184194204214224234244254264274284294304314324334344354364374384394404414424434444454464474484494504514524534544554564574584594604614624634644654664674684694704714724734744754764774784794804814824834844854864874884894904914924934944954964974984995005015025035045055065075085095105115125135145155165175185195205215225235245255265275285295305315325335345355365375385395405415425435445455465475485495505515525535545555565575585595605615625635645655665675685695705715725735745755765775785795805815825835845855865875885895905915925935945955965975985996006016026036046056066076086096106116126136146156166176186196206216226236246256266276286296306316326336346356366376386396406416426436446456466476486496506516526536546556566576586596606616626636646656666676686696706716726736746756766776786796806816826836846856866876886896906916926936946956966976986997007017027037047057067077087097107117127137147157167177187197207217227237247257267277287297307317327337347357367377387397407417427437447457467477487497507517527537547557567577587597607617627637647657667677687697707717727737747757767777787797807817827837847857867877887897907917927937947957967977987998008018028038048058068078088098108118128138148158168178188198208218228238248258268278288298308318328338348358368378388398408418428438448458468478488498508518528538548558568578588598608618628638648658668678688698708718728738748758768778788798808818828838848858868878888898908918928938948958968978988999009019029039049059069079089099109119129139149159169179189199209219229239249259269279289299309319329339349359369379389399409419429439449459469479489499509519529539549559569579589599609619629639649659669679689699709719729739749759769779789799809819829839849859869879889899909919929939949959969979989991000100110021003100410051006100710081009101010111012101310141015101610171018101910201021102210231024102510261027102810291030103110321033103410351036103710381039104010411042104310441045104610471048104910501051105210531054105510561057105810591060106110621063106410651066106710681069107010711072107310741075107610771078107910801081108210831084108510861087108810891090109110921093109410951096109710981099110011011102110311041105110611071108110911101111111211131114111511161117111811191120112111221123112411251126112711281129113011311132113311341135113611371138113911401141114211431144114511461147114811491150115111521153115411551156115711581159116011611162116311641165116611671168116911701171117211731174117511761177117811791180118111821183118411851186118711881189119011911192119311941195119611971198119912001201120212031204120512061207120812091210121112121213121412151216121712181219122012211222122312241225122612271228122912301231123212331234123512361237123812391240124112421243124412451246124712481249125012511252125312541255125612571258125912601261126212631264126512661267126812691270127112721273127412751276127712781279128012811282128312841285128612871288128912901291129212931294129512961297129812991300130113021303130413051306130713081309131013111312131313141315131613171318131913201321132213231324132513261327132813291330133113321333133413351336133713381339134013411342134313441345134613471348134913501351135213531354135513561357135813591360136113621363136413651366136713681369137013711372137313741375137613771378137913801381138213831384138513861387138813891390139113921393139413951396139713981399140014011402140314041405140614071408140914101411141214131414141514161417141814191420142114221423142414251426142714281429143014311432143314341435143614371438143914401441144214431444144514461447144814491450145114521453145414551456145714581459146014611462146314641465146614671468146914701471147214731474147514761477147814791480148114821483148414851486148714881489149014911492149314941495149614971498149915001501150215031504150515061507150815091510151115121513151415151516151715181519152015211522152315241525152615271528152915301531153215331534153515361537153815391540154115421543154415451546154715481549155015511552155315541555155615571558155915601561156215631564156515661567156815691570157115721573157415751576157715781579158015811582158315841585158615871588158915901591159215931594159515961597159815991600160116021603160416051606160716081609161016111612161316141615161616171618161916201621162216231624162516261627162816291630163116321633163416351636163716381639164016411642164316441645164616471648164916501651165216531654165516561657165816591660166116621663166416651666166716681669167016711672167316741675167616771678167916801681168216831684168516861687168816891690169116921693169416951696169716981699170017011702170317041705170617071708170917101711171217131714171517161717171817191720172117221723172417251726172717281729173017311732173317341735173617371738173917401741174217431744174517461747// 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/v7/internal/config" "github.com/router-for-me/CLIProxyAPI/v7/internal/misc" "github.com/router-for-me/CLIProxyAPI/v7/internal/runtime/executor/helps" "github.com/router-for-me/CLIProxyAPI/v7/internal/thinking" "github.com/router-for-me/CLIProxyAPI/v7/internal/util" cliproxyauth "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/auth" cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executor" "github.com/router-for-me/CLIProxyAPI/v7/sdk/proxyutil" sdktranslator "github.com/router-for-me/CLIProxyAPI/v7/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.NewExecutorUsageReporter(ctx, e, 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, body := translateCodexRequestPair(from, to, baseModel, originalPayload, 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.ApplyPayloadConfigWithRequest(e.cfg, baseModel, to.String(), from.String(), "", body, originalTranslated, requestedModel, requestPath, opts.Headers) 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) } body = sanitizeOpenAIResponsesReasoningEncryptedContent(ctx, "codex websockets executor", body)
httpURL := strings.TrimSuffix(baseURL, "/") + "/responses" wsURL, err := buildCodexResponsesWebsocketURL(httpURL) if err != nil { return resp, err }
body, wsHeaders := applyCodexPromptCacheHeaders(from, req, body) 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) applyCodexIdentityConfuseHeaders(wsHeaders, &identityState)
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(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)
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) reporter.StartResponseTTFT() 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(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, 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 } reporter.MarkFirstResponseByte() payload = applyCodexIdentityConfuseResponsePayload(payload, identityState) 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 clientPayload := applyCodexIdentityExposeResponsePayload(payload, identityState) out := sdktranslator.TranslateNonStream(ctx, to, from, req.Model, originalPayload, clientBody, clientPayload, ¶m) 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.NewExecutorUsageReporter(ctx, e, baseModel, auth) defer reporter.TrackFailure(ctx, &err)
from := opts.SourceFormat to := sdktranslator.FromString("codex") body := req.Payload userPayload := req.Payload if len(opts.OriginalRequest) > 0 { userPayload = opts.OriginalRequest }
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.ApplyPayloadConfigWithRequest(e.cfg, baseModel, to.String(), from.String(), "", body, body, requestedModel, requestPath, opts.Headers) body = normalizeCodexInstructions(body) if e.cfg == nil || e.cfg.DisableImageGeneration == config.DisableImageGenerationOff { body = ensureImageGenerationTool(body, baseModel, auth) } body = sanitizeOpenAIResponsesReasoningEncryptedContent(ctx, "codex websockets executor", body)
httpURL := strings.TrimSuffix(baseURL, "/") + "/responses" wsURL, err := buildCodexResponsesWebsocketURL(httpURL) if err != nil { return nil, err }
body, wsHeaders := applyCodexPromptCacheHeaders(from, req, body) clientBody := body var identityState codexIdentityConfuseState upstreamBody, identityState := applyCodexIdentityConfuseBody(e.cfg, auth, userPayload, body) reporter.SetTranslatedReasoningEffort(clientBody, to.String()) wsHeaders = applyCodexWebsocketHeaders(ctx, wsHeaders, auth, apiKey, e.cfg) 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() } }
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)
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) reporter.StartResponseTTFT()
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 } 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, connRetry, wsReqBodyRetry); errSendRetry != nil { helps.RecordAPIWebsocketError(ctx, e.cfg, "send_retry", errSendRetry) e.invalidateUpstreamConn(sess, connRetry, "send_error", errSendRetry) sess.clearActive(readCh) sess.reqMu.Unlock() return nil, errSendRetry } conn = connRetry wsReqBody = wsReqBodyRetry } else { logCodexWebsocketDisconnected(executionSessionID, authID, wsURL, "send_error", errSend) if errClose := conn.Close(); errClose != nil { log.Errorf("codex websockets executor: close websocket error: %v", errClose) } return nil, errSend } }
out := make(chan cliproxyexecutor.StreamChunk) go func() { terminateReason := "completed" var terminateErr error
defer close(out) defer func() { if sess != nil { sess.clearActive(readCh) sess.reqMu.Unlock() return } logCodexWebsocketDisconnected(executionSessionID, authID, wsURL, terminateReason, terminateErr) if errClose := conn.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 } }
var param any 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 } terminateReason = "read_error" terminateErr = errRead helps.RecordAPIWebsocketError(ctx, e.cfg, "read", errRead) reporter.PublishFailure(ctx, errRead) _ = send(cliproxyexecutor.StreamChunk{Err: errRead}) 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)
if wsErr, ok := parseCodexWebsocketError(payload); ok { terminateReason = "upstream_error" terminateErr = wsErr helps.RecordAPIWebsocketError(ctx, e.cfg, "upstream_error", wsErr) reporter.PublishFailure(ctx, wsErr) if sess != nil { e.invalidateUpstreamConn(sess, conn, "upstream_error", wsErr) } _ = send(cliproxyexecutor.StreamChunk{Err: wsErr}) return }
payload = normalizeCodexWebsocketCompletion(payload) eventType := gjson.GetBytes(payload, "type").String() if eventType == "response.completed" || eventType == "response.done" { if detail, ok := helps.ParseCodexUsage(payload); ok { reporter.Publish(ctx, detail) } }
clientPayload := applyCodexIdentityExposeResponsePayload(payload, identityState) line := encodeCodexWebsocketAsSSE(clientPayload) chunks := sdktranslator.TranslateStream(ctx, to, from, req.Model, clientBody, clientBody, line, ¶m) 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}
func (e *CodexWebsocketsExecutor) dialCodexWebsocket(ctx context.Context, auth *cliproxyauth.Auth, wsURL string, headers http.Header) (*websocket.Conn, *http.Response, error) { dialer := newProxyAwareWebsocketDialer(e.cfg, auth) dialer.HandshakeTimeout = codexResponsesWebsocketHandshakeTO dialer.EnableCompression = true if ctx == nil { ctx = context.Background() } conn, resp, err := dialer.DialContext(ctx, wsURL, headers) if conn != nil { // Avoid gorilla/websocket flate tail validation issues on some upstreams/Go versions. // Negotiating permessage-deflate is fine; we just don't compress outbound messages. conn.EnableWriteCompression(false) } return conn, resp, err}
func writeCodexWebsocketMessage(sess *codexWebsocketSession, conn *websocket.Conn, payload []byte) error { if sess != nil { return sess.writeMessage(conn, websocket.TextMessage, payload) } if conn == nil { return fmt.Errorf("codex websockets executor: websocket conn is nil") } return conn.WriteMessage(websocket.TextMessage, payload)}
func buildCodexWebsocketRequestBody(body []byte) []byte { if len(body) == 0 { return nil }
// Match codex-rs websocket v2 semantics: every request is `response.create`. // Incremental follow-up turns continue on the same websocket using // `previous_response_id` + incremental `input`, not `response.append`. wsReqBody, errSet := sjson.SetBytes(bytes.Clone(body), "type", "response.create") if errSet == nil && len(wsReqBody) > 0 { return wsReqBody } fallback := bytes.Clone(body) fallback, _ = sjson.SetBytes(fallback, "type", "response.create") return fallback}
func readCodexWebsocketMessage(ctx context.Context, sess *codexWebsocketSession, conn *websocket.Conn, readCh chan codexWebsocketRead) (int, []byte, error) { if sess == nil { if conn == nil { return 0, nil, fmt.Errorf("codex websockets executor: websocket conn is nil") } _ = conn.SetReadDeadline(time.Now().Add(codexResponsesWebsocketIdleTimeout)) msgType, payload, errRead := conn.ReadMessage() return msgType, payload, errRead } if conn == nil { return 0, nil, fmt.Errorf("codex websockets executor: websocket conn is nil") } if readCh == nil { return 0, nil, fmt.Errorf("codex websockets executor: session read channel is nil") } for { select { case <-ctx.Done(): return 0, nil, ctx.Err() case ev, ok := <-readCh: if !ok { return 0, nil, fmt.Errorf("codex websockets executor: session read channel closed") } if ev.conn != conn { continue } if ev.err != nil { return 0, nil, ev.err } return ev.msgType, ev.payload, nil } }}
func newProxyAwareWebsocketDialer(cfg *config.Config, auth *cliproxyauth.Auth) *websocket.Dialer { dialer := &websocket.Dialer{ Proxy: http.ProxyFromEnvironment, HandshakeTimeout: codexResponsesWebsocketHandshakeTO, EnableCompression: true, NetDialContext: (&net.Dialer{ Timeout: 30 * time.Second, KeepAlive: 30 * time.Second, }).DialContext, }
proxyURL := "" if auth != nil { proxyURL = strings.TrimSpace(auth.ProxyURL) } if proxyURL == "" && cfg != nil { proxyURL = strings.TrimSpace(cfg.ProxyURL) } if proxyURL == "" { return dialer }
setting, errParse := proxyutil.Parse(proxyURL) if errParse != nil { log.Errorf("codex websockets executor: %v", errParse) return dialer }
switch setting.Mode { case proxyutil.ModeDirect: dialer.Proxy = nil return dialer case proxyutil.ModeProxy: default: return dialer }
switch setting.URL.Scheme { case "socks5", "socks5h": var proxyAuth *proxy.Auth if setting.URL.User != nil { username := setting.URL.User.Username() password, _ := setting.URL.User.Password() proxyAuth = &proxy.Auth{User: username, Password: password} } socksDialer, errSOCKS5 := proxy.SOCKS5("tcp", setting.URL.Host, proxyAuth, proxy.Direct) if errSOCKS5 != nil { log.Errorf("codex websockets executor: create SOCKS5 dialer failed: %v", errSOCKS5) return dialer } dialer.Proxy = nil dialer.NetDialContext = func(_ context.Context, network, addr string) (net.Conn, error) { return socksDialer.Dial(network, addr) } case "http", "https": dialer.Proxy = http.ProxyURL(setting.URL) default: log.Errorf("codex websockets executor: unsupported proxy scheme: %s", setting.URL.Scheme) }
return dialer}
func buildCodexResponsesWebsocketURL(httpURL string) (string, error) { parsed, err := url.Parse(strings.TrimSpace(httpURL)) if err != nil { return "", err } switch strings.ToLower(parsed.Scheme) { case "http": parsed.Scheme = "ws" case "https": parsed.Scheme = "wss" default: return "", fmt.Errorf("codex websockets executor: unsupported responses websocket URL scheme %q", parsed.Scheme) } if strings.TrimSpace(parsed.Host) == "" { return "", fmt.Errorf("codex websockets executor: responses websocket URL host is empty") } return parsed.String(), nil}
func applyCodexPromptCacheHeaders(from sdktranslator.Format, req cliproxyexecutor.Request, rawJSON []byte) ([]byte, http.Header) { headers := http.Header{} if len(rawJSON) == 0 { return rawJSON, headers }
var cache helps.CodexCache if sourceFormatEqual(from, sdktranslator.FormatClaude) { if cached, ok := codexClaudeCodePromptCache(req); ok { cache = cached } } else if sourceFormatEqual(from, sdktranslator.FormatOpenAIResponse) { if promptCacheKey := gjson.GetBytes(req.Payload, "prompt_cache_key"); promptCacheKey.Exists() { cache.ID = promptCacheKey.String() } }
if cache.ID != "" { rawJSON, _ = sjson.SetBytes(rawJSON, "prompt_cache_key", cache.ID) setHeaderCasePreserved(headers, "session_id", cache.ID) headers.Set("Conversation_id", cache.ID) }
return rawJSON, headers}
func applyCodexWebsocketHeaders(ctx context.Context, headers http.Header, auth *cliproxyauth.Auth, token string, cfg *config.Config) http.Header { if headers == nil { headers = http.Header{} } if strings.TrimSpace(token) != "" { headers.Set("Authorization", "Bearer "+token) }
var ginHeaders http.Header if ginCtx, ok := ctx.Value("gin").(*gin.Context); ok && ginCtx != nil && ginCtx.Request != nil { ginHeaders = ginCtx.Request.Header.Clone() }
isAPIKey := codexAuthUsesAPIKey(auth) cfgUserAgent, cfgBetaFeatures := codexHeaderDefaults(cfg, auth) ensureHeaderWithPriority(headers, ginHeaders, "x-codex-beta-features", cfgBetaFeatures, "") misc.EnsureHeader(headers, ginHeaders, "x-codex-turn-state", "") misc.EnsureHeader(headers, ginHeaders, "x-codex-turn-metadata", "") misc.EnsureHeader(headers, ginHeaders, "x-client-request-id", "") misc.EnsureHeader(headers, ginHeaders, "x-responsesapi-include-timing-metrics", "") misc.EnsureHeader(headers, ginHeaders, "Version", "") if isAPIKey { ensureHeaderWithPriority(headers, ginHeaders, "User-Agent", "", "") } else { ensureHeaderWithConfigPrecedence(headers, ginHeaders, "User-Agent", cfgUserAgent, codexUserAgent) }
betaHeader := strings.TrimSpace(headers.Get("OpenAI-Beta")) if betaHeader == "" && ginHeaders != nil { betaHeader = strings.TrimSpace(ginHeaders.Get("OpenAI-Beta")) } if betaHeader == "" || !strings.Contains(betaHeader, "responses_websockets=") { betaHeader = codexResponsesWebsocketBetaHeaderValue } headers.Set("OpenAI-Beta", betaHeader) sessionFallback := "" if strings.Contains(headers.Get("User-Agent"), "Mac OS") { sessionFallback = uuid.NewString() } ensureCodexWebsocketSessionHeader(headers, ginHeaders, sessionFallback) if originator := strings.TrimSpace(ginHeaders.Get("Originator")); originator != "" { headers.Set("Originator", originator) } else if !isAPIKey { headers.Set("Originator", codexOriginator) } if !isAPIKey { if auth != nil && auth.Metadata != nil { if accountID, ok := auth.Metadata["account_id"].(string); ok { if trimmed := strings.TrimSpace(accountID); trimmed != "" { setHeaderCasePreserved(headers, "ChatGPT-Account-ID", trimmed) } } } }
var attrs map[string]string if auth != nil { attrs = auth.Attributes } util.ApplyCustomHeadersFromAttrs(&http.Request{Header: headers}, attrs)
return headers}
func ensureCodexWebsocketSessionHeader(target http.Header, source http.Header, fallbackValue string) { if target == nil { return } sessionID := codexSessionHeaderValue(target) if sessionID == "" { sessionID = codexSessionHeaderValue(source) } if sessionID == "" { sessionID = strings.TrimSpace(fallbackValue) } if sessionID != "" { setHeaderCasePreserved(target, "session_id", sessionID) } deleteHeaderCaseInsensitive(target, "Session-Id")}
func codexSessionHeaderValue(headers http.Header) string { for _, key := range []string{"Session-Id", "Session_id", "session_id"} { if value := strings.TrimSpace(headerValueCaseInsensitive(headers, key)); value != "" { return value } } return ""}
func codexAuthUsesAPIKey(auth *cliproxyauth.Auth) bool { if auth == nil || auth.Attributes == nil { return false } return strings.TrimSpace(auth.Attributes["api_key"]) != ""}
func ensureHeaderCasePreserved(target http.Header, source http.Header, key, configValue, fallbackValue string) { if target == nil { return } if strings.TrimSpace(headerValueCaseInsensitive(target, key)) != "" { return } if source != nil { if val := strings.TrimSpace(headerValueCaseInsensitive(source, key)); val != "" { setHeaderCasePreserved(target, key, val) return } } if val := strings.TrimSpace(configValue); val != "" { setHeaderCasePreserved(target, key, val) return } if val := strings.TrimSpace(fallbackValue); val != "" { setHeaderCasePreserved(target, key, val) }}
func setHeaderCasePreserved(headers http.Header, key string, value string) { if headers == nil { return } key = strings.TrimSpace(key) value = strings.TrimSpace(value) if key == "" || value == "" { return } deleteHeaderCaseInsensitive(headers, key) headers[key] = []string{value}}
func setCodexSessionHeaderCasePreserved(headers http.Header, fallbackKey string, value string) { if headers == nil { return } fallbackKey = strings.TrimSpace(fallbackKey) value = strings.TrimSpace(value) if fallbackKey == "" || value == "" { return }
selectedKey := "" if _, ok := headers[fallbackKey]; ok && codexSessionHeaderKeyUsesUnderscore(fallbackKey) { selectedKey = fallbackKey } else { for existingKey := range headers { if codexSessionHeaderKeyUsesUnderscore(existingKey) { selectedKey = existingKey break } } } if selectedKey == "" { selectedKey = fallbackKey } for existingKey := range headers { if codexSessionHeaderKey(existingKey) && existingKey != selectedKey { delete(headers, existingKey) } } headers[selectedKey] = []string{value}}
func codexSessionHeaderKey(key string) bool { normalized := strings.ToLower(strings.TrimSpace(key)) return normalized == "session_id" || normalized == "session-id"}
func codexSessionHeaderKeyUsesUnderscore(key string) bool { return strings.ToLower(strings.TrimSpace(key)) == "session_id"}
func headerValueCaseInsensitive(headers http.Header, key string) string { key = strings.TrimSpace(key) if headers == nil || key == "" { return "" } if val := strings.TrimSpace(headers.Get(key)); val != "" { return val } for existingKey, values := range headers { if !strings.EqualFold(existingKey, key) { continue } for _, value := range values { if trimmed := strings.TrimSpace(value); trimmed != "" { return trimmed } } } return ""}
func deleteHeaderCaseInsensitive(headers http.Header, key string) { for existingKey := range headers { if strings.EqualFold(existingKey, key) { delete(headers, existingKey) } }}
func codexHeaderDefaults(cfg *config.Config, auth *cliproxyauth.Auth) (string, string) { if cfg == nil || auth == nil { return "", "" } if auth.Attributes != nil { if v := strings.TrimSpace(auth.Attributes["api_key"]); v != "" { return "", "" } } return strings.TrimSpace(cfg.CodexHeaderDefaults.UserAgent), strings.TrimSpace(cfg.CodexHeaderDefaults.BetaFeatures)}
func ensureHeaderWithPriority(target http.Header, source http.Header, key, configValue, fallbackValue string) { if target == nil { return } if strings.TrimSpace(target.Get(key)) != "" { return } if source != nil { if val := strings.TrimSpace(source.Get(key)); val != "" { target.Set(key, val) return } } if val := strings.TrimSpace(configValue); val != "" { target.Set(key, val) return } if val := strings.TrimSpace(fallbackValue); val != "" { target.Set(key, val) }}
func ensureHeaderWithConfigPrecedence(target http.Header, source http.Header, key, configValue, fallbackValue string) { if target == nil { return } if strings.TrimSpace(target.Get(key)) != "" { return } if val := strings.TrimSpace(configValue); val != "" { target.Set(key, val) return } if source != nil { if val := strings.TrimSpace(source.Get(key)); val != "" { target.Set(key, val) return } } if val := strings.TrimSpace(fallbackValue); val != "" { target.Set(key, val) }}
type statusErrWithHeaders struct { statusErr headers http.Header}
func (e statusErrWithHeaders) Headers() http.Header { if e.headers == nil { return nil } return e.headers.Clone()}
func parseCodexWebsocketError(payload []byte) (error, bool) { if len(payload) == 0 { return nil, false } if strings.TrimSpace(gjson.GetBytes(payload, "type").String()) != "error" { return nil, false } status := int(gjson.GetBytes(payload, "status").Int()) if status == 0 { status = int(gjson.GetBytes(payload, "status_code").Int()) } if status <= 0 { return nil, false }
out := buildCodexWebsocketErrorPayload(payload, status) headers := parseCodexWebsocketErrorHeaders(payload) statusError := statusErr{code: status, msg: string(out)} if retryAfter := parseCodexRetryAfter(status, out, time.Now()); retryAfter != nil { statusError.retryAfter = retryAfter } else if isCodexWebsocketConnectionLimitError(payload) { retryAfter := time.Duration(0) statusError.retryAfter = &retryAfter } return statusErrWithHeaders{ statusErr: statusError, headers: headers, }, true}
func buildCodexWebsocketErrorPayload(payload []byte, status int) []byte { out := []byte(`{}`) out, _ = sjson.SetBytes(out, "status", status)
if bodyNode := gjson.GetBytes(payload, "body"); bodyNode.Exists() { out, _ = sjson.SetRawBytes(out, "body", []byte(bodyNode.Raw)) if bodyErrorNode := bodyNode.Get("error"); bodyErrorNode.Exists() { out, _ = sjson.SetRawBytes(out, "error", []byte(bodyErrorNode.Raw)) return out } }
if errNode := gjson.GetBytes(payload, "error"); errNode.Exists() { out, _ = sjson.SetRawBytes(out, "error", []byte(errNode.Raw)) return out }
out, _ = sjson.SetBytes(out, "error.type", "server_error") out, _ = sjson.SetBytes(out, "error.message", http.StatusText(status)) return out}
func isCodexWebsocketConnectionLimitError(payload []byte) bool { if len(payload) == 0 { return false } for _, path := range []string{"error.code", "error.type", "body.error.code", "body.error.type", "code", "error"} { if strings.TrimSpace(gjson.GetBytes(payload, path).String()) == "websocket_connection_limit_reached" { return true } } return false}
func parseCodexWebsocketErrorHeaders(payload []byte) http.Header { headersNode := gjson.GetBytes(payload, "headers") if !headersNode.Exists() || !headersNode.IsObject() { return nil } mapped := make(http.Header) headersNode.ForEach(func(key, value gjson.Result) bool { name := strings.TrimSpace(key.String()) if name == "" { return true } switch value.Type { case gjson.String: if v := strings.TrimSpace(value.String()); v != "" { mapped.Set(name, v) } case gjson.Number, gjson.True, gjson.False: if v := strings.TrimSpace(value.Raw); v != "" { mapped.Set(name, v) } default: } return true }) if len(mapped) == 0 { return nil } return mapped}
func normalizeCodexWebsocketCompletion(payload []byte) []byte { if strings.TrimSpace(gjson.GetBytes(payload, "type").String()) == "response.done" { updated, err := sjson.SetBytes(payload, "type", "response.completed") if err == nil && len(updated) > 0 { return updated } } return payload}
func encodeCodexWebsocketAsSSE(payload []byte) []byte { if len(payload) == 0 { return nil } line := make([]byte, 0, len("data: ")+len(payload)) line = append(line, []byte("data: ")...) line = append(line, payload...) return line}
func websocketUpgradeRequestLog(info helps.UpstreamRequestLog) helps.UpstreamRequestLog { upgradeInfo := info upgradeInfo.URL = helps.WebsocketUpgradeRequestURL(info.URL) upgradeInfo.Method = http.MethodGet upgradeInfo.Body = nil upgradeInfo.Headers = info.Headers.Clone() if upgradeInfo.Headers == nil { upgradeInfo.Headers = make(http.Header) } if strings.TrimSpace(upgradeInfo.Headers.Get("Connection")) == "" { upgradeInfo.Headers.Set("Connection", "Upgrade") } if strings.TrimSpace(upgradeInfo.Headers.Get("Upgrade")) == "" { upgradeInfo.Headers.Set("Upgrade", "websocket") } return upgradeInfo}
func recordAPIWebsocketHandshake(ctx context.Context, cfg *config.Config, resp *http.Response) { if resp == nil { return } helps.RecordAPIWebsocketHandshake(ctx, cfg, resp.StatusCode, resp.Header.Clone()) closeHTTPResponseBody(resp, "codex websockets executor: close handshake response body error")}
func websocketHandshakeBody(resp *http.Response) []byte { if resp == nil || resp.Body == nil { return nil } body, _ := io.ReadAll(resp.Body) closeHTTPResponseBody(resp, "codex websockets executor: close handshake response body error") if len(body) == 0 { return nil } return body}
func closeHTTPResponseBody(resp *http.Response, logPrefix string) { if resp == nil || resp.Body == nil { return } if errClose := resp.Body.Close(); errClose != nil { log.Errorf("%s: %v", logPrefix, errClose) }}
func executionSessionIDFromOptions(opts cliproxyexecutor.Options) string { if len(opts.Metadata) == 0 { return "" } raw, ok := opts.Metadata[cliproxyexecutor.ExecutionSessionMetadataKey] if !ok || raw == nil { return "" } switch v := raw.(type) { case string: return strings.TrimSpace(v) case []byte: return strings.TrimSpace(string(v)) default: return "" }}
func (e *CodexWebsocketsExecutor) getOrCreateSession(sessionID string) *codexWebsocketSession { sessionID = strings.TrimSpace(sessionID) if sessionID == "" { return nil } if e == nil { return nil } store := e.store if store == nil { store = globalCodexWebsocketSessionStore } store.mu.Lock() defer store.mu.Unlock() if store.sessions == nil { store.sessions = make(map[string]*codexWebsocketSession) } if sess, ok := store.sessions[sessionID]; ok && sess != nil { return sess } sess := &codexWebsocketSession{ sessionID: sessionID, upstreamDisconnectCh: make(chan error, 1), } store.sessions[sessionID] = sess return sess}
func (e *CodexWebsocketsExecutor) UpstreamDisconnectChan(sessionID string) <-chan error { sess := e.getOrCreateSession(sessionID) if sess == nil { return nil } return sess.upstreamDisconnectCh}
func (e *CodexWebsocketsExecutor) ensureUpstreamConn(ctx context.Context, auth *cliproxyauth.Auth, sess *codexWebsocketSession, authID string, wsURL string, headers http.Header) (*websocket.Conn, *http.Response, error) { if sess == nil { return e.dialCodexWebsocket(ctx, auth, wsURL, headers) }
sess.connMu.Lock() conn := sess.conn readerConn := sess.readerConn sess.connMu.Unlock() if conn != nil { if readerConn != conn { sess.connMu.Lock() sess.readerConn = conn sess.connMu.Unlock() sess.configureConn(conn) go e.readUpstreamLoop(sess, conn) } return conn, nil, nil }
conn, resp, errDial := e.dialCodexWebsocket(ctx, auth, wsURL, headers) if errDial != nil { return nil, resp, errDial }
sess.connMu.Lock() if sess.conn != nil { previous := sess.conn sess.connMu.Unlock() if errClose := conn.Close(); errClose != nil { log.Errorf("codex websockets executor: close websocket error: %v", errClose) } return previous, nil, nil } sess.conn = conn sess.wsURL = wsURL sess.authID = authID sess.readerConn = conn sess.connMu.Unlock()
sess.configureConn(conn) go e.readUpstreamLoop(sess, conn) logCodexWebsocketConnected(sess.sessionID, authID, wsURL) return conn, resp, nil}
func (e *CodexWebsocketsExecutor) readUpstreamLoop(sess *codexWebsocketSession, conn *websocket.Conn) { if e == nil || sess == nil || conn == nil { return } for { _ = conn.SetReadDeadline(time.Now().Add(codexResponsesWebsocketIdleTimeout)) msgType, payload, errRead := conn.ReadMessage() if errRead != nil { sess.activeMu.Lock() ch := sess.activeCh done := sess.activeDone sess.activeMu.Unlock() if ch != nil { select { case ch <- codexWebsocketRead{conn: conn, err: errRead}: case <-done: default: } sess.clearActive(ch) close(ch) } e.invalidateUpstreamConn(sess, conn, "upstream_disconnected", errRead) return }
if msgType != websocket.TextMessage { if msgType == websocket.BinaryMessage { errBinary := fmt.Errorf("codex websockets executor: unexpected binary message") sess.activeMu.Lock() ch := sess.activeCh done := sess.activeDone sess.activeMu.Unlock() if ch != nil { select { case ch <- codexWebsocketRead{conn: conn, err: errBinary}: case <-done: default: } sess.clearActive(ch) close(ch) } e.invalidateUpstreamConn(sess, conn, "unexpected_binary", errBinary) return } continue }
sess.activeMu.Lock() ch := sess.activeCh done := sess.activeDone sess.activeMu.Unlock() if ch == nil { continue } select { case ch <- codexWebsocketRead{conn: conn, msgType: msgType, payload: payload}: case <-done: } }}
func (e *CodexWebsocketsExecutor) invalidateUpstreamConn(sess *codexWebsocketSession, conn *websocket.Conn, reason string, err error) { if sess == nil || conn == nil { return }
sess.connMu.Lock() current := sess.conn authID := sess.authID wsURL := sess.wsURL sessionID := sess.sessionID if current == nil || current != conn { sess.connMu.Unlock() return } sess.conn = nil if sess.readerConn == conn { sess.readerConn = nil } sess.connMu.Unlock()
logCodexWebsocketDisconnected(sessionID, authID, wsURL, reason, err) sess.notifyUpstreamDisconnect(err) if errClose := conn.Close(); errClose != nil { log.Errorf("codex websockets executor: close websocket error: %v", errClose) }}
func (e *CodexWebsocketsExecutor) CloseExecutionSession(sessionID string) { sessionID = strings.TrimSpace(sessionID) if e == nil { return } if sessionID == "" { return } if sessionID == cliproxyauth.CloseAllExecutionSessionsID { // Executor replacement can happen during hot reload (config/credential changes). // Do not force-close upstream websocket sessions here, otherwise in-flight // downstream websocket requests get interrupted. return }
store := e.store if store == nil { store = globalCodexWebsocketSessionStore } store.mu.Lock() sess := store.sessions[sessionID] delete(store.sessions, sessionID) store.mu.Unlock()
e.closeExecutionSession(sess, "session_closed")}
func (e *CodexWebsocketsExecutor) closeAllExecutionSessions(reason string) { if e == nil { return }
store := e.store if store == nil { store = globalCodexWebsocketSessionStore } store.mu.Lock() sessions := make([]*codexWebsocketSession, 0, len(store.sessions)) for sessionID, sess := range store.sessions { delete(store.sessions, sessionID) if sess != nil { sessions = append(sessions, sess) } } store.mu.Unlock()
for i := range sessions { e.closeExecutionSession(sessions[i], reason) }}
func (e *CodexWebsocketsExecutor) closeExecutionSession(sess *codexWebsocketSession, reason string) { closeCodexWebsocketSession(sess, reason)}
func closeCodexWebsocketSession(sess *codexWebsocketSession, reason string) { if sess == nil { return } reason = strings.TrimSpace(reason) if reason == "" { reason = "session_closed" }
sess.connMu.Lock() conn := sess.conn authID := sess.authID wsURL := sess.wsURL sess.conn = nil if sess.readerConn == conn { sess.readerConn = nil } sessionID := sess.sessionID sess.connMu.Unlock()
if conn == nil { return } logCodexWebsocketDisconnected(sessionID, authID, wsURL, reason, nil) if errClose := conn.Close(); errClose != nil { log.Errorf("codex websockets executor: close websocket error: %v", errClose) }}
func logCodexWebsocketConnected(sessionID string, authID string, wsURL string) { log.Infof("codex websockets: upstream connected session=%s auth=%s url=%s", strings.TrimSpace(sessionID), strings.TrimSpace(authID), strings.TrimSpace(wsURL))}
func logCodexWebsocketDisconnected(sessionID string, authID string, wsURL string, reason string, err error) { if err != nil { log.Infof("codex websockets: upstream disconnected session=%s auth=%s url=%s reason=%s err=%v", strings.TrimSpace(sessionID), strings.TrimSpace(authID), strings.TrimSpace(wsURL), strings.TrimSpace(reason), err) return } log.Infof("codex websockets: upstream disconnected session=%s auth=%s url=%s reason=%s", strings.TrimSpace(sessionID), strings.TrimSpace(authID), strings.TrimSpace(wsURL), strings.TrimSpace(reason))}
// CloseCodexWebsocketSessionsForAuthID closes all active Codex upstream websocket sessions// associated with the supplied auth ID.func CloseCodexWebsocketSessionsForAuthID(authID string, reason string) { authID = strings.TrimSpace(authID) if authID == "" { return } reason = strings.TrimSpace(reason) if reason == "" { reason = "auth_removed" }
store := globalCodexWebsocketSessionStore if store == nil { return }
type sessionItem struct { sessionID string sess *codexWebsocketSession }
store.mu.Lock() items := make([]sessionItem, 0, len(store.sessions)) for sessionID, sess := range store.sessions { items = append(items, sessionItem{sessionID: sessionID, sess: sess}) } store.mu.Unlock()
matches := make([]sessionItem, 0) for i := range items { sess := items[i].sess if sess == nil { continue } sess.connMu.Lock() sessAuthID := strings.TrimSpace(sess.authID) sess.connMu.Unlock() if sessAuthID == authID { matches = append(matches, items[i]) } } if len(matches) == 0 { return }
toClose := make([]*codexWebsocketSession, 0, len(matches)) store.mu.Lock() for i := range matches { current, ok := store.sessions[matches[i].sessionID] if !ok || current == nil || current != matches[i].sess { continue } delete(store.sessions, matches[i].sessionID) toClose = append(toClose, current) } store.mu.Unlock()
for i := range toClose { closeCodexWebsocketSession(toClose[i], reason) }}
// CodexAutoExecutor routes Codex requests to the websocket transport only when:// 1. The downstream transport is websocket, and// 2. The selected auth enables websockets.//// For non-websocket downstream requests, it always uses the legacy HTTP implementation.type CodexAutoExecutor struct { httpExec *CodexExecutor wsExec *CodexWebsocketsExecutor}
func NewCodexAutoExecutor(cfg *config.Config) *CodexAutoExecutor { return &CodexAutoExecutor{ httpExec: NewCodexExecutor(cfg), wsExec: NewCodexWebsocketsExecutor(cfg), }}
func (e *CodexAutoExecutor) Identifier() string { return "codex" }
func (e *CodexAutoExecutor) PrepareRequest(req *http.Request, auth *cliproxyauth.Auth) error { if e == nil || e.httpExec == nil { return nil } return e.httpExec.PrepareRequest(req, auth)}
func (e *CodexAutoExecutor) HttpRequest(ctx context.Context, auth *cliproxyauth.Auth, req *http.Request) (*http.Response, error) { if e == nil || e.httpExec == nil { return nil, fmt.Errorf("codex auto executor: http executor is nil") } return e.httpExec.HttpRequest(ctx, auth, req)}
func (e *CodexAutoExecutor) Execute(ctx context.Context, auth *cliproxyauth.Auth, req cliproxyexecutor.Request, opts cliproxyexecutor.Options) (cliproxyexecutor.Response, error) { if e == nil || e.httpExec == nil || e.wsExec == nil { return cliproxyexecutor.Response{}, fmt.Errorf("codex auto executor: executor is nil") } if cliproxyexecutor.DownstreamWebsocket(ctx) && codexWebsocketsEnabled(auth) { return e.wsExec.Execute(ctx, auth, req, opts) } return e.httpExec.Execute(ctx, auth, req, opts)}
func (e *CodexAutoExecutor) ExecuteStream(ctx context.Context, auth *cliproxyauth.Auth, req cliproxyexecutor.Request, opts cliproxyexecutor.Options) (*cliproxyexecutor.StreamResult, error) { if e == nil || e.httpExec == nil || e.wsExec == nil { return nil, fmt.Errorf("codex auto executor: executor is nil") } if cliproxyexecutor.DownstreamWebsocket(ctx) && codexWebsocketsEnabled(auth) { return e.wsExec.ExecuteStream(ctx, auth, req, opts) } return e.httpExec.ExecuteStream(ctx, auth, req, opts)}
func (e *CodexAutoExecutor) Refresh(ctx context.Context, auth *cliproxyauth.Auth) (*cliproxyauth.Auth, error) { if e == nil || e.httpExec == nil { return nil, fmt.Errorf("codex auto executor: http executor is nil") } return e.httpExec.Refresh(ctx, auth)}
func (e *CodexAutoExecutor) CountTokens(ctx context.Context, auth *cliproxyauth.Auth, req cliproxyexecutor.Request, opts cliproxyexecutor.Options) (cliproxyexecutor.Response, error) { if e == nil || e.httpExec == nil { return cliproxyexecutor.Response{}, fmt.Errorf("codex auto executor: http executor is nil") } return e.httpExec.CountTokens(ctx, auth, req, opts)}
func (e *CodexAutoExecutor) CloseExecutionSession(sessionID string) { if e == nil || e.wsExec == nil { return } e.wsExec.CloseExecutionSession(sessionID)}
func (e *CodexAutoExecutor) UpstreamDisconnectChan(sessionID string) <-chan error { if e == nil || e.wsExec == nil { return nil } return e.wsExec.UpstreamDisconnectChan(sessionID)}
func codexWebsocketsEnabled(auth *cliproxyauth.Auth) bool { if auth == nil { return false } if len(auth.Attributes) > 0 { if raw := strings.TrimSpace(auth.Attributes["websockets"]); raw != "" { parsed, errParse := strconv.ParseBool(raw) if errParse == nil { return parsed } } } if len(auth.Metadata) == 0 { return false } raw, ok := auth.Metadata["websockets"] if !ok || raw == nil { return false } switch v := raw.(type) { case bool: return v case string: parsed, errParse := strconv.ParseBool(strings.TrimSpace(v)) if errParse == nil { return parsed } default: } return false}