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 := 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 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")	}

Lines 1–500 of 750. Load more