| package executor |
|
|
| import ( |
| "bytes" |
| "context" |
| "io" |
| "net/http" |
| "net/http/httptest" |
| "testing" |
|
|
| "github.com/router-for-me/CLIProxyAPI/v6/internal/config" |
| cliproxyauth "github.com/router-for-me/CLIProxyAPI/v6/sdk/cliproxy/auth" |
| cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v6/sdk/cliproxy/executor" |
| sdktranslator "github.com/router-for-me/CLIProxyAPI/v6/sdk/translator" |
| "github.com/tidwall/gjson" |
| "github.com/tidwall/sjson" |
| ) |
|
|
| func TestApplyClaudeToolPrefix(t *testing.T) { |
| input := []byte(`{"tools":[{"name":"alpha"},{"name":"proxy_bravo"}],"tool_choice":{"type":"tool","name":"charlie"},"messages":[{"role":"assistant","content":[{"type":"tool_use","name":"delta","id":"t1","input":{}}]}]}`) |
| out := applyClaudeToolPrefix(input, "proxy_") |
|
|
| if got := gjson.GetBytes(out, "tools.0.name").String(); got != "proxy_alpha" { |
| t.Fatalf("tools.0.name = %q, want %q", got, "proxy_alpha") |
| } |
| if got := gjson.GetBytes(out, "tools.1.name").String(); got != "proxy_bravo" { |
| t.Fatalf("tools.1.name = %q, want %q", got, "proxy_bravo") |
| } |
| if got := gjson.GetBytes(out, "tool_choice.name").String(); got != "proxy_charlie" { |
| t.Fatalf("tool_choice.name = %q, want %q", got, "proxy_charlie") |
| } |
| if got := gjson.GetBytes(out, "messages.0.content.0.name").String(); got != "proxy_delta" { |
| t.Fatalf("messages.0.content.0.name = %q, want %q", got, "proxy_delta") |
| } |
| } |
|
|
| func TestApplyClaudeToolPrefix_WithToolReference(t *testing.T) { |
| input := []byte(`{"tools":[{"name":"alpha"}],"messages":[{"role":"user","content":[{"type":"tool_reference","tool_name":"beta"},{"type":"tool_reference","tool_name":"proxy_gamma"}]}]}`) |
| out := applyClaudeToolPrefix(input, "proxy_") |
|
|
| if got := gjson.GetBytes(out, "messages.0.content.0.tool_name").String(); got != "proxy_beta" { |
| t.Fatalf("messages.0.content.0.tool_name = %q, want %q", got, "proxy_beta") |
| } |
| if got := gjson.GetBytes(out, "messages.0.content.1.tool_name").String(); got != "proxy_gamma" { |
| t.Fatalf("messages.0.content.1.tool_name = %q, want %q", got, "proxy_gamma") |
| } |
| } |
|
|
| func TestApplyClaudeToolPrefix_SkipsBuiltinTools(t *testing.T) { |
| input := []byte(`{"tools":[{"type":"web_search_20250305","name":"web_search"},{"name":"my_custom_tool","input_schema":{"type":"object"}}]}`) |
| out := applyClaudeToolPrefix(input, "proxy_") |
|
|
| if got := gjson.GetBytes(out, "tools.0.name").String(); got != "web_search" { |
| t.Fatalf("built-in tool name should not be prefixed: tools.0.name = %q, want %q", got, "web_search") |
| } |
| if got := gjson.GetBytes(out, "tools.1.name").String(); got != "proxy_my_custom_tool" { |
| t.Fatalf("custom tool should be prefixed: tools.1.name = %q, want %q", got, "proxy_my_custom_tool") |
| } |
| } |
|
|
| func TestApplyClaudeToolPrefix_BuiltinToolSkipped(t *testing.T) { |
| body := []byte(`{ |
| "tools": [ |
| {"type": "web_search_20250305", "name": "web_search", "max_uses": 5}, |
| {"name": "Read"} |
| ], |
| "messages": [ |
| {"role": "user", "content": [ |
| {"type": "tool_use", "name": "web_search", "id": "ws1", "input": {}}, |
| {"type": "tool_use", "name": "Read", "id": "r1", "input": {}} |
| ]} |
| ] |
| }`) |
| out := applyClaudeToolPrefix(body, "proxy_") |
|
|
| if got := gjson.GetBytes(out, "tools.0.name").String(); got != "web_search" { |
| t.Fatalf("tools.0.name = %q, want %q", got, "web_search") |
| } |
| if got := gjson.GetBytes(out, "messages.0.content.0.name").String(); got != "web_search" { |
| t.Fatalf("messages.0.content.0.name = %q, want %q", got, "web_search") |
| } |
| if got := gjson.GetBytes(out, "tools.1.name").String(); got != "proxy_Read" { |
| t.Fatalf("tools.1.name = %q, want %q", got, "proxy_Read") |
| } |
| if got := gjson.GetBytes(out, "messages.0.content.1.name").String(); got != "proxy_Read" { |
| t.Fatalf("messages.0.content.1.name = %q, want %q", got, "proxy_Read") |
| } |
| } |
|
|
| func TestApplyClaudeToolPrefix_KnownBuiltinInHistoryOnly(t *testing.T) { |
| body := []byte(`{ |
| "tools": [ |
| {"name": "Read"} |
| ], |
| "messages": [ |
| {"role": "user", "content": [ |
| {"type": "tool_use", "name": "web_search", "id": "ws1", "input": {}} |
| ]} |
| ] |
| }`) |
| out := applyClaudeToolPrefix(body, "proxy_") |
|
|
| if got := gjson.GetBytes(out, "messages.0.content.0.name").String(); got != "web_search" { |
| t.Fatalf("messages.0.content.0.name = %q, want %q", got, "web_search") |
| } |
| if got := gjson.GetBytes(out, "tools.0.name").String(); got != "proxy_Read" { |
| t.Fatalf("tools.0.name = %q, want %q", got, "proxy_Read") |
| } |
| } |
|
|
| func TestApplyClaudeToolPrefix_CustomToolsPrefixed(t *testing.T) { |
| body := []byte(`{ |
| "tools": [{"name": "Read"}, {"name": "Write"}], |
| "messages": [ |
| {"role": "user", "content": [ |
| {"type": "tool_use", "name": "Read", "id": "r1", "input": {}}, |
| {"type": "tool_use", "name": "Write", "id": "w1", "input": {}} |
| ]} |
| ] |
| }`) |
| out := applyClaudeToolPrefix(body, "proxy_") |
|
|
| if got := gjson.GetBytes(out, "tools.0.name").String(); got != "proxy_Read" { |
| t.Fatalf("tools.0.name = %q, want %q", got, "proxy_Read") |
| } |
| if got := gjson.GetBytes(out, "tools.1.name").String(); got != "proxy_Write" { |
| t.Fatalf("tools.1.name = %q, want %q", got, "proxy_Write") |
| } |
| if got := gjson.GetBytes(out, "messages.0.content.0.name").String(); got != "proxy_Read" { |
| t.Fatalf("messages.0.content.0.name = %q, want %q", got, "proxy_Read") |
| } |
| if got := gjson.GetBytes(out, "messages.0.content.1.name").String(); got != "proxy_Write" { |
| t.Fatalf("messages.0.content.1.name = %q, want %q", got, "proxy_Write") |
| } |
| } |
|
|
| func TestApplyClaudeToolPrefix_ToolChoiceBuiltin(t *testing.T) { |
| body := []byte(`{ |
| "tools": [ |
| {"type": "web_search_20250305", "name": "web_search"}, |
| {"name": "Read"} |
| ], |
| "tool_choice": {"type": "tool", "name": "web_search"} |
| }`) |
| out := applyClaudeToolPrefix(body, "proxy_") |
|
|
| if got := gjson.GetBytes(out, "tool_choice.name").String(); got != "web_search" { |
| t.Fatalf("tool_choice.name = %q, want %q", got, "web_search") |
| } |
| } |
|
|
| func TestStripClaudeToolPrefixFromResponse(t *testing.T) { |
| input := []byte(`{"content":[{"type":"tool_use","name":"proxy_alpha","id":"t1","input":{}},{"type":"tool_use","name":"bravo","id":"t2","input":{}}]}`) |
| out := stripClaudeToolPrefixFromResponse(input, "proxy_") |
|
|
| if got := gjson.GetBytes(out, "content.0.name").String(); got != "alpha" { |
| t.Fatalf("content.0.name = %q, want %q", got, "alpha") |
| } |
| if got := gjson.GetBytes(out, "content.1.name").String(); got != "bravo" { |
| t.Fatalf("content.1.name = %q, want %q", got, "bravo") |
| } |
| } |
|
|
| func TestStripClaudeToolPrefixFromResponse_WithToolReference(t *testing.T) { |
| input := []byte(`{"content":[{"type":"tool_reference","tool_name":"proxy_alpha"},{"type":"tool_reference","tool_name":"bravo"}]}`) |
| out := stripClaudeToolPrefixFromResponse(input, "proxy_") |
|
|
| if got := gjson.GetBytes(out, "content.0.tool_name").String(); got != "alpha" { |
| t.Fatalf("content.0.tool_name = %q, want %q", got, "alpha") |
| } |
| if got := gjson.GetBytes(out, "content.1.tool_name").String(); got != "bravo" { |
| t.Fatalf("content.1.tool_name = %q, want %q", got, "bravo") |
| } |
| } |
|
|
| func TestStripClaudeToolPrefixFromStreamLine(t *testing.T) { |
| line := []byte(`data: {"type":"content_block_start","content_block":{"type":"tool_use","name":"proxy_alpha","id":"t1"},"index":0}`) |
| out := stripClaudeToolPrefixFromStreamLine(line, "proxy_") |
|
|
| payload := bytes.TrimSpace(out) |
| if bytes.HasPrefix(payload, []byte("data:")) { |
| payload = bytes.TrimSpace(payload[len("data:"):]) |
| } |
| if got := gjson.GetBytes(payload, "content_block.name").String(); got != "alpha" { |
| t.Fatalf("content_block.name = %q, want %q", got, "alpha") |
| } |
| } |
|
|
| func TestStripClaudeToolPrefixFromStreamLine_WithToolReference(t *testing.T) { |
| line := []byte(`data: {"type":"content_block_start","content_block":{"type":"tool_reference","tool_name":"proxy_beta"},"index":0}`) |
| out := stripClaudeToolPrefixFromStreamLine(line, "proxy_") |
|
|
| payload := bytes.TrimSpace(out) |
| if bytes.HasPrefix(payload, []byte("data:")) { |
| payload = bytes.TrimSpace(payload[len("data:"):]) |
| } |
| if got := gjson.GetBytes(payload, "content_block.tool_name").String(); got != "beta" { |
| t.Fatalf("content_block.tool_name = %q, want %q", got, "beta") |
| } |
| } |
|
|
| func TestApplyClaudeToolPrefix_NestedToolReference(t *testing.T) { |
| input := []byte(`{"messages":[{"role":"user","content":[{"type":"tool_result","tool_use_id":"toolu_123","content":[{"type":"tool_reference","tool_name":"mcp__nia__manage_resource"}]}]}]}`) |
| out := applyClaudeToolPrefix(input, "proxy_") |
| got := gjson.GetBytes(out, "messages.0.content.0.content.0.tool_name").String() |
| if got != "proxy_mcp__nia__manage_resource" { |
| t.Fatalf("nested tool_reference tool_name = %q, want %q", got, "proxy_mcp__nia__manage_resource") |
| } |
| } |
|
|
| func TestClaudeExecutor_ReusesUserIDAcrossModelsWhenCacheEnabled(t *testing.T) { |
| resetUserIDCache() |
|
|
| var userIDs []string |
| var requestModels []string |
| server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { |
| body, _ := io.ReadAll(r.Body) |
| userID := gjson.GetBytes(body, "metadata.user_id").String() |
| model := gjson.GetBytes(body, "model").String() |
| userIDs = append(userIDs, userID) |
| requestModels = append(requestModels, model) |
| t.Logf("HTTP Server received request: model=%s, user_id=%s, url=%s", model, userID, r.URL.String()) |
| w.Header().Set("Content-Type", "application/json") |
| _, _ = w.Write([]byte(`{"id":"msg_1","type":"message","model":"claude-3-5-sonnet","role":"assistant","content":[{"type":"text","text":"ok"}],"usage":{"input_tokens":1,"output_tokens":1}}`)) |
| })) |
| defer server.Close() |
|
|
| t.Logf("End-to-end test: Fake HTTP server started at %s", server.URL) |
|
|
| cacheEnabled := true |
| executor := NewClaudeExecutor(&config.Config{ |
| ClaudeKey: []config.ClaudeKey{ |
| { |
| APIKey: "key-123", |
| BaseURL: server.URL, |
| Cloak: &config.CloakConfig{ |
| CacheUserID: &cacheEnabled, |
| }, |
| }, |
| }, |
| }) |
| auth := &cliproxyauth.Auth{Attributes: map[string]string{ |
| "api_key": "key-123", |
| "base_url": server.URL, |
| }} |
|
|
| payload := []byte(`{"messages":[{"role":"user","content":[{"type":"text","text":"hi"}]}]}`) |
| models := []string{"claude-3-5-sonnet", "claude-3-5-haiku"} |
| for _, model := range models { |
| t.Logf("Sending request for model: %s", model) |
| modelPayload, _ := sjson.SetBytes(payload, "model", model) |
| if _, err := executor.Execute(context.Background(), auth, cliproxyexecutor.Request{ |
| Model: model, |
| Payload: modelPayload, |
| }, cliproxyexecutor.Options{ |
| SourceFormat: sdktranslator.FromString("claude"), |
| }); err != nil { |
| t.Fatalf("Execute(%s) error: %v", model, err) |
| } |
| } |
|
|
| if len(userIDs) != 2 { |
| t.Fatalf("expected 2 requests, got %d", len(userIDs)) |
| } |
| if userIDs[0] == "" || userIDs[1] == "" { |
| t.Fatal("expected user_id to be populated") |
| } |
| t.Logf("user_id[0] (model=%s): %s", requestModels[0], userIDs[0]) |
| t.Logf("user_id[1] (model=%s): %s", requestModels[1], userIDs[1]) |
| if userIDs[0] != userIDs[1] { |
| t.Fatalf("expected user_id to be reused across models, got %q and %q", userIDs[0], userIDs[1]) |
| } |
| if !isValidUserID(userIDs[0]) { |
| t.Fatalf("user_id %q is not valid", userIDs[0]) |
| } |
| t.Logf("✓ End-to-end test passed: Same user_id (%s) was used for both models", userIDs[0]) |
| } |
|
|
| func TestClaudeExecutor_GeneratesNewUserIDByDefault(t *testing.T) { |
| resetUserIDCache() |
|
|
| var userIDs []string |
| server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { |
| body, _ := io.ReadAll(r.Body) |
| userIDs = append(userIDs, gjson.GetBytes(body, "metadata.user_id").String()) |
| w.Header().Set("Content-Type", "application/json") |
| _, _ = w.Write([]byte(`{"id":"msg_1","type":"message","model":"claude-3-5-sonnet","role":"assistant","content":[{"type":"text","text":"ok"}],"usage":{"input_tokens":1,"output_tokens":1}}`)) |
| })) |
| defer server.Close() |
|
|
| executor := NewClaudeExecutor(&config.Config{}) |
| auth := &cliproxyauth.Auth{Attributes: map[string]string{ |
| "api_key": "key-123", |
| "base_url": server.URL, |
| }} |
|
|
| payload := []byte(`{"messages":[{"role":"user","content":[{"type":"text","text":"hi"}]}]}`) |
|
|
| for i := 0; i < 2; i++ { |
| if _, err := executor.Execute(context.Background(), auth, cliproxyexecutor.Request{ |
| Model: "claude-3-5-sonnet", |
| Payload: payload, |
| }, cliproxyexecutor.Options{ |
| SourceFormat: sdktranslator.FromString("claude"), |
| }); err != nil { |
| t.Fatalf("Execute call %d error: %v", i, err) |
| } |
| } |
|
|
| if len(userIDs) != 2 { |
| t.Fatalf("expected 2 requests, got %d", len(userIDs)) |
| } |
| if userIDs[0] == "" || userIDs[1] == "" { |
| t.Fatal("expected user_id to be populated") |
| } |
| if userIDs[0] == userIDs[1] { |
| t.Fatalf("expected user_id to change when caching is not enabled, got identical values %q", userIDs[0]) |
| } |
| if !isValidUserID(userIDs[0]) || !isValidUserID(userIDs[1]) { |
| t.Fatalf("user_ids should be valid, got %q and %q", userIDs[0], userIDs[1]) |
| } |
| } |
|
|
| func TestStripClaudeToolPrefixFromResponse_NestedToolReference(t *testing.T) { |
| input := []byte(`{"content":[{"type":"tool_result","tool_use_id":"toolu_123","content":[{"type":"tool_reference","tool_name":"proxy_mcp__nia__manage_resource"}]}]}`) |
| out := stripClaudeToolPrefixFromResponse(input, "proxy_") |
| got := gjson.GetBytes(out, "content.0.content.0.tool_name").String() |
| if got != "mcp__nia__manage_resource" { |
| t.Fatalf("nested tool_reference tool_name = %q, want %q", got, "mcp__nia__manage_resource") |
| } |
| } |
|
|
| func TestApplyClaudeToolPrefix_NestedToolReferenceWithStringContent(t *testing.T) { |
| |
| input := []byte(`{"messages":[{"role":"user","content":[{"type":"tool_result","tool_use_id":"toolu_123","content":"plain string result"}]}]}`) |
| out := applyClaudeToolPrefix(input, "proxy_") |
| got := gjson.GetBytes(out, "messages.0.content.0.content").String() |
| if got != "plain string result" { |
| t.Fatalf("string content should remain unchanged = %q", got) |
| } |
| } |
|
|
| func TestApplyClaudeToolPrefix_SkipsBuiltinToolReference(t *testing.T) { |
| input := []byte(`{"tools":[{"type":"web_search_20250305","name":"web_search"}],"messages":[{"role":"user","content":[{"type":"tool_result","tool_use_id":"t1","content":[{"type":"tool_reference","tool_name":"web_search"}]}]}]}`) |
| out := applyClaudeToolPrefix(input, "proxy_") |
| got := gjson.GetBytes(out, "messages.0.content.0.content.0.tool_name").String() |
| if got != "web_search" { |
| t.Fatalf("built-in tool_reference should not be prefixed, got %q", got) |
| } |
| } |
|
|
