package executor import ( "bytes" "context" "errors" "net/http" "net/http/httptest" "strings" "testing" "time" "github.com/gin-gonic/gin" "github.com/gorilla/websocket" "github.com/router-for-me/CLIProxyAPI/v6/internal/config" cliproxyauth "github.com/router-for-me/CLIProxyAPI/v6/sdk/cliproxy/auth" cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v6/sdk/cliproxy/executor" sdkconfig "github.com/router-for-me/CLIProxyAPI/v6/sdk/config" sdktranslator "github.com/router-for-me/CLIProxyAPI/v6/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 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 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{}) 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, must not use stale codex-tui prefix", codexUserAgent) } if strings.Contains(codexUserAgent, "(codex-tui;") { t.Fatalf("default Codex User-Agent = %s, must not include stale 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": "sess-client", }) 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 := headerValueCaseInsensitive(headers, "session_id"); got != "sess-client" { t.Fatalf("session_id = %s, want sess-client", got) } if _, ok := headers["session_id"]; !ok { t.Fatalf("expected lowercase session_id header key, got %#v", headers) }} 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 TestApplyCodexPromptCacheHeadersSetsLowercaseSessionAndLegacyConversation(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 := headerValueCaseInsensitive(headers, "session_id"); got != "cache-1" { t.Fatalf("session_id = %s, want cache-1", got) } if _, ok := headers["session_id"]; !ok { t.Fatalf("expected lowercase session_id key, got %#v", headers) } if got := headers.Get("Conversation_id"); got != "cache-1" { t.Fatalf("Conversation_id = %s, want cache-1", 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 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 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)Lines 1–500 of 558. Load more
package executor import ( "bytes" "context" "errors" "net/http" "net/http/httptest" "strings" "testing" "time" "github.com/gin-gonic/gin" "github.com/gorilla/websocket" "github.com/router-for-me/CLIProxyAPI/v6/internal/config" cliproxyauth "github.com/router-for-me/CLIProxyAPI/v6/sdk/cliproxy/auth" cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v6/sdk/cliproxy/executor" sdkconfig "github.com/router-for-me/CLIProxyAPI/v6/sdk/config" sdktranslator "github.com/router-for-me/CLIProxyAPI/v6/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 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 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{}) 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, must not use stale codex-tui prefix", codexUserAgent) } if strings.Contains(codexUserAgent, "(codex-tui;") { t.Fatalf("default Codex User-Agent = %s, must not include stale 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": "sess-client", }) 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 := headerValueCaseInsensitive(headers, "session_id"); got != "sess-client" { t.Fatalf("session_id = %s, want sess-client", got) } if _, ok := headers["session_id"]; !ok { t.Fatalf("expected lowercase session_id header key, got %#v", headers) } } 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 TestApplyCodexPromptCacheHeadersSetsLowercaseSessionAndLegacyConversation(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 := headerValueCaseInsensitive(headers, "session_id"); got != "cache-1" { t.Fatalf("session_id = %s, want cache-1", got) } if _, ok := headers["session_id"]; !ok { t.Fatalf("expected lowercase session_id key, got %#v", headers) } if got := headers.Get("Conversation_id"); got != "cache-1" { t.Fatalf("Conversation_id = %s, want cache-1", 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 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 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)