package executor import ( "bytes" "context" "crypto/sha256" "encoding/base64" "encoding/json" "errors" "fmt" "io" "net/http" "net/http/httptest" "strings" "testing" "time" "github.com/gin-gonic/gin" "github.com/google/uuid" xaiauth "github.com/router-for-me/CLIProxyAPI/v7/internal/auth/xai" 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/translator" cliproxyauth "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/auth" cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executor" "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/usage" sdktranslator "github.com/router-for-me/CLIProxyAPI/v7/sdk/translator" "github.com/tidwall/gjson" "github.com/tidwall/sjson" "github.com/tiktoken-go/tokenizer") func testContextWithAPIKey(apiKey string) context.Context { gin.SetMode(gin.TestMode) rec := httptest.NewRecorder() ginCtx, _ := gin.CreateTestContext(rec) ginCtx.Set("userApiKey", apiKey) return context.WithValue(context.Background(), "gin", ginCtx)} func TestCountXAIInputTokensExcludesRequestStructure(t *testing.T) { enc, err := tokenizer.Get(tokenizer.O200kBase) if err != nil { t.Fatalf("tokenizer.Get() error = %v", err) } semanticBody := []byte(`{ "instructions":"Follow the repository instructions.", "input":[ {"type":"message","role":"user","content":[{"type":"input_text","text":"Review this implementation."}]}, {"type":"function_call","name":"read_file","arguments":"{\"path\":\"main.go\"}"}, {"type":"function_call_output","output":"package main"}, {"type":"reasoning","summary":[{"type":"summary_text","text":"I will inspect the file."}]} ], "tools":[{"type":"function","name":"read_file","description":"Reads a file.","parameters":{"type":"object","properties":{"path":{"type":"string"}}}}], "text":{"format":{"name":"result","schema":{"type":"object"}}} }`) structuralBody := []byte(`{ "model":"grok-4.5", "stream":false, "reasoning":{"effort":"high"}, "metadata":{"large_wrapper":"this metadata must not affect estimated input tokens"}, "prompt_cache_key":"session-123", "max_output_tokens":4096, "instructions":"Follow the repository instructions.", "input":[ {"type":"message","role":"user","content":[{"type":"input_text","text":"Review this implementation."}]}, {"type":"function_call","name":"read_file","arguments":"{\"path\":\"main.go\"}"}, {"type":"function_call_output","output":"package main"}, {"type":"reasoning","summary":[{"type":"summary_text","text":"I will inspect the file."}]} ], "tools":[{"type":"function","name":"read_file","description":"Reads a file.","parameters":{"type":"object","properties":{"path":{"type":"string"}}}}], "text":{"format":{"name":"result","schema":{"type":"object"}}} }`) semanticCount, err := countXAIInputTokens(enc, semanticBody) if err != nil { t.Fatalf("countXAIInputTokens() error = %v", err) } structuralCount, err := countXAIInputTokens(enc, structuralBody) if err != nil { t.Fatalf("countXAIInputTokens() error = %v", err) } if structuralCount != semanticCount { t.Fatalf("structural count = %d, want %d", structuralCount, semanticCount) } for name, tc := range map[string]struct { body []byte expected string }{ "instructions": { body: []byte(`{"instructions":"unique instruction text"}`), expected: "unique instruction text", }, "string input": { body: []byte(`{"input":"unique input text"}`), expected: "unique input text", }, "message content": { body: []byte(`{"input":[{"type":"message","role":"user","content":[{"type":"input_text","text":"unique message text"}]}]}`), expected: "unique message text", }, "refusal": { body: []byte(`{"input":[{"type":"message","content":[{"type":"refusal","refusal":"unique refusal text"}]}]}`), expected: "unique refusal text", }, "input image": { body: []byte(`{"input":[{"type":"message","content":[{"type":"input_image","image_url":"https://example.com/unique.png"}]}]}`), expected: "https://example.com/unique.png", }, "input file": { body: []byte(`{"input":[{"type":"message","content":[{"type":"input_file","file_data":"unique file data","filename":"unique.txt"}]}]}`), expected: "unique file data\nunique.txt", }, "input audio": { body: []byte(`{"input":[{"type":"message","content":[{"type":"input_audio","data":"unique audio data"}]}]}`), expected: "unique audio data", }, "function call": { body: []byte(`{"input":[{"type":"function_call","call_id":"call-1","name":"unique_function","arguments":"{\"value\":\"unique argument\"}"}]}`), expected: "unique_function\n{\"value\":\"unique argument\"}", }, "function call output": { body: []byte(`{"input":[{"type":"function_call_output","call_id":"call-1","output":"unique tool output"}]}`), expected: "unique tool output", }, "reasoning summary": { body: []byte(`{"input":[{"type":"reasoning","summary":[{"type":"summary_text","text":"unique summary text"}]}]}`), expected: "unique summary text", }, "function tool": { body: []byte(`{"tools":[{"type":"function","name":"unique_tool","description":"unique tool description","parameters":{"type":"object","properties":{"value":{"type":"string"}}}}]}`), expected: "unique_tool\nunique tool description\n{\"type\":\"object\",\"properties\":{\"value\":{\"type\":\"string\"}}}", }, "structured text format": { body: []byte(`{"text":{"format":{"name":"unique_format","schema":{"type":"object","properties":{"value":{"type":"string"}}}}}}`), expected: "unique_format\n{\"type\":\"object\",\"properties\":{\"value\":{\"type\":\"string\"}}}", }, } { t.Run(name, func(t *testing.T) { count, errCount := countXAIInputTokens(enc, tc.body) if errCount != nil { t.Fatalf("countXAIInputTokens() error = %v", errCount) } expected, errExpected := enc.Count(tc.expected) if errExpected != nil { t.Fatalf("encoder.Count() error = %v", errExpected) } if count != int64(expected) { t.Fatalf("countXAIInputTokens() = %d, want %d", count, expected) } }) }} func TestXAIExecutorExecuteShapesResponsesRequest(t *testing.T) { var gotPath string var gotAuth string var gotGrokConvID string var gotOriginator string var gotAccountID string var gotBody []byte server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { gotPath = r.URL.Path gotAuth = r.Header.Get("Authorization") gotGrokConvID = r.Header.Get("x-grok-conv-id") gotOriginator = r.Header.Get("Originator") gotAccountID = r.Header.Get("Chatgpt-Account-Id") var errRead error gotBody, errRead = io.ReadAll(r.Body) if errRead != nil { t.Fatalf("read body: %v", errRead) } w.Header().Set("Content-Type", "text/event-stream") _, _ = w.Write([]byte("data: {\"type\":\"response.completed\",\"response\":{\"id\":\"resp_1\",\"object\":\"response\",\"created_at\":0,\"status\":\"completed\",\"model\":\"grok-4.3\",\"output\":[{\"type\":\"message\",\"role\":\"assistant\",\"content\":[{\"type\":\"output_text\",\"text\":\"ok\"}]}],\"usage\":{\"input_tokens\":1,\"output_tokens\":1,\"total_tokens\":2}}}\n\n")) })) defer server.Close() exec := NewXAIExecutor(&config.Config{XAI: config.XAIConfig{InjectXSearch: true}}) auth := &cliproxyauth.Auth{ ID: "xai-auth", Provider: "xai", Attributes: map[string]string{ "base_url": server.URL, "auth_kind": "oauth", }, Metadata: map[string]any{ "access_token": "xai-token", "email": "user@example.com", }, } _, err := exec.Execute(context.Background(), auth, cliproxyexecutor.Request{ Model: "grok-4.3", Payload: []byte(`{"model":"grok-4.3","input":[{"type":"reasoning","summary":[{"type":"summary_text","text":"test"}],"content":null,"encrypted_content":null},{"type":"reasoning","summary":[{"type":"summary_text","text":"second"}]},{"role":"user","content":"hello"}],"include":["reasoning.encrypted_content"],"reasoning":{"effort":"high"},"tools":[{"type":"tool_search"},{"type":"image_generation"},{"type":"custom","name":"apply_patch"},{"type":"custom","name":"custom_lookup"},{"type":"function","name":"lookup"},{"type":"web_search","external_web_access":true,"search_content_types":["text","image"]},{"type":"namespace","name":"codex_app","description":"Tools in the codex_app namespace.","tools":[{"type":"function","name":"automation_update"},{"type":"custom","name":"namespace_custom"},{"type":"tool_search"}]}],"tool_choice":{"type":"allowed_tools","tools":[{"type":"function","name":"automation_update","namespace":"codex_app"},{"type":"function","name":"lookup"},{"type":"web_search"}]}}`), }, cliproxyexecutor.Options{ SourceFormat: sdktranslator.FormatOpenAIResponse, Stream: false, Metadata: map[string]any{ cliproxyexecutor.ExecutionSessionMetadataKey: "conv-xai-1", }, }) if err != nil { t.Fatalf("Execute() error = %v", err) } if gotPath != "/responses" { t.Fatalf("path = %q, want /responses", gotPath) } if gotAuth != "Bearer xai-token" { t.Fatalf("Authorization = %q, want Bearer xai-token", gotAuth) } if gotGrokConvID != "conv-xai-1" { t.Fatalf("x-grok-conv-id = %q, want conv-xai-1", gotGrokConvID) } if gotOriginator != "" { t.Fatalf("Originator = %q, want empty", gotOriginator) } if gotAccountID != "" { t.Fatalf("Chatgpt-Account-Id = %q, want empty", gotAccountID) } if gjson.GetBytes(gotBody, "prompt_cache_key").String() != "conv-xai-1" { t.Fatalf("prompt_cache_key missing from body: %s", string(gotBody)) } if !gjson.GetBytes(gotBody, "stream").Bool() { t.Fatalf("stream = false, want true; body=%s", string(gotBody)) } if gjson.GetBytes(gotBody, "reasoning.effort").String() != "high" { t.Fatalf("reasoning.effort = %q, want high; body=%s", gjson.GetBytes(gotBody, "reasoning.effort").String(), string(gotBody)) } if gjson.GetBytes(gotBody, "input.0.content").Exists() { t.Fatalf("input.0.content exists, want removed; body=%s", string(gotBody)) } if gjson.GetBytes(gotBody, "input.0.encrypted_content").Exists() { t.Fatalf("input.0.encrypted_content exists, want removed; body=%s", string(gotBody)) } if got := gjson.GetBytes(gotBody, "input.0.summary.0.text").String(); got != "test" { t.Fatalf("input.0.summary.0.text = %q, want test; body=%s", got, string(gotBody)) } if got := gjson.GetBytes(gotBody, "input.0.summary.1.text").String(); got != "second" { t.Fatalf("input.0.summary.1.text = %q, want second; body=%s", got, string(gotBody)) } if got := gjson.GetBytes(gotBody, "input.1.role").String(); got != "user" { t.Fatalf("input.1.role = %q, want user; body=%s", got, string(gotBody)) } if gjson.GetBytes(gotBody, "input.2").Exists() { t.Fatalf("input.2 exists, want consecutive reasoning item merged; body=%s", string(gotBody)) } tools := gjson.GetBytes(gotBody, "tools").Array() if len(tools) != 6 { t.Fatalf("tools length = %d, want 6; body=%s", len(tools), string(gotBody)) } foundAutomationUpdate := false foundNamespaceCustom := false foundXSearch := false for i, tool := range tools { toolType := tool.Get("type").String() if toolType == "image_generation" { t.Fatalf("tools.%d.type = image_generation, want removed; body=%s", i, string(gotBody)) } if toolType != "function" && toolType != "web_search" && toolType != "x_search" { t.Fatalf("tools.%d.type = %q, want function, web_search, or x_search; body=%s", i, toolType, string(gotBody)) } if toolType == "x_search" { foundXSearch = true } if toolType == "function" && !tool.Get("parameters").Exists() { t.Fatalf("tools.%d.parameters missing for xAI function tool; body=%s", i, string(gotBody)) } if got := tool.Get("name").String(); got == "apply_patch" { t.Fatalf("tools.%d.name = apply_patch, want removed; body=%s", i, string(gotBody)) } switch tool.Get("name").String() { case "codex_app__automation_update": foundAutomationUpdate = true case "codex_app__namespace_custom": foundNamespaceCustom = true } if toolType == "web_search" { if tool.Get("external_web_access").Exists() { t.Fatalf("tools.%d.external_web_access exists, want removed; body=%s", i, string(gotBody)) } if got := tool.Get("search_content_types.1").String(); got != "image" { t.Fatalf("tools.%d.search_content_types missing image entry; body=%s", i, string(gotBody)) } } } if !foundAutomationUpdate { t.Fatalf("namespace function tool was not moved to top-level tools; body=%s", string(gotBody)) } if !foundNamespaceCustom { t.Fatalf("namespace custom tool was not moved to top-level tools; body=%s", string(gotBody)) } if !foundXSearch { t.Fatalf("native x_search tool was not injected; body=%s", string(gotBody)) } if got := gjson.GetBytes(gotBody, "tool_choice.tools.0.name").String(); got != "codex_app__automation_update" { t.Fatalf("tool_choice.tools.0.name = %q, want codex_app__automation_update; body=%s", got, string(gotBody)) } if gjson.GetBytes(gotBody, "tool_choice.tools.0.namespace").Exists() { t.Fatalf("tool_choice.tools.0.namespace should be removed for xAI upstream: %s", string(gotBody)) } if got := gjson.GetBytes(gotBody, "tool_choice.tools.1.name").String(); got != "lookup" { t.Fatalf("tool_choice.tools.1.name = %q, want lookup; body=%s", got, string(gotBody)) } if got := gjson.GetBytes(gotBody, "tool_choice.tools.2.type").String(); got != "web_search" { t.Fatalf("tool_choice.tools.2.type = %q, want web_search; body=%s", got, string(gotBody)) } if got := gjson.GetBytes(gotBody, "tool_choice.tools.3.type").String(); got != "x_search" { t.Fatalf("tool_choice.tools.3.type = %q, want x_search; body=%s", got, string(gotBody)) } xSearchAllowedCount := 0 for _, tool := range gjson.GetBytes(gotBody, "tool_choice.tools").Array() { if tool.Get("type").String() == "x_search" { xSearchAllowedCount++ } } if xSearchAllowedCount != 1 { t.Fatalf("allowed_tools x_search count = %d, want 1; body=%s", xSearchAllowedCount, string(gotBody)) } foundEncryptedReasoningInclude := false for _, include := range gjson.GetBytes(gotBody, "include").Array() { if include.String() == "reasoning.encrypted_content" { foundEncryptedReasoningInclude = true break } } if !foundEncryptedReasoningInclude { t.Fatalf("xai request must preserve reasoning.encrypted_content include: %s", string(gotBody)) }} func TestXAIExecutorPrepareResponsesRequestRewritesCodexAgentMessage(t *testing.T) { t.Parallel() exec := NewXAIExecutor(&config.Config{Codex: config.CodexConfig{OptimizeMultiAgentV2: true}}) payload := []byte(`{ "model":"grok-4.5", "input":[{ "type":"agent_message", "id":"amsg_019f92c3-6d77-7880-a6e4-f920867dc6a0", "author":"/root", "recipient":"/root/arithmetic_question", "content":[ {"type":"input_text","text":"Message Type: NEW_TASK\nTask name: /root/arithmetic_question\nSender: /root\nPayload:\n"}, {"type":"encrypted_content","encrypted_content":"请出一道四则运算题。只回复题目本身,不要解答;使用中文。"} ], "internal_chat_message_metadata_passthrough":{"turn_id":"019f92c3-6772-7213-8aac-8bd154d528f1"} }] }`) prepared, errPrepare := exec.prepareResponsesRequest(context.Background(), cliproxyexecutor.Request{ Model: "grok-4.5", Payload: payload, }, cliproxyexecutor.Options{ SourceFormat: sdktranslator.FormatOpenAIResponse, Headers: http.Header{"User-Agent": []string{"Codex Desktop/0.146.0-alpha.3.1"}}, }, true) if errPrepare != nil { t.Fatalf("prepareResponsesRequest() error = %v", errPrepare) } message := gjson.GetBytes(prepared.body, "input.0") if message.Get("type").String() != "message" || message.Get("role").String() != "user" { t.Fatalf("agent message was not rewritten: %s", prepared.body) } if message.Get("content.1.type").String() != "input_text" { t.Fatalf("content[1].type = %q, want input_text; body=%s", message.Get("content.1.type").String(), prepared.body) } if text := message.Get("content.1.text").String(); text != "请出一道四则运算题。只回复题目本身,不要解答;使用中文。" { t.Fatalf("content[1].text = %q; body=%s", text, prepared.body) } if message.Get("content.1.encrypted_content").Exists() { t.Fatalf("encrypted_content was preserved: %s", prepared.body) } if message.Get("id").String() != "amsg_019f92c3-6d77-7880-a6e4-f920867dc6a0" || message.Get("author").String() != "/root" || message.Get("recipient").String() != "/root/arithmetic_question" { t.Fatalf("agent message identity fields changed: %s", prepared.body) } if turnID := message.Get("internal_chat_message_metadata_passthrough.turn_id").String(); turnID != "019f92c3-6772-7213-8aac-8bd154d528f1" { t.Fatalf("turn_id = %q; body=%s", turnID, prepared.body) }} func TestXAIExecutorExecuteFoldsNamespacesWhenToolsExceed200(t *testing.T) { var gotBody []byte server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { var errRead error gotBody, errRead = io.ReadAll(r.Body) if errRead != nil { t.Fatalf("read body: %v", errRead) } w.Header().Set("Content-Type", "text/event-stream") _, _ = w.Write([]byte("data: {\"type\":\"response.output_item.done\",\"output_index\":0,\"item\":{\"type\":\"function_call\",\"name\":\"mcp__app_0\",\"call_id\":\"call_1\",\"arguments\":\"{\\\"name\\\":\\\"tool_2\\\",\\\"arguments\\\":{\\\"q\\\":\\\"test\\\"}}\"}}\n\n")) _, _ = w.Write([]byte("data: {\"type\":\"response.completed\",\"response\":{\"id\":\"resp_1\",\"object\":\"response\",\"status\":\"completed\",\"model\":\"grok-4.6\",\"output\":[{\"type\":\"function_call\",\"name\":\"mcp__app_0\",\"call_id\":\"call_1\",\"arguments\":\"{\\\"name\\\":\\\"tool_2\\\",\\\"arguments\\\":{\\\"q\\\":\\\"test\\\"}}\"}],\"usage\":{\"input_tokens\":1,\"output_tokens\":1,\"total_tokens\":2}}}\n\n")) })) defer server.Close() var nsList []string for i := 0; i < 47; i++ { var childTools []string for j := 0; j < 10; j++ { childTools = append(childTools, fmt.Sprintf(`{"type":"function","name":"tool_%d","description":"child tool %d","parameters":{"type":"object","properties":{"q":{"type":"string"}}}}`, j, j)) } nsList = append(nsList, fmt.Sprintf(`{"type":"namespace","name":"mcp__app_%d","description":"App %d tools","tools":[%s]}`, i, i, strings.Join(childTools, ","))) } exec := NewXAIExecutor(&config.Config{}) auth := &cliproxyauth.Auth{ Provider: "xai", Attributes: map[string]string{"base_url": server.URL}, Metadata: map[string]any{"access_token": "xai-token"}, } turn1Payload := fmt.Sprintf(`{ "model":"grok-4.6", "tools":[%s], "input":[{"role":"user","content":"call tool_2"}] }`, strings.Join(nsList, ",")) resp, err := exec.Execute(context.Background(), auth, cliproxyexecutor.Request{ Model: "grok-4.6", Payload: []byte(turn1Payload), }, cliproxyexecutor.Options{ SourceFormat: sdktranslator.FormatOpenAIResponse, ResponseFormat: sdktranslator.FormatOpenAIResponse, Stream: false, }) if err != nil { t.Fatalf("Execute() error = %v", err) } tools := gjson.GetBytes(gotBody, "tools").Array() if len(tools) != 47 { t.Fatalf("upstream tools count = %d, want 47 dispatcher tools; body=%s", len(tools), string(gotBody)) } if got := tools[0].Get("name").String(); got != "mcp__app_0" { t.Fatalf("upstream tools[0].name = %q, want mcp__app_0", got) } if got := tools[0].Get("type").String(); got != "function" { t.Fatalf("upstream tools[0].type = %q, want function", got) } output := gjson.GetBytes(resp.Payload, "output.0") if got := output.Get("name").String(); got != "tool_2" { t.Fatalf("response output name = %q, want tool_2; payload=%s", got, resp.Payload) } if got := output.Get("namespace").String(); got != "mcp__app_0" { t.Fatalf("response output namespace = %q, want mcp__app_0; payload=%s", got, resp.Payload) } if got := output.Get("arguments").String(); got != `{"q":"test"}` { t.Fatalf("response output arguments = %q, want {\"q\":\"test\"}; payload=%s", got, resp.Payload) } // Turn 2: client passes back the restored function_call and function_call_output turn2Payload := fmt.Sprintf(`{ "model":"grok-4.6", "tools":[%s], "input":[ {"role":"user","content":"call tool_2"}, {"type":"function_call","name":"tool_2","namespace":"mcp__app_0","call_id":"call_1","arguments":"{\"q\":\"test\"}"}, {"type":"function_call_output","call_id":"call_1","output":"ok"} ] }`, strings.Join(nsList, ",")) _, err2 := exec.Execute(context.Background(), auth, cliproxyexecutor.Request{ Model: "grok-4.6", Payload: []byte(turn2Payload), }, cliproxyexecutor.Options{ SourceFormat: sdktranslator.FormatOpenAIResponse, ResponseFormat: sdktranslator.FormatOpenAIResponse, Stream: false, }) if err2 != nil { t.Fatalf("Execute() turn 2 error = %v", err2) } input := gjson.GetBytes(gotBody, "input").Array() if len(input) != 3 { t.Fatalf("upstream input count = %d, want 3; body=%s", len(input), string(gotBody)) } fcItem := input[1] if got := fcItem.Get("name").String(); got != "mcp__app_0" { t.Fatalf("turn 2 upstream fc name = %q, want mcp__app_0; body=%s", got, string(gotBody)) } if fcItem.Get("namespace").Exists() { t.Fatalf("turn 2 upstream fc namespace should be removed; body=%s", string(gotBody)) } if gotArgs := fcItem.Get("arguments").String(); gotArgs != `{"arguments":{"q":"test"},"name":"tool_2"}` && gotArgs != `{"name":"tool_2","arguments":{"q":"test"}}` { t.Fatalf("turn 2 upstream fc arguments = %q, want wrapped dispatcher args; body=%s", gotArgs, string(gotBody)) }} func TestXAIExecutorExecuteStreamFoldsNamespacesWhenToolsExceed200(t *testing.T) { var gotBody []byte server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { var errRead error gotBody, errRead = io.ReadAll(r.Body) if errRead != nil { t.Fatalf("read body: %v", errRead) } w.Header().Set("Content-Type", "text/event-stream") _, _ = w.Write([]byte("data: {\"type\":\"response.output_item.done\",\"output_index\":0,\"item\":{\"type\":\"function_call\",\"name\":\"mcp__app_0\",\"call_id\":\"call_1\",\"arguments\":\"{\\\"name\\\":\\\"tool_2\\\",\\\"arguments\\\":{\\\"q\\\":\\\"stream_test\\\"}}\"}}\n\n"))Lines 1–500 of 6318. Load more
package executor import ( "bytes" "context" "crypto/sha256" "encoding/base64" "encoding/json" "errors" "fmt" "io" "net/http" "net/http/httptest" "strings" "testing" "time" "github.com/gin-gonic/gin" "github.com/google/uuid" xaiauth "github.com/router-for-me/CLIProxyAPI/v7/internal/auth/xai" 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/translator" cliproxyauth "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/auth" cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executor" "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/usage" sdktranslator "github.com/router-for-me/CLIProxyAPI/v7/sdk/translator" "github.com/tidwall/gjson" "github.com/tidwall/sjson" "github.com/tiktoken-go/tokenizer" ) func testContextWithAPIKey(apiKey string) context.Context { gin.SetMode(gin.TestMode) rec := httptest.NewRecorder() ginCtx, _ := gin.CreateTestContext(rec) ginCtx.Set("userApiKey", apiKey) return context.WithValue(context.Background(), "gin", ginCtx) } func TestCountXAIInputTokensExcludesRequestStructure(t *testing.T) { enc, err := tokenizer.Get(tokenizer.O200kBase) if err != nil { t.Fatalf("tokenizer.Get() error = %v", err) } semanticBody := []byte(`{ "instructions":"Follow the repository instructions.", "input":[ {"type":"message","role":"user","content":[{"type":"input_text","text":"Review this implementation."}]}, {"type":"function_call","name":"read_file","arguments":"{\"path\":\"main.go\"}"}, {"type":"function_call_output","output":"package main"}, {"type":"reasoning","summary":[{"type":"summary_text","text":"I will inspect the file."}]} ], "tools":[{"type":"function","name":"read_file","description":"Reads a file.","parameters":{"type":"object","properties":{"path":{"type":"string"}}}}], "text":{"format":{"name":"result","schema":{"type":"object"}}} }`) structuralBody := []byte(`{ "model":"grok-4.5", "stream":false, "reasoning":{"effort":"high"}, "metadata":{"large_wrapper":"this metadata must not affect estimated input tokens"}, "prompt_cache_key":"session-123", "max_output_tokens":4096, "instructions":"Follow the repository instructions.", "input":[ {"type":"message","role":"user","content":[{"type":"input_text","text":"Review this implementation."}]}, {"type":"function_call","name":"read_file","arguments":"{\"path\":\"main.go\"}"}, {"type":"function_call_output","output":"package main"}, {"type":"reasoning","summary":[{"type":"summary_text","text":"I will inspect the file."}]} ], "tools":[{"type":"function","name":"read_file","description":"Reads a file.","parameters":{"type":"object","properties":{"path":{"type":"string"}}}}], "text":{"format":{"name":"result","schema":{"type":"object"}}} }`) semanticCount, err := countXAIInputTokens(enc, semanticBody) if err != nil { t.Fatalf("countXAIInputTokens() error = %v", err) } structuralCount, err := countXAIInputTokens(enc, structuralBody) if err != nil { t.Fatalf("countXAIInputTokens() error = %v", err) } if structuralCount != semanticCount { t.Fatalf("structural count = %d, want %d", structuralCount, semanticCount) } for name, tc := range map[string]struct { body []byte expected string }{ "instructions": { body: []byte(`{"instructions":"unique instruction text"}`), expected: "unique instruction text", }, "string input": { body: []byte(`{"input":"unique input text"}`), expected: "unique input text", }, "message content": { body: []byte(`{"input":[{"type":"message","role":"user","content":[{"type":"input_text","text":"unique message text"}]}]}`), expected: "unique message text", }, "refusal": { body: []byte(`{"input":[{"type":"message","content":[{"type":"refusal","refusal":"unique refusal text"}]}]}`), expected: "unique refusal text", }, "input image": { body: []byte(`{"input":[{"type":"message","content":[{"type":"input_image","image_url":"https://example.com/unique.png"}]}]}`), expected: "https://example.com/unique.png", }, "input file": { body: []byte(`{"input":[{"type":"message","content":[{"type":"input_file","file_data":"unique file data","filename":"unique.txt"}]}]}`), expected: "unique file data\nunique.txt", }, "input audio": { body: []byte(`{"input":[{"type":"message","content":[{"type":"input_audio","data":"unique audio data"}]}]}`), expected: "unique audio data", }, "function call": { body: []byte(`{"input":[{"type":"function_call","call_id":"call-1","name":"unique_function","arguments":"{\"value\":\"unique argument\"}"}]}`), expected: "unique_function\n{\"value\":\"unique argument\"}", }, "function call output": { body: []byte(`{"input":[{"type":"function_call_output","call_id":"call-1","output":"unique tool output"}]}`), expected: "unique tool output", }, "reasoning summary": { body: []byte(`{"input":[{"type":"reasoning","summary":[{"type":"summary_text","text":"unique summary text"}]}]}`), expected: "unique summary text", }, "function tool": { body: []byte(`{"tools":[{"type":"function","name":"unique_tool","description":"unique tool description","parameters":{"type":"object","properties":{"value":{"type":"string"}}}}]}`), expected: "unique_tool\nunique tool description\n{\"type\":\"object\",\"properties\":{\"value\":{\"type\":\"string\"}}}", }, "structured text format": { body: []byte(`{"text":{"format":{"name":"unique_format","schema":{"type":"object","properties":{"value":{"type":"string"}}}}}}`), expected: "unique_format\n{\"type\":\"object\",\"properties\":{\"value\":{\"type\":\"string\"}}}", }, } { t.Run(name, func(t *testing.T) { count, errCount := countXAIInputTokens(enc, tc.body) if errCount != nil { t.Fatalf("countXAIInputTokens() error = %v", errCount) } expected, errExpected := enc.Count(tc.expected) if errExpected != nil { t.Fatalf("encoder.Count() error = %v", errExpected) } if count != int64(expected) { t.Fatalf("countXAIInputTokens() = %d, want %d", count, expected) } }) } } func TestXAIExecutorExecuteShapesResponsesRequest(t *testing.T) { var gotPath string var gotAuth string var gotGrokConvID string var gotOriginator string var gotAccountID string var gotBody []byte server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { gotPath = r.URL.Path gotAuth = r.Header.Get("Authorization") gotGrokConvID = r.Header.Get("x-grok-conv-id") gotOriginator = r.Header.Get("Originator") gotAccountID = r.Header.Get("Chatgpt-Account-Id") var errRead error gotBody, errRead = io.ReadAll(r.Body) if errRead != nil { t.Fatalf("read body: %v", errRead) } w.Header().Set("Content-Type", "text/event-stream") _, _ = w.Write([]byte("data: {\"type\":\"response.completed\",\"response\":{\"id\":\"resp_1\",\"object\":\"response\",\"created_at\":0,\"status\":\"completed\",\"model\":\"grok-4.3\",\"output\":[{\"type\":\"message\",\"role\":\"assistant\",\"content\":[{\"type\":\"output_text\",\"text\":\"ok\"}]}],\"usage\":{\"input_tokens\":1,\"output_tokens\":1,\"total_tokens\":2}}}\n\n")) })) defer server.Close() exec := NewXAIExecutor(&config.Config{XAI: config.XAIConfig{InjectXSearch: true}}) auth := &cliproxyauth.Auth{ ID: "xai-auth", Provider: "xai", Attributes: map[string]string{ "base_url": server.URL, "auth_kind": "oauth", }, Metadata: map[string]any{ "access_token": "xai-token", "email": "user@example.com", }, } _, err := exec.Execute(context.Background(), auth, cliproxyexecutor.Request{ Model: "grok-4.3", Payload: []byte(`{"model":"grok-4.3","input":[{"type":"reasoning","summary":[{"type":"summary_text","text":"test"}],"content":null,"encrypted_content":null},{"type":"reasoning","summary":[{"type":"summary_text","text":"second"}]},{"role":"user","content":"hello"}],"include":["reasoning.encrypted_content"],"reasoning":{"effort":"high"},"tools":[{"type":"tool_search"},{"type":"image_generation"},{"type":"custom","name":"apply_patch"},{"type":"custom","name":"custom_lookup"},{"type":"function","name":"lookup"},{"type":"web_search","external_web_access":true,"search_content_types":["text","image"]},{"type":"namespace","name":"codex_app","description":"Tools in the codex_app namespace.","tools":[{"type":"function","name":"automation_update"},{"type":"custom","name":"namespace_custom"},{"type":"tool_search"}]}],"tool_choice":{"type":"allowed_tools","tools":[{"type":"function","name":"automation_update","namespace":"codex_app"},{"type":"function","name":"lookup"},{"type":"web_search"}]}}`), }, cliproxyexecutor.Options{ SourceFormat: sdktranslator.FormatOpenAIResponse, Stream: false, Metadata: map[string]any{ cliproxyexecutor.ExecutionSessionMetadataKey: "conv-xai-1", }, }) if err != nil { t.Fatalf("Execute() error = %v", err) } if gotPath != "/responses" { t.Fatalf("path = %q, want /responses", gotPath) } if gotAuth != "Bearer xai-token" { t.Fatalf("Authorization = %q, want Bearer xai-token", gotAuth) } if gotGrokConvID != "conv-xai-1" { t.Fatalf("x-grok-conv-id = %q, want conv-xai-1", gotGrokConvID) } if gotOriginator != "" { t.Fatalf("Originator = %q, want empty", gotOriginator) } if gotAccountID != "" { t.Fatalf("Chatgpt-Account-Id = %q, want empty", gotAccountID) } if gjson.GetBytes(gotBody, "prompt_cache_key").String() != "conv-xai-1" { t.Fatalf("prompt_cache_key missing from body: %s", string(gotBody)) } if !gjson.GetBytes(gotBody, "stream").Bool() { t.Fatalf("stream = false, want true; body=%s", string(gotBody)) } if gjson.GetBytes(gotBody, "reasoning.effort").String() != "high" { t.Fatalf("reasoning.effort = %q, want high; body=%s", gjson.GetBytes(gotBody, "reasoning.effort").String(), string(gotBody)) } if gjson.GetBytes(gotBody, "input.0.content").Exists() { t.Fatalf("input.0.content exists, want removed; body=%s", string(gotBody)) } if gjson.GetBytes(gotBody, "input.0.encrypted_content").Exists() { t.Fatalf("input.0.encrypted_content exists, want removed; body=%s", string(gotBody)) } if got := gjson.GetBytes(gotBody, "input.0.summary.0.text").String(); got != "test" { t.Fatalf("input.0.summary.0.text = %q, want test; body=%s", got, string(gotBody)) } if got := gjson.GetBytes(gotBody, "input.0.summary.1.text").String(); got != "second" { t.Fatalf("input.0.summary.1.text = %q, want second; body=%s", got, string(gotBody)) } if got := gjson.GetBytes(gotBody, "input.1.role").String(); got != "user" { t.Fatalf("input.1.role = %q, want user; body=%s", got, string(gotBody)) } if gjson.GetBytes(gotBody, "input.2").Exists() { t.Fatalf("input.2 exists, want consecutive reasoning item merged; body=%s", string(gotBody)) } tools := gjson.GetBytes(gotBody, "tools").Array() if len(tools) != 6 { t.Fatalf("tools length = %d, want 6; body=%s", len(tools), string(gotBody)) } foundAutomationUpdate := false foundNamespaceCustom := false foundXSearch := false for i, tool := range tools { toolType := tool.Get("type").String() if toolType == "image_generation" { t.Fatalf("tools.%d.type = image_generation, want removed; body=%s", i, string(gotBody)) } if toolType != "function" && toolType != "web_search" && toolType != "x_search" { t.Fatalf("tools.%d.type = %q, want function, web_search, or x_search; body=%s", i, toolType, string(gotBody)) } if toolType == "x_search" { foundXSearch = true } if toolType == "function" && !tool.Get("parameters").Exists() { t.Fatalf("tools.%d.parameters missing for xAI function tool; body=%s", i, string(gotBody)) } if got := tool.Get("name").String(); got == "apply_patch" { t.Fatalf("tools.%d.name = apply_patch, want removed; body=%s", i, string(gotBody)) } switch tool.Get("name").String() { case "codex_app__automation_update": foundAutomationUpdate = true case "codex_app__namespace_custom": foundNamespaceCustom = true } if toolType == "web_search" { if tool.Get("external_web_access").Exists() { t.Fatalf("tools.%d.external_web_access exists, want removed; body=%s", i, string(gotBody)) } if got := tool.Get("search_content_types.1").String(); got != "image" { t.Fatalf("tools.%d.search_content_types missing image entry; body=%s", i, string(gotBody)) } } } if !foundAutomationUpdate { t.Fatalf("namespace function tool was not moved to top-level tools; body=%s", string(gotBody)) } if !foundNamespaceCustom { t.Fatalf("namespace custom tool was not moved to top-level tools; body=%s", string(gotBody)) } if !foundXSearch { t.Fatalf("native x_search tool was not injected; body=%s", string(gotBody)) } if got := gjson.GetBytes(gotBody, "tool_choice.tools.0.name").String(); got != "codex_app__automation_update" { t.Fatalf("tool_choice.tools.0.name = %q, want codex_app__automation_update; body=%s", got, string(gotBody)) } if gjson.GetBytes(gotBody, "tool_choice.tools.0.namespace").Exists() { t.Fatalf("tool_choice.tools.0.namespace should be removed for xAI upstream: %s", string(gotBody)) } if got := gjson.GetBytes(gotBody, "tool_choice.tools.1.name").String(); got != "lookup" { t.Fatalf("tool_choice.tools.1.name = %q, want lookup; body=%s", got, string(gotBody)) } if got := gjson.GetBytes(gotBody, "tool_choice.tools.2.type").String(); got != "web_search" { t.Fatalf("tool_choice.tools.2.type = %q, want web_search; body=%s", got, string(gotBody)) } if got := gjson.GetBytes(gotBody, "tool_choice.tools.3.type").String(); got != "x_search" { t.Fatalf("tool_choice.tools.3.type = %q, want x_search; body=%s", got, string(gotBody)) } xSearchAllowedCount := 0 for _, tool := range gjson.GetBytes(gotBody, "tool_choice.tools").Array() { if tool.Get("type").String() == "x_search" { xSearchAllowedCount++ } } if xSearchAllowedCount != 1 { t.Fatalf("allowed_tools x_search count = %d, want 1; body=%s", xSearchAllowedCount, string(gotBody)) } foundEncryptedReasoningInclude := false for _, include := range gjson.GetBytes(gotBody, "include").Array() { if include.String() == "reasoning.encrypted_content" { foundEncryptedReasoningInclude = true break } } if !foundEncryptedReasoningInclude { t.Fatalf("xai request must preserve reasoning.encrypted_content include: %s", string(gotBody)) } } func TestXAIExecutorPrepareResponsesRequestRewritesCodexAgentMessage(t *testing.T) { t.Parallel() exec := NewXAIExecutor(&config.Config{Codex: config.CodexConfig{OptimizeMultiAgentV2: true}}) payload := []byte(`{ "model":"grok-4.5", "input":[{ "type":"agent_message", "id":"amsg_019f92c3-6d77-7880-a6e4-f920867dc6a0", "author":"/root", "recipient":"/root/arithmetic_question", "content":[ {"type":"input_text","text":"Message Type: NEW_TASK\nTask name: /root/arithmetic_question\nSender: /root\nPayload:\n"}, {"type":"encrypted_content","encrypted_content":"请出一道四则运算题。只回复题目本身,不要解答;使用中文。"} ], "internal_chat_message_metadata_passthrough":{"turn_id":"019f92c3-6772-7213-8aac-8bd154d528f1"} }] }`) prepared, errPrepare := exec.prepareResponsesRequest(context.Background(), cliproxyexecutor.Request{ Model: "grok-4.5", Payload: payload, }, cliproxyexecutor.Options{ SourceFormat: sdktranslator.FormatOpenAIResponse, Headers: http.Header{"User-Agent": []string{"Codex Desktop/0.146.0-alpha.3.1"}}, }, true) if errPrepare != nil { t.Fatalf("prepareResponsesRequest() error = %v", errPrepare) } message := gjson.GetBytes(prepared.body, "input.0") if message.Get("type").String() != "message" || message.Get("role").String() != "user" { t.Fatalf("agent message was not rewritten: %s", prepared.body) } if message.Get("content.1.type").String() != "input_text" { t.Fatalf("content[1].type = %q, want input_text; body=%s", message.Get("content.1.type").String(), prepared.body) } if text := message.Get("content.1.text").String(); text != "请出一道四则运算题。只回复题目本身,不要解答;使用中文。" { t.Fatalf("content[1].text = %q; body=%s", text, prepared.body) } if message.Get("content.1.encrypted_content").Exists() { t.Fatalf("encrypted_content was preserved: %s", prepared.body) } if message.Get("id").String() != "amsg_019f92c3-6d77-7880-a6e4-f920867dc6a0" || message.Get("author").String() != "/root" || message.Get("recipient").String() != "/root/arithmetic_question" { t.Fatalf("agent message identity fields changed: %s", prepared.body) } if turnID := message.Get("internal_chat_message_metadata_passthrough.turn_id").String(); turnID != "019f92c3-6772-7213-8aac-8bd154d528f1" { t.Fatalf("turn_id = %q; body=%s", turnID, prepared.body) } } func TestXAIExecutorExecuteFoldsNamespacesWhenToolsExceed200(t *testing.T) { var gotBody []byte server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { var errRead error gotBody, errRead = io.ReadAll(r.Body) if errRead != nil { t.Fatalf("read body: %v", errRead) } w.Header().Set("Content-Type", "text/event-stream") _, _ = w.Write([]byte("data: {\"type\":\"response.output_item.done\",\"output_index\":0,\"item\":{\"type\":\"function_call\",\"name\":\"mcp__app_0\",\"call_id\":\"call_1\",\"arguments\":\"{\\\"name\\\":\\\"tool_2\\\",\\\"arguments\\\":{\\\"q\\\":\\\"test\\\"}}\"}}\n\n")) _, _ = w.Write([]byte("data: {\"type\":\"response.completed\",\"response\":{\"id\":\"resp_1\",\"object\":\"response\",\"status\":\"completed\",\"model\":\"grok-4.6\",\"output\":[{\"type\":\"function_call\",\"name\":\"mcp__app_0\",\"call_id\":\"call_1\",\"arguments\":\"{\\\"name\\\":\\\"tool_2\\\",\\\"arguments\\\":{\\\"q\\\":\\\"test\\\"}}\"}],\"usage\":{\"input_tokens\":1,\"output_tokens\":1,\"total_tokens\":2}}}\n\n")) })) defer server.Close() var nsList []string for i := 0; i < 47; i++ { var childTools []string for j := 0; j < 10; j++ { childTools = append(childTools, fmt.Sprintf(`{"type":"function","name":"tool_%d","description":"child tool %d","parameters":{"type":"object","properties":{"q":{"type":"string"}}}}`, j, j)) } nsList = append(nsList, fmt.Sprintf(`{"type":"namespace","name":"mcp__app_%d","description":"App %d tools","tools":[%s]}`, i, i, strings.Join(childTools, ","))) } exec := NewXAIExecutor(&config.Config{}) auth := &cliproxyauth.Auth{ Provider: "xai", Attributes: map[string]string{"base_url": server.URL}, Metadata: map[string]any{"access_token": "xai-token"}, } turn1Payload := fmt.Sprintf(`{ "model":"grok-4.6", "tools":[%s], "input":[{"role":"user","content":"call tool_2"}] }`, strings.Join(nsList, ",")) resp, err := exec.Execute(context.Background(), auth, cliproxyexecutor.Request{ Model: "grok-4.6", Payload: []byte(turn1Payload), }, cliproxyexecutor.Options{ SourceFormat: sdktranslator.FormatOpenAIResponse, ResponseFormat: sdktranslator.FormatOpenAIResponse, Stream: false, }) if err != nil { t.Fatalf("Execute() error = %v", err) } tools := gjson.GetBytes(gotBody, "tools").Array() if len(tools) != 47 { t.Fatalf("upstream tools count = %d, want 47 dispatcher tools; body=%s", len(tools), string(gotBody)) } if got := tools[0].Get("name").String(); got != "mcp__app_0" { t.Fatalf("upstream tools[0].name = %q, want mcp__app_0", got) } if got := tools[0].Get("type").String(); got != "function" { t.Fatalf("upstream tools[0].type = %q, want function", got) } output := gjson.GetBytes(resp.Payload, "output.0") if got := output.Get("name").String(); got != "tool_2" { t.Fatalf("response output name = %q, want tool_2; payload=%s", got, resp.Payload) } if got := output.Get("namespace").String(); got != "mcp__app_0" { t.Fatalf("response output namespace = %q, want mcp__app_0; payload=%s", got, resp.Payload) } if got := output.Get("arguments").String(); got != `{"q":"test"}` { t.Fatalf("response output arguments = %q, want {\"q\":\"test\"}; payload=%s", got, resp.Payload) } // Turn 2: client passes back the restored function_call and function_call_output turn2Payload := fmt.Sprintf(`{ "model":"grok-4.6", "tools":[%s], "input":[ {"role":"user","content":"call tool_2"}, {"type":"function_call","name":"tool_2","namespace":"mcp__app_0","call_id":"call_1","arguments":"{\"q\":\"test\"}"}, {"type":"function_call_output","call_id":"call_1","output":"ok"} ] }`, strings.Join(nsList, ",")) _, err2 := exec.Execute(context.Background(), auth, cliproxyexecutor.Request{ Model: "grok-4.6", Payload: []byte(turn2Payload), }, cliproxyexecutor.Options{ SourceFormat: sdktranslator.FormatOpenAIResponse, ResponseFormat: sdktranslator.FormatOpenAIResponse, Stream: false, }) if err2 != nil { t.Fatalf("Execute() turn 2 error = %v", err2) } input := gjson.GetBytes(gotBody, "input").Array() if len(input) != 3 { t.Fatalf("upstream input count = %d, want 3; body=%s", len(input), string(gotBody)) } fcItem := input[1] if got := fcItem.Get("name").String(); got != "mcp__app_0" { t.Fatalf("turn 2 upstream fc name = %q, want mcp__app_0; body=%s", got, string(gotBody)) } if fcItem.Get("namespace").Exists() { t.Fatalf("turn 2 upstream fc namespace should be removed; body=%s", string(gotBody)) } if gotArgs := fcItem.Get("arguments").String(); gotArgs != `{"arguments":{"q":"test"},"name":"tool_2"}` && gotArgs != `{"name":"tool_2","arguments":{"q":"test"}}` { t.Fatalf("turn 2 upstream fc arguments = %q, want wrapped dispatcher args; body=%s", gotArgs, string(gotBody)) } } func TestXAIExecutorExecuteStreamFoldsNamespacesWhenToolsExceed200(t *testing.T) { var gotBody []byte server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { var errRead error gotBody, errRead = io.ReadAll(r.Body) if errRead != nil { t.Fatalf("read body: %v", errRead) } w.Header().Set("Content-Type", "text/event-stream") _, _ = w.Write([]byte("data: {\"type\":\"response.output_item.done\",\"output_index\":0,\"item\":{\"type\":\"function_call\",\"name\":\"mcp__app_0\",\"call_id\":\"call_1\",\"arguments\":\"{\\\"name\\\":\\\"tool_2\\\",\\\"arguments\\\":{\\\"q\\\":\\\"stream_test\\\"}}\"}}\n\n"))