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/v7/internal/config"	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 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, 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 := headerValueCaseInsensitive(headers, "session_id"); got != "legacy-session" {		t.Fatalf("session_id = %s, want legacy-session", 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 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 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 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 := headerValueCaseInsensitive(headers, "session_id"); gotSession != expectedPromptCacheKey {		t.Fatalf("session_id = %q, want %q", gotSession, expectedPromptCacheKey)	}	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")	}

Lines 1–500 of 669. Load more