package providers import ( "encoding/json" "fmt" "net/http" "net/http/httptest" "testing" "time" "github.com/YspCoder/clawgo/pkg/config" ) func TestGeminiBuildRequestBody(t *testing.T) { p := NewGeminiProvider("gemini", "", "", "gemini-2.5-pro", false, "oauth", 5*time.Second, nil) body := p.buildRequestBody([]Message{ {Role: "system", Content: "You are helpful."}, {Role: "user", Content: "hello"}, { Role: "assistant", Content: "calling tool", ToolCalls: []ToolCall{{ ID: "call_1", Name: "lookup", Function: &FunctionCall{ Name: "lookup", Arguments: `{"q":"weather"}`, }, }}, }, {Role: "tool", ToolCallID: "call_1", Content: `{"ok":true}`}, }, []ToolDefinition{{ Type: "function", Function: ToolFunctionDefinition{ Name: "lookup", Description: "Lookup data", Parameters: map[string]interface{}{"type": "object"}, }, }}, "gemini-2.5-pro", map[string]interface{}{ "max_tokens": 128, "temperature": 0.3, }, false) request := body if system := asString(mapFromAny(request["systemInstruction"])["parts"].([]map[string]any)[0]["text"]); system != "You are helpful." { t.Fatalf("expected system instruction, got %q", system) } if got := len(request["contents"].([]map[string]any)); got != 3 { t.Fatalf("expected 3 content entries, got %d", got) } gen := mapFromAny(request["generationConfig"]) if got := intValue(gen["maxOutputTokens"]); got != 128 { t.Fatalf("expected maxOutputTokens, got %#v", gen["maxOutputTokens"]) } } func TestGeminiProviderCountTokens(t *testing.T) { server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { if r.URL.Path != "/v1beta/models/gemini-2.5-pro:countTokens" { http.NotFound(w, r) return } w.Header().Set("Content-Type", "application/json") _, _ = w.Write([]byte(`{"totalTokens":42}`)) })) defer server.Close() p := NewGeminiProvider("gemini", "token", server.URL, "gemini-2.5-pro", false, "api_key", 5*time.Second, nil) usage, err := p.CountTokens(t.Context(), []Message{{Role: "user", Content: "hello"}}, nil, "gemini-2.5-pro", nil) if err != nil { t.Fatalf("CountTokens error: %v", err) } if usage == nil || usage.PromptTokens != 42 || usage.TotalTokens != 42 { t.Fatalf("usage = %#v, want 42", usage) } } func TestGeminiProviderChat(t *testing.T) { server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { if r.URL.Path != "/v1beta/models/gemini-2.5-pro:generateContent" { http.NotFound(w, r) return } if got := r.Header.Get("x-goog-api-key"); got != "token" { t.Fatalf("x-goog-api-key = %q", got) } w.Header().Set("Content-Type", "application/json") _, _ = w.Write([]byte(`{"candidates":[{"content":{"parts":[{"text":"ok"}]},"finishReason":"STOP"}],"usageMetadata":{"promptTokenCount":1,"candidatesTokenCount":1,"totalTokenCount":2}}`)) })) defer server.Close() p := NewGeminiProvider("gemini", "token", server.URL, "gemini-2.5-pro", false, "api_key", 5*time.Second, nil) resp, err := p.Chat(t.Context(), []Message{{Role: "user", Content: "hello"}}, nil, "gemini-2.5-pro", nil) if err != nil { t.Fatalf("Chat error: %v", err) } if resp.Content != "ok" { t.Fatalf("content = %q, want ok", resp.Content) } if resp.Usage == nil || resp.Usage.TotalTokens != 2 { t.Fatalf("usage = %#v", resp.Usage) } } func TestGeminiBuildRequestBodyFixesImageAspectRatio(t *testing.T) { p := NewGeminiProvider("gemini", "", "", geminiImagePreviewModel, false, "api_key", 5*time.Second, nil) body := p.buildRequestBody([]Message{ {Role: "user", Content: "draw a cat"}, }, nil, geminiImagePreviewModel, map[string]interface{}{ "gemini_generation_config": map[string]interface{}{ "imageConfig": map[string]interface{}{"aspectRatio": "1:1"}, }, }, false) gen := mapFromAny(body["generationConfig"]) if _, ok := gen["imageConfig"]; ok { t.Fatalf("expected imageConfig to be removed, got %#v", gen["imageConfig"]) } if got := len(gen["responseModalities"].([]any)); got != 2 { t.Fatalf("responseModalities len = %d", got) } contents := body["contents"].([]map[string]any) parts := contents[0]["parts"].([]map[string]any) if len(parts) < 3 { t.Fatalf("parts = %#v", parts) } if _, ok := mapFromAny(parts[1]["inlineData"])["data"]; !ok { t.Fatalf("expected inlineData placeholder, got %#v", parts[1]) } } func TestGeminiBuildRequestBodyAppliesBudgetThinkingSuffix(t *testing.T) { p := NewGeminiProvider("gemini", "", "", "gemini-2.5-pro", false, "api_key", 5*time.Second, nil) body := p.buildRequestBody([]Message{{Role: "user", Content: "hi"}}, nil, "gemini-2.5-pro(high)", nil, false) gen := mapFromAny(body["generationConfig"]) thinking := mapFromAny(gen["thinkingConfig"]) if got := intValue(thinking["thinkingBudget"]); got != 24576 { t.Fatalf("thinkingBudget = %v, want 24576", thinking["thinkingBudget"]) } if got := fmt.Sprintf("%v", thinking["includeThoughts"]); got != "true" { t.Fatalf("includeThoughts = %v, want true", thinking["includeThoughts"]) } body = p.buildRequestBody([]Message{{Role: "user", Content: "hi"}}, nil, "gemini-2.5-pro(none)", nil, false) gen = mapFromAny(body["generationConfig"]) thinking = mapFromAny(gen["thinkingConfig"]) if got := intValue(thinking["thinkingBudget"]); got != 128 { t.Fatalf("thinkingBudget = %v, want 128", thinking["thinkingBudget"]) } if got := fmt.Sprintf("%v", thinking["includeThoughts"]); got != "false" { t.Fatalf("includeThoughts = %v, want false", thinking["includeThoughts"]) } body = p.buildRequestBody([]Message{{Role: "user", Content: "hi"}}, nil, "gemini-2.5-pro(auto)", nil, false) gen = mapFromAny(body["generationConfig"]) thinking = mapFromAny(gen["thinkingConfig"]) if got := intValue(thinking["thinkingBudget"]); got != -1 { t.Fatalf("thinkingBudget = %v, want -1", thinking["thinkingBudget"]) } body = p.buildRequestBody([]Message{{Role: "user", Content: "hi"}}, nil, "gemini-2.5-pro(8192)", nil, false) gen = mapFromAny(body["generationConfig"]) thinking = mapFromAny(gen["thinkingConfig"]) if got := intValue(thinking["thinkingBudget"]); got != 8192 { t.Fatalf("thinkingBudget = %v, want 8192", thinking["thinkingBudget"]) } } func TestGeminiBuildRequestBodyAppliesLevelThinkingSuffix(t *testing.T) { p := NewGeminiProvider("gemini", "", "", "gemini-3-pro-preview", false, "api_key", 5*time.Second, nil) body := p.buildRequestBody([]Message{{Role: "user", Content: "hi"}}, nil, "gemini-3-pro-preview(high)", nil, false) gen := mapFromAny(body["generationConfig"]) thinking := mapFromAny(gen["thinkingConfig"]) if got := asString(thinking["thinkingLevel"]); got != "high" { t.Fatalf("thinkingLevel = %q, want high", got) } if got := fmt.Sprintf("%v", thinking["includeThoughts"]); got != "true" { t.Fatalf("includeThoughts = %v, want true", thinking["includeThoughts"]) } body = p.buildRequestBody([]Message{{Role: "user", Content: "hi"}}, nil, "gemini-3-pro-preview(xhigh)", nil, false) gen = mapFromAny(body["generationConfig"]) thinking = mapFromAny(gen["thinkingConfig"]) if got := asString(thinking["thinkingLevel"]); got != "high" { t.Fatalf("thinkingLevel = %q, want high", got) } body = p.buildRequestBody([]Message{{Role: "user", Content: "hi"}}, nil, "gemini-3-pro-preview(none)", nil, false) gen = mapFromAny(body["generationConfig"]) thinking = mapFromAny(gen["thinkingConfig"]) if got := asString(thinking["thinkingLevel"]); got != "low" { t.Fatalf("thinkingLevel = %q, want low", got) } if got := fmt.Sprintf("%v", thinking["includeThoughts"]); got != "false" { t.Fatalf("includeThoughts = %v, want false", thinking["includeThoughts"]) } body = p.buildRequestBody([]Message{{Role: "user", Content: "hi"}}, nil, "gemini-3-pro-preview(auto)", nil, false) gen = mapFromAny(body["generationConfig"]) thinking = mapFromAny(gen["thinkingConfig"]) if got := intValue(thinking["thinkingBudget"]); got != -1 { t.Fatalf("thinkingBudget = %v, want -1", thinking["thinkingBudget"]) } } func TestFilterGeminiSSEUsageMetadataDropsNonTerminalUsage(t *testing.T) { raw := []byte(`{"candidates":[{"content":{"parts":[{"text":"hello"}]}}],"usageMetadata":{"promptTokenCount":1}}`) filtered := filterGeminiSSEUsageMetadata(raw) var payload map[string]any if err := json.Unmarshal(filtered, &payload); err != nil { t.Fatalf("unmarshal filtered: %v", err) } if _, ok := payload["usageMetadata"]; ok { t.Fatalf("expected usageMetadata to be removed: %#v", payload) } } func TestGeminiTextPartsSupportInlineDataAndFileURLs(t *testing.T) { parts := geminiTextParts(Message{ Role: "user", ContentParts: []MessageContentPart{ {Type: "input_image", ImageURL: "data:image/png;base64,AAAA", MIMEType: "image/png"}, {Type: "input_file", FileURL: "https://example.com/doc.pdf", MIMEType: "application/pdf"}, }, }) if len(parts) != 2 { t.Fatalf("parts = %#v", parts) } inline := mapFromAny(parts[0]["inlineData"]) if got := asString(inline["data"]); got != "AAAA" { t.Fatalf("inline data = %q", got) } fileData := mapFromAny(parts[1]["fileData"]) if got := asString(fileData["fileUri"]); got != "https://example.com/doc.pdf" { t.Fatalf("file uri = %q", got) } } func TestGeminiBaseURLForAttemptUsesSessionResourceURL(t *testing.T) { base := NewHTTPProvider("gemini", "token", geminiBaseURL, "gemini-2.5-pro", false, "oauth", 5*time.Second, nil) got := geminiBaseURLForAttempt(base, authAttempt{ kind: "oauth", session: &oauthSession{ ResourceURL: "https://generativelanguage.googleapis.com/v1beta/models", }, }) if got != geminiBaseURL { t.Fatalf("base url = %q, want %q", got, geminiBaseURL) } } func TestCreateProviderByNameRoutesGeminiCLIToGeminiProvider(t *testing.T) { cfg := &config.Config{ Models: config.ModelsConfig{ Providers: map[string]config.ProviderConfig{ "gemini-cli": { APIBase: "", APIKey: "token", TimeoutSec: 30, Models: []string{"gemini-2.5-pro"}, }, }, }, } provider, err := CreateProviderByName(cfg, "gemini-cli") if err != nil { t.Fatalf("CreateProviderByName error: %v", err) } if _, ok := provider.(*GeminiCLIProvider); !ok { t.Fatalf("provider = %T, want *GeminiCLIProvider", provider) } } func TestCreateProviderByNameRoutesAIStudioProviderViaGeminiTests(t *testing.T) { cfg := &config.Config{ Models: config.ModelsConfig{ Providers: map[string]config.ProviderConfig{ "aistudio": { TimeoutSec: 30, Models: []string{"gemini-2.5-pro"}, }, }, }, } provider, err := CreateProviderByName(cfg, "aistudio") if err != nil { t.Fatalf("CreateProviderByName error: %v", err) } if _, ok := provider.(*AistudioProvider); !ok { t.Fatalf("provider = %T, want *AistudioProvider", provider) } } func TestGeminiThinkingSuffixPreservesExplicitIncludeThoughts(t *testing.T) { p := NewGeminiProvider("gemini", "", "", "gemini-3-pro-preview", false, "api_key", 5*time.Second, nil) body := p.buildRequestBody([]Message{{Role: "user", Content: "hi"}}, nil, "gemini-3-pro-preview(high)", map[string]interface{}{ "gemini_generation_config": map[string]interface{}{ "thinkingConfig": map[string]interface{}{ "includeThoughts": false, }, }, }, false) gen := mapFromAny(body["generationConfig"]) thinking := mapFromAny(gen["thinkingConfig"]) if got := asString(thinking["thinkingLevel"]); got != "high" { t.Fatalf("thinkingLevel = %q, want high", got) } if got := fmt.Sprintf("%v", thinking["includeThoughts"]); got != "false" { t.Fatalf("includeThoughts = %v, want false", thinking["includeThoughts"]) } }