Something went wrong. Try again.
sloproxy
Something went wrong. Try again.
100 kB · 2521 lines
Go
at commit 291cfb87
12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970717273747576777879808182838485868788899091929394959697989910010110210310410510610710810911011111211311411511611711811912012112212312412512612712812913013113213313413513613713813914014114214314414514614714814915015115215315415515615715815916016116216316416516616716816917017117217317417517617717817918018118218318418518618718818919019119219319419519619719819920020120220320420520620720820921021121221321421521621721821922022122222322422522622722822923023123223323423523623723823924024124224324424524624724824925025125225325425525625725825926026126226326426526626726826927027127227327427527627727827928028128228328428528628728828929029129229329429529629729829930030130230330430530630730830931031131231331431531631731831932032132232332432532632732832933033133233333433533633733833934034134234334434534634734834935035135235335435535635735835936036136236336436536636736836937037137237337437537637737837938038138238338438538638738838939039139239339439539639739839940040140240340440540640740840941041141241341441541641741841942042142242342442542642742842943043143243343443543643743843944044144244344444544644744844945045145245345445545645745845946046146246346446546646746846947047147247347447547647747847948048148248348448548648748848949049149249349449549649749849950050150250350450550650750850951051151251351451551651751851952052152252352452552652752852953053153253353453553653753853954054154254354454554654754854955055155255355455555655755855956056156256356456556656756856957057157257357457557657757857958058158258358458558658758858959059159259359459559659759859960060160260360460560660760860961061161261361461561661761861962062162262362462562662762862963063163263363463563663763863964064164264364464564664764864965065165265365465565665765865966066166266366466566666766866967067167267367467567667767867968068168268368468568668768868969069169269369469569669769869970070170270370470570670770870971071171271371471571671771871972072172272372472572672772872973073173273373473573673773873974074174274374474574674774874975075175275375475575675775875976076176276376476576676776876977077177277377477577677777877978078178278378478578678778878979079179279379479579679779879980080180280380480580680780880981081181281381481581681781881982082182282382482582682782882983083183283383483583683783883984084184284384484584684784884985085185285385485585685785885986086186286386486586686786886987087187287387487587687787887988088188288388488588688788888989089189289389489589689789889990090190290390490590690790890991091191291391491591691791891992092192292392492592692792892993093193293393493593693793893994094194294394494594694794894995095195295395495595695795895996096196296396496596696796896997097197297397497597697797897998098198298398498598698798898999099199299399499599699799899910001001100210031004100510061007100810091010101110121013101410151016101710181019102010211022102310241025102610271028102910301031103210331034103510361037103810391040104110421043104410451046104710481049105010511052105310541055105610571058105910601061106210631064106510661067106810691070107110721073107410751076107710781079108010811082108310841085108610871088108910901091109210931094109510961097109810991100110111021103110411051106110711081109111011111112111311141115111611171118111911201121112211231124112511261127112811291130113111321133113411351136113711381139114011411142114311441145114611471148114911501151115211531154115511561157115811591160116111621163116411651166116711681169117011711172117311741175117611771178117911801181118211831184118511861187118811891190119111921193119411951196119711981199120012011202120312041205120612071208120912101211121212131214121512161217121812191220122112221223122412251226122712281229123012311232123312341235123612371238123912401241124212431244124512461247124812491250125112521253125412551256125712581259126012611262126312641265126612671268126912701271127212731274127512761277127812791280128112821283128412851286128712881289129012911292129312941295129612971298129913001301130213031304130513061307130813091310131113121313131413151316131713181319132013211322132313241325132613271328132913301331133213331334133513361337133813391340134113421343134413451346134713481349135013511352135313541355135613571358135913601361136213631364136513661367136813691370137113721373137413751376137713781379138013811382138313841385138613871388138913901391139213931394139513961397139813991400140114021403140414051406140714081409141014111412141314141415141614171418141914201421142214231424142514261427142814291430143114321433143414351436143714381439144014411442144314441445144614471448144914501451145214531454145514561457145814591460146114621463146414651466146714681469147014711472147314741475147614771478147914801481148214831484148514861487148814891490149114921493149414951496149714981499150015011502150315041505150615071508150915101511151215131514151515161517151815191520152115221523152415251526152715281529153015311532153315341535153615371538153915401541154215431544154515461547154815491550155115521553155415551556155715581559156015611562156315641565156615671568156915701571157215731574157515761577157815791580158115821583158415851586158715881589159015911592159315941595159615971598159916001601160216031604160516061607160816091610161116121613161416151616161716181619162016211622162316241625162616271628162916301631163216331634163516361637163816391640164116421643164416451646164716481649165016511652165316541655165616571658165916601661166216631664166516661667166816691670167116721673167416751676167716781679168016811682168316841685168616871688168916901691169216931694169516961697169816991700170117021703170417051706170717081709171017111712171317141715171617171718171917201721172217231724172517261727172817291730173117321733173417351736173717381739174017411742174317441745174617471748174917501751175217531754175517561757175817591760176117621763176417651766176717681769177017711772177317741775177617771778177917801781178217831784178517861787178817891790179117921793179417951796179717981799180018011802180318041805180618071808180918101811181218131814181518161817181818191820182118221823182418251826182718281829183018311832183318341835183618371838183918401841184218431844184518461847184818491850185118521853185418551856185718581859186018611862186318641865186618671868186918701871187218731874187518761877187818791880188118821883188418851886188718881889189018911892189318941895189618971898189919001901190219031904190519061907190819091910191119121913191419151916191719181919192019211922192319241925192619271928192919301931193219331934193519361937193819391940194119421943194419451946194719481949195019511952195319541955195619571958195919601961196219631964196519661967196819691970197119721973197419751976197719781979198019811982198319841985198619871988198919901991199219931994199519961997199819992000200120022003200420052006200720082009201020112012201320142015201620172018201920202021202220232024202520262027202820292030203120322033203420352036203720382039204020412042204320442045204620472048204920502051205220532054205520562057205820592060206120622063206420652066206720682069207020712072207320742075207620772078207920802081208220832084208520862087208820892090209120922093209420952096209720982099210021012102210321042105210621072108210921102111211221132114211521162117211821192120212121222123212421252126212721282129213021312132213321342135213621372138213921402141214221432144214521462147214821492150215121522153215421552156215721582159216021612162216321642165216621672168216921702171217221732174217521762177217821792180218121822183218421852186218721882189219021912192219321942195219621972198219922002201220222032204220522062207220822092210221122122213221422152216221722182219222022212222222322242225222622272228222922302231223222332234223522362237223822392240224122422243224422452246224722482249225022512252225322542255225622572258225922602261226222632264226522662267226822692270227122722273227422752276227722782279228022812282228322842285228622872288228922902291229222932294229522962297229822992300230123022303230423052306230723082309231023112312231323142315231623172318231923202321232223232324232523262327232823292330233123322333233423352336233723382339234023412342234323442345234623472348234923502351235223532354235523562357235823592360236123622363236423652366236723682369237023712372237323742375237623772378237923802381238223832384238523862387238823892390239123922393239423952396239723982399240024012402240324042405240624072408240924102411241224132414241524162417241824192420242124222423242424252426242724282429243024312432243324342435243624372438243924402441244224432444244524462447244824492450245124522453245424552456245724582459246024612462246324642465246624672468246924702471247224732474247524762477247824792480248124822483248424852486248724882489249024912492249324942495249624972498249925002501250225032504250525062507250825092510251125122513251425152516251725182519252025212522package executor
import ( "bytes" "context" "errors" "net/http" "net/http/httptest" "strings" "sync/atomic" "testing" "time"
"github.com/gin-gonic/gin" "github.com/google/uuid" "github.com/gorilla/websocket" internalcache "github.com/router-for-me/CLIProxyAPI/v7/internal/cache" "github.com/router-for-me/CLIProxyAPI/v7/internal/config" "github.com/router-for-me/CLIProxyAPI/v7/internal/registry" cliproxyauth "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/auth" cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executor" sdkconfig "github.com/router-for-me/CLIProxyAPI/v7/sdk/config" sdktranslator "github.com/router-for-me/CLIProxyAPI/v7/sdk/translator" "github.com/tidwall/gjson")
var benchmarkBuildCodexWebsocketRequestBodyOutput []byte
func TestBuildCodexWebsocketRequestBodyPreservesPreviousResponseID(t *testing.T) { body := []byte(`{"model":"gpt-5-codex","previous_response_id":"resp-1","input":[{"type":"message","id":"msg-1"}]}`)
wsReqBody := buildCodexWebsocketRequestBody(body)
if got := gjson.GetBytes(wsReqBody, "type").String(); got != "response.create" { t.Fatalf("type = %s, want response.create", got) } if got := gjson.GetBytes(wsReqBody, "previous_response_id").String(); got != "resp-1" { t.Fatalf("previous_response_id = %s, want resp-1", got) } if gjson.GetBytes(wsReqBody, "input.0.id").String() != "msg-1" { t.Fatalf("input item id mismatch") } if got := gjson.GetBytes(wsReqBody, "type").String(); got == "response.append" { t.Fatalf("unexpected websocket request type: %s", got) }}
func BenchmarkBuildCodexWebsocketRequestBodyLargePayload(b *testing.B) { body := []byte(`{"model":"gpt-5.6","input":[{"type":"message","id":"msg_1","role":"user","content":"` + strings.Repeat("x", 8<<20) + `"}]}`) b.ReportAllocs() b.SetBytes(int64(len(body))) b.ResetTimer() for b.Loop() { benchmarkBuildCodexWebsocketRequestBodyOutput = buildCodexWebsocketRequestBody(body) }}
func TestBuildCodexWebsocketRequestBodySanitizesOverlongInputItemIDs(t *testing.T) { longReasoningItemID := "rs_" + strings.Repeat("a", 64) longCallItemID := strings.Repeat("grok-call-item-", 6) longOutputItemID := strings.Repeat("grok-output-item-", 6) body := []byte(`{"model":"gpt-5-codex","input":[{"type":"reasoning","id":"` + longReasoningItemID + `","encrypted_content":"gAAAA-encrypted","summary":[]},{"type":"function_call","id":"` + longCallItemID + `","call_id":"call-1","name":"lookup"},{"type":"function_call_output","id":"` + longOutputItemID + `","call_id":"call-1","output":"ok"},{"type":"message","id":"item_74ec40c883248ebb4885ec84"}]}`)
first := buildCodexWebsocketRequestBody(body) second := buildCodexWebsocketRequestBody(body)
if input := gjson.GetBytes(first, "input").Array(); len(input) != 3 { t.Fatalf("input length = %d, want 3: %s", len(input), first) } if gotType := gjson.GetBytes(first, "input.0.type").String(); gotType != "function_call" { t.Fatalf("input.0.type = %q, want function_call: %s", gotType, first) }
shortCallItemID := gjson.GetBytes(first, "input.0.id").String() shortOutputItemID := gjson.GetBytes(first, "input.1.id").String() if len([]rune(shortCallItemID)) > 64 || shortCallItemID == longCallItemID { t.Fatalf("input.0.id was not shortened to at most 64 characters: %q", shortCallItemID) } if len([]rune(shortOutputItemID)) > 64 || shortOutputItemID == longOutputItemID { t.Fatalf("input.1.id was not shortened to at most 64 characters: %q", shortOutputItemID) } if shortCallItemID == shortOutputItemID { t.Fatalf("distinct long IDs produced the same shortened ID: %q", shortCallItemID) } if got := gjson.GetBytes(second, "input.0.id").String(); got != shortCallItemID { t.Fatalf("input item ID shortening is not deterministic: first=%q second=%q", shortCallItemID, got) } if got := gjson.GetBytes(first, "input.0.call_id").String(); got != "call-1" { t.Fatalf("function call_id = %q, want call-1", got) } if got := gjson.GetBytes(first, "input.1.call_id").String(); got != "call-1" { t.Fatalf("function call output call_id = %q, want call-1", got) } if got := gjson.GetBytes(first, "input.2.id").String(); got != "msg_item_74ec40c883248ebb4885ec84" { t.Fatalf("message input item ID was not normalized: %q", got) }}
func TestCodexWebsocketsExecuteRestoresClaudeAgentReasoningReplay(t *testing.T) { internalcache.ClearCodexReasoningReplayCache() t.Cleanup(internalcache.ClearCodexReasoningReplayCache)
encryptedContent := validCodexReasoningEncryptedContentForTestSeed(31) cacheCodexReasoningReplayFromCompleted(codexReasoningReplayScope{ modelName: "gpt-5.4", sessionKey: "claude:ws-replay-session:agent:agent-a", }, []byte(`{"response":{"output":[`+ `{"type":"reasoning","summary":[],"content":null,"encrypted_content":"`+encryptedContent+`"},`+ `{"type":"message","role":"assistant","content":[{"type":"output_text","text":"previous answer"}]}`+ `]}}`))
upgrader := websocket.Upgrader{CheckOrigin: func(*http.Request) bool { return true }} capturedPayload := make(chan []byte, 1) server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { conn, errUpgrade := upgrader.Upgrade(w, r, nil) if errUpgrade != nil { t.Fatalf("upgrade websocket: %v", errUpgrade) } defer func() { _ = conn.Close() }()
_, payload, errRead := conn.ReadMessage() if errRead != nil { t.Fatalf("read upstream websocket message: %v", errRead) } capturedPayload <- bytes.Clone(payload) completed := []byte(`{"type":"response.completed","response":{"id":"resp-ws-replay","output":[{"type":"message","role":"assistant","content":[{"type":"output_text","text":"next answer"}]}],"usage":{"input_tokens":1,"output_tokens":1,"total_tokens":2}}}`) if errWrite := conn.WriteMessage(websocket.TextMessage, completed); errWrite != nil { t.Fatalf("write completed websocket message: %v", errWrite) } })) defer server.Close()
exec := NewCodexWebsocketsExecutor(&config.Config{SDKConfig: config.SDKConfig{DisableImageGeneration: config.DisableImageGenerationAll}}) auth := &cliproxyauth.Auth{Provider: "codex", Attributes: map[string]string{"api_key": "sk-test", "base_url": server.URL}} req := cliproxyexecutor.Request{ Model: "gpt-5.4", Payload: []byte(`{ "model":"gpt-5.4", "messages":[ {"role":"user","content":"first"}, {"role":"assistant","content":"previous answer"}, {"role":"user","content":"next"} ] }`), } headers := http.Header{} headers.Set("X-Claude-Code-Session-Id", "ws-replay-session") headers.Set("X-Claude-Code-Agent-Id", "agent-a") opts := cliproxyexecutor.Options{SourceFormat: sdktranslator.FromString("claude"), Headers: headers}
if _, errExecute := exec.Execute(context.Background(), auth, req, opts); errExecute != nil { t.Fatalf("Execute() error = %v", errExecute) }
select { case payload := <-capturedPayload: input := gjson.GetBytes(payload, "input").Array() if len(input) != 4 { t.Fatalf("upstream input length = %d, want 4; payload=%s", len(input), payload) } if input[1].Get("type").String() != "reasoning" || input[1].Get("encrypted_content").String() != encryptedContent { t.Fatalf("websocket reasoning replay missing before assistant message: %s", payload) } if input[2].Get("role").String() != "assistant" { t.Fatalf("input.2.role = %q, want assistant; payload=%s", input[2].Get("role").String(), payload) } case <-time.After(5 * time.Second): t.Fatal("timed out waiting for upstream websocket payload") }}
func TestClearCodexReasoningReplayOnWebsocketInvalidSignature(t *testing.T) { internalcache.ClearCodexReasoningReplayCache() t.Cleanup(internalcache.ClearCodexReasoningReplayCache)
scope := codexReasoningReplayScope{modelName: "gpt-5.4", sessionKey: "claude:ws-invalid:agent:main"} encryptedContent := validCodexReasoningEncryptedContentForTestSeed(32) if !internalcache.CacheCodexReasoningReplayItem(scope.modelName, scope.sessionKey, []byte(`{"type":"reasoning","summary":[],"content":null,"encrypted_content":"`+encryptedContent+`"}`)) { t.Fatal("failed to seed websocket replay cache") } payload := []byte(`{"type":"error","status":400,"body":{"error":{"message":"Invalid signature in thinking block","type":"invalid_request_error","code":"invalid_request_error"}}}`) if errClear := clearCodexReasoningReplayOnWebsocketError(context.Background(), scope, payload); errClear != nil { t.Fatalf("clear websocket replay error: %v", errClear) } if _, ok := internalcache.GetCodexReasoningReplayItem(scope.modelName, scope.sessionKey); ok { t.Fatal("websocket invalid signature did not clear replay state") }}
func TestCodexWebsocketsExecuteResponsesLiteDoesNotInjectImageGenerationTool(t *testing.T) { upgrader := websocket.Upgrader{CheckOrigin: func(*http.Request) bool { return true }} capturedPayload := make(chan []byte, 1) server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { conn, err := upgrader.Upgrade(w, r, nil) if err != nil { t.Fatalf("upgrade websocket: %v", err) } defer func() { _ = conn.Close() }()
_, payload, errRead := conn.ReadMessage() if errRead != nil { t.Fatalf("read upstream websocket message: %v", errRead) } capturedPayload <- bytes.Clone(payload)
completed := []byte(`{"type":"response.completed","response":{"id":"resp-1","output":[],"usage":{"input_tokens":0,"output_tokens":0,"total_tokens":0}}}`) if errWrite := conn.WriteMessage(websocket.TextMessage, completed); errWrite != nil { t.Fatalf("write completed websocket message: %v", errWrite) } })) defer server.Close()
exec := NewCodexWebsocketsExecutor(&config.Config{}) auth := &cliproxyauth.Auth{ Provider: "codex", Attributes: map[string]string{ "api_key": "sk-test", "base_url": server.URL, "plan_type": "pro", }, } req := cliproxyexecutor.Request{ Model: "gpt-5.6-sol", Payload: []byte(`{"model":"gpt-5.6-sol","input":[{"type":"additional_tools","role":"developer","tools":[{"type":"custom","name":"exec"}]},{"role":"user","content":"hello"}],"parallel_tool_calls":true,"client_metadata":{"ws_request_header_x_openai_internal_codex_responses_lite":"true"}}`), } opts := cliproxyexecutor.Options{SourceFormat: sdktranslator.FromString("codex")}
if _, err := exec.Execute(context.Background(), auth, req, opts); err != nil { t.Fatalf("Execute() error = %v", err) }
select { case payload := <-capturedPayload: if tools := gjson.GetBytes(payload, "tools"); tools.Exists() { t.Fatalf("unexpected tools in responses-lite upstream payload: %s", tools.Raw) } if got := gjson.GetBytes(payload, "input.0.type").String(); got != "additional_tools" { t.Fatalf("input.0.type = %q, want additional_tools; payload=%s", got, payload) } if got := gjson.GetBytes(payload, "client_metadata.ws_request_header_x_openai_internal_codex_responses_lite").String(); got != "true" { t.Fatalf("responses-lite metadata = %q, want true; payload=%s", got, payload) } parallelToolCalls := gjson.GetBytes(payload, "parallel_tool_calls") if !parallelToolCalls.Exists() || parallelToolCalls.Bool() { t.Fatalf("responses-lite parallel_tool_calls should be false: %s", payload) } case <-time.After(5 * time.Second): t.Fatal("timed out waiting for upstream websocket payload") }}
func TestCodexWebsocketsExecuteStreamResponsesLiteForcesParallelToolCallsFalse(t *testing.T) { upgrader := websocket.Upgrader{CheckOrigin: func(*http.Request) bool { return true }} capturedPayload := make(chan []byte, 1) server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { conn, errUpgrade := upgrader.Upgrade(w, r, nil) if errUpgrade != nil { t.Errorf("upgrade websocket: %v", errUpgrade) return } defer func() { _ = conn.Close() }()
_, payload, errRead := conn.ReadMessage() if errRead != nil { t.Errorf("read upstream websocket message: %v", errRead) return } capturedPayload <- bytes.Clone(payload)
completed := []byte(`{"type":"response.completed","response":{"id":"resp-1","output":[],"usage":{"input_tokens":0,"output_tokens":0,"total_tokens":0}}}`) if errWrite := conn.WriteMessage(websocket.TextMessage, completed); errWrite != nil { t.Errorf("write completed websocket message: %v", errWrite) } })) defer server.Close()
exec := NewCodexWebsocketsExecutor(&config.Config{}) auth := &cliproxyauth.Auth{ Provider: "codex", Attributes: map[string]string{ "api_key": "sk-test", "base_url": server.URL, "plan_type": "pro", }, } req := cliproxyexecutor.Request{ Model: "gpt-5.6-luna", Payload: []byte(`{"model":"gpt-5.6-luna","input":[{"type":"additional_tools","role":"developer","tools":[{"type":"custom","name":"exec"}]},{"role":"user","content":"hello"}],"parallel_tool_calls":true,"client_metadata":{"ws_request_header_x_openai_internal_codex_responses_lite":"true"}}`), } opts := cliproxyexecutor.Options{SourceFormat: sdktranslator.FromString("codex")}
result, errExecute := exec.ExecuteStream(context.Background(), auth, req, opts) if errExecute != nil { t.Fatalf("ExecuteStream() error = %v", errExecute) } streamComplete := false for !streamComplete { select { case chunk, ok := <-result.Chunks: if !ok { streamComplete = true continue } if chunk.Err != nil { t.Fatalf("stream chunk error = %v", chunk.Err) } case <-time.After(5 * time.Second): t.Fatal("timed out waiting for websocket stream completion") } }
select { case payload := <-capturedPayload: parallelToolCalls := gjson.GetBytes(payload, "parallel_tool_calls") if !parallelToolCalls.Exists() || parallelToolCalls.Bool() { t.Fatalf("responses-lite parallel_tool_calls should be false: %s", payload) } case <-time.After(5 * time.Second): t.Fatal("timed out waiting for upstream websocket payload") }}
func TestCodexWebsocketsExecutePreservesPreviousResponseIDUpstream(t *testing.T) { upgrader := websocket.Upgrader{CheckOrigin: func(*http.Request) bool { return true }} capturedPayload := make(chan []byte, 1) server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { if r.URL.Path != "/responses" { t.Fatalf("request path = %s, want /responses", r.URL.Path) } conn, err := upgrader.Upgrade(w, r, nil) if err != nil { t.Fatalf("upgrade websocket: %v", err) } defer func() { _ = conn.Close() }()
msgType, payload, err := conn.ReadMessage() if err != nil { t.Fatalf("read upstream websocket message: %v", err) } if msgType != websocket.TextMessage { t.Fatalf("message type = %d, want text", msgType) } capturedPayload <- bytes.Clone(payload)
completed := []byte(`{"type":"response.completed","response":{"id":"resp-2","output":[],"usage":{"input_tokens":0,"output_tokens":0,"total_tokens":0}}}`) if errWrite := conn.WriteMessage(websocket.TextMessage, completed); errWrite != nil { t.Fatalf("write completed websocket message: %v", errWrite) } })) defer server.Close()
exec := NewCodexWebsocketsExecutor(&config.Config{SDKConfig: config.SDKConfig{DisableImageGeneration: config.DisableImageGenerationAll}}) auth := &cliproxyauth.Auth{Attributes: map[string]string{"api_key": "sk-test", "base_url": server.URL}} req := cliproxyexecutor.Request{ Model: "gpt-5-codex", Payload: []byte(`{"model":"gpt-5-codex","previous_response_id":"resp-1","input":[{"type":"message","id":"msg-1"}]}`), } opts := cliproxyexecutor.Options{SourceFormat: sdktranslator.FromString("codex")}
if _, err := exec.Execute(context.Background(), auth, req, opts); err != nil { t.Fatalf("Execute() error = %v", err) }
select { case payload := <-capturedPayload: if got := gjson.GetBytes(payload, "type").String(); got != "response.create" { t.Fatalf("upstream type = %s, want response.create; payload=%s", got, payload) } if got := gjson.GetBytes(payload, "previous_response_id").String(); got != "resp-1" { t.Fatalf("upstream previous_response_id = %s, want resp-1; payload=%s", got, payload) } case <-time.After(5 * time.Second): t.Fatal("timed out waiting for upstream websocket payload") }}
func TestCodexWebsocketsExecuteStreamUpgradeRequiredReturnsWithoutLockingSession(t *testing.T) { upgradeAttempts := make(chan struct{}, 2) server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { if !strings.EqualFold(r.Header.Get("Upgrade"), "websocket") { t.Errorf("unexpected HTTP fallback request: %s %s", r.Method, r.URL.Path) w.WriteHeader(http.StatusInternalServerError) return } upgradeAttempts <- struct{}{} w.WriteHeader(http.StatusUpgradeRequired) _, _ = w.Write([]byte(`{"error":{"message":"websocket unavailable"}}`)) })) defer server.Close()
exec := NewCodexWebsocketsExecutor(&config.Config{SDKConfig: config.SDKConfig{DisableImageGeneration: config.DisableImageGenerationAll}}) const executionSessionID = "ws-upgrade-required-session" t.Cleanup(func() { exec.CloseExecutionSession(executionSessionID) }) auth := &cliproxyauth.Auth{ ID: "codex-test", Provider: "codex", Attributes: map[string]string{ "api_key": "sk-test", "base_url": server.URL, }, } opts := cliproxyexecutor.Options{ SourceFormat: sdktranslator.FromString("openai-response"), ResponseFormat: sdktranslator.FromString("openai-response"), Metadata: map[string]any{ cliproxyexecutor.ExecutionSessionMetadataKey: executionSessionID, }, } ctx := cliproxyexecutor.WithDownstreamWebsocket(context.Background())
execute := func(payload string) { t.Helper() done := make(chan error, 1) go func() { _, errExecute := exec.ExecuteStream(ctx, auth, cliproxyexecutor.Request{ Model: "gpt-5.4", Payload: []byte(payload), }, opts) done <- errExecute }()
select { case errExecute := <-done: if errExecute == nil { t.Fatal("upgrade-required error = nil") } statusErr, ok := errExecute.(interface{ StatusCode() int }) if !ok || statusErr.StatusCode() != http.StatusUpgradeRequired { t.Fatalf("upgrade-required error = %T %v, want status 426", errExecute, errExecute) } case <-time.After(5 * time.Second): t.Fatal("timed out waiting for upgrade-required error; execution session may still be locked") } }
execute(`{"model":"gpt-5.4","generate":false,"input":[]}`) execute(`{"model":"gpt-5.4","previous_response_id":"resp-1","input":[{"type":"message","id":"msg-2"}]}`)
if got := len(upgradeAttempts); got != 2 { t.Fatalf("websocket upgrade attempts = %d, want 2", got) }}
func TestCodexWebsocketsExecuteStreamHandshakeErrorReturnsWithoutLockingSession(t *testing.T) { server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { w.WriteHeader(http.StatusUnauthorized) _, _ = w.Write([]byte(`{"error":{"message":"unauthorized"}}`)) })) defer server.Close()
exec := NewCodexWebsocketsExecutor(&config.Config{SDKConfig: config.SDKConfig{DisableImageGeneration: config.DisableImageGenerationAll}}) const executionSessionID = "ws-handshake-error-session" t.Cleanup(func() { exec.CloseExecutionSession(executionSessionID) }) auth := &cliproxyauth.Auth{ ID: "codex-test", Provider: "codex", Attributes: map[string]string{ "api_key": "sk-test", "base_url": server.URL, }, } opts := cliproxyexecutor.Options{ SourceFormat: sdktranslator.FromString("openai-response"), Metadata: map[string]any{ cliproxyexecutor.ExecutionSessionMetadataKey: executionSessionID, }, }
for i := 0; i < 2; i++ { attemptCtx := cliproxyexecutor.WithUpstreamAttemptTracker(context.Background()) done := make(chan error, 1) go func(ctx context.Context) { _, errExecute := exec.ExecuteStream(ctx, auth, cliproxyexecutor.Request{ Model: "gpt-5.4", Payload: []byte(`{"model":"gpt-5.4","input":[{"type":"message","id":"msg-1"}]}`), }, opts) done <- errExecute }(attemptCtx) select { case errExecute := <-done: statusErr, ok := errExecute.(interface{ StatusCode() int }) if !ok || statusErr.StatusCode() != http.StatusUnauthorized { t.Fatalf("attempt %d error = %T %v, want status 401", i+1, errExecute, errExecute) } if !cliproxyexecutor.UpstreamAttempted(attemptCtx) { t.Fatalf("attempt %d websocket handshake was not marked as upstream", i+1) } case <-time.After(5 * time.Second): t.Fatalf("attempt %d timed out; execution session remained locked", i+1) } }}
func TestExistingWebsocketSessionConnRequiresMatchingHealthyConnection(t *testing.T) { conn := &websocket.Conn{} closer := newWebsocketConnectionCloser(conn) sess := &codexWebsocketSession{ conn: conn, connCloser: closer, authID: "auth-a", wsURL: "ws://example.test/responses", } sess.resetUpstreamDisconnectError(conn) if gotConn, gotCloser := existingWebsocketSessionConn(sess, "auth-a", "ws://example.test/responses"); gotConn != conn || gotCloser != closer { t.Fatal("matching healthy websocket session was not reusable") } if got, _ := existingWebsocketSessionConn(sess, "auth-b", "ws://example.test/responses"); got != nil { t.Fatal("websocket session matched a different auth") } if got, _ := existingWebsocketSessionConn(sess, "auth-a", "ws://other.test/responses"); got != nil { t.Fatal("websocket session matched a different URL") } sess.setUpstreamDisconnectError(conn, errors.New("upstream disconnected")) if got, _ := existingWebsocketSessionConn(sess, "auth-a", "ws://example.test/responses"); got != nil { t.Fatal("disconnected websocket session remained reusable") }}
func TestCodexAutoExecutorRequiredUpstreamWebsocketRejectsHTTPFallback(t *testing.T) { exec := NewCodexAutoExecutor(&config.Config{SDKConfig: config.SDKConfig{DisableImageGeneration: config.DisableImageGenerationAll}}) auth := &cliproxyauth.Auth{ ID: "codex-http-only", Provider: "codex", Attributes: map[string]string{ "api_key": "sk-test", }, } ctx := cliproxyexecutor.WithRequiredUpstreamWebsocket( cliproxyexecutor.WithDownstreamWebsocket(context.Background()), ) _, errExecute := exec.ExecuteStream(ctx, auth, cliproxyexecutor.Request{ Model: "gpt-5.4", Payload: []byte(`{"model":"gpt-5.4","previous_response_id":"resp-1","input":[{"type":"message","id":"msg-2"}]}`), }, cliproxyexecutor.Options{SourceFormat: sdktranslator.FromString("openai-response")}) if errExecute == nil { t.Fatal("ExecuteStream() error = nil, want replay-required error") } statusErr, ok := errExecute.(interface{ StatusCode() int }) if !ok || statusErr.StatusCode() != http.StatusUpgradeRequired { t.Fatalf("ExecuteStream() error = %T %v, want status 426", errExecute, errExecute) } if got := gjson.Get(errExecute.Error(), "error.code").String(); got != "upstream_http_replay_required" { t.Fatalf("ExecuteStream() error code = %q, want upstream_http_replay_required", got) } requestScoped, ok := errExecute.(cliproxyexecutor.RequestScopedError) if !ok || !requestScoped.IsRequestScoped() { t.Fatalf("ExecuteStream() error = %T, want request-scoped replay signal", errExecute) }}
func TestCodexWebsocketUpgradeFallbackLocalErrorDoesNotMarkUpstreamAttempt(t *testing.T) { tests := []struct { name string execute func(*CodexWebsocketsExecutor, context.Context, *cliproxyauth.Auth, cliproxyexecutor.Request, cliproxyexecutor.Options) error }{ { name: "non-stream", execute: func(exec *CodexWebsocketsExecutor, ctx context.Context, auth *cliproxyauth.Auth, req cliproxyexecutor.Request, opts cliproxyexecutor.Options) error { _, errExecute := exec.Execute(ctx, auth, req, opts) return errExecute }, }, { name: "stream", execute: func(exec *CodexWebsocketsExecutor, ctx context.Context, auth *cliproxyauth.Auth, req cliproxyexecutor.Request, opts cliproxyexecutor.Options) error { _, errExecute := exec.ExecuteStream(ctx, auth, req, opts) return errExecute }, }, }
for _, tc := range tests { t.Run(tc.name, func(t *testing.T) { var upgradeAttempts atomic.Int32 server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { if !strings.EqualFold(r.Header.Get("Upgrade"), "websocket") { t.Errorf("unexpected HTTP fallback request: %s %s", r.Method, r.URL.Path) w.WriteHeader(http.StatusInternalServerError) return } upgradeAttempts.Add(1) w.WriteHeader(http.StatusUpgradeRequired) _, _ = w.Write([]byte(`{"error":{"message":"websocket unavailable"}}`)) })) defer server.Close()
exec := NewCodexWebsocketsExecutor(&config.Config{SDKConfig: config.SDKConfig{DisableImageGeneration: config.DisableImageGenerationAll}}) auth := &cliproxyauth.Auth{ ID: "codex-fallback-local-error", Provider: "codex", Attributes: map[string]string{ "api_key": "sk-test", "base_url": server.URL, }, } ctx := cliproxyexecutor.WithUpstreamAttemptTracker(context.Background()) opts := codexOpenAIImageTestOptions(codexImagesGenerationsPath, tc.name == "stream") errExecute := tc.execute(exec, ctx, auth, cliproxyexecutor.Request{ Model: "gpt-5.4", Payload: []byte("not-json"), }, opts) if errExecute == nil || !strings.Contains(errExecute.Error(), "invalid OpenAI image generation request JSON") { t.Fatalf("Execute() error = %v, want local image request validation error", errExecute) } if got := upgradeAttempts.Load(); got != 1 { t.Fatalf("websocket upgrade attempts = %d, want 1", got) } if cliproxyexecutor.UpstreamAttempted(ctx) { t.Fatal("transparent 426 fallback marked a local HTTP preparation error as an upstream attempt") } }) }}
func TestCodexWebsocketMissingRequiredSessionDoesNotMarkUpstreamAttempt(t *testing.T) { exec := NewCodexWebsocketsExecutor(&config.Config{SDKConfig: config.SDKConfig{DisableImageGeneration: config.DisableImageGenerationAll}}) exec.store = &codexWebsocketSessionStore{sessions: make(map[string]*codexWebsocketSession)} auth := &cliproxyauth.Auth{ ID: "codex-required-session", Provider: "codex", Attributes: map[string]string{ "api_key": "sk-test", }, } ctx := cliproxyexecutor.WithUpstreamAttemptTracker( cliproxyexecutor.WithRequiredUpstreamWebsocket(context.Background()), ) _, errExecute := exec.Execute(ctx, auth, cliproxyexecutor.Request{ Model: "gpt-5.4", Payload: []byte(`{"model":"gpt-5.4","previous_response_id":"resp-1","input":[{"type":"message","role":"user","content":"hello"}]}`), }, cliproxyexecutor.Options{ SourceFormat: sdktranslator.FormatOpenAIResponse, Metadata: map[string]any{ cliproxyexecutor.ExecutionSessionMetadataKey: "missing-codex-session", }, }) if !cliproxyexecutor.IsUpstreamWebsocketReplayRequired(errExecute) { t.Fatalf("Execute() error = %T %v, want replay-required", errExecute, errExecute) } if cliproxyexecutor.UpstreamAttempted(ctx) { t.Fatal("missing retained websocket connection was marked as an upstream attempt") }}
func TestCodexWebsocketsExecuteStreamPassesThroughUpstreamWebsocketPayloadForDownstreamWebsocket(t *testing.T) { upgrader := websocket.Upgrader{CheckOrigin: func(*http.Request) bool { return true }} capturedPayload := make(chan []byte, 1) delta := []byte(`{"type":"response.output_text.delta","delta":"hello"}`) completed := []byte(`{"type":"response.completed","response":{"id":"resp-1","output":[],"usage":{"input_tokens":0,"output_tokens":0,"total_tokens":0}}}`) server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { conn, err := upgrader.Upgrade(w, r, nil) if err != nil { t.Errorf("upgrade websocket: %v", err) return } defer func() { _ = conn.Close() }()
_, payload, errRead := conn.ReadMessage() if errRead != nil { t.Errorf("read upstream websocket message: %v", errRead) return } capturedPayload <- bytes.Clone(payload) if errWrite := conn.WriteMessage(websocket.TextMessage, delta); errWrite != nil { t.Errorf("write delta websocket message: %v", errWrite) return } if errWrite := conn.WriteMessage(websocket.TextMessage, completed); errWrite != nil { t.Errorf("write completed websocket message: %v", errWrite) return } })) defer server.Close()
exec := NewCodexWebsocketsExecutor(&config.Config{SDKConfig: config.SDKConfig{DisableImageGeneration: config.DisableImageGenerationAll}}) auth := &cliproxyauth.Auth{Attributes: map[string]string{"api_key": "sk-test", "base_url": server.URL}} req := cliproxyexecutor.Request{ Model: "gpt-5-codex", Payload: []byte(`{"model":"prolite/gpt-5-codex","input":[{"type":"additional_tools","role":"developer","tools":[{"type":"custom","name":"exec"}]},{"type":"message","role":"user","content":"hello"}],"parallel_tool_calls":true}`), } opts := cliproxyexecutor.Options{ SourceFormat: sdktranslator.FromString("openai-response"), ResponseFormat: sdktranslator.FromString("openai-response"), } ctx := cliproxyexecutor.WithUpstreamAttemptTracker( cliproxyexecutor.WithDownstreamWebsocket(context.Background()), )
result, err := exec.ExecuteStream(ctx, auth, req, opts) if err != nil { t.Fatalf("ExecuteStream() error = %v", err) } if !cliproxyexecutor.UpstreamAttempted(ctx) { t.Fatal("websocket write did not mark an upstream attempt") }
select { case chunk, ok := <-result.Chunks: if !ok { t.Fatal("stream closed before first chunk") } if chunk.Err != nil { t.Fatalf("first chunk error = %v", chunk.Err) } if !bytes.Equal(bytes.TrimSpace(chunk.Payload), delta) { t.Fatalf("first chunk = %q, want raw upstream websocket payload %q", chunk.Payload, delta) } case <-time.After(5 * time.Second): t.Fatal("timed out waiting for first stream chunk") }
select { case payload := <-capturedPayload: if got := gjson.GetBytes(payload, "model").String(); got != "gpt-5-codex" { t.Fatalf("upstream model = %s, want gpt-5-codex; payload=%s", got, payload) } parallelToolCalls := gjson.GetBytes(payload, "parallel_tool_calls") if !parallelToolCalls.Exists() || !parallelToolCalls.Bool() { t.Fatalf("non-lite parallel_tool_calls should be preserved: %s", payload) } case <-time.After(5 * time.Second): t.Fatal("timed out waiting for upstream websocket payload") }}
func TestCodexWebsocketsExecuteStreamPropagatesUpstreamErrorForDownstreamWebsocket(t *testing.T) { upgrader := websocket.Upgrader{CheckOrigin: func(*http.Request) bool { return true }} errorPayload := []byte(`{"type":"error","status":429,"error":{"code":"websocket_connection_limit_reached","message":"too many websockets"}}`) server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { conn, err := upgrader.Upgrade(w, r, nil) if err != nil { t.Errorf("upgrade websocket: %v", err) return } defer func() { _ = conn.Close() }()
if _, _, errRead := conn.ReadMessage(); errRead != nil { t.Errorf("read upstream websocket message: %v", errRead) return } if errWrite := conn.WriteMessage(websocket.TextMessage, errorPayload); errWrite != nil { t.Errorf("write error websocket message: %v", errWrite) return } })) defer server.Close()
exec := NewCodexWebsocketsExecutor(&config.Config{SDKConfig: config.SDKConfig{DisableImageGeneration: config.DisableImageGenerationAll}}) auth := &cliproxyauth.Auth{Attributes: map[string]string{"api_key": "sk-test", "base_url": server.URL}} req := cliproxyexecutor.Request{ Model: "gpt-5-codex", Payload: []byte(`{"model":"gpt-5-codex","input":[{"type":"message","role":"user","content":"hello"}]}`), } opts := cliproxyexecutor.Options{ SourceFormat: sdktranslator.FromString("openai-response"), ResponseFormat: sdktranslator.FromString("openai-response"), } ctx := cliproxyexecutor.WithDownstreamWebsocket(context.Background())
result, err := exec.ExecuteStream(ctx, auth, req, opts) if err != nil { t.Fatalf("ExecuteStream() error = %v", err) }
select { case chunk, ok := <-result.Chunks: if !ok { t.Fatal("stream closed before error chunk") } if len(bytes.TrimSpace(chunk.Payload)) != 0 { t.Fatalf("error chunk payload = %q, want empty", chunk.Payload) } if chunk.Err == nil { t.Fatal("error chunk Err = nil, want upstream error") } statusErr, ok := chunk.Err.(interface{ StatusCode() int }) if !ok { t.Fatalf("error type %T does not expose StatusCode", chunk.Err) } if got := statusErr.StatusCode(); got != http.StatusTooManyRequests { t.Fatalf("status = %d, want %d", got, http.StatusTooManyRequests) } case <-time.After(5 * time.Second): t.Fatal("timed out waiting for error stream chunk") }}
func TestSendTerminalWebsocketReadInvalidatesBeforeWaitingForCapacity(t *testing.T) { terminalErr := &websocket.CloseError{Code: websocket.CloseMessageTooBig}
t.Run("available channel keeps fast path ordering", func(t *testing.T) { ch := make(chan codexWebsocketRead, 1) done := make(chan struct{}) invalidateCalls := 0 invalidated := sendTerminalWebsocketRead(ch, done, codexWebsocketRead{err: terminalErr}, func() { invalidateCalls++ }) if invalidated { t.Fatal("available channel should not invalidate before delivery") } if invalidateCalls != 0 { t.Fatalf("invalidate calls = %d, want 0", invalidateCalls) } event := <-ch if !errors.Is(event.err, terminalErr) { t.Fatalf("terminal error = %v, want %v", event.err, terminalErr) } })
t.Run("full channel invalidates before waiting", func(t *testing.T) { ch := make(chan codexWebsocketRead, 1) ch <- codexWebsocketRead{payload: []byte("queued")} done := make(chan struct{}) invalidateCalled := make(chan struct{}) result := make(chan bool, 1)
go func() { result <- sendTerminalWebsocketRead(ch, done, codexWebsocketRead{err: terminalErr}, func() { close(invalidateCalled) }) }()
select { case <-invalidateCalled: case <-time.After(time.Second): t.Fatal("invalidation did not happen before waiting for channel capacity") } select { case <-result: t.Fatal("terminal sender returned before capacity was released") default: }
<-ch select { case event := <-ch: if !errors.Is(event.err, terminalErr) { t.Fatalf("terminal error = %v, want %v", event.err, terminalErr) } case <-time.After(time.Second): t.Fatal("timed out waiting for terminal read") } select { case invalidated := <-result: if !invalidated { t.Fatal("full channel should report early invalidation") } case <-time.After(time.Second): t.Fatal("terminal sender did not finish") } })
t.Run("full channel stops when invalidation cancels active read", func(t *testing.T) { ch := make(chan codexWebsocketRead, 1) ch <- codexWebsocketRead{payload: []byte("queued")} done := make(chan struct{}) invalidated := sendTerminalWebsocketRead(ch, done, codexWebsocketRead{err: terminalErr}, func() { close(done) }) if !invalidated { t.Fatal("full channel should report early invalidation") } if len(ch) != 1 { t.Fatalf("channel length = %d, want queued payload only", len(ch)) } })}
func TestMapCodexWebsocketWriteErrorStopsRetryForMessageTooBig(t *testing.T) { networkWriteErr := errors.New("write: broken pipe") tests := []struct { name string closeCode int writeErr error wantStatus int wantRetry bool }{ { name: "close sent after message too big is request scoped", closeCode: websocket.CloseMessageTooBig, writeErr: websocket.ErrCloseSent, wantStatus: http.StatusRequestEntityTooLarge, wantRetry: false, }, { name: "network write error after message too big is request scoped", closeCode: websocket.CloseMessageTooBig, writeErr: networkWriteErr, wantStatus: http.StatusRequestEntityTooLarge, wantRetry: false, }, { name: "other close keeps stale connection retry", closeCode: websocket.CloseNormalClosure, writeErr: websocket.ErrCloseSent, wantRetry: true, }, }
for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { sess := &codexWebsocketSession{} conn := &websocket.Conn{} sess.resetUpstreamDisconnectError(conn) sess.setUpstreamDisconnectError(conn, &websocket.CloseError{Code: tt.closeCode})
mappedErr := mapCodexWebsocketWriteError(sess, conn, tt.writeErr) if got := shouldRetryCodexWebsocketSend(mappedErr); got != tt.wantRetry { t.Fatalf("shouldRetryCodexWebsocketSend() = %v, want %v; err=%v", got, tt.wantRetry, mappedErr) } if tt.wantStatus == 0 { if !errors.Is(mappedErr, tt.writeErr) { t.Fatalf("mapped error = %v, want %v", mappedErr, tt.writeErr) } return } statusErr, ok := mappedErr.(interface{ StatusCode() int }) if !ok || statusErr.StatusCode() != tt.wantStatus { t.Fatalf("mapped status = %v, want %d; err=%v", statusErr, tt.wantStatus, mappedErr) } requestErr, ok := mappedErr.(interface{ IsRequestScoped() bool }) if !ok || !requestErr.IsRequestScoped() { t.Fatalf("mapped error should be request scoped, got %T", mappedErr) } }) }}
func TestMapCodexWebsocketWriteErrorDoesNotReusePriorConnectionClose(t *testing.T) { sess := &codexWebsocketSession{} priorConn := &websocket.Conn{} replacementConn := &websocket.Conn{}
sess.resetUpstreamDisconnectError(priorConn) sess.setUpstreamDisconnectError(priorConn, &websocket.CloseError{Code: websocket.CloseMessageTooBig}) priorErr := mapCodexWebsocketWriteError(sess, priorConn, websocket.ErrCloseSent) if shouldRetryCodexWebsocketSend(priorErr) { t.Fatalf("prior connection 1009 should not retry, got %v", priorErr) }
sess.resetUpstreamDisconnectError(replacementConn) // A late close callback from the prior connection must not overwrite the // replacement connection's close state. sess.setUpstreamDisconnectError(priorConn, &websocket.CloseError{Code: websocket.CloseMessageTooBig}) sess.setUpstreamDisconnectError(replacementConn, &websocket.CloseError{Code: websocket.CloseNormalClosure}) replacementErr := mapCodexWebsocketWriteError(sess, replacementConn, websocket.ErrCloseSent) if !errors.Is(replacementErr, websocket.ErrCloseSent) { t.Fatalf("replacement connection error = %v, want %v", replacementErr, websocket.ErrCloseSent) } if !shouldRetryCodexWebsocketSend(replacementErr) { t.Fatalf("replacement connection should keep stale-connection retry, got %v", replacementErr) }}
func TestCodexWebsocketsExecuteStreamMapsMessageTooBigClose(t *testing.T) { upgrader := websocket.Upgrader{CheckOrigin: func(*http.Request) bool { return true }} server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { conn, err := upgrader.Upgrade(w, r, nil) if err != nil { t.Errorf("upgrade websocket: %v", err) return } defer func() { _ = conn.Close() }()
if _, _, errRead := conn.ReadMessage(); errRead != nil { t.Errorf("read upstream websocket message: %v", errRead) return } deadline := time.Now().Add(time.Second) closeMessage := websocket.FormatCloseMessage(websocket.CloseMessageTooBig, "message too big") if errWrite := conn.WriteControl(websocket.CloseMessage, closeMessage, deadline); errWrite != nil { t.Errorf("write close websocket message: %v", errWrite) return } })) defer server.Close()
exec := NewCodexWebsocketsExecutor(&config.Config{SDKConfig: config.SDKConfig{DisableImageGeneration: config.DisableImageGenerationAll}}) auth := &cliproxyauth.Auth{Attributes: map[string]string{"api_key": "sk-test", "base_url": server.URL}} req := cliproxyexecutor.Request{ Model: "gpt-5-codex", Payload: []byte(`{"model":"gpt-5-codex","input":[{"type":"message","role":"user","content":"hello"}]}`), } opts := cliproxyexecutor.Options{ SourceFormat: sdktranslator.FromString("openai-response"), ResponseFormat: sdktranslator.FromString("openai-response"), }
result, err := exec.ExecuteStream(context.Background(), auth, req, opts) if err != nil { t.Fatalf("ExecuteStream() error = %v", err) }
select { case chunk, ok := <-result.Chunks: if !ok { t.Fatal("stream closed before error chunk") } if chunk.Err == nil { t.Fatal("error chunk Err = nil, want message-too-big error") } statusErr, ok := chunk.Err.(interface{ StatusCode() int }) if !ok { t.Fatalf("error type %T does not expose StatusCode", chunk.Err) } if got := statusErr.StatusCode(); got != http.StatusRequestEntityTooLarge { t.Fatalf("status = %d, want %d", got, http.StatusRequestEntityTooLarge) } if got := gjson.Get(chunk.Err.Error(), "error.code").String(); got != "message_too_big" { t.Fatalf("error code = %q, want message_too_big; err=%v", got, chunk.Err) } requestErr, ok := chunk.Err.(interface{ IsRequestScoped() bool }) if !ok || !requestErr.IsRequestScoped() { t.Fatalf("message-too-big error should be request scoped, got %T", chunk.Err) } case <-time.After(5 * time.Second): t.Fatal("timed out waiting for error stream chunk") }}
func TestCodexWebsocketsUpstreamDisconnectChanSignalsOnInvalidate(t *testing.T) { upgrader := websocket.Upgrader{CheckOrigin: func(*http.Request) bool { return true }} server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { conn, err := upgrader.Upgrade(w, r, nil) if err != nil { t.Errorf("upgrade websocket: %v", err) return } defer func() { _ = conn.Close() }() for { if _, _, errRead := conn.ReadMessage(); errRead != nil { return } } })) defer server.Close()
wsURL := "ws" + strings.TrimPrefix(server.URL, "http") conn, _, err := websocket.DefaultDialer.Dial(wsURL, nil) if err != nil { t.Fatalf("dial websocket: %v", err) } defer func() { _ = conn.Close() }()
exec := NewCodexWebsocketsExecutor(&config.Config{}) exec.store = &codexWebsocketSessionStore{sessions: make(map[string]*codexWebsocketSession)} sessionID := "sess-1" disconnectCh := exec.UpstreamDisconnectChan(sessionID) if disconnectCh == nil { t.Fatal("expected disconnect channel") }
sess := exec.getOrCreateSession(sessionID) if sess == nil { t.Fatal("expected session") } sess.connMu.Lock() sess.conn = conn sess.authID = "auth-1" sess.wsURL = "ws://example.test/responses" sess.readerConn = conn sess.connMu.Unlock()
upstreamErr := errors.New("upstream gone") exec.invalidateUpstreamConn(sess, conn, "test_invalidate", upstreamErr)
select { case errRead, ok := <-disconnectCh: if !ok { t.Fatal("expected disconnect channel to deliver error before closing") } if errRead == nil || errRead.Error() != upstreamErr.Error() { t.Fatalf("disconnect error = %v, want %v", errRead, upstreamErr) } case <-time.After(5 * time.Second): t.Fatal("timed out waiting for disconnect signal") }}
func TestApplyCodexWebsocketHeadersDefaultsToCurrentResponsesBeta(t *testing.T) { headers := applyCodexWebsocketHeaders(context.Background(), http.Header{}, nil, "", nil)
if got := headers.Get("OpenAI-Beta"); got != codexResponsesWebsocketBetaHeaderValue { t.Fatalf("OpenAI-Beta = %s, want %s", got, codexResponsesWebsocketBetaHeaderValue) } if got := headers.Get("User-Agent"); got != codexUserAgent { t.Fatalf("User-Agent = %s, want %s", got, codexUserAgent) } if !strings.HasPrefix(codexUserAgent, codexOriginator+"/") { t.Fatalf("default Codex User-Agent = %s, want prefix %s/", codexUserAgent, codexOriginator) } if !strings.HasPrefix(codexUserAgent, "codex-tui/") { t.Fatalf("default Codex User-Agent = %s, want codex-tui prefix", codexUserAgent) } if !strings.Contains(codexUserAgent, "(codex-tui;") { t.Fatalf("default Codex User-Agent = %s, want codex-tui suffix", codexUserAgent) } if got := headers.Get("Originator"); got != codexOriginator { t.Fatalf("Originator = %s, want %s", got, codexOriginator) } if got := headers.Get("Version"); got != "" { t.Fatalf("Version = %q, want empty", got) } if got := headers.Get("x-codex-beta-features"); got != "" { t.Fatalf("x-codex-beta-features = %q, want empty", got) } if got := headers.Get("X-Codex-Turn-Metadata"); got != "" { t.Fatalf("X-Codex-Turn-Metadata = %q, want empty", got) } if got := headers.Get("X-Client-Request-Id"); got != "" { t.Fatalf("X-Client-Request-Id = %q, want empty", got) }}
func TestApplyCodexWebsocketHeadersDefaultsToCodexCloaking(t *testing.T) { tests := []struct { name string auth *cliproxyauth.Auth token string }{ { name: "OAuth", auth: &cliproxyauth.Auth{ Provider: "codex", Attributes: map[string]string{ "header:User-Agent": "custom-ua", "header:Originator": "custom-origin", }, }, }, { name: "API key", auth: &cliproxyauth.Auth{ Provider: "codex", Attributes: map[string]string{ "api_key": "sk-test", "header:User-Agent": "custom-ua", "header:Originator": "custom-origin", }, }, token: "sk-test", }, }
for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { cfg := &config.Config{ CodexHeaderDefaults: config.CodexHeaderDefaults{UserAgent: "config-ua"}, } ctx := contextWithGinHeaders(map[string]string{ "User-Agent": "client-ua", "Originator": "client-origin", }) headers := http.Header{} headers.Set("User-Agent", "existing-ua") headers.Set("Originator", "existing-origin")
headers = applyCodexWebsocketHeaders(ctx, headers, tt.auth, tt.token, cfg)
if got := headers.Get("User-Agent"); got != codexUserAgent { t.Fatalf("User-Agent = %q, want %q", got, codexUserAgent) } if got := headers.Get("Originator"); got != codexOriginator { t.Fatalf("Originator = %q, want %q", got, codexOriginator) } }) }}
func TestApplyCodexWebsocketHeadersPassesThroughClientIdentityHeadersWhenCloakingDisabled(t *testing.T) { cfg := &config.Config{Codex: config.CodexConfig{DisableCodexCloaking: true}} auth := &cliproxyauth.Auth{ Provider: "codex", Metadata: map[string]any{"email": "user@example.com"}, } ctx := contextWithGinHeaders(map[string]string{ "Originator": "Codex Desktop", "User-Agent": "codex_cli_rs/0.1.0", "Version": "0.115.0-alpha.27", "X-Codex-Turn-Metadata": `{"turn_id":"turn-1"}`, "X-Client-Request-Id": "019d2233-e240-7162-992d-38df0a2a0e0d", "session-id": "legacy-session", })
headers := applyCodexWebsocketHeaders(ctx, http.Header{}, auth, "", cfg)
if got := headers.Get("Originator"); got != "Codex Desktop" { t.Fatalf("Originator = %s, want %s", got, "Codex Desktop") } if got := headers.Get("User-Agent"); got != "codex_cli_rs/0.1.0" { t.Fatalf("User-Agent = %s, want %s", got, "codex_cli_rs/0.1.0") } if got := headers.Get("Version"); got != "0.115.0-alpha.27" { t.Fatalf("Version = %s, want %s", got, "0.115.0-alpha.27") } if got := headers.Get("X-Codex-Turn-Metadata"); got != `{"turn_id":"turn-1"}` { t.Fatalf("X-Codex-Turn-Metadata = %s, want %s", got, `{"turn_id":"turn-1"}`) } if got := headers.Get("X-Client-Request-Id"); got != "019d2233-e240-7162-992d-38df0a2a0e0d" { t.Fatalf("X-Client-Request-Id = %s, want %s", got, "019d2233-e240-7162-992d-38df0a2a0e0d") } if got := headers["session_id"]; len(got) != 1 || got[0] != "legacy-session" { t.Fatalf("session_id = %#v, want [legacy-session]", got) } if got := headers.Get("Session-Id"); got != "" { t.Fatalf("Session-Id = %s, want empty", got) }}
func TestApplyCodexWebsocketHeadersCanonicalizesLegacyUnderscoreSessionHeader(t *testing.T) { auth := &cliproxyauth.Auth{ Provider: "codex", Metadata: map[string]any{"email": "user@example.com"}, } ctx := contextWithGinHeaders(map[string]string{ "Originator": "Codex Desktop", "User-Agent": "codex_cli_rs/0.1.0", "Session_id": "legacy-underscore-session", })
headers := applyCodexWebsocketHeaders(ctx, http.Header{}, auth, "", nil)
if got := headers["session_id"]; len(got) != 1 || got[0] != "legacy-underscore-session" { t.Fatalf("session_id = %#v, want [legacy-underscore-session]", got) } if got := headers.Get("Session-Id"); got != "" { t.Fatalf("Session-Id = %s, want empty", got) }}
func TestApplyCodexWebsocketHeadersUsesConfigDefaultsForOAuth(t *testing.T) { cfg := &config.Config{ Codex: config.CodexConfig{DisableCodexCloaking: true}, CodexHeaderDefaults: config.CodexHeaderDefaults{ UserAgent: "my-codex-client/1.0", BetaFeatures: "feature-a,feature-b", }, } auth := &cliproxyauth.Auth{ Provider: "codex", Metadata: map[string]any{"email": "user@example.com"}, }
headers := applyCodexWebsocketHeaders(context.Background(), http.Header{}, auth, "", cfg)
if got := headers.Get("User-Agent"); got != "my-codex-client/1.0" { t.Fatalf("User-Agent = %s, want %s", got, "my-codex-client/1.0") } if got := headers.Get("x-codex-beta-features"); got != "feature-a,feature-b" { t.Fatalf("x-codex-beta-features = %s, want %s", got, "feature-a,feature-b") } if got := headers.Get("OpenAI-Beta"); got != codexResponsesWebsocketBetaHeaderValue { t.Fatalf("OpenAI-Beta = %s, want %s", got, codexResponsesWebsocketBetaHeaderValue) }}
func TestApplyCodexWebsocketHeadersPrefersExistingHeadersOverClientAndConfig(t *testing.T) { cfg := &config.Config{ Codex: config.CodexConfig{DisableCodexCloaking: true}, CodexHeaderDefaults: config.CodexHeaderDefaults{ UserAgent: "config-ua", BetaFeatures: "config-beta", }, } auth := &cliproxyauth.Auth{ Provider: "codex", Metadata: map[string]any{"email": "user@example.com"}, } ctx := contextWithGinHeaders(map[string]string{ "User-Agent": "client-ua", "X-Codex-Beta-Features": "client-beta", }) headers := http.Header{} headers.Set("User-Agent", "existing-ua") headers.Set("X-Codex-Beta-Features", "existing-beta")
got := applyCodexWebsocketHeaders(ctx, headers, auth, "", cfg)
if gotVal := got.Get("User-Agent"); gotVal != "existing-ua" { t.Fatalf("User-Agent = %s, want %s", gotVal, "existing-ua") } if gotVal := got.Get("x-codex-beta-features"); gotVal != "existing-beta" { t.Fatalf("x-codex-beta-features = %s, want %s", gotVal, "existing-beta") }}
func TestApplyCodexWebsocketHeadersConfigUserAgentOverridesClientHeader(t *testing.T) { cfg := &config.Config{ Codex: config.CodexConfig{DisableCodexCloaking: true}, CodexHeaderDefaults: config.CodexHeaderDefaults{ UserAgent: "config-ua", BetaFeatures: "config-beta", }, } auth := &cliproxyauth.Auth{ Provider: "codex", Metadata: map[string]any{"email": "user@example.com"}, } ctx := contextWithGinHeaders(map[string]string{ "User-Agent": "client-ua", "X-Codex-Beta-Features": "client-beta", })
headers := applyCodexWebsocketHeaders(ctx, http.Header{}, auth, "", cfg)
if got := headers.Get("User-Agent"); got != "config-ua" { t.Fatalf("User-Agent = %s, want %s", got, "config-ua") } if got := headers.Get("x-codex-beta-features"); got != "client-beta" { t.Fatalf("x-codex-beta-features = %s, want %s", got, "client-beta") }}
func TestApplyCodexWebsocketHeadersIgnoresConfigForAPIKeyAuth(t *testing.T) { cfg := &config.Config{ Codex: config.CodexConfig{DisableCodexCloaking: true}, CodexHeaderDefaults: config.CodexHeaderDefaults{ UserAgent: "config-ua", BetaFeatures: "config-beta", }, } auth := &cliproxyauth.Auth{ Provider: "codex", Attributes: map[string]string{"api_key": "sk-test"}, }
headers := applyCodexWebsocketHeaders(context.Background(), http.Header{}, auth, "sk-test", cfg)
if got := headers.Get("User-Agent"); got != "" { t.Fatalf("User-Agent = %s, want empty", got) } if got := headers.Get("x-codex-beta-features"); got != "" { t.Fatalf("x-codex-beta-features = %q, want empty", got) } if got := headers.Get("Originator"); got != "" { t.Fatalf("Originator = %s, want empty", got) }}
func TestApplyCodexWebsocketHeadersPreservesExplicitAPIKeyUserAgent(t *testing.T) { auth := &cliproxyauth.Auth{Provider: "codex", Attributes: map[string]string{"api_key": "sk-test"}} ctx := contextWithGinHeaders(map[string]string{"User-Agent": "api-key-client/1.0", "Originator": "explicit-origin"})
headers := applyCodexWebsocketHeaders(ctx, http.Header{}, auth, "sk-test", nil)
if got := headers.Get("User-Agent"); got != "api-key-client/1.0" { t.Fatalf("User-Agent = %s, want api-key-client/1.0", got) } if got := headers.Get("Originator"); got != "explicit-origin" { t.Fatalf("Originator = %s, want explicit-origin", got) }}
func TestApplyCodexWebsocketHeadersUsesCanonicalAccountHeader(t *testing.T) { auth := &cliproxyauth.Auth{Provider: "codex", Metadata: map[string]any{"account_id": "acct-1"}}
headers := applyCodexWebsocketHeaders(context.Background(), http.Header{}, auth, "", nil)
if got := headerValueCaseInsensitive(headers, "ChatGPT-Account-ID"); got != "acct-1" { t.Fatalf("ChatGPT-Account-ID = %s, want acct-1", got) } values, ok := headers["ChatGPT-Account-ID"] if !ok { t.Fatalf("expected exact ChatGPT-Account-ID key, got %#v", headers) } if len(values) != 1 || values[0] != "acct-1" { t.Fatalf("ChatGPT-Account-ID values = %#v, want [acct-1]", values) }}
func TestApplyCodexPromptCacheHeadersSetsSessionIDAndLegacyConversation(t *testing.T) { req := cliproxyexecutor.Request{Model: "gpt-5-codex", Payload: []byte(`{"prompt_cache_key":"cache-1"}`)}
_, headers := applyCodexPromptCacheHeaders("openai-response", req, []byte(`{"model":"gpt-5-codex"}`))
if got := headers["session_id"]; len(got) != 1 || got[0] != "cache-1" { t.Fatalf("session_id = %#v, want [cache-1]", got) } if got := headers.Get("Session-Id"); got != "" { t.Fatalf("Session-Id = %s, want empty", got) } if got := headers.Get("Conversation_id"); got != "cache-1" { t.Fatalf("Conversation_id = %s, want cache-1", got) }}
func TestApplyCodexPromptCacheHeadersUsesDerivedSessionUUID(t *testing.T) { t.Parallel()
req := cliproxyexecutor.Request{ Model: "gpt-5-codex", Payload: []byte(`{"input":"hello"}`), Metadata: map[string]any{cliproxyexecutor.DerivedSessionIDMetadataKey: "ctx:v1:derived-root"}, } body, headers := applyCodexPromptCacheHeaders(sdktranslator.FormatInteractions, req, []byte(`{"model":"gpt-5-codex"}`)) cacheKey := gjson.GetBytes(body, "prompt_cache_key").String() if _, errParse := uuid.Parse(cacheKey); errParse != nil { t.Fatalf("prompt_cache_key %q is not a UUID: %v", cacheKey, errParse) } if got := headers["session_id"]; len(got) != 1 || got[0] != cacheKey { t.Fatalf("session_id = %#v, want [%q]", got, cacheKey) } if got := headers.Get("Conversation_id"); got != cacheKey { t.Fatalf("Conversation_id = %q, want %q", got, cacheKey) }}
func TestApplyCodexPromptCacheHeadersKeepsExecutionSessionAcrossIncrementalRoots(t *testing.T) { t.Parallel()
firstReq := cliproxyexecutor.Request{ Model: "gpt-5-codex", Payload: []byte(`{"input":"first"}`), Metadata: map[string]any{ cliproxyexecutor.ExecutionSessionMetadataKey: "connection-1", cliproxyexecutor.DerivedSessionIDMetadataKey: "ctx:v1:first-root", }, } secondReq := cliproxyexecutor.Request{ Model: "gpt-5-codex", Payload: []byte(`{"input":"second"}`), Metadata: map[string]any{ cliproxyexecutor.ExecutionSessionMetadataKey: "connection-1", cliproxyexecutor.DerivedSessionIDMetadataKey: "ctx:v1:second-root", }, } firstBody, _ := applyCodexPromptCacheHeaders(sdktranslator.FormatOpenAIResponse, firstReq, []byte(`{"model":"gpt-5-codex"}`)) secondBody, _ := applyCodexPromptCacheHeaders(sdktranslator.FormatOpenAIResponse, secondReq, []byte(`{"model":"gpt-5-codex"}`)) firstKey := gjson.GetBytes(firstBody, "prompt_cache_key").String() secondKey := gjson.GetBytes(secondBody, "prompt_cache_key").String() if firstKey == "" || firstKey != secondKey { t.Fatalf("incremental websocket roots changed prompt cache key: first=%q second=%q", firstKey, secondKey) }}
func TestApplyCodexPromptCacheHeadersClaudeUsesClaudeCodeSessionID(t *testing.T) { firstReq := cliproxyexecutor.Request{ Model: "gpt-5-codex-claude-ws-cache-session", Payload: []byte(`{ "metadata":{"user_id":"{\"device_id\":\"device-a\",\"account_uuid\":\"\",\"session_id\":\"ws-cache-session-1\"}"}, "messages":[{"role":"user","content":[{"type":"text","text":"first"}]}] }`), } secondReq := cliproxyexecutor.Request{ Model: "gpt-5-codex-claude-ws-cache-session", Payload: []byte(`{ "metadata":{"user_id":"{\"device_id\":\"device-b\",\"account_uuid\":\"\",\"session_id\":\"ws-cache-session-1\"}"}, "messages":[{"role":"user","content":[{"type":"text","text":"next"}]}] }`), }
firstBody, firstHeaders := applyCodexPromptCacheHeaders("claude", firstReq, []byte(`{"model":"gpt-5-codex"}`)) secondBody, secondHeaders := applyCodexPromptCacheHeaders("claude", secondReq, []byte(`{"model":"gpt-5-codex"}`))
firstKey := gjson.GetBytes(firstBody, "prompt_cache_key").String() secondKey := gjson.GetBytes(secondBody, "prompt_cache_key").String() if firstKey == "" { t.Fatalf("first prompt_cache_key is empty; body=%s", string(firstBody)) } if secondKey != firstKey { t.Fatalf("same Claude Code session_id produced different websocket prompt_cache_key: first=%q second=%q", firstKey, secondKey) } if got := firstHeaders["session_id"]; len(got) != 1 || got[0] != firstKey { t.Fatalf("first session_id = %#v, want [%q]", got, firstKey) } if got := secondHeaders["session_id"]; len(got) != 1 || got[0] != firstKey { t.Fatalf("second session_id = %#v, want [%q]", got, firstKey) }}
func TestApplyCodexPromptCacheHeadersClaudeRejectsBareUserID(t *testing.T) { req := cliproxyexecutor.Request{ Model: "gpt-5-codex-claude-ws-cache-bare-user", Payload: []byte(`{"metadata":{"user_id":"same-user-across-chats"},"messages":[{"role":"user","content":[{"type":"text","text":"first"}]}]}`), }
body, headers := applyCodexPromptCacheHeaders("claude", req, []byte(`{"model":"gpt-5-codex"}`))
if got := gjson.GetBytes(body, "prompt_cache_key").String(); got != "" { t.Fatalf("bare metadata.user_id must not create websocket prompt_cache_key, got %q; body=%s", got, string(body)) } if got := headers["session_id"]; len(got) != 0 { t.Fatalf("bare metadata.user_id must not create websocket session_id, got %#v", got) } if got := headers.Get("Session-Id"); got != "" { t.Fatalf("bare metadata.user_id must not create websocket Session-Id, got %q", got) } if got := headers.Get("Conversation_id"); got != "" { t.Fatalf("bare metadata.user_id must not create websocket Conversation_id, got %q", got) }}
func TestApplyCodexWebsocketHeadersIdentityConfuseRemapsPromptCacheKey(t *testing.T) { cfg := &config.Config{ Routing: config.RoutingConfig{SessionAffinity: true}, Codex: config.CodexConfig{IdentityConfuse: true}, } auth := &cliproxyauth.Auth{ID: "auth-ws-1", Provider: "codex"} req := cliproxyexecutor.Request{ Model: "gpt-5-codex", Payload: []byte(`{"prompt_cache_key":"cache-ws-1","client_metadata":{"x-codex-installation-id":"install-ws-1"}}`), }
body, headers := applyCodexPromptCacheHeaders("openai-response", req, []byte(`{"model":"gpt-5-codex"}`)) body, identityState := applyCodexIdentityConfuseBody(cfg, auth, req.Payload, body) ctx := contextWithGinHeaders(map[string]string{ "X-Codex-Turn-Metadata": `{"prompt_cache_key":"cache-ws-1","turn_id":"turn-ws-1","window_id":"cache-ws-1:0"}`, "X-Client-Request-Id": "client-request-1", }) headers = applyCodexWebsocketHeaders(ctx, headers, auth, "oauth-token", cfg) applyCodexIdentityConfuseHeaders(headers, &identityState)
expectedPromptCacheKey := codexIdentityConfuseUUID("auth-ws-1", "prompt-cache", "cache-ws-1") expectedTurnID := codexIdentityConfuseUUID("auth-ws-1", "turn", "turn-ws-1") if gotKey := gjson.GetBytes(body, "prompt_cache_key").String(); gotKey != expectedPromptCacheKey { t.Fatalf("prompt_cache_key = %q, want %q", gotKey, expectedPromptCacheKey) } if gotSession := headers["session_id"]; len(gotSession) != 1 || gotSession[0] != expectedPromptCacheKey { t.Fatalf("session_id = %#v, want [%q]", gotSession, expectedPromptCacheKey) } if gotCanonicalSession := headers.Get("Session-Id"); gotCanonicalSession != "" { t.Fatalf("Session-Id = %q, want empty", gotCanonicalSession) } if gotRequestID := headers.Get("X-Client-Request-Id"); gotRequestID != expectedPromptCacheKey { t.Fatalf("X-Client-Request-Id = %q, want %q", gotRequestID, expectedPromptCacheKey) } if gotThreadID := headers.Get("Thread-Id"); gotThreadID != expectedPromptCacheKey { t.Fatalf("Thread-Id = %q, want %q", gotThreadID, expectedPromptCacheKey) } if gotConversation := headers.Get("Conversation_id"); gotConversation != expectedPromptCacheKey { t.Fatalf("Conversation_id = %q, want %q", gotConversation, expectedPromptCacheKey) } if gotWindowID := headers.Get("X-Codex-Window-Id"); gotWindowID != expectedPromptCacheKey+":0" { t.Fatalf("X-Codex-Window-Id = %q, want %q", gotWindowID, expectedPromptCacheKey+":0") } gotMetadata := headers.Get("X-Codex-Turn-Metadata") if gotMetadataPromptCacheKey := gjson.Get(gotMetadata, "prompt_cache_key").String(); gotMetadataPromptCacheKey != expectedPromptCacheKey { t.Fatalf("X-Codex-Turn-Metadata.prompt_cache_key = %q, want %q", gotMetadataPromptCacheKey, expectedPromptCacheKey) } if gotMetadataTurnID := gjson.Get(gotMetadata, "turn_id").String(); gotMetadataTurnID != expectedTurnID { t.Fatalf("X-Codex-Turn-Metadata.turn_id = %q, want %q", gotMetadataTurnID, expectedTurnID) } if gotMetadataWindowID := gjson.Get(gotMetadata, "window_id").String(); gotMetadataWindowID != expectedPromptCacheKey+":0" { t.Fatalf("X-Codex-Turn-Metadata.window_id = %q, want %q", gotMetadataWindowID, expectedPromptCacheKey+":0") } expectedInstallationID := codexIdentityConfuseUUID("auth-ws-1", "installation", "install-ws-1") if gotInstallationID := gjson.GetBytes(body, "client_metadata.x-codex-installation-id").String(); gotInstallationID != expectedInstallationID { t.Fatalf("installation id = %q, want %q", gotInstallationID, expectedInstallationID) }}
func TestCodexIdentityConfuseResponsePayloadHidesUpstreamAndRestoresClient(t *testing.T) { state := codexIdentityConfuseState{ enabled: true, authID: "auth-ws-1", originalPromptCacheKey: "cache-ws-1", promptCacheKey: codexIdentityConfuseUUID("auth-ws-1", "prompt-cache", "cache-ws-1"), } expectedTurnID := state.confuseTurnID("turn-ws-1") rawPayload := []byte(`{"type":"response.completed","response":{"prompt_cache_key":"cache-ws-1","turn_id":"turn-ws-1"},"prompt_cache_key":"cache-ws-1","turn_id":"turn-ws-1"}`)
upstreamPayload := applyCodexIdentityConfuseResponsePayload(rawPayload, state) if bytes.Contains(upstreamPayload, []byte(`cache-ws-1`)) { t.Fatalf("upstream payload still contains original prompt_cache_key: %s", string(upstreamPayload)) } if bytes.Contains(upstreamPayload, []byte(`turn-ws-1`)) { t.Fatalf("upstream payload still contains original turn_id: %s", string(upstreamPayload)) } if !bytes.Contains(upstreamPayload, []byte(state.promptCacheKey)) { t.Fatalf("upstream payload missing confused prompt_cache_key: %s", string(upstreamPayload)) } if !bytes.Contains(upstreamPayload, []byte(expectedTurnID)) { t.Fatalf("upstream payload missing confused turn_id: %s", string(upstreamPayload)) }
clientPayload := applyCodexIdentityExposeResponsePayload(upstreamPayload, state) if bytes.Contains(clientPayload, []byte(state.promptCacheKey)) { t.Fatalf("client payload still contains confused prompt_cache_key: %s", string(clientPayload)) } if bytes.Contains(clientPayload, []byte(expectedTurnID)) { t.Fatalf("client payload still contains confused turn_id: %s", string(clientPayload)) } if !bytes.Contains(clientPayload, []byte(`cache-ws-1`)) { t.Fatalf("client payload missing original prompt_cache_key: %s", string(clientPayload)) } if !bytes.Contains(clientPayload, []byte(`turn-ws-1`)) { t.Fatalf("client payload missing original turn_id: %s", string(clientPayload)) }
rawSSE := []byte(`data: {"type":"response.completed","response":{"prompt_cache_key":"cache-ws-1","turn_id":"turn-ws-1"}}`) upstreamSSE := applyCodexIdentityConfuseResponsePayload(rawSSE, state) if bytes.Contains(upstreamSSE, []byte(`cache-ws-1`)) { t.Fatalf("upstream SSE still contains original prompt_cache_key: %s", string(upstreamSSE)) } if bytes.Contains(upstreamSSE, []byte(`turn-ws-1`)) { t.Fatalf("upstream SSE still contains original turn_id: %s", string(upstreamSSE)) } clientSSE := applyCodexIdentityExposeResponsePayload(upstreamSSE, state) if !bytes.Contains(clientSSE, []byte(`cache-ws-1`)) || bytes.Contains(clientSSE, []byte(state.promptCacheKey)) { t.Fatalf("client SSE prompt_cache_key was not restored: %s", string(clientSSE)) } if !bytes.Contains(clientSSE, []byte(`turn-ws-1`)) || bytes.Contains(clientSSE, []byte(expectedTurnID)) { t.Fatalf("client SSE turn_id was not restored: %s", string(clientSSE)) }}
func TestBuildCodexResponsesWebsocketURLRequiresHTTPURL(t *testing.T) { if got, err := buildCodexResponsesWebsocketURL("https://example.com/backend/responses"); err != nil || got != "wss://example.com/backend/responses" { t.Fatalf("https URL = %q, %v; want wss URL", got, err) } if _, err := buildCodexResponsesWebsocketURL("ftp://example.com/responses"); err == nil { t.Fatalf("expected unsupported scheme error") } if _, err := buildCodexResponsesWebsocketURL("https:///responses"); err == nil { t.Fatalf("expected empty host error") }}
func TestParseCodexWebsocketErrorMarksConnectionLimitRetryable(t *testing.T) { err, ok := parseCodexWebsocketError([]byte(`{"type":"error","status":429,"error":{"code":"websocket_connection_limit_reached","message":"too many websockets"},"headers":{"retry-after":"1"}}`)) if !ok { t.Fatalf("expected websocket error") } status, ok := err.(interface{ StatusCode() int }) if !ok || status.StatusCode() != http.StatusTooManyRequests { t.Fatalf("status = %#v, want 429", err) } retryable, ok := err.(interface{ RetryAfter() *time.Duration }) if !ok || retryable.RetryAfter() == nil { t.Fatalf("expected retryable websocket connection limit error") } if got := *retryable.RetryAfter(); got != 0 { t.Fatalf("retryAfter = %v, want connection-limit fallback 0", got) } withHeaders, ok := err.(interface{ Headers() http.Header }) if !ok || withHeaders.Headers().Get("retry-after") != "1" { t.Fatalf("headers = %#v, want retry-after", err) }}
func TestParseCodexWebsocketErrorUsesUsageLimitRetryMetadata(t *testing.T) { err, ok := parseCodexWebsocketError([]byte(`{"type":"error","status":429,"body":{"error":{"type":"usage_limit_reached","message":"usage limit reached","resets_in_seconds":7}}}`)) if !ok { t.Fatalf("expected websocket error") }
retryable, ok := err.(interface{ RetryAfter() *time.Duration }) if !ok || retryable.RetryAfter() == nil { t.Fatalf("expected retryable usage limit websocket error") } if got := *retryable.RetryAfter(); got != 7*time.Second { t.Fatalf("retryAfter = %v, want 7s", got) }}
func TestParseCodexWebsocketErrorPreservesWrappedBodyAndHeaders(t *testing.T) { err, ok := parseCodexWebsocketError([]byte(`{"type":"error","status":429,"body":{"error":{"code":"websocket_connection_limit_reached","type":"server_error","message":"too many websocket connections"}},"headers":{"x-request-id":"req-1"}}`)) if !ok { t.Fatalf("expected websocket error") }
parsed := gjson.Parse(err.Error()) if got := parsed.Get("status").Int(); got != http.StatusTooManyRequests { t.Fatalf("wrapped status = %d, want 429; payload=%s", got, err.Error()) } if got := parsed.Get("body.error.code").String(); got != "websocket_connection_limit_reached" { t.Fatalf("wrapped body error code = %s, want websocket_connection_limit_reached; payload=%s", got, err.Error()) } if got := parsed.Get("error.code").String(); got != "websocket_connection_limit_reached" { t.Fatalf("surface error code = %s, want websocket_connection_limit_reached; payload=%s", got, err.Error()) } retryable, ok := err.(interface{ RetryAfter() *time.Duration }) if !ok || retryable.RetryAfter() == nil { t.Fatalf("expected body.error.code websocket connection limit to be retryable") } withHeaders, ok := err.(interface{ Headers() http.Header }) if !ok || withHeaders.Headers().Get("x-request-id") != "req-1" { t.Fatalf("headers = %#v, want x-request-id", err) }}
func TestApplyCodexHeadersUsesConfigUserAgentForOAuth(t *testing.T) { req, err := http.NewRequest(http.MethodPost, "https://example.com/responses", nil) if err != nil { t.Fatalf("NewRequest() error = %v", err) } cfg := &config.Config{ Codex: config.CodexConfig{DisableCodexCloaking: true}, CodexHeaderDefaults: config.CodexHeaderDefaults{ UserAgent: "config-ua", BetaFeatures: "config-beta", }, } auth := &cliproxyauth.Auth{ Provider: "codex", Metadata: map[string]any{"email": "user@example.com"}, } req = req.WithContext(contextWithGinHeaders(map[string]string{ "User-Agent": "client-ua", }))
applyCodexHeaders(req, auth, "oauth-token", true, cfg)
if got := req.Header.Get("User-Agent"); got != "config-ua" { t.Fatalf("User-Agent = %s, want %s", got, "config-ua") } if got := req.Header.Get("x-codex-beta-features"); got != "" { t.Fatalf("x-codex-beta-features = %q, want empty", got) }}
func TestApplyCodexHeadersDefaultsToCodexCloaking(t *testing.T) { req, err := http.NewRequest(http.MethodPost, "https://example.com/responses", nil) if err != nil { t.Fatalf("NewRequest() error = %v", err) } req.Header.Set("User-Agent", "existing-ua") req.Header.Set("Originator", "existing-origin") cfg := &config.Config{ CodexHeaderDefaults: config.CodexHeaderDefaults{ UserAgent: "config-ua", }, } auth := &cliproxyauth.Auth{ Provider: "codex", Attributes: map[string]string{ "api_key": "api-key", "header:User-Agent": "custom-ua", "header:Originator": "custom-origin", }, } ginHeaders := http.Header{ "User-Agent": []string{"client-ua"}, "Originator": []string{"client-origin"}, }
applyCodexHeadersFromSources(req, auth, "api-key", false, cfg, ginHeaders)
if got := req.Header.Get("User-Agent"); got != codexUserAgent { t.Fatalf("User-Agent = %q, want %q", got, codexUserAgent) } if got := req.Header.Get("Originator"); got != codexOriginator { t.Fatalf("Originator = %q, want %q", got, codexOriginator) }}
func TestApplyCodexHeaders_EmptyAPIKey_OmitsAuthorizationAndOAuthHeaders(t *testing.T) { req, err := http.NewRequest(http.MethodPost, "https://example.com/responses", nil) if err != nil { t.Fatalf("NewRequest() error = %v", err) } auth := &cliproxyauth.Auth{ Provider: "codex", Attributes: map[string]string{ "auth_kind": "apikey", "base_url": "https://custom-codex.example.com", }, Metadata: map[string]any{ "account_id": "acc-12345", }, } cfg := &config.Config{ Codex: config.CodexConfig{ DisableCodexCloaking: true, }, CodexHeaderDefaults: config.CodexHeaderDefaults{ UserAgent: "oauth-default-ua", }, } applyCodexHeaders(req, auth, "", false, cfg)
if got := req.Header.Get("Authorization"); got != "" { t.Fatalf("Authorization = %q, want empty for empty API key", got) } if got := req.Header.Get("Chatgpt-Account-Id"); got != "" { t.Fatalf("Chatgpt-Account-Id = %q, want empty for API key auth_kind", got) } if got := req.Header.Get("Originator"); got != "" { t.Fatalf("Originator = %q, want empty for API key auth_kind when client originator omitted", got) } if got := req.Header.Get("User-Agent"); got == "oauth-default-ua" { t.Fatalf("User-Agent unexpectedly used OAuth default UA %q for API key auth_kind", got) }}
func TestApplyCodexWebsocketHeaders_EmptyAPIKey_OmitsAuthorizationAndOAuthHeaders(t *testing.T) { auth := &cliproxyauth.Auth{ Provider: "codex", Attributes: map[string]string{ "auth_kind": "apikey", "base_url": "https://custom-codex.example.com", }, Metadata: map[string]any{ "account_id": "acc-ws-123", }, } cfg := &config.Config{ Codex: config.CodexConfig{ DisableCodexCloaking: true, }, CodexHeaderDefaults: config.CodexHeaderDefaults{ UserAgent: "oauth-default-ua", BetaFeatures: "oauth-beta", }, } headers := applyCodexWebsocketHeaders(context.Background(), nil, auth, "", cfg) if got := headers.Get("Authorization"); got != "" { t.Fatalf("Authorization = %q, want empty for empty API key", got) } if got := headers.Get("ChatGPT-Account-ID"); got != "" { t.Fatalf("ChatGPT-Account-ID = %q, want empty for API key auth_kind", got) } if got := headers.Get("Originator"); got != "" { t.Fatalf("Originator = %q, want empty for API key auth_kind", got) } if got := headers.Get("x-codex-beta-features"); got != "" { t.Fatalf("x-codex-beta-features = %q, want empty for API key auth_kind", got) } if got := headers.Get("User-Agent"); got == "oauth-default-ua" { t.Fatalf("User-Agent unexpectedly used OAuth default UA %q for API key auth_kind", got) }}
func TestApplyModelHeaderOverridesFromModelConfig(t *testing.T) { const wantUA = "codex-tui/0.144.0 (Mac OS 26.5.1; arm64) iTerm.app/3.6.11 (codex-tui; 0.144.0)" req, err := http.NewRequest(http.MethodPost, "https://example.com/responses", nil) if err != nil { t.Fatalf("NewRequest() error = %v", err) } cfg := &config.Config{ CodexHeaderDefaults: config.CodexHeaderDefaults{ UserAgent: "config-ua", }, } auth := &cliproxyauth.Auth{ Provider: "codex", Metadata: map[string]any{"email": "user@example.com"}, }
applyCodexHeaders(req, auth, "oauth-token", true, cfg) applyModelHeaderOverrides(req.Header, "gpt-5.6-luna")
if got := req.Header.Get("User-Agent"); got != wantUA { t.Fatalf("User-Agent = %q, want %q", got, wantUA) } if got := codexSessionHeaderValue(req.Header); got == "" { t.Fatal("expected Session_id to be set for Mac OS User-Agent override") }
applyModelHeaderOverrides(req.Header, "gpt-5.4") if got := req.Header.Get("User-Agent"); got != wantUA { t.Fatalf("User-Agent after no-op override = %q, want %q", got, wantUA) }}
func TestApplyModelHeaderOverridesMultipleHeaders(t *testing.T) { reg := registry.GetGlobalRegistry() clientID := "test-model-header-override" reg.RegisterClient(clientID, "codex", []*registry.ModelInfo{{ ID: "test-override-headers-model", Config: ®istry.ModelConfig{ OverrideHeader: map[string]string{ "user-agent": "custom-ua/1.0", "originator": "custom-origin", "x-test-header": "forced-value", }, }, }}) t.Cleanup(func() { reg.UnregisterClient(clientID) })
headers := http.Header{} headers.Set("User-Agent", "old-ua") headers.Set("Originator", "old-origin") headers.Set("X-Test-Header", "old-value")
applyModelHeaderOverrides(headers, "test-override-headers-model")
if got := headers.Get("User-Agent"); got != "custom-ua/1.0" { t.Fatalf("User-Agent = %q, want custom-ua/1.0", got) } if got := headers.Get("Originator"); got != "custom-origin" { t.Fatalf("Originator = %q, want custom-origin", got) } if got := headers.Get("X-Test-Header"); got != "forced-value" { t.Fatalf("X-Test-Header = %q, want forced-value", got) }}
func TestApplyCodexHeadersPassesThroughClientIdentityHeaders(t *testing.T) { req, err := http.NewRequest(http.MethodPost, "https://example.com/responses", nil) if err != nil { t.Fatalf("NewRequest() error = %v", err) } auth := &cliproxyauth.Auth{ Provider: "codex", Metadata: map[string]any{"email": "user@example.com"}, } req = req.WithContext(contextWithGinHeaders(map[string]string{ "Originator": "Codex Desktop", "Version": "0.115.0-alpha.27", "X-Codex-Turn-Metadata": `{"turn_id":"turn-1"}`, "X-Client-Request-Id": "019d2233-e240-7162-992d-38df0a2a0e0d", }))
cfg := &config.Config{Codex: config.CodexConfig{DisableCodexCloaking: true}} applyCodexHeaders(req, auth, "oauth-token", true, cfg)
if got := req.Header.Get("Originator"); got != "Codex Desktop" { t.Fatalf("Originator = %s, want %s", got, "Codex Desktop") } if got := req.Header.Get("Version"); got != "0.115.0-alpha.27" { t.Fatalf("Version = %s, want %s", got, "0.115.0-alpha.27") } if got := req.Header.Get("X-Codex-Turn-Metadata"); got != `{"turn_id":"turn-1"}` { t.Fatalf("X-Codex-Turn-Metadata = %s, want %s", got, `{"turn_id":"turn-1"}`) } if got := req.Header.Get("X-Client-Request-Id"); got != "019d2233-e240-7162-992d-38df0a2a0e0d" { t.Fatalf("X-Client-Request-Id = %s, want %s", got, "019d2233-e240-7162-992d-38df0a2a0e0d") }}
func TestApplyCodexHeadersDoesNotInjectClientOnlyHeadersByDefault(t *testing.T) { req, err := http.NewRequest(http.MethodPost, "https://example.com/responses", nil) if err != nil { t.Fatalf("NewRequest() error = %v", err) }
applyCodexHeaders(req, nil, "oauth-token", true, nil)
if got := req.Header.Get("Version"); got != "" { t.Fatalf("Version = %q, want empty", got) } if got := req.Header.Get("X-Codex-Turn-Metadata"); got != "" { t.Fatalf("X-Codex-Turn-Metadata = %q, want empty", got) } if got := req.Header.Get("X-Client-Request-Id"); got != "" { t.Fatalf("X-Client-Request-Id = %q, want empty", got) }}
func contextWithGinHeaders(headers map[string]string) context.Context { gin.SetMode(gin.TestMode) recorder := httptest.NewRecorder() ginCtx, _ := gin.CreateTestContext(recorder) ginCtx.Request = httptest.NewRequest(http.MethodPost, "/", nil) ginCtx.Request.Header = make(http.Header, len(headers)) for key, value := range headers { ginCtx.Request.Header.Set(key, value) } return context.WithValue(context.Background(), "gin", ginCtx)}
func TestNewProxyAwareWebsocketDialerDirectDisablesProxy(t *testing.T) { t.Parallel()
dialer := newProxyAwareWebsocketDialer( &config.Config{SDKConfig: sdkconfig.SDKConfig{ProxyURL: "http://global-proxy.example.com:8080"}}, &cliproxyauth.Auth{ProxyURL: "direct"}, )
if dialer.Proxy != nil { t.Fatal("expected websocket proxy function to be nil for direct mode") }}
func TestCodexWebsocketUpgradeRequiredDoesNotFallbackToHTTPWithLifecycle(t *testing.T) { var httpFallbackCalls atomic.Int32 server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { if r.Method == http.MethodPost { httpFallbackCalls.Add(1) http.Error(w, "unexpected HTTP fallback", http.StatusInternalServerError) return } http.Error(w, "websocket upgrade required", http.StatusUpgradeRequired) })) defer server.Close()
exec := NewCodexWebsocketsExecutor(&config.Config{SDKConfig: config.SDKConfig{DisableImageGeneration: config.DisableImageGenerationAll}}) auth := &cliproxyauth.Auth{ID: "auth-a", Provider: "codex", Attributes: map[string]string{"api_key": "sk-test", "base_url": server.URL}} req := cliproxyexecutor.Request{Model: "gpt-5-codex", Payload: []byte(`{"model":"gpt-5-codex","input":[{"type":"message","role":"user","content":"hello"}]}`)} opts := cliproxyexecutor.Options{ SourceFormat: sdktranslator.FromString("openai-response"), ResponseFormat: sdktranslator.FromString("openai-response"), ExecutionLifecycle: newTerminalFailureLifecycle(), }
if _, errExecute := exec.ExecuteStream(context.Background(), auth, req, opts); errExecute == nil { t.Fatal("ExecuteStream() error = nil, want failed Home lifecycle attempt") } if got := httpFallbackCalls.Load(); got != 0 { t.Fatalf("HTTP fallback calls = %d, want 0 with an execution lifecycle", got) }}
func TestCodexWebsocketHandshakeFailureReleasesSessionRequestLock(t *testing.T) { for _, statusCode := range []int{http.StatusUpgradeRequired, http.StatusBadGateway} { t.Run(http.StatusText(statusCode), func(t *testing.T) { server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { http.Error(w, "upstream rejected websocket", statusCode) })) defer server.Close()
exec := NewCodexWebsocketsExecutor(&config.Config{SDKConfig: config.SDKConfig{DisableImageGeneration: config.DisableImageGenerationAll}}) exec.store = &codexWebsocketSessionStore{sessions: make(map[string]*codexWebsocketSession)} auth := &cliproxyauth.Auth{ID: "auth-a", Provider: "codex", Attributes: map[string]string{"api_key": "sk-test", "base_url": server.URL}} req := cliproxyexecutor.Request{Model: "gpt-5-codex", Payload: []byte(`{"model":"gpt-5-codex","input":[{"type":"message","role":"user","content":"hello"}]}`)} opts := cliproxyexecutor.Options{ SourceFormat: sdktranslator.FromString("openai-response"), ResponseFormat: sdktranslator.FromString("openai-response"), Metadata: map[string]any{ cliproxyexecutor.ExecutionSessionMetadataKey: "failed-handshake", }, }
_, _ = exec.ExecuteStream(context.Background(), auth, req, opts) sess := exec.getOrCreateSession("failed-handshake") acquired := make(chan struct{}) go func() { sess.reqMu.Lock() close(acquired) sess.reqMu.Unlock() }() select { case <-acquired: case <-time.After(time.Second): t.Fatal("websocket handshake failure left the session request lock held") } }) }}
type terminalFailureLifecycle struct { active atomic.Bool ends atomic.Int32}
func newTerminalFailureLifecycle() *terminalFailureLifecycle { lifecycle := &terminalFailureLifecycle{} lifecycle.active.Store(true) return lifecycle}
func (*terminalFailureLifecycle) Bind(func() error) error { return nil }func (l *terminalFailureLifecycle) End(string) { l.ends.Add(1) l.active.Store(false)}func (*terminalFailureLifecycle) Retain() {}
func TestCodexWebsocketTerminalFailureInvalidatesRetainedLifecycle(t *testing.T) { upgrader := websocket.Upgrader{CheckOrigin: func(*http.Request) bool { return true }} var connections atomic.Int32 firstRelease := make(chan struct{}) server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { conn, errUpgrade := upgrader.Upgrade(w, r, nil) if errUpgrade != nil { t.Errorf("upgrade websocket: %v", errUpgrade) return } defer func() { _ = conn.Close() }() connection := connections.Add(1) if _, _, errRead := conn.ReadMessage(); errRead != nil { return } terminal := []byte(`{"type":"response.failed","response":{"error":{"type":"authentication_error","code":"invalid_api_key","message":"Invalid token."}}}`) if errWrite := conn.WriteMessage(websocket.TextMessage, terminal); errWrite != nil { t.Errorf("write terminal response: %v", errWrite) } if connection == 1 { <-firstRelease } })) defer server.Close()
exec := NewCodexWebsocketsExecutor(&config.Config{SDKConfig: config.SDKConfig{DisableImageGeneration: config.DisableImageGenerationAll}}) exec.store = &codexWebsocketSessionStore{sessions: make(map[string]*codexWebsocketSession)} auth := &cliproxyauth.Auth{ID: "auth-a", Provider: "codex", Attributes: map[string]string{"api_key": "sk-test", "base_url": server.URL}} req := cliproxyexecutor.Request{Model: "gpt-5-codex", Payload: []byte(`{"model":"gpt-5-codex","input":[{"type":"message","role":"user","content":"hello"}]}`)} opts := cliproxyexecutor.Options{ SourceFormat: sdktranslator.FromString("openai-response"), ResponseFormat: sdktranslator.FromString("openai-response"), ExecutionLifecycle: newTerminalFailureLifecycle(), Metadata: map[string]any{ cliproxyexecutor.ExecutionSessionMetadataKey: "terminal-failure", }, }
result, errExecute := exec.ExecuteStream(context.Background(), auth, req, opts) if errExecute != nil { t.Fatalf("first ExecuteStream() error = %v", errExecute) } for chunk := range result.Chunks { if chunk.Err == nil { continue } } lifecycle := opts.ExecutionLifecycle.(*terminalFailureLifecycle) if lifecycle.active.Load() { t.Fatal("terminal failure left the retained lifecycle active") } if got := lifecycle.ends.Load(); got != 1 { t.Fatalf("retained lifecycle End calls = %d, want 1", got) } sess := exec.getOrCreateSession("terminal-failure") sess.connMu.Lock() connected := sess.conn != nil sess.connMu.Unlock() if connected { t.Fatal("terminal failure left the upstream session connection cached") } close(firstRelease)
opts.ExecutionLifecycle = newTerminalFailureLifecycle() result, errExecute = exec.ExecuteStream(context.Background(), auth, req, opts) if errExecute != nil { t.Fatalf("second ExecuteStream() error = %v", errExecute) } for range result.Chunks { } if got := connections.Load(); got != 2 { t.Fatalf("websocket connections = %d, want 2 after terminal invalidation", got) }}
type rejectingExecutionLifecycle struct{}
func (rejectingExecutionLifecycle) Bind(func() error) error { return errors.New("lifecycle bind rejected")}func (rejectingExecutionLifecycle) End(string) {}
func TestCodexWebsocketNonstreamLifecycleBindFailureDetachesConnection(t *testing.T) { upgrader := websocket.Upgrader{CheckOrigin: func(*http.Request) bool { return true }} var connections atomic.Int32 closed := make(chan struct{}, 1) server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { conn, errUpgrade := upgrader.Upgrade(w, r, nil) if errUpgrade != nil { t.Errorf("upgrade websocket: %v", errUpgrade) return } connection := connections.Add(1) defer func() { _ = conn.Close() if connection == 1 { closed <- struct{}{} } }() if _, _, errRead := conn.ReadMessage(); errRead != nil { return } completed := []byte(`{"type":"response.completed","response":{"id":"resp-1","output":[],"usage":{"input_tokens":0,"output_tokens":0,"total_tokens":0}}}`) if errWrite := conn.WriteMessage(websocket.TextMessage, completed); errWrite != nil { t.Errorf("write completed response: %v", errWrite) } })) defer server.Close()
exec := NewCodexWebsocketsExecutor(&config.Config{SDKConfig: config.SDKConfig{DisableImageGeneration: config.DisableImageGenerationAll}}) exec.store = &codexWebsocketSessionStore{sessions: make(map[string]*codexWebsocketSession)} auth := &cliproxyauth.Auth{ID: "auth-a", Provider: "codex", Attributes: map[string]string{"api_key": "sk-test", "base_url": server.URL}} req := cliproxyexecutor.Request{Model: "gpt-5-codex", Payload: []byte(`{"model":"gpt-5-codex","input":[{"type":"message","role":"user","content":"hello"}]}`)} opts := cliproxyexecutor.Options{ SourceFormat: sdktranslator.FromString("openai-response"), ResponseFormat: sdktranslator.FromString("openai-response"), ExecutionLifecycle: rejectingExecutionLifecycle{}, Metadata: map[string]any{ cliproxyexecutor.ExecutionSessionMetadataKey: "nonstream-bind-failed", }, } attemptCtx := cliproxyexecutor.WithUpstreamAttemptTracker(context.Background()) if _, errExecute := exec.Execute(attemptCtx, auth, req, opts); errExecute == nil { t.Fatal("Execute() error = nil, want lifecycle bind failure") } if cliproxyexecutor.UpstreamAttempted(attemptCtx) { t.Fatal("successful handshake marked a lifecycle bind failure as an upstream request attempt") } select { case <-closed: case <-time.After(time.Second): t.Fatal("nonstream lifecycle bind failure did not close the upstream websocket") } sess := exec.getOrCreateSession("nonstream-bind-failed") sess.connMu.Lock() connected := sess.conn != nil sess.connMu.Unlock() if connected { t.Fatal("nonstream lifecycle bind failure left the closed connection attached to the session") }
opts.ExecutionLifecycle = nil if _, errExecute := exec.Execute(context.Background(), auth, req, opts); errExecute != nil { t.Fatalf("second Execute() error = %v", errExecute) } if got := connections.Load(); got != 2 { t.Fatalf("websocket connections = %d, want 2 after bind failure", got) }}
func TestCodexWebsocketLifecycleBindFailureReleasesSessionRequestLock(t *testing.T) { upgrader := websocket.Upgrader{CheckOrigin: func(*http.Request) bool { return true }} closed := make(chan struct{}, 1) server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { conn, errUpgrade := upgrader.Upgrade(w, r, nil) if errUpgrade != nil { t.Errorf("upgrade websocket: %v", errUpgrade) return } defer func() { _ = conn.Close() closed <- struct{}{} }() for { if _, _, errRead := conn.ReadMessage(); errRead != nil { return } } })) defer server.Close()
exec := NewCodexWebsocketsExecutor(&config.Config{SDKConfig: config.SDKConfig{DisableImageGeneration: config.DisableImageGenerationAll}}) exec.store = &codexWebsocketSessionStore{sessions: make(map[string]*codexWebsocketSession)} auth := &cliproxyauth.Auth{ID: "auth-a", Provider: "codex", Attributes: map[string]string{"api_key": "sk-test", "base_url": server.URL}} req := cliproxyexecutor.Request{Model: "gpt-5-codex", Payload: []byte(`{"model":"gpt-5-codex","input":[{"type":"message","role":"user","content":"hello"}]}`)} opts := cliproxyexecutor.Options{ SourceFormat: sdktranslator.FromString("openai-response"), ResponseFormat: sdktranslator.FromString("openai-response"), ExecutionLifecycle: rejectingExecutionLifecycle{}, Metadata: map[string]any{ cliproxyexecutor.ExecutionSessionMetadataKey: "bind-failed", }, } attemptCtx := cliproxyexecutor.WithUpstreamAttemptTracker(context.Background()) if _, errExecute := exec.ExecuteStream(attemptCtx, auth, req, opts); errExecute == nil { t.Fatal("ExecuteStream() error = nil, want lifecycle bind failure") } if cliproxyexecutor.UpstreamAttempted(attemptCtx) { t.Fatal("successful handshake marked a lifecycle bind failure as an upstream request attempt") } select { case <-closed: case <-time.After(time.Second): t.Fatal("lifecycle bind failure did not close the upstream websocket") }
sess := exec.getOrCreateSession("bind-failed") acquired := make(chan struct{}) go func() { sess.reqMu.Lock() close(acquired) sess.reqMu.Unlock() }() select { case <-acquired: case <-time.After(time.Second): t.Fatal("lifecycle bind failure left the session request lock held") }}
func TestCodexWebsocketsExecuteObservesWebSocketResponseEvents(t *testing.T) { upgrader := websocket.Upgrader{CheckOrigin: func(*http.Request) bool { return true }} server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { conn, err := upgrader.Upgrade(w, r, nil) if err != nil { t.Fatalf("upgrade websocket: %v", err) } defer func() { _ = conn.Close() }()
if _, _, errRead := conn.ReadMessage(); errRead != nil { t.Fatalf("read upstream websocket message: %v", errRead) }
rateLimitMsg := []byte(`{"type":"codex.rate_limits","rate_limits":{"primary":{"used_percent":42}}}`) if errWrite := conn.WriteMessage(websocket.TextMessage, rateLimitMsg); errWrite != nil { t.Fatalf("write rate limit websocket message: %v", errWrite) }
completed := []byte(`{"type":"response.completed","response":{"id":"resp-1","output":[],"usage":{"input_tokens":0,"output_tokens":0,"total_tokens":0}}}`) if errWrite := conn.WriteMessage(websocket.TextMessage, completed); errWrite != nil { t.Fatalf("write completed websocket message: %v", errWrite) } })) defer server.Close()
exec := NewCodexWebsocketsExecutor(&config.Config{}) auth := &cliproxyauth.Auth{ ID: "auth-codex-1", Label: "codex-account", Provider: "codex", Attributes: map[string]string{ "api_key": "sk-test", "base_url": server.URL, "auth_kind": "oauth", }, Metadata: map[string]any{ "email": "user@example.com", }, } req := cliproxyexecutor.Request{ Model: "gpt-5.6-sol", Payload: []byte(`{"model":"gpt-5.6-sol","input":[{"role":"user","content":"hello"}]}`), }
var observedEvents []cliproxyexecutor.WebSocketResponseEvent opts := cliproxyexecutor.Options{ SourceFormat: sdktranslator.FromString("codex"), Metadata: map[string]any{ "request_id": "test-req-codex", }, WebSocketResponseObserver: func(_ context.Context, ev cliproxyexecutor.WebSocketResponseEvent) { observedEvents = append(observedEvents, ev) }, }
if _, err := exec.Execute(context.Background(), auth, req, opts); err != nil { t.Fatalf("Execute() error = %v", err) }
if len(observedEvents) < 2 { t.Fatalf("observed %d events, want at least 2", len(observedEvents)) }
rateLimitEvent := observedEvents[0] if rateLimitEvent.RequestID != "test-req-codex" { t.Fatalf("RequestID = %q, want test-req-codex", rateLimitEvent.RequestID) } if rateLimitEvent.AuthID != "auth-codex-1" || rateLimitEvent.AuthLabel != "codex-account" { t.Fatalf("Auth = (%q, %q), want (auth-codex-1, codex-account)", rateLimitEvent.AuthID, rateLimitEvent.AuthLabel) } if rateLimitEvent.EventType != "codex.rate_limits" { t.Fatalf("EventType = %q, want codex.rate_limits", rateLimitEvent.EventType) } if !bytes.Contains(rateLimitEvent.Payload, []byte(`"used_percent":42`)) { t.Fatalf("Payload = %s, want used_percent 42", rateLimitEvent.Payload) }}
func TestCodexWebsocketsExecuteStreamObservesWebSocketResponseEvents(t *testing.T) { upgrader := websocket.Upgrader{CheckOrigin: func(*http.Request) bool { return true }} server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { conn, err := upgrader.Upgrade(w, r, nil) if err != nil { t.Fatalf("upgrade websocket: %v", err) } defer func() { _ = conn.Close() }()
if _, _, errRead := conn.ReadMessage(); errRead != nil { t.Fatalf("read upstream websocket message: %v", errRead) }
rateLimitMsg := []byte(`{"type":"codex.rate_limits","rate_limits":{"primary":{"used_percent":75}}}`) if errWrite := conn.WriteMessage(websocket.TextMessage, rateLimitMsg); errWrite != nil { t.Fatalf("write rate limit websocket message: %v", errWrite) }
completed := []byte(`{"type":"response.completed","response":{"id":"resp-1","output":[],"usage":{"input_tokens":0,"output_tokens":0,"total_tokens":0}}}`) if errWrite := conn.WriteMessage(websocket.TextMessage, completed); errWrite != nil { t.Fatalf("write completed websocket message: %v", errWrite) } })) defer server.Close()
exec := NewCodexWebsocketsExecutor(&config.Config{}) auth := &cliproxyauth.Auth{ ID: "auth-stream-1", Label: "codex-stream-account", Provider: "codex", Attributes: map[string]string{ "api_key": "sk-test", "base_url": server.URL, "auth_kind": "oauth", }, } req := cliproxyexecutor.Request{ Model: "gpt-5.6-sol", Payload: []byte(`{"model":"gpt-5.6-sol","input":[{"role":"user","content":"hello"}]}`), }
var observedEvents []cliproxyexecutor.WebSocketResponseEvent opts := cliproxyexecutor.Options{ SourceFormat: sdktranslator.FromString("codex"), Metadata: map[string]any{ "request_id": "test-req-stream", }, WebSocketResponseObserver: func(_ context.Context, ev cliproxyexecutor.WebSocketResponseEvent) { observedEvents = append(observedEvents, ev) }, }
streamResult, err := exec.ExecuteStream(context.Background(), auth, req, opts) if err != nil { t.Fatalf("ExecuteStream() error = %v", err) }
// Drain stream for chunk := range streamResult.Chunks { if chunk.Err != nil { t.Fatalf("stream chunk error: %v", chunk.Err) } }
if len(observedEvents) < 2 { t.Fatalf("observed %d events, want at least 2", len(observedEvents)) }
rateLimitEvent := observedEvents[0] if rateLimitEvent.RequestID != "test-req-stream" { t.Fatalf("RequestID = %q, want test-req-stream", rateLimitEvent.RequestID) } if rateLimitEvent.AuthID != "auth-stream-1" || rateLimitEvent.AuthLabel != "codex-stream-account" { t.Fatalf("Auth = (%q, %q), want (auth-stream-1, codex-stream-account)", rateLimitEvent.AuthID, rateLimitEvent.AuthLabel) } if rateLimitEvent.EventType != "codex.rate_limits" { t.Fatalf("EventType = %q, want codex.rate_limits", rateLimitEvent.EventType) } if !bytes.Contains(rateLimitEvent.Payload, []byte(`"used_percent":75`)) { t.Fatalf("Payload = %s, want used_percent 75", rateLimitEvent.Payload) }}
func TestCodexWebsocketsExecuteHandshakeUsageLimitReachedSetsRetryAfter(t *testing.T) { body := []byte(`{"error":{"type":"usage_limit_reached","message":"The usage limit has been reached","resets_in_seconds":120}}`) server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { w.Header().Set("Content-Type", "application/json") w.WriteHeader(http.StatusTooManyRequests) if _, errWrite := w.Write(body); errWrite != nil { t.Errorf("write handshake rejection: %v", errWrite) } })) defer server.Close()
exec := NewCodexWebsocketsExecutor(&config.Config{}) auth := &cliproxyauth.Auth{ ID: "codex-auth-quota-exhausted", Provider: "codex", Attributes: map[string]string{ "base_url": server.URL, "websockets": "true", "api_key": "sk-test", }, } req := cliproxyexecutor.Request{ Model: "gpt-5.6-luna", Payload: []byte(`{"model":"gpt-5.6-luna","input":[{"type":"message","id":"msg-1"}]}`), } opts := cliproxyexecutor.Options{ SourceFormat: sdktranslator.FromString("openai-response"), ResponseFormat: sdktranslator.FromString("openai-response"), }
ctx := cliproxyexecutor.WithUpstreamAttemptTracker(context.Background()) _, errExecute := exec.Execute(ctx, auth, req, opts) if errExecute == nil { t.Fatal("Execute() error = nil, want handshake rejection") } if !cliproxyexecutor.UpstreamAttempted(ctx) { t.Fatal("429 websocket handshake was not marked as an upstream attempt") } statusErr, ok := errExecute.(interface{ StatusCode() int }) if !ok || statusErr.StatusCode() != http.StatusTooManyRequests { t.Fatalf("status = %#v, want 429", errExecute) } retryable, ok := errExecute.(interface{ RetryAfter() *time.Duration }) if !ok || retryable.RetryAfter() == nil { t.Fatalf("expected RetryAfter for usage_limit_reached handshake error: %#v", errExecute) } if got := *retryable.RetryAfter(); got != 120*time.Second { t.Fatalf("RetryAfter = %v, want 120s", got) }}
func TestCodexWebsocketsExecuteStreamHandshakeUsageLimitReachedSetsRetryAfter(t *testing.T) { body := []byte(`{"error":{"type":"usage_limit_reached","message":"The usage limit has been reached","resets_in_seconds":120}}`) server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { w.Header().Set("Content-Type", "application/json") w.WriteHeader(http.StatusTooManyRequests) if _, errWrite := w.Write(body); errWrite != nil { t.Errorf("write handshake rejection: %v", errWrite) } })) defer server.Close()
exec := NewCodexWebsocketsExecutor(&config.Config{}) auth := &cliproxyauth.Auth{ ID: "codex-auth-quota-exhausted-stream", Provider: "codex", Attributes: map[string]string{ "base_url": server.URL, "websockets": "true", "api_key": "sk-test", }, } req := cliproxyexecutor.Request{ Model: "gpt-5.6-luna", Payload: []byte(`{"model":"gpt-5.6-luna","input":[{"type":"message","id":"msg-1"}]}`), } opts := cliproxyexecutor.Options{ SourceFormat: sdktranslator.FromString("openai-response"), ResponseFormat: sdktranslator.FromString("openai-response"), }
ctx := cliproxyexecutor.WithUpstreamAttemptTracker(context.Background()) _, errExecuteStream := exec.ExecuteStream(ctx, auth, req, opts) if errExecuteStream == nil { t.Fatal("ExecuteStream() error = nil, want handshake rejection") } if !cliproxyexecutor.UpstreamAttempted(ctx) { t.Fatal("429 streaming websocket handshake was not marked as an upstream attempt") } statusErr, ok := errExecuteStream.(interface{ StatusCode() int }) if !ok || statusErr.StatusCode() != http.StatusTooManyRequests { t.Fatalf("status = %#v, want 429", errExecuteStream) } retryable, ok := errExecuteStream.(interface{ RetryAfter() *time.Duration }) if !ok || retryable.RetryAfter() == nil { t.Fatalf("expected RetryAfter for usage_limit_reached handshake error: %#v", errExecuteStream) } if got := *retryable.RetryAfter(); got != 120*time.Second { t.Fatalf("RetryAfter = %v, want 120s", got) }}