diff options
| author | Claude Sonnet 5 <noreply@anthropic.com> | 2026-07-03 22:33:22 +0000 |
|---|---|---|
| committer | Claude Sonnet 5 <noreply@anthropic.com> | 2026-07-03 22:33:22 +0000 |
| commit | 1f203a7ac0efad15ec3fc0a4c5b335ad7073a52f (patch) | |
| tree | 8047b9f63d0b1e80faf3e549b3ad21b9c1f71f1f /internal | |
| parent | 767ddade57f189827fa956ff8081ca47404a4798 (diff) | |
feat(provider): add native Google Gemini API adapter (Phase 4)
Adds internal/provider/google, the second native cloud adapter (following
internal/provider/anthropic's pattern) on top of Phase 1's provider-neutral
tool-use loop, wired to a Docker-sandboxed NativeRunner under agent.type:
"google" -- a separate execution path and budget bucket from the existing
CLI-subprocess "gemini" ContainerRunner, which is untouched.
Wire-format research (the highest-risk part of this adapter): Gemini's
multi-turn function-calling shape was resolved by cross-referencing the
REST API reference's own generateContent example against the go-genai SDK's
struct tags on GitHub -- both agree on functionCall/functionResponse parts
keyed by "name" (with an optional "id" for round-tripping ToolCall.ID),
with the response fed back inside a "user"-role Content (Gemini has no
tool/function role, mirroring Anthropic's lack of one). A separate fetched
source (the function-calling guide page) was deliberately discarded as a
reference for this shape -- it documents a different, newer "Interactions
API" whose call_id/type:"function_result" structure doesn't fit the
contents/parts/candidates shape used everywhere else.
- internal/provider/google: request/response translation, systemInstruction
handling, role mapping (assistant->model, tool-results->user role),
per-model-prefix pricing table (2.5 Pro/Flash/Flash-Lite, 2.0, 1.5 tiers)
- internal/retry: IsRateLimitError additively extended for RESOURCE_EXHAUSTED
- internal/config: RunnersConfig.Google/GoogleEnabled()
- internal/cli/serve.go, run.go: runners["google"] construction mirroring
the Anthropic wiring exactly (Docker sandbox default)
- docs/api-keys-setup.md: Google marked wired-up, budget-bucket/disable/
verify guidance added matching the Anthropic section
go build/vet/test -race all pass. No live Gemini API key available in this
environment; verified via fake-httptest-server adapter tests (plain text,
tool-use round-trip, multi-turn tool-result, rate-limit error matching) plus
a Pool/NativeRunner routing test. Live E2E is a follow-up once a key is
configured, same as Phase 2's Anthropic adapter.
Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01V1moSNCJRcP6kykA4tyUSs
Diffstat (limited to 'internal')
| -rw-r--r-- | internal/cli/run.go | 18 | ||||
| -rw-r--r-- | internal/cli/serve.go | 21 | ||||
| -rw-r--r-- | internal/config/config.go | 9 | ||||
| -rw-r--r-- | internal/config/config_test.go | 27 | ||||
| -rw-r--r-- | internal/executor/google_routing_test.go | 89 | ||||
| -rw-r--r-- | internal/provider/google/google.go | 467 | ||||
| -rw-r--r-- | internal/provider/google/google_test.go | 413 | ||||
| -rw-r--r-- | internal/provider/google/pricing.go | 60 | ||||
| -rw-r--r-- | internal/retry/backoff.go | 11 | ||||
| -rw-r--r-- | internal/retry/backoff_test.go | 7 |
10 files changed, 1115 insertions, 7 deletions
diff --git a/internal/cli/run.go b/internal/cli/run.go index 497ba6b..34c3d42 100644 --- a/internal/cli/run.go +++ b/internal/cli/run.go @@ -11,6 +11,7 @@ import ( "github.com/spf13/cobra" "github.com/thepeterstone/claudomator/internal/executor" "github.com/thepeterstone/claudomator/internal/provider/anthropic" + "github.com/thepeterstone/claudomator/internal/provider/google" "github.com/thepeterstone/claudomator/internal/provider/openaicompat" "github.com/thepeterstone/claudomator/internal/storage" "github.com/thepeterstone/claudomator/internal/task" @@ -135,6 +136,23 @@ func runTasks(file string, parallel int, dryRun bool) error { } } + if pc, ok := cfg.Providers["google"]; ok && pc.APIKey != "" && cfg.Runners.GoogleEnabled() { + timeout := time.Duration(0) + if pc.TimeoutSeconds > 0 { + timeout = time.Duration(pc.TimeoutSeconds) * time.Second + } + runners["google"] = &executor.NativeRunner{ + Provider: google.New(pc.APIKey, pc.Endpoint, timeout), + Logger: logger, + LogDir: cfg.LogDir, + DefaultModel: pc.DefaultModel, + // See internal/cli/serve.go: native cloud providers get + // sandbox.DockerSandbox rather than the weaker HostSandbox. + SandboxKind: "docker", + SandboxImage: cfg.SandboxImage, + } + } + pool := executor.NewPool(parallel, runners, store, logger) pool.Classifier = &executor.Classifier{ LLM: localClient, diff --git a/internal/cli/serve.go b/internal/cli/serve.go index 8b6cc6a..57c2cbd 100644 --- a/internal/cli/serve.go +++ b/internal/cli/serve.go @@ -17,6 +17,7 @@ import ( "github.com/thepeterstone/claudomator/internal/executor" "github.com/thepeterstone/claudomator/internal/notify" "github.com/thepeterstone/claudomator/internal/provider/anthropic" + "github.com/thepeterstone/claudomator/internal/provider/google" "github.com/thepeterstone/claudomator/internal/provider/openaicompat" "github.com/thepeterstone/claudomator/internal/storage" "github.com/thepeterstone/claudomator/internal/version" @@ -174,6 +175,26 @@ func serve(addr, basePath string) error { logger.Info("anthropic runner registered", "default_model", pc.DefaultModel, "sandbox", "docker") } + if pc, ok := cfg.Providers["google"]; ok && pc.APIKey != "" && cfg.Runners.GoogleEnabled() { + timeout := time.Duration(0) + if pc.TimeoutSeconds > 0 { + timeout = time.Duration(pc.TimeoutSeconds) * time.Second + } + runners["google"] = &executor.NativeRunner{ + Provider: google.New(pc.APIKey, pc.Endpoint, timeout), + Logger: logger, + LogDir: cfg.LogDir, + DefaultModel: pc.DefaultModel, + // Native cloud providers get real container isolation + // (sandbox.DockerSandbox) rather than HostSandbox's host-side + // path-prefix confinement — same reasoning as the "anthropic" + // runner above. + SandboxKind: "docker", + SandboxImage: cfg.SandboxImage, + } + logger.Info("google runner registered", "default_model", pc.DefaultModel, "sandbox", "docker") + } + pool := executor.NewPool(cfg.MaxConcurrent, runners, store, logger) pool.Classifier = &executor.Classifier{ LLM: localClient, diff --git a/internal/config/config.go b/internal/config/config.go index afbd12e..b9454d1 100644 --- a/internal/config/config.go +++ b/internal/config/config.go @@ -64,6 +64,14 @@ type RunnersConfig struct { // the same Anthropic account. Also requires [providers.anthropic].api_key // to be set; Anthropic = true with no configured key has no effect. Anthropic *bool `toml:"anthropic"` + // Google gates the native provider.Provider-backed NativeRunner + // registered under agent type "google" (internal/provider/google). It is + // distinct from Gemini, which gates the Docker/CLI-subprocess + // ContainerRunner for agent type "gemini" — the two are separate + // execution paths (and separate budget buckets), same reasoning as + // Anthropic vs. Claude above. Also requires [providers.google].api_key to + // be set; Google = true with no configured key has no effect. + Google *bool `toml:"google"` } func (r RunnersConfig) ClaudeEnabled() bool { return r.Claude == nil || *r.Claude } @@ -71,6 +79,7 @@ func (r RunnersConfig) GeminiEnabled() bool { return r.Gemini == nil || *r.Ge func (r RunnersConfig) LocalEnabled() bool { return r.Local == nil || *r.Local } func (r RunnersConfig) ContainerEnabled() bool { return r.Container == nil || *r.Container } func (r RunnersConfig) AnthropicEnabled() bool { return r.Anthropic == nil || *r.Anthropic } +func (r RunnersConfig) GoogleEnabled() bool { return r.Google == nil || *r.Google } // ProviderConfig configures a native (or OpenAI-compatible) LLM provider // backend for the multi-provider harness — see docs/api-keys-setup.md. Phase diff --git a/internal/config/config_test.go b/internal/config/config_test.go index f71b9cc..4167ff7 100644 --- a/internal/config/config_test.go +++ b/internal/config/config_test.go @@ -71,11 +71,14 @@ func TestRunnersConfig_DefaultsAllEnabled(t *testing.T) { if !r.AnthropicEnabled() { t.Error("anthropic should be enabled by default (nil)") } + if !r.GoogleEnabled() { + t.Error("google should be enabled by default (nil)") + } } func TestRunnersConfig_ExplicitFalseDisables(t *testing.T) { f := false - r := RunnersConfig{Claude: &f, Gemini: &f, Local: &f, Container: &f, Anthropic: &f} + r := RunnersConfig{Claude: &f, Gemini: &f, Local: &f, Container: &f, Anthropic: &f, Google: &f} if r.ClaudeEnabled() { t.Error("explicit false should disable claude") } @@ -91,12 +94,15 @@ func TestRunnersConfig_ExplicitFalseDisables(t *testing.T) { if r.AnthropicEnabled() { t.Error("explicit false should disable anthropic") } + if r.GoogleEnabled() { + t.Error("explicit false should disable google") + } } func TestRunnersConfig_ExplicitTrueEnables(t *testing.T) { tr := true - r := RunnersConfig{Claude: &tr, Gemini: &tr, Local: &tr, Container: &tr, Anthropic: &tr} - if !r.ClaudeEnabled() || !r.GeminiEnabled() || !r.LocalEnabled() || !r.ContainerEnabled() || !r.AnthropicEnabled() { + r := RunnersConfig{Claude: &tr, Gemini: &tr, Local: &tr, Container: &tr, Anthropic: &tr, Google: &tr} + if !r.ClaudeEnabled() || !r.GeminiEnabled() || !r.LocalEnabled() || !r.ContainerEnabled() || !r.AnthropicEnabled() || !r.GoogleEnabled() { t.Error("explicit true should keep all runners enabled") } } @@ -185,6 +191,10 @@ timeout_seconds = 30 [providers.anthropic] api_key = "sk-ant-xyz" default_model = "claude-sonnet-4-5" + +[providers.google] +api_key = "AIza-abc" +default_model = "gemini-2.5-flash" ` if err := os.WriteFile(path, []byte(toml), 0600); err != nil { t.Fatal(err) @@ -195,8 +205,8 @@ default_model = "claude-sonnet-4-5" if err != nil { t.Fatalf("unexpected error: %v", err) } - if len(cfg.Providers) != 2 { - t.Fatalf("expected 2 providers, got %d: %+v", len(cfg.Providers), cfg.Providers) + if len(cfg.Providers) != 3 { + t.Fatalf("expected 3 providers, got %d: %+v", len(cfg.Providers), cfg.Providers) } groq, ok := cfg.Providers["groq"] if !ok { @@ -213,4 +223,11 @@ default_model = "claude-sonnet-4-5" if anthropic.Endpoint != "" || anthropic.APIKey != "sk-ant-xyz" || anthropic.DefaultModel != "claude-sonnet-4-5" { t.Errorf("anthropic provider config: %+v", anthropic) } + google, ok := cfg.Providers["google"] + if !ok { + t.Fatal("expected [providers.google] table") + } + if google.Endpoint != "" || google.APIKey != "AIza-abc" || google.DefaultModel != "gemini-2.5-flash" { + t.Errorf("google provider config: %+v", google) + } } diff --git a/internal/executor/google_routing_test.go b/internal/executor/google_routing_test.go new file mode 100644 index 0000000..9690a4b --- /dev/null +++ b/internal/executor/google_routing_test.go @@ -0,0 +1,89 @@ +package executor + +import ( + "context" + "log/slog" + "net/http" + "net/http/httptest" + "os" + "strings" + "sync/atomic" + "testing" + + "github.com/thepeterstone/claudomator/internal/provider/google" +) + +// TestPool_Submit_GoogleAgentType_RoutesThroughGoogleProvider is the +// NativeRunner-level confirmation (Phase 4 verification item 5) that a task +// with agent.type: "google" is actually routed through the new +// internal/provider/google.Provider when the pool is wired with a "google" +// runner — exactly how internal/cli/serve.go and internal/cli/run.go +// register it from cfg.Providers["google"] when a non-empty APIKey is +// configured. This exercises the full path: Pool.Submit -> execute() -> +// NativeRunner.Run() -> agentloop.Loop.Run() -> google.Provider.Chat() -> +// HTTP POST to the (fake) Gemini generateContent API endpoint. Mirrors +// internal/executor/anthropic_routing_test.go. +func TestPool_Submit_GoogleAgentType_RoutesThroughGoogleProvider(t *testing.T) { + var calls int32 + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + atomic.AddInt32(&calls, 1) + if !strings.Contains(r.URL.Path, ":generateContent") { + t.Errorf("unexpected path %q", r.URL.Path) + } + if r.URL.Query().Get("key") != "test-google-key" { + t.Errorf("wrong key query param: %q", r.URL.Query().Get("key")) + } + w.Header().Set("Content-Type", "application/json") + w.Write([]byte(`{ + "candidates": [{ + "content": {"role": "model", "parts": [{"text": "## Summary\ndone"}]}, + "finishReason": "STOP" + }], + "usageMetadata": {"promptTokenCount": 12, "candidatesTokenCount": 4, "totalTokenCount": 16} + }`)) + })) + defer srv.Close() + + // Mirrors the exact construction in internal/cli/serve.go's + // cfg.Providers["google"] wiring block: google.New(apiKey, endpoint, + // timeout) wrapped in a NativeRunner registered under the "google" + // runner-map key. + runners := map[string]Runner{ + "google": &NativeRunner{ + Provider: google.New("test-google-key", srv.URL, 0), + Logger: slog.New(slog.NewTextHandler(os.Stderr, &slog.HandlerOptions{Level: slog.LevelError})), + LogDir: t.TempDir(), + }, + } + + store := testStore(t) + logger := slog.New(slog.NewTextHandler(os.Stderr, &slog.HandlerOptions{Level: slog.LevelError})) + pool := NewPool(2, runners, store, logger) + + tk := makeTask("google-routing-1") + tk.Agent.Type = "google" + tk.Agent.Model = "gemini-2.5-flash" + tk.Agent.SkipPlanning = true + if err := store.CreateTask(tk); err != nil { + t.Fatalf("CreateTask: %v", err) + } + + if err := pool.Submit(context.Background(), tk); err != nil { + t.Fatalf("Submit: %v", err) + } + + result := <-pool.Results() + if result.Err != nil { + t.Fatalf("expected no error, got: %v", result.Err) + } + if got := atomic.LoadInt32(&calls); got != 1 { + t.Fatalf("expected exactly 1 call to the fake Gemini server, got %d", got) + } + if result.Execution.Agent != "google" { + t.Errorf("execution.Agent: want %q, got %q", "google", result.Execution.Agent) + } + if result.Execution.TokensIn != 12 || result.Execution.TokensOut != 4 { + t.Errorf("tokens should come from the google provider's usage: got in=%d out=%d", + result.Execution.TokensIn, result.Execution.TokensOut) + } +} diff --git a/internal/provider/google/google.go b/internal/provider/google/google.go new file mode 100644 index 0000000..45fcfa9 --- /dev/null +++ b/internal/provider/google/google.go @@ -0,0 +1,467 @@ +// Package google implements provider.Provider against Google's native Gemini +// API generateContent REST endpoint +// (https://generativelanguage.googleapis.com/v1beta/models/{model}:generateContent) +// — the second native cloud adapter added on top of Phase 1's provider- +// neutral tool-use loop (internal/agentloop), following the same pattern as +// internal/provider/anthropic. +// +// Gemini's wire shape is structurally different from both Anthropic's +// (typed content blocks under a flat "content" array) and OpenAI's (flat +// tool_calls array): +// +// - The system prompt is a top-level "systemInstruction" field (itself a +// Content{parts:[...]} object), not a role:"system" message and not a +// bare string like Anthropic's "system". +// - Message content is "contents[]", each a {"role":..., "parts":[...]} +// object. Valid role values are only "user" and "model" — there is no +// "system" or "tool"/"function" role. +// - A model turn that calls a function is a "model"-role Content whose +// parts include a {"functionCall": {"name":..., "args": {...}}} part +// (alongside an optional text part for any accompanying prose). +// - Feeding a function's result back is done via a {"functionResponse": +// {"name":..., "response": {...}}} part **inside a "user"-role** +// Content — Gemini has no "tool"/"function" role, mirroring (for +// different reasons) Anthropic's lack of one. This was the single +// highest-risk detail in this adapter; see the wire-shape research note +// below for how it was confirmed. +// - Tools are declared as "tools": [{"functionDeclarations": [...]}] +// (one Tool wrapper holding all declarations), each declaration using +// "parameters" (not "input_schema" like Anthropic, and not nested under +// a "function" wrapper like OpenAI's flat tool array). +// +// # Wire-format research note (function-response shape) +// +// The REST API reference (https://ai.google.dev/api/generate-content) is +// the canonical source and documents the full multi-turn shape directly: +// a "model"-role Content with a functionCall part, followed by a +// "user"-role Content with a functionResponse part carrying {"name":..., +// "response": {...}}. This was cross-checked against the Go SDK +// (github.com/googleapis/go-genai)'s type definitions on GitHub, which +// confirm the exact JSON struct tags: Part has `FunctionCall +// *FunctionCall `json:"functionCall,omitempty"`` and `FunctionResponse +// *FunctionResponse `json:"functionResponse,omitempty"``; FunctionCall has +// `Name string `json:"name,omitempty"``, `Args map[string]any +// `json:"args,omitempty"``, and (for parallel-function-calling +// disambiguation) `ID string `json:"id,omitempty"``; FunctionResponse +// mirrors this with `Name`, `Response map[string]any`, and `ID`. Both +// sources agree, so this adapter uses functionCall/functionResponse parts +// keyed primarily by "name" with an optional "id" for round-tripping +// provider.ToolCall.ID. +// +// One fetched source (the function-calling guide, +// https://ai.google.dev/gemini-api/docs/function-calling) was NOT used to +// confirm this shape: as of this research it defaults to documenting a +// newer, structurally different "Interactions API" +// (/v1beta/interactions, snake_case fields like "call_id" and +// "type":"function_result") rather than the classic generateContent API +// this adapter targets. That shape does not fit the +// contents/parts/candidates structure seen everywhere else in the +// generateContent reference and in the SDK, which is exactly the kind of +// red-flag mismatch called out before writing this adapter — it was +// discarded in favor of the REST reference's own generateContent-specific +// multi-turn example and the SDK struct tags, which agree with each other. +// +// The API key is passed as the "?key=" query parameter on the endpoint URL, +// per the REST reference's own curl example for this specific endpoint +// (some other Gemini docs pages show an "x-goog-api-key" header instead, +// but the query-param form is what the generateContent reference itself +// documents, so that's what this adapter uses). +package google + +import ( + "bytes" + "context" + "encoding/json" + "fmt" + "io" + "net/http" + "net/url" + "strings" + "time" + + "github.com/thepeterstone/claudomator/internal/provider" +) + +// DefaultEndpoint is the Gemini API base, used when Provider.Endpoint is +// empty. +const DefaultEndpoint = "https://generativelanguage.googleapis.com" + +// APIVersion is the Gemini REST API version path segment this adapter +// speaks ("v1beta" — the version used throughout the current generateContent +// REST reference; there is no stable "v1" for generateContent as of this +// writing). +const APIVersion = "v1beta" + +// DefaultTimeout is used when New is called with a non-positive timeout. +const DefaultTimeout = 120 * time.Second + +// Provider implements provider.Provider against the Gemini generateContent +// REST API. +type Provider struct { + APIKey string + Endpoint string + HTTPClient *http.Client +} + +var _ provider.Provider = (*Provider)(nil) + +// New returns a Provider configured with apiKey. endpoint defaults to +// DefaultEndpoint when empty (the native Gemini API base — most callers +// should leave this blank; an override exists only for testing or for +// pointing at a compatible proxy). timeout configures the underlying HTTP +// client and defaults to DefaultTimeout when non-positive. +func New(apiKey, endpoint string, timeout time.Duration) *Provider { + if endpoint == "" { + endpoint = DefaultEndpoint + } + if timeout <= 0 { + timeout = DefaultTimeout + } + return &Provider{ + APIKey: apiKey, + Endpoint: endpoint, + HTTPClient: &http.Client{Timeout: timeout}, + } +} + +// Name implements provider.Provider. +func (p *Provider) Name() string { return "google" } + +// Chat implements provider.Provider by translating req into a Gemini +// generateContent request, performing the HTTP call, and translating the +// response (or error) back into provider-neutral shape. +func (p *Provider) Chat(ctx context.Context, req provider.ChatRequest) (*provider.ChatResponse, error) { + if p == nil || p.APIKey == "" { + return nil, fmt.Errorf("google: no API key configured") + } + if req.Model == "" { + return nil, fmt.Errorf("google: no model configured") + } + + wireReq := toWireRequest(req) + body, err := json.Marshal(wireReq) + if err != nil { + return nil, fmt.Errorf("google: marshal request: %w", err) + } + + endpoint := fmt.Sprintf("%s/%s/models/%s:generateContent?key=%s", + strings.TrimRight(p.Endpoint, "/"), APIVersion, req.Model, url.QueryEscape(p.APIKey)) + httpReq, err := http.NewRequestWithContext(ctx, http.MethodPost, endpoint, bytes.NewReader(body)) + if err != nil { + return nil, fmt.Errorf("google: build http request: %w", err) + } + httpReq.Header.Set("content-type", "application/json") + + client := p.HTTPClient + if client == nil { + client = &http.Client{Timeout: DefaultTimeout} + } + httpResp, err := client.Do(httpReq) + if err != nil { + return nil, fmt.Errorf("google: http: %w", err) + } + defer httpResp.Body.Close() + + raw, err := io.ReadAll(httpResp.Body) + if err != nil { + return nil, fmt.Errorf("google: read response: %w", err) + } + if httpResp.StatusCode >= 400 { + return nil, errFromStatus(httpResp.StatusCode, raw) + } + + var wireResp generateContentResponse + if err := json.Unmarshal(raw, &wireResp); err != nil { + return nil, fmt.Errorf("google: decode response: %w", err) + } + return fromWireResponse(&wireResp, req.Model), nil +} + +// --- request/response wire types (Gemini generateContent API) --- + +type wireRequest struct { + Contents []wireContent `json:"contents"` + Tools []wireTool `json:"tools,omitempty"` + SystemInstruction *wireContent `json:"systemInstruction,omitempty"` + GenerationConfig *wireGenerationConfig `json:"generationConfig,omitempty"` +} + +// wireContent is Gemini's "Content" object: a role ("user" or "model"; left +// empty for systemInstruction, which needs none) plus an ordered list of +// parts. +type wireContent struct { + Role string `json:"role,omitempty"` + Parts []wirePart `json:"parts"` +} + +// wirePart is a single Gemini "Part". Only one of Text/FunctionCall/ +// FunctionResponse is set per part on the way out; on the way in +// (fromWireResponse), only text and functionCall parts are ever produced by +// a non-streaming generateContent response, so those are the only two +// branches handled there. +type wirePart struct { + Text string `json:"text,omitempty"` + FunctionCall *wireFunctionCall `json:"functionCall,omitempty"` + FunctionResponse *wireFunctionResponse `json:"functionResponse,omitempty"` +} + +// wireFunctionCall mirrors the go-genai SDK's FunctionCall struct tags +// (functionCall/name/args/id). Args is a raw JSON object rather than a +// decoded map so provider.ToolCall.ArgsJSON round-trips byte-for-byte +// without an unnecessary unmarshal/remarshal, same trick +// internal/provider/anthropic uses for its "input" field. +type wireFunctionCall struct { + ID string `json:"id,omitempty"` + Name string `json:"name"` + Args json.RawMessage `json:"args,omitempty"` +} + +// wireFunctionResponse mirrors the go-genai SDK's FunctionResponse struct +// tags (functionResponse/name/response/id). Unlike Anthropic's tool_result +// block, Gemini imposes no schema on "response" beyond "some JSON object" — +// there is no dedicated is_error field. This adapter's convention (a +// judgment call, not something the wire spec dictates) is +// {"content": "..."} for a successful ToolResult and {"error": "..."} for +// one with IsError set, which keeps the object shape trivially valid JSON +// while still surfacing success/failure to the model. +type wireFunctionResponse struct { + ID string `json:"id,omitempty"` + Name string `json:"name"` + Response map[string]any `json:"response"` +} + +type wireTool struct { + FunctionDeclarations []wireFunctionDeclaration `json:"functionDeclarations"` +} + +type wireFunctionDeclaration struct { + Name string `json:"name"` + Description string `json:"description,omitempty"` + Parameters map[string]any `json:"parameters,omitempty"` +} + +// wireGenerationConfig carries the subset of Gemini's GenerationConfig this +// adapter translates. Unlike Anthropic, Gemini does not require +// maxOutputTokens on every request (it has a model-specific default), so +// (unlike anthropic.defaultMaxTokens) this adapter leaves it unset rather +// than synthesizing a default when ChatRequest.MaxTokens is zero. +type wireGenerationConfig struct { + Temperature *float64 `json:"temperature,omitempty"` + MaxOutputTokens int `json:"maxOutputTokens,omitempty"` +} + +type generateContentResponse struct { + Candidates []wireCandidate `json:"candidates"` + UsageMetadata wireUsageMetadata `json:"usageMetadata"` +} + +type wireCandidate struct { + Content wireContent `json:"content"` + FinishReason string `json:"finishReason"` +} + +type wireUsageMetadata struct { + PromptTokenCount int `json:"promptTokenCount"` + CandidatesTokenCount int `json:"candidatesTokenCount"` + TotalTokenCount int `json:"totalTokenCount"` +} + +// wireErrorEnvelope is Gemini's error response shape: +// +// {"error": {"code": 429, "message": "...", "status": "RESOURCE_EXHAUSTED"}} +type wireErrorEnvelope struct { + Error struct { + Code int `json:"code"` + Message string `json:"message"` + Status string `json:"status"` + } `json:"error"` +} + +// --- translation: provider-neutral -> Gemini wire shape --- + +// toWireRequest translates a provider-neutral ChatRequest into the Gemini +// generateContent request shape. +func toWireRequest(req provider.ChatRequest) *wireRequest { + contents := make([]wireContent, 0, len(req.Messages)) + for _, m := range req.Messages { + if wc := toWireContent(m); wc != nil { + contents = append(contents, *wc) + } + } + + var tools []wireTool + if len(req.Tools) > 0 { + decls := make([]wireFunctionDeclaration, 0, len(req.Tools)) + for _, ts := range req.Tools { + decls = append(decls, wireFunctionDeclaration{ + Name: ts.Name, + Description: ts.Description, + Parameters: ts.ParametersJSONSchema, + }) + } + tools = []wireTool{{FunctionDeclarations: decls}} + } + + var sysInstr *wireContent + if req.System != "" { + sysInstr = &wireContent{Parts: []wirePart{{Text: req.System}}} + } + + var genConfig *wireGenerationConfig + if req.Temperature != nil || req.MaxTokens > 0 { + genConfig = &wireGenerationConfig{ + Temperature: req.Temperature, + MaxOutputTokens: req.MaxTokens, + } + } + + return &wireRequest{ + Contents: contents, + Tools: tools, + SystemInstruction: sysInstr, + GenerationConfig: genConfig, + } +} + +// toWireContent translates a single provider-neutral Message into a Gemini +// wire Content, or nil if it carries no parts to send (e.g. an assistant +// turn with neither text nor tool calls — Gemini rejects a Content with an +// empty parts array). +// +// Role mapping: provider-neutral "user" -> Gemini "user"; "assistant" -> +// Gemini "model" (Gemini's own name for the assistant role); "tool" (a +// Message carrying ToolResults, as agentloop.Loop emits after dispatching a +// tool call) has no Gemini role of its own, so — mirroring how Anthropic +// has no "tool" role either — it becomes a **user**-role Content whose +// parts are functionResponse blocks, the standard Gemini pattern for +// feeding a function's result back to the model (see the package doc's +// wire-format research note). +func toWireContent(m provider.Message) *wireContent { + if len(m.ToolResults) > 0 { + parts := make([]wirePart, 0, len(m.ToolResults)) + for _, tr := range m.ToolResults { + respObj := map[string]any{"content": tr.Content} + if tr.IsError { + respObj = map[string]any{"error": tr.Content} + } + parts = append(parts, wirePart{ + FunctionResponse: &wireFunctionResponse{ + ID: tr.ToolCallID, + Name: tr.Name, + Response: respObj, + }, + }) + } + return &wireContent{Role: "user", Parts: parts} + } + + role := "user" + if m.Role == "assistant" { + role = "model" + } else if m.Role != "user" { + // agentloop never emits a Message with role "system" (the system + // prompt travels in ChatRequest.System instead) or any other role, + // but fall back to "user" defensively rather than silently dropping + // content if some future caller does. + role = "user" + } + + var parts []wirePart + if m.Text != "" { + parts = append(parts, wirePart{Text: m.Text}) + } + for _, tc := range m.ToolCalls { + args := json.RawMessage(tc.ArgsJSON) + if len(args) == 0 || !json.Valid(args) { + args = json.RawMessage("{}") + } + parts = append(parts, wirePart{ + FunctionCall: &wireFunctionCall{ + ID: tc.ID, + Name: tc.Name, + Args: args, + }, + }) + } + if len(parts) == 0 { + return nil + } + return &wireContent{Role: role, Parts: parts} +} + +// --- translation: Gemini wire shape -> provider-neutral --- + +// fromWireResponse translates a Gemini generateContent response into a +// provider-neutral ChatResponse, concatenating any text parts and +// collecting functionCall parts into ToolCalls from the first candidate +// (Gemini can return multiple candidates when candidateCount > 1 is +// requested; this adapter, like the Anthropic one, never requests more than +// the default single candidate, so only candidates[0] is consulted). +// model is the model name that was requested (ChatRequest.Model) — used for +// pricing lookup instead of trying to parse a response-echoed model field, +// since generateContent responses do not reliably include one. +func fromWireResponse(r *generateContentResponse, model string) *provider.ChatResponse { + var text strings.Builder + var calls []provider.ToolCall + var finishReason string + + if len(r.Candidates) > 0 { + cand := r.Candidates[0] + finishReason = cand.FinishReason + for _, part := range cand.Content.Parts { + switch { + case part.FunctionCall != nil: + args := string(part.FunctionCall.Args) + if args == "" { + args = "{}" + } + calls = append(calls, provider.ToolCall{ + ID: part.FunctionCall.ID, + Name: part.FunctionCall.Name, + ArgsJSON: args, + }) + case part.Text != "": + text.WriteString(part.Text) + } + } + } + + inTok, outTok := r.UsageMetadata.PromptTokenCount, r.UsageMetadata.CandidatesTokenCount + inPerMTok, outPerMTok := pricingFor(model) + cost := float64(inTok)/1_000_000*inPerMTok + float64(outTok)/1_000_000*outPerMTok + + return &provider.ChatResponse{ + Text: text.String(), + ToolCalls: calls, + StopReason: finishReason, + Usage: provider.Usage{ + InputTokens: inTok, + OutputTokens: outTok, + CostUSD: cost, + }, + } +} + +// errFromStatus builds an error from a non-2xx Gemini API response, +// embedding the HTTP status code and Gemini's own error "status" (e.g. +// "RESOURCE_EXHAUSTED", "INVALID_ARGUMENT", "PERMISSION_DENIED") in the +// message text so retry.IsRateLimitError can pattern-match quota/rate-limit +// conditions. +func errFromStatus(status int, body []byte) error { + var env wireErrorEnvelope + _ = json.Unmarshal(body, &env) + + errStatus := env.Error.Status + if errStatus == "" { + errStatus = "UNKNOWN" + } + msg := env.Error.Message + if msg == "" { + snippet := strings.TrimSpace(string(body)) + if len(snippet) > 500 { + snippet = snippet[:500] + "..." + } + msg = snippet + } + return fmt.Errorf("google: http %d %s: %s", status, errStatus, msg) +} diff --git a/internal/provider/google/google_test.go b/internal/provider/google/google_test.go new file mode 100644 index 0000000..5b67839 --- /dev/null +++ b/internal/provider/google/google_test.go @@ -0,0 +1,413 @@ +package google + +import ( + "context" + "encoding/json" + "net/http" + "net/http/httptest" + "strings" + "testing" + "time" + + "github.com/thepeterstone/claudomator/internal/provider" + "github.com/thepeterstone/claudomator/internal/retry" +) + +// fakeGoogleServer replies to POST /v1beta/models/{model}:generateContent +// with the given status and raw JSON body, capturing the last decoded +// request for assertions. Modeled on +// internal/provider/anthropic/anthropic_test.go's fakeAnthropicServer, +// adapted to the Gemini generateContent wire shape (model in the URL path, +// key as a query param rather than a header). +func fakeGoogleServer(t *testing.T, status int, body string, capture *wireRequest) *httptest.Server { + t.Helper() + return httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if !strings.Contains(r.URL.Path, ":generateContent") { + t.Errorf("unexpected path %q", r.URL.Path) + } + if !strings.HasPrefix(r.URL.Path, "/v1beta/models/") { + t.Errorf("expected path to start with /v1beta/models/, got %q", r.URL.Path) + } + if r.URL.Query().Get("key") != "test-key" { + t.Errorf("missing/wrong key query param: %q", r.URL.Query().Get("key")) + } + if r.Header.Get("content-type") != "application/json" { + t.Errorf("missing/wrong content-type header: %q", r.Header.Get("content-type")) + } + if capture != nil { + if err := json.NewDecoder(r.Body).Decode(capture); err != nil { + t.Fatalf("decode request body: %v", err) + } + } + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(status) + w.Write([]byte(body)) + })) +} + +func newTestProvider(url string) *Provider { + return New("test-key", url, 5*time.Second) +} + +// TestChat_PlainTextResponse covers a plain model text reply: no function +// calls, finishReason "STOP", usage translated to provider.Usage with a +// non-zero CostUSD computed from the pricing table. +func TestChat_PlainTextResponse(t *testing.T) { + var got wireRequest + srv := fakeGoogleServer(t, http.StatusOK, `{ + "candidates": [{ + "content": {"role": "model", "parts": [{"text": "Hello there!"}]}, + "finishReason": "STOP" + }], + "usageMetadata": {"promptTokenCount": 10, "candidatesTokenCount": 5, "totalTokenCount": 15} + }`, &got) + defer srv.Close() + + p := newTestProvider(srv.URL) + resp, err := p.Chat(context.Background(), provider.ChatRequest{ + Model: "gemini-2.5-flash", + System: "You are helpful.", + Messages: []provider.Message{ + {Role: "user", Text: "hi"}, + }, + }) + if err != nil { + t.Fatalf("Chat: %v", err) + } + if resp.Text != "Hello there!" { + t.Errorf("Text: want %q got %q", "Hello there!", resp.Text) + } + if len(resp.ToolCalls) != 0 { + t.Errorf("expected no tool calls, got %+v", resp.ToolCalls) + } + if resp.StopReason != "STOP" { + t.Errorf("StopReason: want STOP got %q", resp.StopReason) + } + if resp.Usage.InputTokens != 10 || resp.Usage.OutputTokens != 5 { + t.Errorf("usage tokens: got %+v", resp.Usage) + } + wantCost := float64(10)/1_000_000*0.30 + float64(5)/1_000_000*2.50 + if resp.Usage.CostUSD != wantCost { + t.Errorf("CostUSD: want %v got %v", wantCost, resp.Usage.CostUSD) + } + + // Verify the outgoing request shape: system is a top-level + // systemInstruction Content, not a message in "contents". + if got.SystemInstruction == nil || len(got.SystemInstruction.Parts) != 1 || got.SystemInstruction.Parts[0].Text != "You are helpful." { + t.Errorf("systemInstruction: got %+v", got.SystemInstruction) + } + if len(got.Contents) != 1 || got.Contents[0].Role != "user" { + t.Fatalf("contents: %+v", got.Contents) + } + if len(got.Contents[0].Parts) != 1 || got.Contents[0].Parts[0].Text != "hi" { + t.Errorf("user content part: %+v", got.Contents[0].Parts) + } +} + +// TestChat_ToolUseResponse verifies that a response whose parts mix text and +// a functionCall round-trips into ChatResponse.ToolCalls correctly +// (id/name/args), and that tools declared on the request are translated to +// Gemini's {tools: [{functionDeclarations: [...]}]} shape with "parameters" +// (not "input_schema" or a "function" wrapper). +func TestChat_ToolUseResponse(t *testing.T) { + var got wireRequest + srv := fakeGoogleServer(t, http.StatusOK, `{ + "candidates": [{ + "content": { + "role": "model", + "parts": [ + {"text": "Let me check that."}, + {"functionCall": {"id": "call_01", "name": "read_file", "args": {"path": "a.txt"}}} + ] + }, + "finishReason": "STOP" + }], + "usageMetadata": {"promptTokenCount": 20, "candidatesTokenCount": 8, "totalTokenCount": 28} + }`, &got) + defer srv.Close() + + p := newTestProvider(srv.URL) + resp, err := p.Chat(context.Background(), provider.ChatRequest{ + Model: "gemini-2.5-pro", + Messages: []provider.Message{{Role: "user", Text: "read a.txt"}}, + Tools: []provider.ToolSpec{ + { + Name: "read_file", + Description: "Read a file.", + ParametersJSONSchema: map[string]any{ + "type": "object", + "properties": map[string]any{"path": map[string]any{"type": "string"}}, + "required": []string{"path"}, + }, + }, + }, + }) + if err != nil { + t.Fatalf("Chat: %v", err) + } + if resp.Text != "Let me check that." { + t.Errorf("Text: got %q", resp.Text) + } + if resp.StopReason != "STOP" { + t.Errorf("StopReason: want STOP got %q", resp.StopReason) + } + if len(resp.ToolCalls) != 1 { + t.Fatalf("expected 1 tool call, got %d: %+v", len(resp.ToolCalls), resp.ToolCalls) + } + tc := resp.ToolCalls[0] + if tc.ID != "call_01" || tc.Name != "read_file" { + t.Errorf("tool call id/name: got %+v", tc) + } + var args struct { + Path string `json:"path"` + } + if err := json.Unmarshal([]byte(tc.ArgsJSON), &args); err != nil { + t.Fatalf("tool call args not valid JSON: %v (%q)", err, tc.ArgsJSON) + } + if args.Path != "a.txt" { + t.Errorf("tool call args.path: got %q", args.Path) + } + + // Verify tools were declared as tools:[{functionDeclarations:[...]}]. + if len(got.Tools) != 1 || len(got.Tools[0].FunctionDeclarations) != 1 { + t.Fatalf("expected 1 tool wrapper with 1 declaration, got %+v", got.Tools) + } + decl := got.Tools[0].FunctionDeclarations[0] + if decl.Name != "read_file" || decl.Description != "Read a file." { + t.Errorf("function declaration: %+v", decl) + } + if decl.Parameters == nil || decl.Parameters["type"] != "object" { + t.Errorf("parameters not forwarded: %+v", decl.Parameters) + } +} + +// TestChat_MultiTurnWithPriorToolResult exercises the functionResponse-as- +// a-user-content translation: a provider.Message carrying ToolResults (as +// agentloop.Loop emits with Role "tool" after dispatching a tool call) must +// serialize as a **user**-role Content with functionResponse parts — Gemini +// has no "tool"/"function" role. Also verifies a prior "assistant" turn +// with a tool call round-trips into a single "model"-role Content with a +// functionCall part carrying the same id used in the subsequent +// functionResponse. +func TestChat_MultiTurnWithPriorToolResult(t *testing.T) { + var got wireRequest + srv := fakeGoogleServer(t, http.StatusOK, `{ + "candidates": [{ + "content": {"role": "model", "parts": [{"text": "The file says hello."}]}, + "finishReason": "STOP" + }], + "usageMetadata": {"promptTokenCount": 30, "candidatesTokenCount": 6, "totalTokenCount": 36} + }`, &got) + defer srv.Close() + + p := newTestProvider(srv.URL) + _, err := p.Chat(context.Background(), provider.ChatRequest{ + Model: "gemini-2.5-flash", + Messages: []provider.Message{ + {Role: "user", Text: "read a.txt"}, + { + Role: "assistant", + Text: "", + ToolCalls: []provider.ToolCall{ + {ID: "call_01", Name: "read_file", ArgsJSON: `{"path":"a.txt"}`}, + }, + }, + { + Role: "tool", + ToolResults: []provider.ToolResult{ + {ToolCallID: "call_01", Name: "read_file", Content: "hello", IsError: false}, + }, + }, + }, + }) + if err != nil { + t.Fatalf("Chat: %v", err) + } + + if len(got.Contents) != 3 { + t.Fatalf("expected 3 wire contents, got %d: %+v", len(got.Contents), got.Contents) + } + + // Turn 2: model functionCall, no text part since Text was empty. + asst := got.Contents[1] + if asst.Role != "model" { + t.Errorf("turn 2 role: want model got %q", asst.Role) + } + if len(asst.Parts) != 1 || asst.Parts[0].FunctionCall == nil { + t.Fatalf("turn 2 parts: %+v", asst.Parts) + } + if asst.Parts[0].FunctionCall.ID != "call_01" || asst.Parts[0].FunctionCall.Name != "read_file" { + t.Errorf("turn 2 functionCall part: %+v", asst.Parts[0].FunctionCall) + } + + // Turn 3: the tool result must be a USER-role Content with a + // functionResponse part (Gemini has no "tool"/"function" role). + toolTurn := got.Contents[2] + if toolTurn.Role != "user" { + t.Errorf("tool result turn role: want user (Gemini has no tool role) got %q", toolTurn.Role) + } + if len(toolTurn.Parts) != 1 || toolTurn.Parts[0].FunctionResponse == nil { + t.Fatalf("tool result parts: %+v", toolTurn.Parts) + } + fr := toolTurn.Parts[0].FunctionResponse + if fr.ID != "call_01" { + t.Errorf("functionResponse id: want call_01 got %q", fr.ID) + } + if fr.Name != "read_file" { + t.Errorf("functionResponse name: want read_file got %q", fr.Name) + } + if fr.Response["content"] != "hello" { + t.Errorf("functionResponse.response.content: want hello got %+v", fr.Response) + } + if _, hasErr := fr.Response["error"]; hasErr { + t.Errorf("response should not carry an error key when IsError is false: %+v", fr.Response) + } +} + +// TestChat_ToolResultIsError verifies IsError round-trips onto the +// functionResponse's response object as an "error" key (this adapter's +// convention — Gemini imposes no schema on the response object). +func TestChat_ToolResultIsError(t *testing.T) { + var got wireRequest + srv := fakeGoogleServer(t, http.StatusOK, `{ + "candidates": [{"content": {"role": "model", "parts": [{"text": "ok"}]}, "finishReason": "STOP"}], + "usageMetadata": {"promptTokenCount": 1, "candidatesTokenCount": 1, "totalTokenCount": 2} + }`, &got) + defer srv.Close() + + p := newTestProvider(srv.URL) + _, err := p.Chat(context.Background(), provider.ChatRequest{ + Model: "gemini-2.5-flash", + Messages: []provider.Message{ + {Role: "tool", ToolResults: []provider.ToolResult{ + {ToolCallID: "call_x", Name: "run_bash", Content: "boom", IsError: true}, + }}, + }, + }) + if err != nil { + t.Fatalf("Chat: %v", err) + } + if len(got.Contents) != 1 || len(got.Contents[0].Parts) != 1 { + t.Fatalf("unexpected contents: %+v", got.Contents) + } + fr := got.Contents[0].Parts[0].FunctionResponse + if fr == nil { + t.Fatalf("expected a functionResponse part") + } + if fr.Response["error"] != "boom" { + t.Errorf("expected response.error = boom, got %+v", fr.Response) + } +} + +// TestChat_RateLimitError429 verifies a 429 RESOURCE_EXHAUSTED response +// produces a Go error whose text retry.IsRateLimitError recognizes. +func TestChat_RateLimitError429(t *testing.T) { + srv := fakeGoogleServer(t, http.StatusTooManyRequests, `{ + "error": {"code": 429, "message": "You exceeded your current quota.", "status": "RESOURCE_EXHAUSTED"} + }`, nil) + defer srv.Close() + + p := newTestProvider(srv.URL) + _, err := p.Chat(context.Background(), provider.ChatRequest{ + Model: "gemini-2.5-flash", + Messages: []provider.Message{{Role: "user", Text: "hi"}}, + }) + if err == nil { + t.Fatal("expected an error for HTTP 429") + } + if !strings.Contains(err.Error(), "429") || !strings.Contains(err.Error(), "RESOURCE_EXHAUSTED") { + t.Errorf("error should mention 429 and RESOURCE_EXHAUSTED, got: %v", err) + } + if !retry.IsRateLimitError(err) { + t.Errorf("retry.IsRateLimitError should recognize: %v", err) + } +} + +// TestChat_InvalidArgumentError400 is a non-retryable error case: confirms +// the error text carries the INVALID_ARGUMENT status and is NOT +// misclassified as a rate-limit error. +func TestChat_InvalidArgumentError400(t *testing.T) { + srv := fakeGoogleServer(t, http.StatusBadRequest, `{ + "error": {"code": 400, "message": "Request contains an invalid argument.", "status": "INVALID_ARGUMENT"} + }`, nil) + defer srv.Close() + + p := newTestProvider(srv.URL) + _, err := p.Chat(context.Background(), provider.ChatRequest{ + Model: "gemini-2.5-flash", + Messages: []provider.Message{{Role: "user", Text: "hi"}}, + }) + if err == nil { + t.Fatal("expected an error for HTTP 400") + } + if !strings.Contains(err.Error(), "INVALID_ARGUMENT") { + t.Errorf("error should mention INVALID_ARGUMENT, got: %v", err) + } + if retry.IsRateLimitError(err) { + t.Errorf("400 INVALID_ARGUMENT should NOT be classified as a rate-limit error: %v", err) + } +} + +// TestChat_NoAPIKey_Errors confirms Chat fails fast without ever making an +// HTTP request when no API key is configured. +func TestChat_NoAPIKey_Errors(t *testing.T) { + p := New("", "http://127.0.0.1:1", time.Second) + _, err := p.Chat(context.Background(), provider.ChatRequest{Model: "gemini-2.5-flash"}) + if err == nil || !strings.Contains(err.Error(), "no API key") { + t.Errorf("expected 'no API key' error, got %v", err) + } +} + +// TestChat_NoModel_Errors confirms Chat fails fast when no model is set, +// since the model name is part of the URL path (unlike Anthropic, where a +// missing model just gets sent through and rejected server-side). +func TestChat_NoModel_Errors(t *testing.T) { + p := New("test-key", "http://127.0.0.1:1", time.Second) + _, err := p.Chat(context.Background(), provider.ChatRequest{}) + if err == nil || !strings.Contains(err.Error(), "no model") { + t.Errorf("expected 'no model' error, got %v", err) + } +} + +// TestNew_DefaultsEndpointAndTimeout verifies New falls back to +// DefaultEndpoint/DefaultTimeout for zero-value arguments. +func TestNew_DefaultsEndpointAndTimeout(t *testing.T) { + p := New("key", "", 0) + if p.Endpoint != DefaultEndpoint { + t.Errorf("Endpoint: want %q got %q", DefaultEndpoint, p.Endpoint) + } + if p.HTTPClient.Timeout != DefaultTimeout { + t.Errorf("Timeout: want %v got %v", DefaultTimeout, p.HTTPClient.Timeout) + } +} + +// TestName returns the provider's identifier, used as the agent-type/budget +// key ("google") once wired into a NativeRunner. +func TestName(t *testing.T) { + p := New("key", "", 0) + if p.Name() != "google" { + t.Errorf("Name: got %q", p.Name()) + } +} + +// TestPricingFor_KnownAndUnknownModels sanity-checks the pricing table +// lookup: exact/prefix match for known models, and a non-erroring fallback +// for an unrecognized model string. +func TestPricingFor_KnownAndUnknownModels(t *testing.T) { + in, out := pricingFor("gemini-2.5-pro") + if in != 1.25 || out != 10.00 { + t.Errorf("2.5 pro pricing: got %v/%v", in, out) + } + in, out = pricingFor("gemini-2.5-flash-lite-preview-09-2025") + if in != 0.10 || out != 0.40 { + t.Errorf("2.5 flash-lite dated snapshot pricing: got %v/%v", in, out) + } + in, out = pricingFor("gemini-2.5-flash-preview") + if in != 0.30 || out != 2.50 { + t.Errorf("2.5 flash (not lite) pricing: got %v/%v", in, out) + } + in, out = pricingFor("some-future-unknown-model") + if in != defaultPricing.InputPerMTok || out != defaultPricing.OutputPerMTok { + t.Errorf("unknown model should fall back to default pricing, got %v/%v", in, out) + } +} diff --git a/internal/provider/google/pricing.go b/internal/provider/google/pricing.go new file mode 100644 index 0000000..55ea66f --- /dev/null +++ b/internal/provider/google/pricing.go @@ -0,0 +1,60 @@ +package google + +import "strings" + +// modelPrice is USD cost per million tokens for one model (or model family +// prefix). +type modelPrice struct { + InputPerMTok float64 + OutputPerMTok float64 +} + +// pricingTable maps a model-name *prefix* to its per-million-token pricing. +// Keyed by prefix (rather than exact model string) so dated/versioned model +// IDs (e.g. "gemini-2.5-flash-preview-09-2025") still resolve without +// needing an entry per snapshot. Edit this table as Google's pricing or +// model lineup changes — it intentionally has no other logic attached. +// +// Prices below are cached from training/research as of mid-2026 and are not +// fetched live; verify against https://ai.google.dev/gemini-api/docs/pricing +// before relying on them for real billing decisions. Two caveats specific to +// Gemini pricing that this flat per-model table does NOT model: +// - Some tiers (notably 1.5 Pro and some 2.5 Pro contexts) charge a higher +// rate once the prompt exceeds 128K/200K tokens; the price below is the +// lower/base-context rate. +// - Some Flash-tier models charge a different (higher) output rate for +// "thinking"/reasoning tokens vs. plain output tokens; the price below +// is the standard output rate, not the thinking-token rate. +var pricingTable = map[string]modelPrice{ + // Pro tier (most capable, most expensive). + "gemini-2.5-pro": {InputPerMTok: 1.25, OutputPerMTok: 10.00}, + "gemini-1.5-pro": {InputPerMTok: 1.25, OutputPerMTok: 5.00}, + + // Flash tier (balanced cost/capability; default for most tasks). + "gemini-2.5-flash-lite": {InputPerMTok: 0.10, OutputPerMTok: 0.40}, + "gemini-2.5-flash": {InputPerMTok: 0.30, OutputPerMTok: 2.50}, + "gemini-2.0-flash-lite": {InputPerMTok: 0.075, OutputPerMTok: 0.30}, + "gemini-2.0-flash": {InputPerMTok: 0.10, OutputPerMTok: 0.40}, + "gemini-1.5-flash-8b": {InputPerMTok: 0.0375, OutputPerMTok: 0.15}, + "gemini-1.5-flash": {InputPerMTok: 0.075, OutputPerMTok: 0.30}, +} + +// defaultPricing is used for unrecognized models rather than erroring — +// Flash-tier pricing is a reasonable default since it's the most commonly +// used tier for the harness's escalation ladder. +var defaultPricing = modelPrice{InputPerMTok: 0.30, OutputPerMTok: 2.50} + +// pricingFor returns the per-million-token input/output pricing for model, +// matching the longest pricingTable prefix that model starts with. Falls +// back to defaultPricing for unrecognized models. +func pricingFor(model string) (inputPerMTok, outputPerMTok float64) { + price := defaultPricing + bestLen := -1 + for prefix, p := range pricingTable { + if len(prefix) > bestLen && strings.HasPrefix(model, prefix) { + bestLen = len(prefix) + price = p + } + } + return price.InputPerMTok, price.OutputPerMTok +} diff --git a/internal/retry/backoff.go b/internal/retry/backoff.go index a372b37..72b923c 100644 --- a/internal/retry/backoff.go +++ b/internal/retry/backoff.go @@ -22,7 +22,13 @@ const maxBackoffDelay = 5 * time.Minute // OpenAI-compatible error text (internal/llm); "rate_limit_error", // "overloaded_error", and "529" additionally match Anthropic's Messages API // error shape (HTTP 429 with error.type "rate_limit_error", or HTTP 529 -// with error.type "overloaded_error" — see internal/provider/anthropic). +// with error.type "overloaded_error" — see internal/provider/anthropic); +// "resource_exhausted" additionally matches Gemini's generateContent API +// error shape (HTTP 429 with error.status "RESOURCE_EXHAUSTED" — see +// internal/provider/google). The plain "429" check above already covers +// Gemini's HTTP status code, but the status string is matched too for +// parity with how the Anthropic error types are matched by name, not just +// by status code. func IsRateLimitError(err error) bool { if err == nil { return false @@ -34,7 +40,8 @@ func IsRateLimitError(err error) bool { strings.Contains(msg, "overloaded") || strings.Contains(msg, "rate_limit_error") || strings.Contains(msg, "overloaded_error") || - strings.Contains(msg, "529") + strings.Contains(msg, "529") || + strings.Contains(msg, "resource_exhausted") } // ParseRetryAfter extracts a Retry-After duration from an error message. diff --git a/internal/retry/backoff_test.go b/internal/retry/backoff_test.go index d605f29..4d893c3 100644 --- a/internal/retry/backoff_test.go +++ b/internal/retry/backoff_test.go @@ -59,6 +59,13 @@ func TestIsRateLimitError_HTTP529(t *testing.T) { } } +func TestIsRateLimitError_GoogleResourceExhausted(t *testing.T) { + err := errors.New(`google: http 429 RESOURCE_EXHAUSTED: You exceeded your current quota.`) + if !IsRateLimitError(err) { + t.Error("want true for Gemini 'RESOURCE_EXHAUSTED' shape, got false") + } +} + func TestIsRateLimitError_NonRateLimitError(t *testing.T) { err := errors.New("claude exited with error: exit status 1") if IsRateLimitError(err) { |
