Something went wrong. Try again.
sloproxy
Something went wrong. Try again.
57 kB · 1706 lines
Go
at main
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818819820821822823824825826827828829830831832833834835836837838839840841842843844845846847848849850851852853854855856857858859860861862863864865866867868869870871872873874875876877878879880881882883884885886887888889890891892893894895896897898899900901902903904905906907908909910911912913914915916917918919920921922923924925926927928929930931932933934935936937938939940941942943944945946947948949950951952953954955956957958959960961962963964965966967968969970971972973974975976977978979980981982983984985986987988989990991992993994995996997998999100010011002100310041005100610071008100910101011101210131014101510161017101810191020102110221023102410251026102710281029103010311032103310341035103610371038103910401041104210431044104510461047104810491050105110521053105410551056105710581059106010611062106310641065106610671068106910701071107210731074107510761077107810791080108110821083108410851086108710881089109010911092109310941095109610971098109911001101110211031104110511061107110811091110111111121113111411151116111711181119112011211122112311241125112611271128112911301131113211331134113511361137113811391140114111421143114411451146114711481149115011511152115311541155115611571158115911601161116211631164116511661167116811691170117111721173117411751176117711781179118011811182118311841185118611871188118911901191119211931194119511961197119811991200120112021203120412051206120712081209121012111212121312141215121612171218121912201221122212231224122512261227122812291230123112321233123412351236123712381239124012411242124312441245124612471248124912501251125212531254125512561257125812591260126112621263126412651266126712681269127012711272127312741275127612771278127912801281128212831284128512861287128812891290129112921293129412951296129712981299130013011302130313041305130613071308130913101311131213131314131513161317131813191320132113221323132413251326132713281329133013311332133313341335133613371338133913401341134213431344134513461347134813491350135113521353135413551356135713581359136013611362136313641365136613671368136913701371137213731374137513761377137813791380138113821383138413851386138713881389139013911392139313941395139613971398139914001401140214031404140514061407140814091410141114121413141414151416141714181419142014211422142314241425142614271428142914301431143214331434143514361437143814391440144114421443144414451446144714481449145014511452145314541455145614571458145914601461146214631464146514661467146814691470147114721473147414751476147714781479148014811482148314841485148614871488148914901491149214931494149514961497149814991500150115021503150415051506150715081509151015111512151315141515151615171518151915201521152215231524152515261527152815291530153115321533153415351536153715381539154015411542154315441545154615471548154915501551155215531554155515561557155815591560156115621563156415651566156715681569157015711572157315741575157615771578157915801581158215831584158515861587158815891590159115921593159415951596159715981599160016011602160316041605160616071608160916101611161216131614161516161617161816191620162116221623162416251626162716281629163016311632163316341635163616371638163916401641164216431644164516461647164816491650165116521653165416551656165716581659166016611662166316641665166616671668166916701671167216731674167516761677167816791680168116821683168416851686168716881689169016911692169316941695169616971698169917001701170217031704170517061707// Package executor provides runtime execution capabilities for various AI service providers.// This file implements an xAI executor that uses the Responses API WebSocket transport.package executor
import ( "bytes" "context" "encoding/json" "fmt" "net/http" "net/url" "strconv" "strings" "sync" "time"
"github.com/gorilla/websocket" xaiauth "github.com/router-for-me/CLIProxyAPI/v7/internal/auth/xai" "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/util" cliproxyauth "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/auth" cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executor" log "github.com/sirupsen/logrus" "github.com/tidwall/gjson" "github.com/tidwall/sjson")
// XAIWebsocketsExecutor executes xAI Responses requests using a WebSocket transport.type XAIWebsocketsExecutor struct { *XAIExecutor
store *codexWebsocketSessionStore idStore *xaiWebsocketIDStateStore}
var globalXAIWebsocketSessionStore = &codexWebsocketSessionStore{ sessions: make(map[string]*codexWebsocketSession),}
var globalXAIWebsocketIDStates = &xaiWebsocketIDStateStore{ sessions: make(map[string]*xaiWebsocketIDState),}
type xaiWebsocketIDStateStore struct { mu sync.Mutex sessions map[string]*xaiWebsocketIDState}
type xaiWebsocketIDState struct { requestMu sync.Mutex mu sync.Mutex downstreamToUpstream map[string]string sequence int transcriptInput []json.RawMessage replayCompactedTranscriptOnReset bool}
type xaiWebsocketRequestIDMapper struct { state *xaiWebsocketIDState downstreamPreviousID string upstreamPreviousID string upstreamResponseID string downstreamResponseID string replayedCompactedTranscript bool}
func NewXAIWebsocketsExecutor(cfg *config.Config) *XAIWebsocketsExecutor { return &XAIWebsocketsExecutor{ XAIExecutor: NewXAIExecutor(cfg), store: globalXAIWebsocketSessionStore, idStore: globalXAIWebsocketIDStates, }}
func getXAIWebsocketIDState(store *xaiWebsocketIDStateStore, sessionID string) *xaiWebsocketIDState { sessionID = strings.TrimSpace(sessionID) if sessionID == "" || store == nil { return nil } store.mu.Lock() defer store.mu.Unlock() if store.sessions == nil { store.sessions = make(map[string]*xaiWebsocketIDState) } if state := store.sessions[sessionID]; state != nil { return state } state := &xaiWebsocketIDState{ downstreamToUpstream: make(map[string]string), } store.sessions[sessionID] = state return state}
func deleteXAIWebsocketIDState(store *xaiWebsocketIDStateStore, sessionID string) { sessionID = strings.TrimSpace(sessionID) if sessionID == "" || store == nil { return } store.mu.Lock() delete(store.sessions, sessionID) store.mu.Unlock()}
func newXAIWebsocketRequestIDMapper(store *xaiWebsocketIDStateStore, sessionID string, downstreamRequest []byte) *xaiWebsocketRequestIDMapper { state := getXAIWebsocketIDState(store, sessionID) if state == nil { return nil } downstreamPreviousID := strings.TrimSpace(gjson.GetBytes(downstreamRequest, "previous_response_id").String()) upstreamPreviousID := downstreamPreviousID if downstreamPreviousID != "" { upstreamPreviousID = state.upstreamIDForDownstream(downstreamPreviousID) } return &xaiWebsocketRequestIDMapper{ state: state, downstreamPreviousID: downstreamPreviousID, upstreamPreviousID: upstreamPreviousID, }}
func (s *xaiWebsocketIDState) upstreamIDForDownstream(downstreamID string) string { downstreamID = strings.TrimSpace(downstreamID) if s == nil || downstreamID == "" { return downstreamID } s.mu.Lock() defer s.mu.Unlock() if upstreamID, ok := s.downstreamToUpstream[downstreamID]; ok { return strings.TrimSpace(upstreamID) } return downstreamID}
func (s *xaiWebsocketIDState) mapDownstreamToUpstream(downstreamID string, upstreamID string) { downstreamID = strings.TrimSpace(downstreamID) if s == nil || downstreamID == "" { return } s.mu.Lock() if s.downstreamToUpstream == nil { s.downstreamToUpstream = make(map[string]string) } s.downstreamToUpstream[downstreamID] = strings.TrimSpace(upstreamID) s.mu.Unlock()}
func (s *xaiWebsocketIDState) snapshotTranscriptInput() []byte { if s == nil { return nil } s.mu.Lock() defer s.mu.Unlock() if len(s.transcriptInput) == 0 { return nil } return xaiMarshalRawMessages(s.transcriptInput)}
func (s *xaiWebsocketIDState) prependTranscriptInput(payload []byte) []byte { if s == nil || len(payload) == 0 { return payload } s.mu.Lock() prefix := make([]json.RawMessage, 0, len(s.transcriptInput)) for _, item := range s.transcriptInput { prefix = append(prefix, bytes.Clone(item)) } s.mu.Unlock() if len(prefix) == 0 { return payload } current := xaiJSONRawMessages(gjson.GetBytes(payload, "input")) merged := append(prefix, current...) out, errSet := sjson.SetRawBytes(payload, "input", xaiMarshalRawMessages(merged)) if errSet != nil { return payload } return out}
func (s *xaiWebsocketIDState) recordTranscriptTurn(requestPayload []byte, completedPayload []byte, reset bool) { if s == nil || len(requestPayload) == 0 || len(completedPayload) == 0 { return } inputItems := xaiJSONRawMessages(gjson.GetBytes(requestPayload, "input")) outputItems := xaiJSONRawMessages(gjson.GetBytes(completedPayload, "response.output"))
s.mu.Lock() defer s.mu.Unlock() if reset { s.transcriptInput = nil s.replayCompactedTranscriptOnReset = false } if len(inputItems) == 0 && len(outputItems) == 0 { return } s.transcriptInput = append(s.transcriptInput, inputItems...) s.transcriptInput = append(s.transcriptInput, outputItems...)}
func (s *xaiWebsocketIDState) replaceTranscriptWithItems(items ...[]byte) { if s == nil { return } next := make([]json.RawMessage, 0, len(items)) for _, item := range items { item = bytes.TrimSpace(item) if len(item) == 0 || !json.Valid(item) { continue } next = append(next, bytes.Clone(item)) } s.mu.Lock() s.transcriptInput = next s.replayCompactedTranscriptOnReset = len(next) > 0 s.mu.Unlock()}
func (s *xaiWebsocketIDState) prependCompactedTranscriptOnReset(payload []byte) ([]byte, bool) { if s == nil || len(payload) == 0 { return payload, false } s.mu.Lock() if !s.replayCompactedTranscriptOnReset || len(s.transcriptInput) == 0 { s.mu.Unlock() return payload, false } prefix := make([]json.RawMessage, 0, len(s.transcriptInput)) for _, item := range s.transcriptInput { prefix = append(prefix, bytes.Clone(item)) } s.mu.Unlock()
current := xaiJSONRawMessages(gjson.GetBytes(payload, "input")) merged := append(prefix, current...) out, errSet := sjson.SetRawBytes(payload, "input", xaiMarshalRawMessages(merged)) if errSet != nil { return payload, false } return out, true}
func xaiJSONRawMessages(result gjson.Result) []json.RawMessage { if !result.Exists() || !result.IsArray() { return nil } items := result.Array() out := make([]json.RawMessage, 0, len(items)) for _, item := range items { raw := bytes.TrimSpace([]byte(item.Raw)) if len(raw) == 0 || !json.Valid(raw) { continue } out = append(out, bytes.Clone(raw)) } return out}
func xaiMarshalRawMessages(items []json.RawMessage) []byte { var buf bytes.Buffer buf.WriteByte('[') for i, item := range items { if i > 0 { buf.WriteByte(',') } buf.Write(bytes.TrimSpace(item)) } buf.WriteByte(']') return buf.Bytes()}
func (m *xaiWebsocketRequestIDMapper) upstreamRequestPayload(payload []byte) []byte { if m == nil || len(payload) == 0 { return payload } if m.downstreamPreviousID == m.upstreamPreviousID { requestType := strings.TrimSpace(gjson.GetBytes(payload, "type").String()) if m.downstreamPreviousID == "" && requestType == "response.append" && m.state != nil { out, replayed := m.state.prependCompactedTranscriptOnReset(payload) m.replayedCompactedTranscript = replayed return out } return payload } if m.upstreamPreviousID == "" { out, errDelete := sjson.DeleteBytes(payload, "previous_response_id") if errDelete == nil { if m.downstreamPreviousID != "" && m.state != nil { out = m.state.prependTranscriptInput(out) m.replayedCompactedTranscript = true } return out } return payload } out, errSet := sjson.SetBytes(payload, "previous_response_id", m.upstreamPreviousID) if errSet != nil { return payload } return out}
func (m *xaiWebsocketRequestIDMapper) downstreamResponsePayload(payload []byte) []byte { if m == nil || len(payload) == 0 { return payload } upstreamResponseID := strings.TrimSpace(gjson.GetBytes(payload, "response.id").String()) downstreamResponseID := m.downstreamIDForUpstreamResponse(upstreamResponseID) if downstreamResponseID == "" { return payload } return rewriteXAIWebsocketDownstreamIDs(payload, m.upstreamResponseID, downstreamResponseID, m.upstreamPreviousID, m.downstreamPreviousID)}
func (m *xaiWebsocketRequestIDMapper) downstreamIDForUpstreamResponse(upstreamResponseID string) string { upstreamResponseID = strings.TrimSpace(upstreamResponseID) if m == nil || m.state == nil { return upstreamResponseID } if m.upstreamResponseID != "" { return m.downstreamResponseID } if upstreamResponseID == "" { return "" }
m.state.mu.Lock() defer m.state.mu.Unlock() m.upstreamResponseID = upstreamResponseID m.downstreamResponseID = upstreamResponseID if m.state.downstreamToUpstream == nil { m.state.downstreamToUpstream = make(map[string]string) } _, upstreamResponseIDSeen := m.state.downstreamToUpstream[upstreamResponseID] if (m.downstreamPreviousID != "" && m.upstreamPreviousID != "" && upstreamResponseID == m.upstreamPreviousID) || upstreamResponseIDSeen { m.state.sequence++ m.downstreamResponseID = fmt.Sprintf("%s-xai-%d", upstreamResponseID, m.state.sequence) } m.state.downstreamToUpstream[upstreamResponseID] = upstreamResponseID m.state.downstreamToUpstream[m.downstreamResponseID] = upstreamResponseID return m.downstreamResponseID}
func rewriteXAIWebsocketDownstreamIDs(payload []byte, upstreamResponseID string, downstreamResponseID string, upstreamPreviousID string, downstreamPreviousID string) []byte { upstreamResponseID = strings.TrimSpace(upstreamResponseID) downstreamResponseID = strings.TrimSpace(downstreamResponseID) upstreamPreviousID = strings.TrimSpace(upstreamPreviousID) downstreamPreviousID = strings.TrimSpace(downstreamPreviousID) if len(payload) == 0 || (upstreamResponseID == downstreamResponseID && upstreamPreviousID == downstreamPreviousID) { return payload }
var value any decoder := json.NewDecoder(bytes.NewReader(payload)) decoder.UseNumber() if errDecode := decoder.Decode(&value); errDecode != nil { return payload } if !rewriteXAIWebsocketDownstreamIDValue(value, upstreamResponseID, downstreamResponseID, upstreamPreviousID, downstreamPreviousID, "") { return payload } out, errMarshal := json.Marshal(value) if errMarshal != nil { return payload } return out}
func rewriteXAIWebsocketDownstreamIDValue(value any, upstreamResponseID string, downstreamResponseID string, upstreamPreviousID string, downstreamPreviousID string, key string) bool { switch typed := value.(type) { case map[string]any: changed := false for childKey, childValue := range typed { if childString, ok := childValue.(string); ok { replaced := rewriteXAIWebsocketDownstreamIDString(childString, childKey, upstreamResponseID, downstreamResponseID, upstreamPreviousID, downstreamPreviousID) if replaced != childString { typed[childKey] = replaced changed = true } continue } if rewriteXAIWebsocketDownstreamIDValue(childValue, upstreamResponseID, downstreamResponseID, upstreamPreviousID, downstreamPreviousID, childKey) { changed = true } } return changed case []any: changed := false for i := range typed { if rewriteXAIWebsocketDownstreamIDValue(typed[i], upstreamResponseID, downstreamResponseID, upstreamPreviousID, downstreamPreviousID, key) { changed = true } } return changed default: return false }}
func rewriteXAIWebsocketDownstreamIDString(value string, key string, upstreamResponseID string, downstreamResponseID string, upstreamPreviousID string, downstreamPreviousID string) string { switch key { case "id", "item_id": if upstreamResponseID != "" && downstreamResponseID != "" && downstreamResponseID != upstreamResponseID && strings.Contains(value, upstreamResponseID) { return strings.ReplaceAll(value, upstreamResponseID, downstreamResponseID) } case "previous_response_id": if upstreamPreviousID != "" && downstreamPreviousID != "" && value == upstreamPreviousID { return downstreamPreviousID } } return value}
func (e *XAIWebsocketsExecutor) Execute(ctx context.Context, auth *cliproxyauth.Auth, req cliproxyexecutor.Request, opts cliproxyexecutor.Options) (cliproxyexecutor.Response, error) { if e == nil || e.XAIExecutor == nil { return cliproxyexecutor.Response{}, fmt.Errorf("xai websockets executor: executor is nil") } return e.XAIExecutor.Execute(ctx, auth, req, opts)}
func (e *XAIWebsocketsExecutor) ExecuteStream(ctx context.Context, auth *cliproxyauth.Auth, req cliproxyexecutor.Request, opts cliproxyexecutor.Options) (_ *cliproxyexecutor.StreamResult, err error) { if e == nil || e.XAIExecutor == nil { return nil, fmt.Errorf("xai websockets executor: executor is nil") } if ctx == nil { ctx = context.Background() } if opts.Alt == "responses/compact" { return nil, statusErr{code: http.StatusBadRequest, msg: "streaming not supported for /responses/compact"} } executionSessionID := executionSessionIDFromOptions(opts) stateSessionID := xaiExecutionSessionID(req, opts) if stateSessionID == "" { stateSessionID = executionSessionID } state := getXAIWebsocketIDState(e.idStore, stateSessionID) stateRequestLocked := false stateRequestLockTransferred := false if executionSessionID == "" && state != nil { state.requestMu.Lock() stateRequestLocked = true } defer func() { if stateRequestLocked && !stateRequestLockTransferred { state.requestMu.Unlock() } }() if xaiInputHasItemType(req.Payload, "compaction_trigger") { if cliproxyexecutor.RequiredUpstreamWebsocket(ctx) { return nil, cliproxyexecutor.NewUpstreamWebsocketReplayRequiredError() } if executionSessionID != "" { sess := e.getOrCreateSession(executionSessionID) if sess != nil { sess.reqMu.Lock() defer sess.reqMu.Unlock() } } idMapper := newXAIWebsocketRequestIDMapper(e.idStore, stateSessionID, req.Payload) return e.executeCompactionTriggerFromWebsocketContext(ctx, auth, req, opts, idMapper) }
// Keep websocket on the official API base URL (or an explicit non-default // base_url). Do not reuse xaiChatBaseURL: cli-chat-proxy only accepts HTTP // POST and returns 405 for websocket upgrades. token, baseURL := xaiCreds(auth) if baseURL == "" { baseURL = xaiauth.DefaultAPIBaseURL }
prepared, err := e.prepareResponsesWebsocketRequest(ctx, req, opts) if err != nil { return nil, err }
reporter := helps.NewExecutorUsageReporter(ctx, e, prepared.baseModel, auth) defer reporter.TrackFailure(ctx, &err)
httpURL := strings.TrimSuffix(baseURL, "/") + "/responses" wsURL, err := buildXAIResponsesWebsocketURL(httpURL) if err != nil { return nil, err }
var authID, authLabel, authType, authValue string if auth != nil { authID = auth.ID authLabel = auth.Label authType, authValue = auth.AccountInfo() }
var sess *codexWebsocketSession if executionSessionID != "" { sess = e.getOrCreateSession(executionSessionID) if sess != nil { sess.reqMu.Lock() } } idMapper := newXAIWebsocketRequestIDMapper(e.idStore, stateSessionID, req.Payload) if idMapper != nil { if websocketSessionTargetChanged(sess, authID, wsURL) { idMapper.upstreamPreviousID = "" } prepared.body = idMapper.upstreamRequestPayload(prepared.body) } reporter.SetTranslatedReasoningEffort(prepared.body, e.Identifier())
wsHeaders := applyXAIWebsocketHeaders(http.Header{}, auth, token, prepared.sessionID, opts.Headers) wsReqBody := buildXAIWebsocketRequestBody(prepared.body) requestType := strings.TrimSpace(gjson.GetBytes(req.Payload, "type").String()) transcriptReset := strings.TrimSpace(gjson.GetBytes(wsReqBody, "previous_response_id").String()) == "" && (requestType != "response.append" || (idMapper != nil && idMapper.replayedCompactedTranscript)) warmupRequest := xaiWebsocketGenerateFalse(wsReqBody)
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) logXAIWebsocketRequest(executionSessionID, authID, wsURL, wsReqBody)
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 > 0 { if sess != nil { sess.reqMu.Unlock() } return nil, xaiStatusErr(respHS.StatusCode, 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 { logXAIWebsocketConnected(executionSessionID, authID, wsURL) }
var readCh chan codexWebsocketRead if sess != nil { readCh = sess.activate(conn) }
cliproxyexecutor.MarkUpstreamAttempt(ctx) if errSend := writeCodexWebsocketMessage(sess, conn, wsReqBody); errSend != nil { errSend = mapXAIWebsocketWriteError(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 !shouldRetryXAIWebsocketSend(errSend) { return nil, errSend } return nil, cliproxyexecutor.NewUpstreamWebsocketReplayRequiredError() } e.invalidateUpstreamConn(sess, conn, "send_error", errSend) if !shouldRetryXAIWebsocketSend(errSend) { sess.clearActive(conn, readCh) sess.reqMu.Unlock() return nil, errSend } connRetry, closerRetry, respHSRetry, errDialRetry := e.ensureUpstreamConn(ctx, auth, sess, authID, wsURL, wsHeaders) if errDialRetry != nil || connRetry == nil { bodyErrRetry := websocketHandshakeBody(respHSRetry) closeHTTPResponseBody(respHSRetry, "xai websockets executor: close handshake response body error") helps.RecordAPIWebsocketError(ctx, e.cfg, "dial_retry", errDialRetry) sess.clearActive(conn, readCh) sess.reqMu.Unlock() if respHSRetry != nil && respHSRetry.StatusCode > 0 { return nil, xaiStatusErr(respHSRetry.StatusCode, bodyErrRetry) } 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) wsReqBodyRetry := buildXAIWebsocketRequestBody(prepared.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, }) logXAIWebsocketRequest(executionSessionID, authID, wsURL, wsReqBodyRetry) recordAPIWebsocketHandshake(ctx, e.cfg, respHSRetry) reporter.StartResponseTTFT() cliproxyexecutor.MarkUpstreamAttempt(ctx) if errSendRetry := writeCodexWebsocketMessage(sess, conn, wsReqBodyRetry); errSendRetry != nil { errSendRetry = mapXAIWebsocketWriteError(sess, connRetry, errSendRetry) helps.RecordAPIWebsocketError(ctx, e.cfg, "send_retry", errSendRetry) e.invalidateUpstreamConn(sess, connRetry, "send_error", errSendRetry) sess.clearActive(conn, readCh) sess.reqMu.Unlock() return nil, errSendRetry } wsReqBody = wsReqBodyRetry } else { logXAIWebsocketDisconnected(executionSessionID, authID, wsURL, "send_error", errSend) if errClose := closer.Close(); errClose != nil { log.Errorf("xai websockets executor: close websocket error: %v", errClose) } return nil, errSend } }
out := make(chan cliproxyexecutor.StreamChunk) if stateRequestLocked { stateRequestLockTransferred = true } go func() { if stateRequestLocked { defer state.requestMu.Unlock() } terminateReason := "completed" var terminateErr error
defer close(out) defer func() { if sess != nil { sess.clearActive(conn, readCh) sess.reqMu.Unlock() return } logXAIWebsocketDisconnected(executionSessionID, authID, wsURL, terminateReason, terminateErr) if errClose := closer.Close(); errClose != nil { log.Errorf("xai 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 } }
claudeInputTokens := helps.NewClaudeInputTokenState(prepared.from, prepared.to, prepared.responseFormat, prepared.originalPayload) var param any outputItemsByIndex := make(map[int64][]byte) var outputItemsFallback [][]byte responseFilter := newXAIInternalXSearchResponseFilter(prepared.filterInternalXSearch, prepared.clientDeclaredTools) namespaceRestorer := newXAINamespaceRestorer(prepared.namespaceTools) recordedTranscript := false for { if ctx != nil && ctx.Err() != nil { terminateReason = "context_done" terminateErr = ctx.Err() _ = send(cliproxyexecutor.StreamChunk{Err: ctx.Err()}) return } msgType, payload, errRead := readXAIWebsocketMessage(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 := mapXAIWebsocketReadError(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 { errBinary := fmt.Errorf("xai websockets executor: unexpected binary message") terminateReason = "unexpected_binary" terminateErr = errBinary helps.RecordAPIWebsocketError(ctx, e.cfg, "unexpected_binary", errBinary) reporter.PublishFailure(ctx, errBinary) if sess != nil { e.invalidateUpstreamConn(sess, conn, "unexpected_binary", errBinary) } _ = send(cliproxyexecutor.StreamChunk{Err: errBinary}) return } continue }
payload = bytes.TrimSpace(payload) if len(payload) == 0 { continue } reporter.MarkFirstResponseByte() helps.AppendAPIWebsocketResponse(ctx, e.cfg, payload) helps.EmitWebSocketResponseEvent(ctx, opts, auth, e.Identifier(), req.Model, payload)
if wsErr, ok := parseXAIWebsocketError(payload); ok { terminateReason = "upstream_error" terminateErr = wsErr helps.RecordAPIWebsocketError(ctx, e.cfg, "upstream_error", wsErr) reporter.PublishFailure(ctx, wsErr) if sess != nil { e.invalidateUpstreamConnWithoutDisconnectNotify(sess, conn, "upstream_error", wsErr) } _ = send(cliproxyexecutor.StreamChunk{Err: wsErr}) return }
for _, payload := range xaiNormalizeReasoningSummaryDataEvents(payload) { payload = namespaceRestorer.restore(payload) payload = responseFilter.apply(payload) if len(payload) == 0 { continue } eventType := gjson.GetBytes(payload, "type").String() isTerminalEvent := eventType == "response.completed" || eventType == "response.done" || eventType == "error" warmupCompletedPayload := []byte(nil) switch eventType { case "response.created": if warmupRequest { warmupCompletedPayload = buildXAIWebsocketWarmupCompletedPayload(payload) if idMapper != nil && idMapper.state != nil && !recordedTranscript { idMapper.state.recordTranscriptTurn(wsReqBody, warmupCompletedPayload, transcriptReset) recordedTranscript = true } logXAIWebsocketWarmupCompleted(executionSessionID, authID, wsURL, payload) } case "response.output_item.done": xaiCollectOutputItemDone(payload, outputItemsByIndex, &outputItemsFallback) case "response.completed": logXAIWebsocketTerminalResponse(executionSessionID, authID, wsURL, eventType, payload) if detail, ok := helps.ParseCodexUsage(payload); ok { reporter.Publish(ctx, detail) } payload = xaiPatchCompletedOutput(payload, outputItemsByIndex, outputItemsFallback) payload = xaiNormalizeReasoningSummaryData(payload) cacheXAIReasoningReplayFromCompleted(ctx, prepared.replayScope, payload) if !warmupRequest && idMapper != nil && idMapper.state != nil && !recordedTranscript { idMapper.state.recordTranscriptTurn(wsReqBody, payload, transcriptReset) recordedTranscript = true } case "response.done": logXAIWebsocketTerminalResponse(executionSessionID, authID, wsURL, eventType, payload) if detail, ok := helps.ParseCodexUsage(payload); ok { reporter.Publish(ctx, detail) } if !warmupRequest && idMapper != nil && idMapper.state != nil && !recordedTranscript { idMapper.state.recordTranscriptTurn(wsReqBody, payload, transcriptReset) recordedTranscript = true } }
if cliproxyexecutor.DownstreamWebsocket(ctx) { downstreamPayload := helps.EnsureResponsesUsageDetails(payload) downstreamWarmupCompletedPayload := helps.EnsureResponsesUsageDetails(warmupCompletedPayload) if idMapper != nil { downstreamPayload = idMapper.downstreamResponsePayload(downstreamPayload) if len(warmupCompletedPayload) > 0 { downstreamWarmupCompletedPayload = idMapper.downstreamResponsePayload(downstreamWarmupCompletedPayload) } } if !send(cliproxyexecutor.StreamChunk{Payload: downstreamPayload}) { terminateReason = "context_done" terminateErr = ctx.Err() return } if len(downstreamWarmupCompletedPayload) > 0 { if !send(cliproxyexecutor.StreamChunk{Payload: downstreamWarmupCompletedPayload}) { terminateReason = "context_done" terminateErr = ctx.Err() return } return } if isTerminalEvent { return } continue }
payload = normalizeCodexWebsocketCompletion(payload) line := encodeCodexWebsocketAsSSE(payload) chunks := helps.TranslateStreamWithClaudeInputTokens(ctx, prepared.to, prepared.responseFormat, req.Model, prepared.originalPayload, prepared.body, line, ¶m, claudeInputTokens) for i := range chunks { if !send(cliproxyexecutor.StreamChunk{Payload: chunks[i]}) { terminateReason = "context_done" terminateErr = ctx.Err() return } } if len(warmupCompletedPayload) > 0 { line = encodeCodexWebsocketAsSSE(warmupCompletedPayload) chunks = helps.TranslateStreamWithClaudeInputTokens(ctx, prepared.to, prepared.responseFormat, req.Model, prepared.originalPayload, prepared.body, line, ¶m, claudeInputTokens) for i := range chunks { if !send(cliproxyexecutor.StreamChunk{Payload: chunks[i]}) { terminateReason = "context_done" terminateErr = ctx.Err() return } } return } if eventType == "response.completed" || eventType == "response.done" { return } } } }() return &cliproxyexecutor.StreamResult{Headers: upstreamHeaders, Chunks: out}, nil}
func (e *XAIWebsocketsExecutor) executeCompactionTriggerFromWebsocketContext(ctx context.Context, auth *cliproxyauth.Auth, req cliproxyexecutor.Request, opts cliproxyexecutor.Options, idMapper *xaiWebsocketRequestIDMapper) (*cliproxyexecutor.StreamResult, error) { if idMapper == nil || idMapper.state == nil { return nil, statusErr{code: http.StatusBadRequest, msg: "xai websocket compaction context is unavailable"} } transcriptInput := idMapper.state.snapshotTranscriptInput() if len(transcriptInput) == 0 { return nil, statusErr{code: http.StatusBadRequest, msg: "xai websocket compaction context is empty"} } authID := "" if auth != nil { authID = auth.ID } log.Infof( "xai websockets: compact fallback session=%s auth=%s input_items=%d", xaiExecutionSessionID(req, opts), strings.TrimSpace(authID), len(gjson.ParseBytes(transcriptInput).Array()), ) compactPayload, err := buildXAIWebsocketCompactionPayload(req.Payload, transcriptInput) if err != nil { return nil, err } compactReq := req compactReq.Payload = compactPayload
prepared, data, headers, err := e.XAIExecutor.executeCompactRequest(ctx, auth, compactReq, opts) if err != nil { return nil, err }
responseID, compactionItem, errValidate := validateXAIWebsocketCompactionResponse(data) if errValidate != nil { return nil, errValidate } idMapper.state.replaceTranscriptWithItems(compactionItem) idMapper.state.mapDownstreamToUpstream(responseID, "")
headers = headers.Clone() if headers == nil { headers = make(http.Header) } headers.Set("Content-Type", "text/event-stream")
chunks := xaiBuildCompactionTriggerStreamChunks(prepared, data) out := make(chan cliproxyexecutor.StreamChunk, len(chunks)) for _, chunk := range chunks { out <- cliproxyexecutor.StreamChunk{Payload: chunk} } close(out) return &cliproxyexecutor.StreamResult{Headers: headers, Chunks: out}, nil}
func validateXAIWebsocketCompactionResponse(data []byte) (string, []byte, error) { if len(data) == 0 || !json.Valid(data) { return "", nil, statusErr{code: http.StatusBadGateway, msg: "xai websocket compaction returned invalid JSON"} } responseIDResult := gjson.GetBytes(data, "id") output := gjson.GetBytes(data, "output") if responseIDResult.Type != gjson.String || strings.TrimSpace(responseIDResult.String()) == "" || !output.Exists() || !output.IsArray() { return "", nil, statusErr{code: http.StatusBadGateway, msg: "xai websocket compaction response is missing compacted state"} } items := output.Array() if len(items) == 0 { return "", nil, statusErr{code: http.StatusBadGateway, msg: "xai websocket compaction response is missing compacted state"} } item := items[0] itemType := item.Get("type") encryptedContent := item.Get("encrypted_content") if item.Type != gjson.JSON || itemType.Type != gjson.String || strings.TrimSpace(itemType.String()) != "compaction" || encryptedContent.Type != gjson.String || strings.TrimSpace(encryptedContent.String()) == "" { return "", nil, statusErr{code: http.StatusBadGateway, msg: "xai websocket compaction response is missing compacted state"} } normalizedResponseID := xaiCompactionResponseID(data) return normalizedResponseID, xaiCompactionOutputItem(data, normalizedResponseID), nil}
func buildXAIWebsocketCompactionPayload(payload []byte, transcriptInput []byte) ([]byte, error) { if len(payload) == 0 { payload = []byte(`{}`) } if len(transcriptInput) == 0 { transcriptInput = []byte("[]") } out := bytes.Clone(payload) var err error out, err = sjson.SetRawBytes(out, "input", transcriptInput) if err != nil { return nil, err } out, _ = sjson.DeleteBytes(out, "previous_response_id") return out, nil}
func xaiWebsocketGenerateFalse(payload []byte) bool { generate := gjson.GetBytes(payload, "generate") return generate.Exists() && !generate.Bool()}
func buildXAIWebsocketWarmupCompletedPayload(createdPayload []byte) []byte { completed := []byte(`{"type":"response.completed","response":{"output":[],"usage":{"input_tokens":0,"input_tokens_details":{"cached_tokens":0},"output_tokens":0,"output_tokens_details":{"reasoning_tokens":0},"total_tokens":0}}}`) if sequence := gjson.GetBytes(createdPayload, "sequence_number"); sequence.Exists() { completed, _ = sjson.SetBytes(completed, "sequence_number", sequence.Int()+1) } if response := gjson.GetBytes(createdPayload, "response"); response.Exists() && response.IsObject() { responsePayload := []byte(response.Raw) responsePayload, _ = sjson.SetBytes(responsePayload, "status", "completed") if !gjson.GetBytes(responsePayload, "output").Exists() { responsePayload, _ = sjson.SetRawBytes(responsePayload, "output", []byte("[]")) } if !gjson.GetBytes(responsePayload, "usage").Exists() { responsePayload, _ = sjson.SetRawBytes(responsePayload, "usage", []byte(`{"input_tokens":0,"input_tokens_details":{"cached_tokens":0},"output_tokens":0,"output_tokens_details":{"reasoning_tokens":0},"total_tokens":0}`)) } completed, _ = sjson.SetRawBytes(completed, "response", responsePayload) } return helps.EnsureResponsesUsageDetails(completed)}
func parseXAIWebsocketError(payload []byte) (error, bool) { if wsErr, ok := parseCodexWebsocketError(payload); ok { if statusError, okStatus := wsErr.(statusErrWithHeaders); okStatus { xaiError := xaiStatusErr(statusError.code, payload) // Apply normalized status (e.g. 403 bad-credentials -> 401) and any // provider-specific retry hint while preserving websocket headers. statusError.code = xaiError.code if xaiError.retryAfter != nil { statusError.retryAfter = xaiError.retryAfter } return statusError, true } return wsErr, true } if len(payload) == 0 || !gjson.GetBytes(payload, "error").Exists() { return nil, false } status := int(gjson.GetBytes(payload, "status").Int()) if status <= 0 { status = int(gjson.GetBytes(payload, "status_code").Int()) } if status <= 0 { status = xaiBareWebsocketErrorStatus(payload) } out := []byte(`{}`) out, _ = sjson.SetBytes(out, "type", "error") out, _ = sjson.SetBytes(out, "status", status) if errNode := gjson.GetBytes(payload, "error"); errNode.Exists() { out, _ = sjson.SetRawBytes(out, "error", []byte(errNode.Raw)) } return xaiStatusErr(status, out), true}
func xaiBareWebsocketErrorStatus(payload []byte) int { for _, path := range []string{"error.code", "error.status", "code"} { raw := strings.TrimSpace(gjson.GetBytes(payload, path).String()) if raw == "" { continue } status, errAtoi := strconv.Atoi(raw) if errAtoi == nil && status > 0 { return status } } message := strings.TrimSpace(gjson.GetBytes(payload, "error.message").String()) if strings.Contains(message, `"code":"400"`) || strings.Contains(message, "Request validation error") { return http.StatusBadRequest } return http.StatusInternalServerError}
func (e *XAIWebsocketsExecutor) prepareResponsesWebsocketRequest(ctx context.Context, req cliproxyexecutor.Request, opts cliproxyexecutor.Options) (*xaiPreparedRequest, error) { prepared, err := e.prepareResponsesRequest(ctx, req, opts, true) if err != nil { return nil, err } if previousResponseID := strings.TrimSpace(gjson.GetBytes(req.Payload, "previous_response_id").String()); previousResponseID != "" { prepared.body, _ = sjson.SetBytes(prepared.body, "previous_response_id", previousResponseID) } return prepared, nil}
func (e *XAIWebsocketsExecutor) dialXAIWebsocket(ctx context.Context, auth *cliproxyauth.Auth, wsURL string, headers http.Header) (*websocket.Conn, *websocketConnectionCloser, *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 err != nil { cliproxyexecutor.MarkUpstreamAttempt(ctx) } closer := newWebsocketConnectionCloser(conn) if conn != nil { // Avoid gorilla/websocket flate tail validation issues on some upstreams/Go versions. conn.EnableWriteCompression(false) } return conn, closer, resp, err}
func (e *XAIWebsocketsExecutor) getOrCreateSession(sessionID string) *codexWebsocketSession { sessionID = strings.TrimSpace(sessionID) if sessionID == "" || e == nil { return nil } store := e.store if store == nil { store = globalXAIWebsocketSessionStore } 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 *XAIWebsocketsExecutor) UpstreamDisconnectChan(sessionID string) <-chan error { sess := e.getOrCreateSession(sessionID) if sess == nil { return nil } return sess.upstreamDisconnectCh}
func (e *XAIWebsocketsExecutor) ensureUpstreamConn(ctx context.Context, auth *cliproxyauth.Auth, sess *codexWebsocketSession, authID string, wsURL string, headers http.Header) (*websocket.Conn, *websocketConnectionCloser, *http.Response, error) { if sess == nil { return e.dialXAIWebsocket(ctx, auth, wsURL, headers) }
if staleConn, staleCloser, staleAuthID, staleWSURL, staleLifecycle := detachMismatchedWebsocketSessionConn(sess, authID, wsURL); staleConn != nil { logXAIWebsocketDisconnected(sess.sessionID, staleAuthID, staleWSURL, "target_changed", nil) if staleCloser != nil { if errClose := staleCloser.Close(); errClose != nil { log.Errorf("xai websockets executor: close stale websocket error: %v", errClose) } } if staleLifecycle != nil { staleLifecycle.End("target_changed") } }
sess.connMu.Lock() conn := sess.conn closer := sess.connCloser readerConn := sess.readerConn sess.connMu.Unlock() if conn != nil { if readerConn != conn { sess.connMu.Lock() sess.readerConn = conn sess.connMu.Unlock() configureXAIWebsocketConn(sess, conn) go e.readUpstreamLoop(sess, conn) } return conn, closer, nil, nil }
conn, closer, resp, errDial := e.dialXAIWebsocket(ctx, auth, wsURL, headers) if errDial != nil { return nil, closer, resp, errDial }
sess.connMu.Lock() if sess.conn != nil { previous := sess.conn previousCloser := sess.connCloser sess.connMu.Unlock() if errClose := closer.Close(); errClose != nil { log.Errorf("xai websockets executor: close websocket error: %v", errClose) } return previous, previousCloser, nil, nil } sess.conn = conn sess.connCloser = closer sess.wsURL = wsURL sess.authID = authID sess.readerConn = conn sess.connMu.Unlock()
configureXAIWebsocketConn(sess, conn) go e.readUpstreamLoop(sess, conn) logXAIWebsocketConnected(sess.sessionID, authID, wsURL) return conn, closer, resp, nil}
func configureXAIWebsocketConn(sess *codexWebsocketSession, conn *websocket.Conn) { if sess == nil || conn == nil { return } sess.resetUpstreamDisconnectError(conn) conn.SetPingHandler(func(appData string) error { sess.writeMu.Lock() defer sess.writeMu.Unlock() return conn.WriteControl(websocket.PongMessage, []byte(appData), time.Time{}) }) defaultCloseHandler := conn.CloseHandler() conn.SetCloseHandler(func(code int, text string) error { sess.setUpstreamDisconnectError(conn, &websocket.CloseError{Code: code, Text: text}) return defaultCloseHandler(code, text) })}
func mapXAIWebsocketReadError(err error) error { return mapCodexWebsocketReadError(err)}
func mapXAIWebsocketWriteError(sess *codexWebsocketSession, conn *websocket.Conn, err error) error { return mapCodexWebsocketWriteError(sess, conn, err)}
func shouldRetryXAIWebsocketSend(err error) bool { return shouldRetryCodexWebsocketSend(err)}
func readXAIWebsocketMessage(ctx context.Context, sess *codexWebsocketSession, conn *websocket.Conn, readCh chan codexWebsocketRead) (int, []byte, error) { if ctx == nil { ctx = context.Background() } if sess == nil { if conn == nil { return 0, nil, fmt.Errorf("xai websockets executor: websocket conn is nil") } msgType, payload, errRead := conn.ReadMessage() return msgType, payload, errRead } if conn == nil { return 0, nil, fmt.Errorf("xai websockets executor: websocket conn is nil") } if readCh == nil { return 0, nil, fmt.Errorf("xai 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("xai 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 (e *XAIWebsocketsExecutor) readUpstreamLoop(sess *codexWebsocketSession, conn *websocket.Conn) { if e == nil || sess == nil || conn == nil { return } for { msgType, payload, errRead := conn.ReadMessage() if errRead != nil { invalidate := func() { e.invalidateUpstreamConn(sess, conn, "upstream_disconnected", errRead) } invalidated := false ch, done := sess.activeForConn(conn) if ch != nil { invalidated = sendTerminalWebsocketRead(ch, done, codexWebsocketRead{conn: conn, err: errRead}, invalidate) if sess.clearActive(conn, ch) { close(ch) } } if !invalidated { invalidate() } return }
if msgType != websocket.TextMessage { if msgType == websocket.BinaryMessage { errBinary := fmt.Errorf("xai websockets executor: unexpected binary message") invalidate := func() { e.invalidateUpstreamConn(sess, conn, "unexpected_binary", errBinary) } invalidated := false ch, done := sess.activeForConn(conn) if ch != nil { invalidated = sendTerminalWebsocketRead(ch, done, codexWebsocketRead{conn: conn, err: errBinary}, invalidate) if sess.clearActive(conn, ch) { close(ch) } } if !invalidated { invalidate() } return } continue }
ch, done := sess.activeForConn(conn) if ch == nil { continue } select { case ch <- codexWebsocketRead{conn: conn, msgType: msgType, payload: payload}: case <-done: } }}
func (e *XAIWebsocketsExecutor) invalidateUpstreamConn(sess *codexWebsocketSession, conn *websocket.Conn, reason string, err error) { e.invalidateUpstreamConnWithNotify(sess, conn, reason, err, true)}
func (e *XAIWebsocketsExecutor) invalidateUpstreamConnWithoutDisconnectNotify(sess *codexWebsocketSession, conn *websocket.Conn, reason string, err error) { e.invalidateUpstreamConnWithNotify(sess, conn, reason, err, false)}
func (e *XAIWebsocketsExecutor) invalidateUpstreamConnWithNotify(sess *codexWebsocketSession, conn *websocket.Conn, reason string, err error, notify bool) { 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 } lifecycle := sess.lifecycle closer := sess.connCloser sess.lifecycle = nil sess.lifecycleModel = "" sess.conn = nil sess.connCloser = nil if sess.readerConn == conn { sess.readerConn = nil } sess.connMu.Unlock()
logXAIWebsocketDisconnected(sessionID, authID, wsURL, reason, err) if notify { sess.notifyUpstreamDisconnect(err) } if closer != nil { if errClose := closer.Close(); errClose != nil { log.Errorf("xai websockets executor: close websocket error: %v", errClose) } } if lifecycle != nil { lifecycle.End(reason) }}
func (e *XAIWebsocketsExecutor) CloseExecutionSession(sessionID string) { sessionID = strings.TrimSpace(sessionID) if e == nil || sessionID == "" { return } if sessionID == cliproxyauth.CloseAllExecutionSessionsID { e.closeAllExecutionSessions("executor_shutdown") return }
store := e.store if store == nil { store = globalXAIWebsocketSessionStore } store.mu.Lock() sess := store.sessions[sessionID] delete(store.sessions, sessionID) store.mu.Unlock() deleteXAIWebsocketIDState(e.idStore, sessionID)
e.closeExecutionSession(sess, "session_closed")}
func (e *XAIWebsocketsExecutor) closeExecutionSession(sess *codexWebsocketSession, reason string) { closeXAIWebsocketSession(sess, reason)}
func (e *XAIWebsocketsExecutor) closeAllExecutionSessions(reason string) { if e == nil { return } store := e.store if store == nil { store = globalXAIWebsocketSessionStore } 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 _, sess := range sessions { closeXAIWebsocketSession(sess, reason) }}
func closeXAIWebsocketSession(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 lifecycle := sess.lifecycle closer := sess.connCloser sess.lifecycle = nil sess.lifecycleModel = "" sess.conn = nil sess.connCloser = nil if sess.readerConn == conn { sess.readerConn = nil } sessionID := sess.sessionID sess.connMu.Unlock()
if conn != nil { logXAIWebsocketDisconnected(sessionID, authID, wsURL, reason, nil) if closer != nil { if errClose := closer.Close(); errClose != nil { log.Errorf("xai websockets executor: close websocket error: %v", errClose) } } } if lifecycle != nil { lifecycle.End(reason) }}
func buildXAIWebsocketRequestBody(body []byte) []byte { if len(body) == 0 { return nil } wsReqBody := bytes.Clone(body) wsReqBody, _ = sjson.SetBytes(wsReqBody, "type", "response.create") wsReqBody, _ = sjson.DeleteBytes(wsReqBody, "stream") wsReqBody, _ = sjson.DeleteBytes(wsReqBody, "stream_options") wsReqBody, _ = sjson.DeleteBytes(wsReqBody, "background") wsReqBody, _ = sjson.SetBytes(wsReqBody, "store", true) if strings.TrimSpace(gjson.GetBytes(wsReqBody, "previous_response_id").String()) != "" { wsReqBody, _ = sjson.DeleteBytes(wsReqBody, "instructions") } return wsReqBody}
func buildXAIResponsesWebsocketURL(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" case "ws", "wss": default: return "", fmt.Errorf("xai websockets executor: unsupported responses websocket URL scheme %q", parsed.Scheme) } if strings.TrimSpace(parsed.Host) == "" { return "", fmt.Errorf("xai websockets executor: responses websocket URL host is empty") } return parsed.String(), nil}
func applyXAIWebsocketHeaders(headers http.Header, auth *cliproxyauth.Auth, token string, sessionID string, clientHeaders ...http.Header) http.Header { if headers == nil { headers = http.Header{} } headers.Set("Content-Type", "application/json") if strings.TrimSpace(token) != "" { headers.Set("Authorization", "Bearer "+token) } else { headers.Del("Authorization") } if sessionID != "" { headers.Set("x-grok-conv-id", sessionID) } var attrs map[string]string if auth != nil { attrs = auth.Attributes } util.ApplyCustomHeadersFromAttrs(&http.Request{Header: headers}, attrs, clientHeaders...) return headers}
func logXAIWebsocketConnected(sessionID string, authID string, wsURL string) { log.Infof("xai websockets: upstream connected session=%s auth=%s url=%s", strings.TrimSpace(sessionID), strings.TrimSpace(authID), strings.TrimSpace(wsURL))}
func logXAIWebsocketRequest(sessionID string, authID string, wsURL string, payload []byte) { if len(payload) == 0 { log.Infof("xai websockets: upstream request sent session=%s auth=%s url=%s", strings.TrimSpace(sessionID), strings.TrimSpace(authID), strings.TrimSpace(wsURL)) return } generateValue := "default" if generate := gjson.GetBytes(payload, "generate"); generate.Exists() { generateValue = strings.TrimSpace(generate.Raw) } log.Infof( "xai websockets: upstream request sent session=%s auth=%s url=%s event=%s previous_response_id=%s generate=%s input_items=%d", strings.TrimSpace(sessionID), strings.TrimSpace(authID), strings.TrimSpace(wsURL), strings.TrimSpace(gjson.GetBytes(payload, "type").String()), strings.TrimSpace(gjson.GetBytes(payload, "previous_response_id").String()), generateValue, len(gjson.GetBytes(payload, "input").Array()), )}
func logXAIWebsocketWarmupCompleted(sessionID string, authID string, wsURL string, payload []byte) { log.Infof( "xai websockets: upstream warmup completed session=%s auth=%s url=%s response_id=%s", strings.TrimSpace(sessionID), strings.TrimSpace(authID), strings.TrimSpace(wsURL), strings.TrimSpace(gjson.GetBytes(payload, "response.id").String()), )}
func logXAIWebsocketTerminalResponse(sessionID string, authID string, wsURL string, eventType string, payload []byte) { log.Infof( "xai websockets: upstream terminal response session=%s auth=%s url=%s event=%s response_id=%s previous_response_id=%s", strings.TrimSpace(sessionID), strings.TrimSpace(authID), strings.TrimSpace(wsURL), strings.TrimSpace(eventType), strings.TrimSpace(gjson.GetBytes(payload, "response.id").String()), strings.TrimSpace(gjson.GetBytes(payload, "response.previous_response_id").String()), )}
func logXAIWebsocketDisconnected(sessionID string, authID string, wsURL string, reason string, err error) { if err != nil { log.Infof("xai 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("xai websockets: upstream disconnected session=%s auth=%s url=%s reason=%s", strings.TrimSpace(sessionID), strings.TrimSpace(authID), strings.TrimSpace(wsURL), strings.TrimSpace(reason))}
// CloseXAIWebsocketSessionsForAuthID closes all active xAI upstream websocket sessions// associated with the supplied auth ID.func CloseXAIWebsocketSessionsForAuthID(authID string, reason string) { authID = strings.TrimSpace(authID) if authID == "" { return } reason = strings.TrimSpace(reason) if reason == "" { reason = "auth_removed" }
store := globalXAIWebsocketSessionStore 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) deleteXAIWebsocketIDState(globalXAIWebsocketIDStates, matches[i].sessionID) toClose = append(toClose, current) } store.mu.Unlock()
for i := range toClose { closeXAIWebsocketSession(toClose[i], reason) }}
// XAIAutoExecutor routes xAI stream requests to the websocket transport only// when the downstream transport is websocket and the selected auth enables// websockets. Non-stream requests keep using the HTTP implementation.type XAIAutoExecutor struct { httpExec *XAIExecutor wsExec *XAIWebsocketsExecutor}
func NewXAIAutoExecutor(cfg *config.Config) *XAIAutoExecutor { return &XAIAutoExecutor{ httpExec: NewXAIExecutor(cfg), wsExec: NewXAIWebsocketsExecutor(cfg), }}
func (e *XAIAutoExecutor) Identifier() string { return "xai" }
// UsesConfig reports whether the executor was created for cfg.func (e *XAIAutoExecutor) UsesConfig(cfg *config.Config) bool { return e != nil && e.httpExec != nil && e.httpExec.cfg == cfg}
func (e *XAIAutoExecutor) PrepareRequest(req *http.Request, auth *cliproxyauth.Auth) error { if e == nil || e.httpExec == nil { return nil } return e.httpExec.PrepareRequest(req, auth)}
func (e *XAIAutoExecutor) HttpRequest(ctx context.Context, auth *cliproxyauth.Auth, req *http.Request) (*http.Response, error) { if e == nil || e.httpExec == nil { return nil, fmt.Errorf("xai auto executor: http executor is nil") } return e.httpExec.HttpRequest(ctx, auth, req)}
func (e *XAIAutoExecutor) Execute(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("xai auto executor: executor is nil") } return e.httpExec.Execute(ctx, auth, req, opts)}
func (e *XAIAutoExecutor) 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("xai auto executor: executor is nil") } if cliproxyexecutor.DownstreamWebsocket(ctx) && xaiWebsocketsEnabled(auth) { return e.wsExec.ExecuteStream(ctx, auth, req, opts) } if cliproxyexecutor.RequiredUpstreamWebsocket(ctx) { return nil, cliproxyexecutor.NewUpstreamWebsocketReplayRequiredError() } return e.httpExec.ExecuteStream(ctx, auth, req, opts)}
func (e *XAIAutoExecutor) Refresh(ctx context.Context, auth *cliproxyauth.Auth) (*cliproxyauth.Auth, error) { if e == nil || e.httpExec == nil { return nil, fmt.Errorf("xai auto executor: http executor is nil") } return e.httpExec.Refresh(ctx, auth)}
func (e *XAIAutoExecutor) 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("xai auto executor: http executor is nil") } return e.httpExec.CountTokens(ctx, auth, req, opts)}
func (e *XAIAutoExecutor) CloseExecutionSession(sessionID string) { if e == nil || e.wsExec == nil { return } e.wsExec.CloseExecutionSession(sessionID)}
func (e *XAIAutoExecutor) UpstreamDisconnectChan(sessionID string) <-chan error { if e == nil || e.wsExec == nil { return nil } return e.wsExec.UpstreamDisconnectChan(sessionID)}
func xaiWebsocketsEnabled(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}