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 TestCodexWebsocketsExecuteStreamPassesThroughUpstreamWebsocketPayloadForDownstreamWebsocket(t *testing.T) {	upgrader := websocket.Upgrader{CheckOrigin: func(*http.Request) bool { return true }}	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() }() 		if _, _, errRead := conn.ReadMessage(); errRead != nil {			t.Errorf("read upstream websocket message: %v", errRead)			return		}		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":"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 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")	}} 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 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)

Lines 1–500 of 872. Load more