Something went wrong. Try again.
sloproxy
Something went wrong. Try again.
79 kB · 1961 lines
Go
at commit e8e39526
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818819820821822823824825826827828829830831832833834835836837838839840841842843844845846847848849850851852853854855856857858859860861862863864865866867868869870871872873874875876877878879880881882883884885886887888889890891892893894895896897898899900901902903904905906907908909910911912913914915916917918919920921922923924925926927928929930931932933934935936937938939940941942943944945946947948949950951952953954955956957958959960961962963964965966967968969970971972973974975976977978979980981982983984985986987988989990991992993994995996997998999100010011002100310041005100610071008100910101011101210131014101510161017101810191020102110221023102410251026102710281029103010311032103310341035103610371038103910401041104210431044104510461047104810491050105110521053105410551056105710581059106010611062106310641065106610671068106910701071107210731074107510761077107810791080108110821083108410851086108710881089109010911092109310941095109610971098109911001101110211031104110511061107110811091110111111121113111411151116111711181119112011211122112311241125112611271128112911301131113211331134113511361137113811391140114111421143114411451146114711481149115011511152115311541155115611571158115911601161116211631164116511661167116811691170117111721173117411751176117711781179118011811182118311841185118611871188118911901191119211931194119511961197119811991200120112021203120412051206120712081209121012111212121312141215121612171218121912201221122212231224122512261227122812291230123112321233123412351236123712381239124012411242124312441245124612471248124912501251125212531254125512561257125812591260126112621263126412651266126712681269127012711272127312741275127612771278127912801281128212831284128512861287128812891290129112921293129412951296129712981299130013011302130313041305130613071308130913101311131213131314131513161317131813191320132113221323132413251326132713281329133013311332133313341335133613371338133913401341134213431344134513461347134813491350135113521353135413551356135713581359136013611362136313641365136613671368136913701371137213731374137513761377137813791380138113821383138413851386138713881389139013911392139313941395139613971398139914001401140214031404140514061407140814091410141114121413141414151416141714181419142014211422142314241425142614271428142914301431143214331434143514361437143814391440144114421443144414451446144714481449145014511452145314541455145614571458145914601461146214631464146514661467146814691470147114721473147414751476147714781479148014811482148314841485148614871488148914901491149214931494149514961497149814991500150115021503150415051506150715081509151015111512151315141515151615171518151915201521152215231524152515261527152815291530153115321533153415351536153715381539154015411542154315441545154615471548154915501551155215531554155515561557155815591560156115621563156415651566156715681569157015711572157315741575157615771578157915801581158215831584158515861587158815891590159115921593159415951596159715981599160016011602160316041605160616071608160916101611161216131614161516161617161816191620162116221623162416251626162716281629163016311632163316341635163616371638163916401641164216431644164516461647164816491650165116521653165416551656165716581659166016611662166316641665166616671668166916701671167216731674167516761677167816791680168116821683168416851686168716881689169016911692169316941695169616971698169917001701170217031704170517061707170817091710171117121713171417151716171717181719172017211722172317241725172617271728172917301731173217331734173517361737173817391740174117421743174417451746174717481749175017511752175317541755175617571758175917601761176217631764176517661767176817691770177117721773177417751776177717781779178017811782178317841785178617871788178917901791179217931794179517961797179817991800180118021803180418051806180718081809181018111812181318141815181618171818181918201821182218231824182518261827182818291830183118321833183418351836183718381839184018411842184318441845184618471848184918501851185218531854185518561857185818591860186118621863186418651866186718681869187018711872187318741875187618771878187918801881188218831884188518861887188818891890189118921893189418951896189718981899190019011902190319041905190619071908190919101911191219131914191519161917191819191920192119221923192419251926192719281929193019311932193319341935193619371938193919401941194219431944194519461947194819491950195119521953195419551956195719581959196019611962package 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")
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 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":"msg-1"}]}`)
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-1" { t.Fatalf("valid input item ID changed: %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++ { done := make(chan error, 1) go func() { _, errExecute := exec.ExecuteStream(context.Background(), auth, cliproxyexecutor.Request{ Model: "gpt-5.4", Payload: []byte(`{"model":"gpt-5.4","input":[{"type":"message","id":"msg-1"}]}`), }, opts) done <- errExecute }() 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) } 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 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.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 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 TestApplyCodexWebsocketHeadersPassesThroughClientIdentityHeaders(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", "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, "", nil)
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{ 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{ 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{ 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{ 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{ 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 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", }))
applyCodexHeaders(req, auth, "oauth-token", true, nil)
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", }, } if _, errExecute := exec.Execute(context.Background(), auth, req, opts); errExecute == nil { t.Fatal("Execute() error = nil, want lifecycle bind failure") } 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", }, } if _, errExecute := exec.ExecuteStream(context.Background(), auth, req, opts); errExecute == nil { t.Fatal("ExecuteStream() error = nil, want lifecycle bind failure") } 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") }}