package executor import (	"fmt"	"net/http"	"net/http/httptest"	"strings"	"sync/atomic"	"testing" 	"github.com/gorilla/websocket"	"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"	sdktranslator "github.com/router-for-me/CLIProxyAPI/v7/sdk/translator"	"github.com/tidwall/gjson") func TestCodexWebsocketsExecutorRestoresMultiAgentV2NamespaceAcrossIncrementalTurns(t *testing.T) {	for _, tt := range []struct {		name   string		stream bool	}{		{name: "execute"},		{name: "stream", stream: true},	} {		t.Run(tt.name, func(t *testing.T) {			upgrader := websocket.Upgrader{CheckOrigin: func(*http.Request) bool { return true }}			capturedPayload := make(chan []byte, 6)			var connectionCount atomic.Int32			var requestCount atomic.Int32			server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, request *http.Request) {				connectionCount.Add(1)				conn, errUpgrade := upgrader.Upgrade(w, request, nil)				if errUpgrade != nil {					t.Errorf("upgrade websocket: %v", errUpgrade)					return				}				defer func() { _ = conn.Close() }() 				for {					_, payload, errRead := conn.ReadMessage()					if errRead != nil {						return					}					capturedPayload <- append([]byte(nil), payload...)					turn := requestCount.Add(1)					completed := []byte(fmt.Sprintf(`{"type":"response.completed","response":{"id":"resp_%d","object":"response","status":"completed","output":[{"type":"function_call","name":"spawn_agent","namespace":"collaboration-optimize","arguments":"{}","call_id":"call_%d"}]}}`, turn, turn))					if errWrite := conn.WriteMessage(websocket.TextMessage, completed); errWrite != nil {						t.Errorf("write websocket response: %v", errWrite)						return					}					if turn == 6 {						return					}				}			}))			t.Cleanup(server.Close) 			executor := NewCodexWebsocketsExecutor(&config.Config{Codex: config.CodexConfig{OptimizeMultiAgentV2: true}})			const executionSessionID = "multi-agent-v2-incremental"			t.Cleanup(func() { executor.CloseExecutionSession(executionSessionID) })			auth := &cliproxyauth.Auth{				ID: "codex-test",				Attributes: map[string]string{					"base_url": server.URL,					"api_key":  "test",				},			}			opts := cliproxyexecutor.Options{				SourceFormat:   sdktranslator.FromString("openai-response"),				ResponseFormat: sdktranslator.FromString("openai-response"),				Headers:        http.Header{"User-Agent": []string{"overridden-client/1.0"}},				Metadata: map[string]any{					cliproxyexecutor.ExecutionSessionMetadataKey: executionSessionID,				},			}			execute := func(payload []byte) []byte {				t.Helper()				req := cliproxyexecutor.Request{Model: "gpt-5.4", Payload: payload}				if !tt.stream {					response, errExecute := executor.Execute(codexSpawnAgentTestContext(), auth, req, opts)					if errExecute != nil {						t.Fatalf("Execute() error = %v", errExecute)					}					return response.Payload				} 				result, errExecute := executor.ExecuteStream(codexSpawnAgentTestContext(), auth, req, opts)				if errExecute != nil {					t.Fatalf("ExecuteStream() error = %v", errExecute)				}				var responsePayload []byte				for chunk := range result.Chunks {					if chunk.Err != nil {						t.Fatalf("stream chunk error = %v", chunk.Err)					}					responsePayload = append(responsePayload, chunk.Payload...)				}				return responsePayload			} 			firstClientPayload := execute(codexSpawnAgentTestPayload())			firstUpstreamPayload := <-capturedPayload			if namespace := gjson.GetBytes(firstUpstreamPayload, "input.0.tools.0.name").String(); namespace != "collaboration-optimize" {				t.Fatalf("first upstream namespace = %q, want collaboration-optimize", namespace)			}			assertCodexSpawnAgentClientNamespace(t, firstClientPayload) 			secondRequest := []byte(`{"model":"gpt-5.4","previous_response_id":"resp_1","input":[{"type":"function_call_output","call_id":"call_1","output":"done"}]}`)			secondClientPayload := execute(secondRequest)			secondUpstreamPayload := <-capturedPayload			if strings.Contains(string(secondUpstreamPayload), "collaboration") || strings.Contains(string(secondUpstreamPayload), "spawn_agent") {				t.Fatalf("incremental upstream request unexpectedly contains collaboration tools: %s", secondUpstreamPayload)			}			assertCodexSpawnAgentClientNamespace(t, secondClientPayload) 			conflictingRequest := []byte(`{"model":"gpt-5.4","tools":[{"type":"namespace","name":"collaboration-optimize","tools":[{"type":"function","name":"spawn_agent","description":"User-defined tool."}]}],"input":[{"type":"message","role":"user","content":"use the user-defined namespace"}]}`)			conflictingClientPayload := execute(conflictingRequest)			conflictingUpstreamPayload := <-capturedPayload			if namespace := gjson.GetBytes(conflictingUpstreamPayload, "tools.0.name").String(); namespace != "collaboration-optimize" {				t.Fatalf("conflicting upstream namespace = %q, want collaboration-optimize", namespace)			}			if !strings.Contains(string(conflictingClientPayload), `"namespace":"collaboration-optimize"`) {				t.Fatalf("user-defined collaboration-optimize namespace was rewritten: %s", conflictingClientPayload)			} 			fourthRequest := []byte(`{"model":"gpt-5.4","previous_response_id":"resp_3","input":[{"type":"function_call_output","call_id":"call_3","output":"done"}]}`)			fourthClientPayload := execute(fourthRequest)			fourthUpstreamPayload := <-capturedPayload			if strings.Contains(string(fourthUpstreamPayload), "collaboration") || strings.Contains(string(fourthUpstreamPayload), "spawn_agent") {				t.Fatalf("post-conflict incremental upstream request unexpectedly contains collaboration tools: %s", fourthUpstreamPayload)			}			if !strings.Contains(string(fourthClientPayload), `"namespace":"collaboration-optimize"`) {				t.Fatalf("user-defined namespace was rewritten on the post-conflict incremental turn: %s", fourthClientPayload)			} 			fifthClientPayload := execute(codexSpawnAgentTestPayload())			fifthUpstreamPayload := <-capturedPayload			if namespace := gjson.GetBytes(fifthUpstreamPayload, "input.0.tools.0.name").String(); namespace != "collaboration-optimize" {				t.Fatalf("re-enabled upstream namespace = %q, want collaboration-optimize", namespace)			}			assertCodexSpawnAgentClientNamespace(t, fifthClientPayload) 			sixthRequest := []byte(`{"model":"gpt-5.4","previous_response_id":"resp_5","input":[{"type":"function_call_output","call_id":"call_5","output":"done"}]}`)			sixthClientPayload := execute(sixthRequest)			sixthUpstreamPayload := <-capturedPayload			if strings.Contains(string(sixthUpstreamPayload), "collaboration") || strings.Contains(string(sixthUpstreamPayload), "spawn_agent") {				t.Fatalf("re-enabled incremental upstream request unexpectedly contains collaboration tools: %s", sixthUpstreamPayload)			}			assertCodexSpawnAgentClientNamespace(t, sixthClientPayload) 			if got := connectionCount.Load(); got != 1 {				t.Fatalf("upstream websocket connections = %d, want 1", got)			}		})	}} func TestCodexWebsocketsExecutorOptimizeMultiAgentV2(t *testing.T) {	modelID := "codex-websocket-spawn-agent-test-model"	clientID := "codex-websocket-spawn-agent-test-client"	modelRegistry := registry.GetGlobalRegistry()	modelRegistry.RegisterClient(clientID, "codex", []*registry.ModelInfo{{		ID:          modelID,		Description: "Executor test model.",		Thinking: &registry.ThinkingSupport{			Levels: []string{"low", "medium", "high"},		},	}})	defer modelRegistry.UnregisterClient(clientID) 	upgrader := websocket.Upgrader{CheckOrigin: func(*http.Request) bool { return true }}	capturedPayload := make(chan []byte, 2)	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, request *http.Request) {		conn, errUpgrade := upgrader.Upgrade(w, request, 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 websocket request: %v", errRead)			return		}		capturedPayload <- payload		namespace := gjson.GetBytes(payload, "input.0.tools.0.name").String()		completed := []byte(fmt.Sprintf(`{"type":"response.completed","response":{"id":"resp_1","object":"response","status":"completed","output":[{"type":"function_call","name":"spawn_agent","namespace":%q,"arguments":"{}","call_id":"call_1"}]}}`, namespace))		if errWrite := conn.WriteMessage(websocket.TextMessage, completed); errWrite != nil {			t.Errorf("write websocket response: %v", errWrite)		}	}))	defer server.Close() 	auth := &cliproxyauth.Auth{Attributes: map[string]string{		"base_url": server.URL,		"api_key":  "test",	}}	req := cliproxyexecutor.Request{Model: "gpt-5.4", Payload: codexSpawnAgentTestPayload()}	opts := cliproxyexecutor.Options{		SourceFormat: sdktranslator.FromString("openai-response"),		Headers:      http.Header{"User-Agent": []string{"overridden-client/1.0"}},	} 	for _, tt := range []struct {		name    string		enabled bool		stream  bool	}{		{name: "execute enabled", enabled: true},		{name: "execute disabled", enabled: false},		{name: "stream enabled", enabled: true, stream: true},		{name: "stream disabled", enabled: false, stream: true},	} {		t.Run(tt.name, func(t *testing.T) {			executor := NewCodexWebsocketsExecutor(&config.Config{Codex: config.CodexConfig{OptimizeMultiAgentV2: tt.enabled}})			var clientPayload []byte			if tt.stream {				result, errExecute := executor.ExecuteStream(codexSpawnAgentTestContext(), auth, req, opts)				if errExecute != nil {					t.Fatalf("ExecuteStream() error = %v", errExecute)				}				for chunk := range result.Chunks {					clientPayload = append(clientPayload, chunk.Payload...)				}			} else {				response, errExecute := executor.Execute(codexSpawnAgentTestContext(), auth, req, opts)				if errExecute != nil {					t.Fatalf("Execute() error = %v", errExecute)				}				clientPayload = response.Payload			}			upstreamPayload := <-capturedPayload			assertCodexSpawnAgentOptimization(t, upstreamPayload, modelID, tt.enabled)			assertCodexSpawnAgentRequestMessage(t, upstreamPayload, tt.enabled)			assertCodexSpawnAgentClientNamespace(t, clientPayload)		})	}}