package executor import (	"bytes"	"context"	"errors"	"fmt"	"io"	"net/http"	"net/http/httptest"	"strings"	"sync/atomic"	"testing"	"time" 	"github.com/gin-gonic/gin"	"github.com/google/uuid"	"github.com/gorilla/websocket"	internalcache "github.com/router-for-me/CLIProxyAPI/v7/internal/cache"	"github.com/router-for-me/CLIProxyAPI/v7/internal/config"	"github.com/router-for-me/CLIProxyAPI/v7/internal/registry"	cliproxyauth "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/auth"	cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executor"	sdkconfig "github.com/router-for-me/CLIProxyAPI/v7/sdk/config"	sdktranslator "github.com/router-for-me/CLIProxyAPI/v7/sdk/translator"	log "github.com/sirupsen/logrus"	"github.com/tidwall/gjson") var benchmarkBuildCodexWebsocketRequestBodyOutput []byte 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 BenchmarkBuildCodexWebsocketRequestBodyLargePayload(b *testing.B) {	body := []byte(`{"model":"gpt-5.6","input":[{"type":"message","id":"msg_1","role":"user","content":"` + strings.Repeat("x", 8<<20) + `"}]}`)	b.ReportAllocs()	b.SetBytes(int64(len(body)))	b.ResetTimer()	for b.Loop() {		benchmarkBuildCodexWebsocketRequestBodyOutput = buildCodexWebsocketRequestBody(body)	}} func TestBuildCodexWebsocketRequestBodySanitizesOverlongInputItemIDs(t *testing.T) {	longReasoningItemID := "rs_" + strings.Repeat("a", 64)	longCallItemID := strings.Repeat("grok-call-item-", 6)	longOutputItemID := strings.Repeat("grok-output-item-", 6)	body := []byte(`{"model":"gpt-5-codex","input":[{"type":"reasoning","id":"` + longReasoningItemID + `","encrypted_content":"gAAAA-encrypted","summary":[]},{"type":"function_call","id":"` + longCallItemID + `","call_id":"call-1","name":"lookup"},{"type":"function_call_output","id":"` + longOutputItemID + `","call_id":"call-1","output":"ok"},{"type":"message","id":"item_74ec40c883248ebb4885ec84"}]}`) 	first := buildCodexWebsocketRequestBody(body)	second := buildCodexWebsocketRequestBody(body) 	if input := gjson.GetBytes(first, "input").Array(); len(input) != 3 {		t.Fatalf("input length = %d, want 3: %s", len(input), first)	}	if gotType := gjson.GetBytes(first, "input.0.type").String(); gotType != "function_call" {		t.Fatalf("input.0.type = %q, want function_call: %s", gotType, first)	} 	shortCallItemID := gjson.GetBytes(first, "input.0.id").String()	shortOutputItemID := gjson.GetBytes(first, "input.1.id").String()	if len([]rune(shortCallItemID)) > 64 || shortCallItemID == longCallItemID {		t.Fatalf("input.0.id was not shortened to at most 64 characters: %q", shortCallItemID)	}	if len([]rune(shortOutputItemID)) > 64 || shortOutputItemID == longOutputItemID {		t.Fatalf("input.1.id was not shortened to at most 64 characters: %q", shortOutputItemID)	}	if shortCallItemID == shortOutputItemID {		t.Fatalf("distinct long IDs produced the same shortened ID: %q", shortCallItemID)	}	if got := gjson.GetBytes(second, "input.0.id").String(); got != shortCallItemID {		t.Fatalf("input item ID shortening is not deterministic: first=%q second=%q", shortCallItemID, got)	}	if got := gjson.GetBytes(first, "input.0.call_id").String(); got != "call-1" {		t.Fatalf("function call_id = %q, want call-1", got)	}	if got := gjson.GetBytes(first, "input.1.call_id").String(); got != "call-1" {		t.Fatalf("function call output call_id = %q, want call-1", got)	}	if got := gjson.GetBytes(first, "input.2.id").String(); got != "msg_item_74ec40c883248ebb4885ec84" {		t.Fatalf("message input item ID was not normalized: %q", got)	}} func TestCodexWebsocketsExecuteRestoresClaudeAgentReasoningReplay(t *testing.T) {	internalcache.ClearCodexReasoningReplayCache()	t.Cleanup(internalcache.ClearCodexReasoningReplayCache) 	encryptedContent := validCodexReasoningEncryptedContentForTestSeed(31)	cacheCodexReasoningReplayFromCompleted(codexReasoningReplayScope{		modelName:  "gpt-5.4",		sessionKey: "claude:ws-replay-session:agent:agent-a",	}, []byte(`{"response":{"output":[`+		`{"type":"reasoning","summary":[],"content":null,"encrypted_content":"`+encryptedContent+`"},`+		`{"type":"message","role":"assistant","content":[{"type":"output_text","text":"previous answer"}]}`+		`]}}`)) 	upgrader := websocket.Upgrader{CheckOrigin: func(*http.Request) bool { return true }}	capturedPayload := make(chan []byte, 1)	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {		conn, errUpgrade := upgrader.Upgrade(w, r, nil)		if errUpgrade != nil {			t.Fatalf("upgrade websocket: %v", errUpgrade)		}		defer func() { _ = conn.Close() }() 		_, payload, errRead := conn.ReadMessage()		if errRead != nil {			t.Fatalf("read upstream websocket message: %v", errRead)		}		capturedPayload <- bytes.Clone(payload)		completed := []byte(`{"type":"response.completed","response":{"id":"resp-ws-replay","output":[{"type":"message","role":"assistant","content":[{"type":"output_text","text":"next answer"}]}],"usage":{"input_tokens":1,"output_tokens":1,"total_tokens":2}}}`)		if errWrite := conn.WriteMessage(websocket.TextMessage, completed); errWrite != nil {			t.Fatalf("write completed websocket message: %v", errWrite)		}	}))	defer server.Close() 	exec := NewCodexWebsocketsExecutor(&config.Config{SDKConfig: config.SDKConfig{DisableImageGeneration: config.DisableImageGenerationAll}})	auth := &cliproxyauth.Auth{Provider: "codex", Attributes: map[string]string{"api_key": "sk-test", "base_url": server.URL}}	req := cliproxyexecutor.Request{		Model: "gpt-5.4",		Payload: []byte(`{			"model":"gpt-5.4",			"messages":[				{"role":"user","content":"first"},				{"role":"assistant","content":"previous answer"},				{"role":"user","content":"next"}			]		}`),	}	headers := http.Header{}	headers.Set("X-Claude-Code-Session-Id", "ws-replay-session")	headers.Set("X-Claude-Code-Agent-Id", "agent-a")	opts := cliproxyexecutor.Options{SourceFormat: sdktranslator.FromString("claude"), Headers: headers} 	if _, errExecute := exec.Execute(context.Background(), auth, req, opts); errExecute != nil {		t.Fatalf("Execute() error = %v", errExecute)	} 	select {	case payload := <-capturedPayload:		input := gjson.GetBytes(payload, "input").Array()		if len(input) != 4 {			t.Fatalf("upstream input length = %d, want 4; payload=%s", len(input), payload)		}		if input[1].Get("type").String() != "reasoning" || input[1].Get("encrypted_content").String() != encryptedContent {			t.Fatalf("websocket reasoning replay missing before assistant message: %s", payload)		}		if input[2].Get("role").String() != "assistant" {			t.Fatalf("input.2.role = %q, want assistant; payload=%s", input[2].Get("role").String(), payload)		}	case <-time.After(5 * time.Second):		t.Fatal("timed out waiting for upstream websocket payload")	}} func TestClearCodexReasoningReplayOnWebsocketInvalidSignature(t *testing.T) {	internalcache.ClearCodexReasoningReplayCache()	t.Cleanup(internalcache.ClearCodexReasoningReplayCache) 	scope := codexReasoningReplayScope{modelName: "gpt-5.4", sessionKey: "claude:ws-invalid:agent:main"}	encryptedContent := validCodexReasoningEncryptedContentForTestSeed(32)	if !internalcache.CacheCodexReasoningReplayItem(scope.modelName, scope.sessionKey, []byte(`{"type":"reasoning","summary":[],"content":null,"encrypted_content":"`+encryptedContent+`"}`)) {		t.Fatal("failed to seed websocket replay cache")	}	payload := []byte(`{"type":"error","status":400,"body":{"error":{"message":"Invalid signature in thinking block","type":"invalid_request_error","code":"invalid_request_error"}}}`)	if errClear := clearCodexReasoningReplayOnWebsocketError(context.Background(), scope, payload); errClear != nil {		t.Fatalf("clear websocket replay error: %v", errClear)	}	if _, ok := internalcache.GetCodexReasoningReplayItem(scope.modelName, scope.sessionKey); ok {		t.Fatal("websocket invalid signature did not clear replay state")	}} func TestCodexWebsocketsExecuteResponsesLiteDoesNotInjectImageGenerationTool(t *testing.T) {	upgrader := websocket.Upgrader{CheckOrigin: func(*http.Request) bool { return true }}	capturedPayload := make(chan []byte, 1)	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {		conn, err := upgrader.Upgrade(w, r, nil)		if err != nil {			t.Fatalf("upgrade websocket: %v", err)		}		defer func() { _ = conn.Close() }() 		_, payload, errRead := conn.ReadMessage()		if errRead != nil {			t.Fatalf("read upstream websocket message: %v", errRead)		}		capturedPayload <- bytes.Clone(payload) 		completed := []byte(`{"type":"response.completed","response":{"id":"resp-1","output":[],"usage":{"input_tokens":0,"output_tokens":0,"total_tokens":0}}}`)		if errWrite := conn.WriteMessage(websocket.TextMessage, completed); errWrite != nil {			t.Fatalf("write completed websocket message: %v", errWrite)		}	}))	defer server.Close() 	exec := NewCodexWebsocketsExecutor(&config.Config{})	auth := &cliproxyauth.Auth{		Provider: "codex",		Attributes: map[string]string{			"api_key":   "sk-test",			"base_url":  server.URL,			"plan_type": "pro",		},	}	req := cliproxyexecutor.Request{		Model:   "gpt-5.6-sol",		Payload: []byte(`{"model":"gpt-5.6-sol","input":[{"type":"additional_tools","role":"developer","tools":[{"type":"custom","name":"exec"}]},{"role":"user","content":"hello"}],"parallel_tool_calls":true,"client_metadata":{"ws_request_header_x_openai_internal_codex_responses_lite":"true"}}`),	}	opts := cliproxyexecutor.Options{SourceFormat: sdktranslator.FromString("codex")} 	if _, err := exec.Execute(context.Background(), auth, req, opts); err != nil {		t.Fatalf("Execute() error = %v", err)	} 	select {	case payload := <-capturedPayload:		if instructions := gjson.GetBytes(payload, "instructions"); instructions.Exists() {			t.Errorf("unexpected instructions in responses-lite upstream payload: %s", payload)		}		if tools := gjson.GetBytes(payload, "tools"); tools.Exists() {			t.Fatalf("unexpected tools in responses-lite upstream payload: %s", tools.Raw)		}		if got := gjson.GetBytes(payload, "input.0.type").String(); got != "additional_tools" {			t.Fatalf("input.0.type = %q, want additional_tools; payload=%s", got, payload)		}		if got := gjson.GetBytes(payload, "client_metadata.ws_request_header_x_openai_internal_codex_responses_lite").String(); got != "true" {			t.Fatalf("responses-lite metadata = %q, want true; payload=%s", got, payload)		}		parallelToolCalls := gjson.GetBytes(payload, "parallel_tool_calls")		if !parallelToolCalls.Exists() || parallelToolCalls.Bool() {			t.Fatalf("responses-lite parallel_tool_calls should be false: %s", payload)		}	case <-time.After(5 * time.Second):		t.Fatal("timed out waiting for upstream websocket payload")	}} func TestCodexWebsocketsExecuteStreamResponsesLiteForcesParallelToolCallsFalse(t *testing.T) {	upgrader := websocket.Upgrader{CheckOrigin: func(*http.Request) bool { return true }}	capturedPayload := make(chan []byte, 1)	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {		conn, errUpgrade := upgrader.Upgrade(w, r, nil)		if errUpgrade != nil {			t.Errorf("upgrade websocket: %v", errUpgrade)			return		}		defer func() { _ = conn.Close() }() 		_, payload, errRead := conn.ReadMessage()		if errRead != nil {			t.Errorf("read upstream websocket message: %v", errRead)			return		}		capturedPayload <- bytes.Clone(payload) 		completed := []byte(`{"type":"response.completed","response":{"id":"resp-1","output":[],"usage":{"input_tokens":0,"output_tokens":0,"total_tokens":0}}}`)		if errWrite := conn.WriteMessage(websocket.TextMessage, completed); errWrite != nil {			t.Errorf("write completed websocket message: %v", errWrite)		}	}))	defer server.Close() 	exec := NewCodexWebsocketsExecutor(&config.Config{})	auth := &cliproxyauth.Auth{		Provider: "codex",		Attributes: map[string]string{			"api_key":   "sk-test",			"base_url":  server.URL,			"plan_type": "pro",		},	}	req := cliproxyexecutor.Request{		Model:   "gpt-5.6-luna",		Payload: []byte(`{"model":"gpt-5.6-luna","input":[{"type":"additional_tools","role":"developer","tools":[{"type":"custom","name":"exec"}]},{"role":"user","content":"hello"}],"parallel_tool_calls":true,"client_metadata":{"ws_request_header_x_openai_internal_codex_responses_lite":"true"}}`),	}	opts := cliproxyexecutor.Options{SourceFormat: sdktranslator.FromString("codex")} 	result, errExecute := exec.ExecuteStream(context.Background(), auth, req, opts)	if errExecute != nil {		t.Fatalf("ExecuteStream() error = %v", errExecute)	}	streamComplete := false	for !streamComplete {		select {		case chunk, ok := <-result.Chunks:			if !ok {				streamComplete = true				continue			}			if chunk.Err != nil {				t.Fatalf("stream chunk error = %v", chunk.Err)			}		case <-time.After(5 * time.Second):			t.Fatal("timed out waiting for websocket stream completion")		}	} 	select {	case payload := <-capturedPayload:		parallelToolCalls := gjson.GetBytes(payload, "parallel_tool_calls")		if !parallelToolCalls.Exists() || parallelToolCalls.Bool() {			t.Fatalf("responses-lite parallel_tool_calls should be false: %s", payload)		}	case <-time.After(5 * time.Second):		t.Fatal("timed out waiting for upstream websocket payload")	}} func TestCodexWebsocketsExecutePreservesPreviousResponseIDUpstream(t *testing.T) {	upgrader := websocket.Upgrader{CheckOrigin: func(*http.Request) bool { return true }}	capturedPayload := make(chan []byte, 1)	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {		if r.URL.Path != "/responses" {			t.Fatalf("request path = %s, want /responses", r.URL.Path)		}		conn, err := upgrader.Upgrade(w, r, nil)		if err != nil {			t.Fatalf("upgrade websocket: %v", err)		}		defer func() { _ = conn.Close() }() 		msgType, payload, err := conn.ReadMessage()		if err != nil {			t.Fatalf("read upstream websocket message: %v", err)		}		if msgType != websocket.TextMessage {			t.Fatalf("message type = %d, want text", msgType)		}		capturedPayload <- bytes.Clone(payload) 		completed := []byte(`{"type":"response.completed","response":{"id":"resp-2","output":[],"usage":{"input_tokens":0,"output_tokens":0,"total_tokens":0}}}`)		if errWrite := conn.WriteMessage(websocket.TextMessage, completed); errWrite != nil {			t.Fatalf("write completed websocket message: %v", errWrite)		}	}))	defer server.Close() 	exec := NewCodexWebsocketsExecutor(&config.Config{SDKConfig: config.SDKConfig{DisableImageGeneration: config.DisableImageGenerationAll}})	auth := &cliproxyauth.Auth{Attributes: map[string]string{"api_key": "sk-test", "base_url": server.URL}}	req := cliproxyexecutor.Request{		Model:   "gpt-5-codex",		Payload: []byte(`{"model":"gpt-5-codex","previous_response_id":"resp-1","input":[{"type":"message","id":"msg-1"}]}`),	}	opts := cliproxyexecutor.Options{SourceFormat: sdktranslator.FromString("codex")} 	if _, err := exec.Execute(context.Background(), auth, req, opts); err != nil {		t.Fatalf("Execute() error = %v", err)	} 	select {	case payload := <-capturedPayload:		if got := gjson.GetBytes(payload, "type").String(); got != "response.create" {			t.Fatalf("upstream type = %s, want response.create; payload=%s", got, payload)		}		if got := gjson.GetBytes(payload, "previous_response_id").String(); got != "resp-1" {			t.Fatalf("upstream previous_response_id = %s, want resp-1; payload=%s", got, payload)		}	case <-time.After(5 * time.Second):		t.Fatal("timed out waiting for upstream websocket payload")	}} func TestCodexWebsocketsExecuteStreamUpgradeRequiredReturnsWithoutLockingSession(t *testing.T) {	upgradeAttempts := make(chan struct{}, 2)	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {		if !strings.EqualFold(r.Header.Get("Upgrade"), "websocket") {			t.Errorf("unexpected HTTP fallback request: %s %s", r.Method, r.URL.Path)			w.WriteHeader(http.StatusInternalServerError)			return		}		upgradeAttempts <- struct{}{}		w.WriteHeader(http.StatusUpgradeRequired)		_, _ = w.Write([]byte(`{"error":{"message":"websocket unavailable"}}`))	}))	defer server.Close() 	exec := NewCodexWebsocketsExecutor(&config.Config{SDKConfig: config.SDKConfig{DisableImageGeneration: config.DisableImageGenerationAll}})	const executionSessionID = "ws-upgrade-required-session"	t.Cleanup(func() { exec.CloseExecutionSession(executionSessionID) })	auth := &cliproxyauth.Auth{		ID:       "codex-test",		Provider: "codex",		Attributes: map[string]string{			"api_key":  "sk-test",			"base_url": server.URL,		},	}	opts := cliproxyexecutor.Options{		SourceFormat:   sdktranslator.FromString("openai-response"),		ResponseFormat: sdktranslator.FromString("openai-response"),		Metadata: map[string]any{			cliproxyexecutor.ExecutionSessionMetadataKey: executionSessionID,		},	}	ctx := cliproxyexecutor.WithDownstreamWebsocket(context.Background()) 	execute := func(payload string) {		t.Helper()		done := make(chan error, 1)		go func() {			_, errExecute := exec.ExecuteStream(ctx, auth, cliproxyexecutor.Request{				Model:   "gpt-5.4",				Payload: []byte(payload),			}, opts)			done <- errExecute		}() 		select {		case errExecute := <-done:			if errExecute == nil {				t.Fatal("upgrade-required error = nil")			}			statusErr, ok := errExecute.(interface{ StatusCode() int })			if !ok || statusErr.StatusCode() != http.StatusUpgradeRequired {				t.Fatalf("upgrade-required error = %T %v, want status 426", errExecute, errExecute)			}		case <-time.After(5 * time.Second):			t.Fatal("timed out waiting for upgrade-required error; execution session may still be locked")		}	} 	execute(`{"model":"gpt-5.4","generate":false,"input":[]}`)	execute(`{"model":"gpt-5.4","previous_response_id":"resp-1","input":[{"type":"message","id":"msg-2"}]}`) 	if got := len(upgradeAttempts); got != 2 {		t.Fatalf("websocket upgrade attempts = %d, want 2", got)	}} func TestCodexWebsocketsExecuteStreamHandshakeErrorReturnsWithoutLockingSession(t *testing.T) {	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {		w.WriteHeader(http.StatusUnauthorized)		_, _ = w.Write([]byte(`{"error":{"message":"unauthorized"}}`))	}))	defer server.Close() 	exec := NewCodexWebsocketsExecutor(&config.Config{SDKConfig: config.SDKConfig{DisableImageGeneration: config.DisableImageGenerationAll}})	const executionSessionID = "ws-handshake-error-session"	t.Cleanup(func() { exec.CloseExecutionSession(executionSessionID) })	auth := &cliproxyauth.Auth{		ID:       "codex-test",		Provider: "codex",		Attributes: map[string]string{			"api_key":  "sk-test",			"base_url": server.URL,		},	}	opts := cliproxyexecutor.Options{		SourceFormat: sdktranslator.FromString("openai-response"),		Metadata: map[string]any{			cliproxyexecutor.ExecutionSessionMetadataKey: executionSessionID,		},	} 	for i := 0; i < 2; i++ {		attemptCtx := cliproxyexecutor.WithUpstreamAttemptTracker(context.Background())		done := make(chan error, 1)		go func(ctx context.Context) {			_, errExecute := exec.ExecuteStream(ctx, auth, cliproxyexecutor.Request{				Model:   "gpt-5.4",				Payload: []byte(`{"model":"gpt-5.4","input":[{"type":"message","id":"msg-1"}]}`),			}, opts)			done <- errExecute		}(attemptCtx)		select {		case errExecute := <-done:			statusErr, ok := errExecute.(interface{ StatusCode() int })			if !ok || statusErr.StatusCode() != http.StatusUnauthorized {				t.Fatalf("attempt %d error = %T %v, want status 401", i+1, errExecute, errExecute)			}			if !cliproxyexecutor.UpstreamAttempted(attemptCtx) {				t.Fatalf("attempt %d websocket handshake was not marked as upstream", i+1)			}		case <-time.After(5 * time.Second):			t.Fatalf("attempt %d timed out; execution session remained locked", i+1)		}	}} func TestExistingWebsocketSessionConnRequiresMatchingHealthyConnection(t *testing.T) {

Lines 1–500 of 3405. Load more