package executor import (	"bytes"	"context"	"errors"	"fmt"	"io"	"net/http"	"net/http/httptest"	"strings"	"testing"	"time" 	"github.com/gorilla/websocket"	"github.com/router-for-me/CLIProxyAPI/v7/internal/config"	_ "github.com/router-for-me/CLIProxyAPI/v7/internal/translator"	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 TestXAIWebsocketsEnabledForConfigAPIKey(t *testing.T) {	auth := &cliproxyauth.Auth{		Provider: "xai",		Attributes: map[string]string{			"api_key":    "xai-key",			"websockets": "true",		},	}	if !xaiWebsocketsEnabled(auth) {		t.Fatal("xaiWebsocketsEnabled() = false, want true")	}} func TestXAIAutoExecutorRequiredUpstreamWebsocketRejectsHTTPFallback(t *testing.T) {	exec := NewXAIAutoExecutor(&config.Config{})	auth := &cliproxyauth.Auth{		ID:       "xai-http-only",		Provider: "xai",		Attributes: map[string]string{			"api_key": "xai-key",		},	}	ctx := cliproxyexecutor.WithRequiredUpstreamWebsocket(		cliproxyexecutor.WithDownstreamWebsocket(context.Background()),	)	_, errExecute := exec.ExecuteStream(ctx, auth, cliproxyexecutor.Request{		Model:   "grok-4",		Payload: []byte(`{"model":"grok-4","previous_response_id":"resp-1","input":[{"type":"message","id":"msg-2"}]}`),	}, cliproxyexecutor.Options{SourceFormat: sdktranslator.FromString("openai-response")})	if errExecute == nil {		t.Fatal("ExecuteStream() error = nil, want replay-required error")	}	statusErr, ok := errExecute.(interface{ StatusCode() int })	if !ok || statusErr.StatusCode() != http.StatusUpgradeRequired {		t.Fatalf("ExecuteStream() error = %T %v, want status 426", errExecute, errExecute)	}	if got := gjson.Get(errExecute.Error(), "error.code").String(); got != "upstream_http_replay_required" {		t.Fatalf("ExecuteStream() error code = %q, want upstream_http_replay_required", got)	}	requestScoped, ok := errExecute.(cliproxyexecutor.RequestScopedError)	if !ok || !requestScoped.IsRequestScoped() {		t.Fatalf("ExecuteStream() error = %T, want request-scoped replay signal", errExecute)	}} func TestXAIWebsocketsRequiredUpstreamRejectsCompactionHTTPFallback(t *testing.T) {	exec := NewXAIWebsocketsExecutor(&config.Config{})	ctx := cliproxyexecutor.WithRequiredUpstreamWebsocket(context.Background())	_, errExecute := exec.ExecuteStream(ctx, &cliproxyauth.Auth{}, cliproxyexecutor.Request{		Model:   "grok-4",		Payload: []byte(`{"model":"grok-4","input":[{"type":"compaction_trigger"}]}`),	}, cliproxyexecutor.Options{})	if !cliproxyexecutor.IsUpstreamWebsocketReplayRequired(errExecute) {		t.Fatalf("ExecuteStream() error = %T %v, want replay-required", errExecute, errExecute)	}} func TestXAIWebsocketMissingRequiredSessionDoesNotMarkUpstreamAttempt(t *testing.T) {	exec := NewXAIWebsocketsExecutor(&config.Config{})	exec.store = &codexWebsocketSessionStore{sessions: make(map[string]*codexWebsocketSession)}	auth := &cliproxyauth.Auth{		ID:       "xai-required-session",		Provider: "xai",		Attributes: map[string]string{			"api_key": "xai-key",		},	}	ctx := cliproxyexecutor.WithUpstreamAttemptTracker(		cliproxyexecutor.WithRequiredUpstreamWebsocket(context.Background()),	)	_, errExecute := exec.ExecuteStream(ctx, auth, cliproxyexecutor.Request{		Model:   "grok-4",		Payload: []byte(`{"model":"grok-4","previous_response_id":"resp-1","input":[{"type":"message","role":"user","content":"hello"}]}`),	}, cliproxyexecutor.Options{		SourceFormat: sdktranslator.FormatOpenAIResponse,		Metadata: map[string]any{			cliproxyexecutor.ExecutionSessionMetadataKey: "missing-xai-session",		},	})	if !cliproxyexecutor.IsUpstreamWebsocketReplayRequired(errExecute) {		t.Fatalf("ExecuteStream() error = %T %v, want replay-required", errExecute, errExecute)	}	if cliproxyexecutor.UpstreamAttempted(ctx) {		t.Fatal("missing retained websocket connection was marked as an upstream attempt")	}} func TestXAIWebsocketSuccessfulHandshakeDoesNotMarkRequestAttempt(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, errUpgrade := upgrader.Upgrade(w, r, nil)		if errUpgrade != nil {			t.Errorf("upgrade websocket: %v", errUpgrade)			return		}		defer func() { _ = conn.Close() }()		_, _, _ = conn.ReadMessage()	}))	defer server.Close() 	exec := NewXAIWebsocketsExecutor(&config.Config{})	ctx := cliproxyexecutor.WithUpstreamAttemptTracker(context.Background())	conn, closer, resp, errDial := exec.dialXAIWebsocket(ctx, &cliproxyauth.Auth{}, strings.Replace(server.URL, "http", "ws", 1), http.Header{})	if errDial != nil || conn == nil {		t.Fatalf("dialXAIWebsocket() = (%p, %v), want successful connection", conn, errDial)	}	if resp != nil && resp.Body != nil {		defer func() { _ = resp.Body.Close() }()	}	defer func() { _ = closer.Close() }()	if cliproxyexecutor.UpstreamAttempted(ctx) {		t.Fatal("successful handshake was marked before an upstream request was sent")	}} func TestMapXAIWebsocketWriteErrorStopsRetryForMessageTooBig(t *testing.T) {	networkWriteErr := errors.New("write: broken pipe")	tests := []struct {		name       string		closeCode  int		writeErr   error		wantStatus int		wantRetry  bool	}{		{			name:       "close sent after message too big is request scoped",			closeCode:  websocket.CloseMessageTooBig,			writeErr:   websocket.ErrCloseSent,			wantStatus: http.StatusRequestEntityTooLarge,			wantRetry:  false,		},		{			name:       "network write error after message too big is request scoped",			closeCode:  websocket.CloseMessageTooBig,			writeErr:   networkWriteErr,			wantStatus: http.StatusRequestEntityTooLarge,			wantRetry:  false,		},		{			name:      "other close keeps stale connection retry",			closeCode: websocket.CloseNormalClosure,			writeErr:  websocket.ErrCloseSent,			wantRetry: true,		},	} 	for _, tt := range tests {		t.Run(tt.name, func(t *testing.T) {			sess := &codexWebsocketSession{}			conn := &websocket.Conn{}			sess.resetUpstreamDisconnectError(conn)			sess.setUpstreamDisconnectError(conn, &websocket.CloseError{Code: tt.closeCode}) 			mappedErr := mapXAIWebsocketWriteError(sess, conn, tt.writeErr)			if got := shouldRetryXAIWebsocketSend(mappedErr); got != tt.wantRetry {				t.Fatalf("shouldRetryXAIWebsocketSend() = %v, want %v; err=%v", got, tt.wantRetry, mappedErr)			}			if tt.wantStatus == 0 {				if !errors.Is(mappedErr, tt.writeErr) {					t.Fatalf("mapped error = %v, want %v", mappedErr, tt.writeErr)				}				return			}			statusErr, ok := mappedErr.(interface{ StatusCode() int })			if !ok || statusErr.StatusCode() != tt.wantStatus {				t.Fatalf("mapped status = %v, want %d; err=%v", statusErr, tt.wantStatus, mappedErr)			}			requestErr, ok := mappedErr.(interface{ IsRequestScoped() bool })			if !ok || !requestErr.IsRequestScoped() {				t.Fatalf("mapped error should be request scoped, got %T", mappedErr)			}		})	}} func TestXAIWebsocketsExecuteStreamMapsMessageTooBigClose(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() }() 		if _, _, errRead := conn.ReadMessage(); errRead != nil {			t.Errorf("read upstream websocket message: %v", errRead)			return		}		deadline := time.Now().Add(time.Second)		closeMessage := websocket.FormatCloseMessage(websocket.CloseMessageTooBig, "message too big")		if errWrite := conn.WriteControl(websocket.CloseMessage, closeMessage, deadline); errWrite != nil {			t.Errorf("write close websocket message: %v", errWrite)		}	}))	defer server.Close() 	exec := NewXAIWebsocketsExecutor(&config.Config{})	auth := &cliproxyauth.Auth{		Provider: "xai",		Attributes: map[string]string{			"base_url":   server.URL,			"websockets": "true",		},		Metadata: map[string]any{"access_token": "xai-token"},	}	req := cliproxyexecutor.Request{		Model:   "grok-4.5",		Payload: []byte(`{"model":"grok-4.5","input":[{"type":"message","role":"user","content":"hello"}]}`),	}	opts := cliproxyexecutor.Options{		SourceFormat:   sdktranslator.FormatOpenAIResponse,		ResponseFormat: sdktranslator.FormatOpenAIResponse,	} 	result, err := exec.ExecuteStream(context.Background(), 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 chunk.Err == nil {			t.Fatal("error chunk Err = nil, want message-too-big error")		}		statusErr, ok := chunk.Err.(interface{ StatusCode() int })		if !ok || statusErr.StatusCode() != http.StatusRequestEntityTooLarge {			t.Fatalf("status error = %v, want %d; err=%v", statusErr, http.StatusRequestEntityTooLarge, chunk.Err)		}		if got := gjson.Get(chunk.Err.Error(), "error.code").String(); got != "message_too_big" {			t.Fatalf("error code = %q, want message_too_big; err=%v", got, chunk.Err)		}		requestErr, ok := chunk.Err.(interface{ IsRequestScoped() bool })		if !ok || !requestErr.IsRequestScoped() {			t.Fatalf("message-too-big error should be request scoped, got %T", chunk.Err)		}	case <-time.After(5 * time.Second):		t.Fatal("timed out waiting for error stream chunk")	}} func TestXAIWebsocketsExecuteStreamSendsResponseCreateWithPreviousResponseID(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.Errorf("path = %q, want /responses", r.URL.Path)		}		if got := r.Header.Get("Authorization"); got != "Bearer xai-token" {			t.Errorf("Authorization = %q, want Bearer xai-token", got)		}		if got := r.Header.Get("x-grok-conv-id"); got != "execution-session-1" {			t.Errorf("x-grok-conv-id = %q, want execution-session-1", got)		}		conn, err := upgrader.Upgrade(w, r, nil)		if err != nil {			t.Errorf("upgrade websocket: %v", err)			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-xai-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 := NewXAIWebsocketsExecutor(&config.Config{})	auth := &cliproxyauth.Auth{		ID:       "xai-auth",		Provider: "xai",		Attributes: map[string]string{			"base_url":   server.URL,			"websockets": "true",		},		Metadata: map[string]any{"access_token": "xai-token"},	}	req := cliproxyexecutor.Request{		Model:   "grok-4.3",		Payload: []byte(`{"model":"grok-4.3","stream":true,"previous_response_id":"resp-prev","instructions":"system prompt","input":[{"type":"message","role":"user","content":"hello"}]}`),	}	opts := cliproxyexecutor.Options{		SourceFormat:   sdktranslator.FormatOpenAIResponse,		ResponseFormat: sdktranslator.FormatOpenAIResponse,		Metadata: map[string]any{			cliproxyexecutor.ExecutionSessionMetadataKey: "execution-session-1",		},	}	ctx := cliproxyexecutor.WithUpstreamAttemptTracker(		cliproxyexecutor.WithDownstreamWebsocket(context.Background()),	) 	result, err := exec.ExecuteStream(ctx, auth, req, opts)	if err != nil {		t.Fatalf("ExecuteStream() error = %v", err)	}	if !cliproxyexecutor.UpstreamAttempted(ctx) {		t.Fatal("websocket write did not mark an upstream attempt")	} 	select {	case payload := <-capturedPayload:		if got := gjson.GetBytes(payload, "type").String(); got != "response.create" {			t.Fatalf("type = %q, want response.create; payload=%s", got, payload)		}		if got := gjson.GetBytes(payload, "previous_response_id").String(); got != "resp-prev" {			t.Fatalf("previous_response_id = %q, want resp-prev; payload=%s", got, payload)		}		if gjson.GetBytes(payload, "stream").Exists() {			t.Fatalf("stream must be omitted for xAI websocket payload: %s", payload)		}		if gjson.GetBytes(payload, "instructions").Exists() {			t.Fatalf("instructions must be omitted when previous_response_id is set: %s", payload)		}		if got := gjson.GetBytes(payload, "prompt_cache_key").String(); got != "execution-session-1" {			t.Fatalf("prompt_cache_key = %q, want execution-session-1; payload=%s", got, payload)		}		if got := gjson.GetBytes(payload, "store").Bool(); !got {			t.Fatalf("store = false, want true; payload=%s", payload)		}	case <-time.After(5 * time.Second):		t.Fatal("timed out waiting for upstream websocket payload")	} 	select {	case chunk, ok := <-result.Chunks:		if !ok {			t.Fatal("stream closed before completed chunk")		}		if chunk.Err != nil {			t.Fatalf("chunk error = %v", chunk.Err)		}		if got := gjson.GetBytes(bytes.TrimSpace(chunk.Payload), "type").String(); got != "response.completed" {			t.Fatalf("chunk type = %q, want response.completed; payload=%s", got, chunk.Payload)		}	case <-time.After(5 * time.Second):		t.Fatal("timed out waiting for completed chunk")	}} func TestXAIWebsocketsExecuteStreamRestoresNamespaceToolCalls(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.Errorf("upgrade websocket: %v", err)			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) 		events := [][]byte{			[]byte(`{"type":"response.output_item.done","output_index":0,"item":{"type":"function_call","name":"mcp__exa__web_search_exa","call_id":"call_1","arguments":"{}"}}`),			[]byte(`{"type":"response.completed","response":{"id":"resp_1","output":[],"usage":{"input_tokens":1,"output_tokens":1,"total_tokens":2}}}`),		}		for _, event := range events {			if errWrite := conn.WriteMessage(websocket.TextMessage, event); errWrite != nil {				t.Errorf("write websocket event: %v", errWrite)				return			}		}	}))	defer server.Close() 	exec := NewXAIWebsocketsExecutor(&config.Config{})	auth := &cliproxyauth.Auth{		Provider: "xai",		Attributes: map[string]string{			"base_url":   server.URL,			"websockets": "true",		},		Metadata: map[string]any{"access_token": "xai-token"},	}	req := cliproxyexecutor.Request{		Model: "grok-4.3",		Payload: []byte(`{			"model":"grok-4.3",			"input":[				{"type":"additional_tools","role":"developer","tools":[{					"type":"namespace",					"name":"mcp__exa",					"tools":[{"type":"function","name":"web_search_exa","parameters":{"type":"object"}}]				}]},				{"role":"user","content":"use Exa"}			]		}`),	}	opts := cliproxyexecutor.Options{		SourceFormat:   sdktranslator.FormatOpenAIResponse,		ResponseFormat: sdktranslator.FormatOpenAIResponse,		Stream:         true,	}	ctx := cliproxyexecutor.WithDownstreamWebsocket(context.Background()) 	result, err := exec.ExecuteStream(ctx, auth, req, opts)	if err != nil {		t.Fatalf("ExecuteStream() error = %v", err)	} 	select {	case payload := <-capturedPayload:		for _, item := range gjson.GetBytes(payload, "input").Array() {			if got := item.Get("type").String(); got == "additional_tools" {				t.Fatalf("upstream input contains unsupported additional_tools item: %s", payload)			}		}		if got := gjson.GetBytes(payload, "input.0.role").String(); got != "user" {			t.Fatalf("input.0.role = %q, want user; payload=%s", got, payload)		}		tool := gjson.GetBytes(payload, "tools.0")		if got := tool.Get("name").String(); got != "mcp__exa__web_search_exa" {			t.Fatalf("upstream tool name = %q, want qualified name; payload=%s", got, payload)		}		if tool.Get("tools").Exists() {			t.Fatalf("upstream tool should not contain namespace children: %s", payload)		}	case <-time.After(5 * time.Second):		t.Fatal("timed out waiting for upstream websocket payload")	} 	var outputItemDone, completed gjson.Result	for chunk := range result.Chunks {		if chunk.Err != nil {			t.Fatalf("stream chunk error = %v", chunk.Err)		}		payload := gjson.ParseBytes(bytes.TrimSpace(chunk.Payload))		switch payload.Get("type").String() {		case "response.output_item.done":			outputItemDone = payload		case "response.completed":			completed = payload		}	} 	for label, item := range map[string]gjson.Result{		"output_item.done": outputItemDone.Get("item"),		"completed":        completed.Get("response.output.0"),	} {		if got := item.Get("name").String(); got != "web_search_exa" {			t.Fatalf("%s name = %q, want child name; item=%s", label, got, item.Raw)		}		if got := item.Get("namespace").String(); got != "mcp__exa" {			t.Fatalf("%s namespace = %q, want mcp__exa; item=%s", label, got, item.Raw)		}	}} func TestXAIWebsocketsExecuteStreamPreservesClientSameNameToolsWithXSearch(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() }() 		if _, _, errRead := conn.ReadMessage(); errRead != nil {			t.Errorf("read upstream websocket message: %v", errRead)			return

Lines 1–500 of 2177. Load more