diff --git a/.env.template b/.env.template index 68dd2048..3fa474e1 100644 --- a/.env.template +++ b/.env.template @@ -358,9 +358,16 @@ # Add more instances with __*. # Example: OPENAI_EAST_API_KEY and OPENAI_EAST_BASE_URL register provider openai-east. # Underscores in the suffix become hyphens in the provider name. +# +# Give ONE provider several API keys with _API_KEY_, numbered from 2. +# Requests then rotate across the keys round robin, which lifts the per-key rate +# limit. Works for any provider that authenticates with an API key. +# Caveat: rotation defeats provider prompt caching, because providers scope the +# cache to the key that filled it. Rotate to raise rate limits, not to cut cost. # OpenAI # OPENAI_API_KEY=sk-... +# OPENAI_API_KEY_2=sk-... # OPENAI_BASE_URL=https://api.openai.com/v1 # Anthropic diff --git a/CLAUDE.md b/CLAUDE.md index b8ccd159..876186c1 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -132,5 +132,6 @@ Full reference: `.env.template` and `config/config.yaml` - **Resilience:** Configured via `config/config.yaml` - global `resilience.retry.*` and `resilience.circuit_breaker.*` defaults with optional per-provider overrides under `providers..resilience.retry.*` and `providers..resilience.circuit_breaker.*`. Retry defaults: `max_retries` (3), `initial_backoff` (1s), `max_backoff` (30s), `backoff_factor` (2.0), `jitter_factor` (0.1). Circuit breaker defaults: `failure_threshold` (5), `success_threshold` (2), `timeout` (30s). Breaker state is per-process, not shown in the dashboard, and exported as the `gomodel_circuit_breaker_state` gauge when metrics are enabled. - **Metrics:** `METRICS_ENABLED` (false), `METRICS_ENDPOINT` (/metrics) - **Guardrails:** Definitions are persisted in the `guardrail_definitions` store and managed via the admin API/dashboard; `config/config.yaml` entries are validated and upserted into that store at startup (a seed, not the source of truth). `GUARDRAILS_ENABLED` env var gates the feature. +- **Provider API key rotation:** Any API-key provider accepts several keys: `[_SUFFIX]_API_KEY_` env vars (numbered from 2; `_1` is accepted as a synonym for the unsuffixed key) or `providers..api_keys` in `config.yaml` (merged after `api_key`, de-duplicated, unresolved `${...}` entries dropped; env replaces the whole YAML list). Two or more keys turn on round-robin rotation, drawn per outbound HTTP request — including retries, so a 429'd request retries under the next key. Realtime websocket sessions pick a key per session. Counters are in-memory per instance. The trailing number names a key, not a provider: `OPENAI_API_KEY_2` is key 2 of `openai`, while `OPENAI_REGION_2_API_KEY` is the sole key of provider `openai-region-2`. **Rotation defeats provider prompt caching** (providers scope the cache to the key that filled it); use it to lift per-key rate limits, not to save cost. Keyless (Ollama, vLLM) and non-API-key providers (Vertex, Bedrock) are unaffected. - **Providers:** `OPENAI_API_KEY`, `ANTHROPIC_API_KEY`, `ANTHROPIC_DEFAULT_MAX_TOKENS` (optional default `max_tokens` for Anthropic-translated requests that omit it; default 4096), `GEMINI_API_KEY`, `USE_GOOGLE_GEMINI_NATIVE_API` (true by default; false uses Gemini's OpenAI-compatible chat API), `XAI_API_KEY`, `GROQ_API_KEY`, `FIREWORKS_API_KEY`, `FIREWORKS_BASE_URL` (optional Fireworks AI endpoint override; default `https://api.fireworks.ai/inference/v1`), `OPENROUTER_API_KEY`, `OPENROUTER_SITE_URL`/`OPENROUTER_APP_NAME` (optional OpenRouter attribution headers), `ZAI_API_KEY`, `ZAI_BASE_URL` (optional Z.ai endpoint override), `MINIMAX_API_KEY`, `MINIMAX_BASE_URL` (optional MiniMax endpoint override), `XIAOMI_API_KEY`, `XIAOMI_BASE_URL` (optional Xiaomi MiMo endpoint override), `OPENCODE_GO_API_KEY`, `OPENCODE_GO_BASE_URL` (optional OpenCode Go/Zen endpoint override; default `https://opencode.ai/zen/go/v1`), `OPENCODE_GO_MESSAGES_MODELS` (optional comma-separated model IDs routed to the Anthropic-native `/messages` endpoint instead of `/chat/completions`; default `qwen3.7-max`), `BAILIAN_API_KEY`, `BAILIAN_BASE_URL` (optional Bailian base URL for region switching; default `https://dashscope.aliyuncs.com/compatible-mode/v1`), `AZURE_API_KEY`, `AZURE_BASE_URL` (Azure OpenAI deployment base URL), `AZURE_API_VERSION` (optional Azure API version), `ORACLE_API_KEY` (Oracle API key), `ORACLE_BASE_URL` (Oracle OpenAI-compatible base URL), `[_SUFFIX]_MODELS` (comma-separated configured model list for any provider type), `OLLAMA_BASE_URL`, `VLLM_BASE_URL`, `VLLM_API_KEY` (optional upstream vLLM bearer token) - **Provider model metadata:** `providers..models` accepts either model IDs (strings) or `{id, metadata}` objects. When `metadata` is supplied (`display_name`, `context_window`, `max_output_tokens`, `modes`, `capabilities`, `pricing`, …) it is merged onto the remote ai-model-list entry during enrichment, with operator values winning per-field. Primary use case: advertising context windows, capabilities, and pricing for local models (Ollama) and other custom endpoints whose IDs are not in the upstream registry. diff --git a/config/config.example.yaml b/config/config.example.yaml index dd7c462f..b4fa933a 100644 --- a/config/config.example.yaml +++ b/config/config.example.yaml @@ -266,6 +266,13 @@ providers: openai: type: openai api_key: "sk-..." + # Several keys for one provider rotate round robin, raising the per-key rate + # limit. Equivalent env vars: OPENAI_API_KEY_2, OPENAI_API_KEY_3, ... + # Caveat: rotation defeats provider prompt caching, since providers scope the + # cache to the key that filled it. + # api_keys: + # - "${OPENAI_API_KEY_2}" + # - "${OPENAI_API_KEY_3}" # Per-provider resilience overrides (optional). # Only specified fields override the global defaults above. # resilience: diff --git a/config/providers.go b/config/providers.go index 20876024..07428b1b 100644 --- a/config/providers.go +++ b/config/providers.go @@ -4,8 +4,12 @@ package config // overrides, credential filtering, or resilience merging. Exported so the // providers package can resolve it into a fully-configured ProviderConfig. type RawProviderConfig struct { - Type string `yaml:"type"` - APIKey string `yaml:"api_key"` + Type string `yaml:"type"` + APIKey string `yaml:"api_key"` + // APIKeys lists additional API keys for this provider. When more than one + // key is resolved (counting APIKey), requests rotate across them round + // robin. Set it via `api_keys:` or the `_API_KEY_` env vars. + APIKeys []string `yaml:"api_keys"` BaseURL string `yaml:"base_url"` APIVersion string `yaml:"api_version"` Backend string `yaml:"backend"` diff --git a/docs/docs.json b/docs/docs.json index 5b1437fd..8036bff7 100644 --- a/docs/docs.json +++ b/docs/docs.json @@ -123,6 +123,7 @@ "icon": "plug", "pages": [ "providers/overview", + "providers/key-rotation", "providers/anthropic", "providers/gemini", "providers/deepseek", diff --git a/docs/providers/key-rotation.mdx b/docs/providers/key-rotation.mdx new file mode 100644 index 00000000..5f516887 --- /dev/null +++ b/docs/providers/key-rotation.mdx @@ -0,0 +1,130 @@ +--- +title: "API Key Rotation" +description: "Give one provider several API keys and spread requests across them round robin to lift per-key rate limits." +icon: "key-round" +--- + +A provider normally holds one API key, and every request to it uses that key. +Configure more than one and GoModel spreads requests across them round robin, so +each key carries a fraction of the traffic. This raises the per-key rate limits +you hit before the provider starts returning `429`. + +Rotation is off until a second key is configured. One key behaves exactly as it +always has. + + + **Rotation defeats provider prompt caching.** Providers scope their prompt + cache to the API key that created it, so consecutive requests sent under + different keys never hit a warm cache. Expect higher cost and slower + time-to-first-token on workloads with long, repeated prefixes — system + prompts, few-shot examples, long documents, or coding agents replaying a + conversation. + + Rotate when you are rate limited. Do not rotate to save money on cached + prompts; it does the opposite. + + +## Configure + +Add keys with numbered environment variables. `_API_KEY` is the first +key; `_API_KEY_2`, `_3`, and so on add the rest. + +```bash +OPENAI_API_KEY=sk-first +OPENAI_API_KEY_2=sk-second +OPENAI_API_KEY_3=sk-third +``` + +Three keys, so requests to `openai` cycle `sk-first` → `sk-second` → +`sk-third` → `sk-first` → … + +Or in `config.yaml`: + +```yaml +providers: + openai: + type: openai + api_keys: + - "${OPENAI_API_KEY}" + - "${OPENAI_API_KEY_2}" + - "${OPENAI_API_KEY_3}" +``` + +`api_key` and `api_keys` may be combined; `api_key` is used first. + +Any provider that authenticates with an API key can rotate — `ANTHROPIC_API_KEY_2`, +`GEMINI_API_KEY_2`, `GROQ_API_KEY_2`, and so on. Providers that authenticate +another way (Vertex AI and Bedrock, which use Google and AWS credentials) ignore +these variables. + +## Rules + +- Numbering starts at `2`, since `_API_KEY` is key 1. Spelling out + `_API_KEY_1` instead of the unsuffixed form also works. +- Gaps are fine. Setting only `_API_KEY` and `_API_KEY_3` gives two keys. +- Duplicate keys are collapsed, so a key repeated across two variables does not + take a double share of the traffic. +- Keys are used in the order configured: the unsuffixed key first, then + ascending by number. +- Environment variables replace the provider's whole `api_keys` list from + `config.yaml` rather than merging into it. +- A key that does not resolve — an `${UNSET_VAR}` placeholder — is dropped. A + provider left with no keys at all is not registered. + + + Keys rotate per outbound HTTP request, including retries. A request that a + provider rejects with `429` is retried under the next key rather than the one + that was just throttled. + + +## Rotation on suffixed providers + +Suffixed environment variables register a *separate provider* of the same type. +The two mechanisms compose, and the trailing number is what tells them apart: + +```bash +OPENAI_EU_API_KEY=sk-eu-first +OPENAI_EU_API_KEY_2=sk-eu-second +``` + +This is provider `openai-eu` with two keys. The trailing `_2` names a key. + +```bash +OPENAI_REGION_2_API_KEY=sk-only +``` + +This is provider `openai-region-2` with one key. Here the `2` is part of the +provider suffix, not a key number. + +## What rotation covers + +Every call GoModel makes to the provider draws the next key: chat completions, +responses, embeddings, audio, images, batches, and provider-native +[passthrough routes](/features/passthrough-api). Realtime websocket sessions pick a +key when the session opens and keep it for the life of that session. + +Rotation is per gateway process. Counters live in memory and reset on restart, +so `N` replicas each rotate independently — which is fine, since the goal is to +spread load rather than to sequence it exactly. + +## Choosing between rotation and separate providers + +Both let you use several keys of one provider type. They solve different problems. + +| | Key rotation | Suffixed providers | +| --- | --- | --- | +| Config | `OPENAI_API_KEY_2` | `OPENAI_EAST_API_KEY` | +| Model IDs | unchanged (`gpt-4o`) | provider-qualified (`openai-east/gpt-4o`) | +| Routing | automatic, round robin | caller picks, or a [virtual model](/features/virtual-models) balances | +| Prompt caching | lost | preserved per provider | +| Use it for | lifting rate limits on one account | separate accounts, regions, or base URLs | + +If you want several keys *and* prompt caching, register them as suffixed +providers and put a `round_robin` [virtual model](/features/virtual-models) in +front. Each provider keeps its own key, so each keeps its own warm cache, at the +cost of qualified model IDs and more configuration. + +## Verify + +Rotation is not reported by the API. To confirm it, watch the provider's own +usage dashboard: traffic should divide roughly evenly across the keys. diff --git a/docs/providers/overview.mdx b/docs/providers/overview.mdx index a874a57f..a0ef1f68 100644 --- a/docs/providers/overview.mdx +++ b/docs/providers/overview.mdx @@ -75,6 +75,10 @@ support, not every individual model capability exposed by an upstream provider. for providers that define a list, skipping their upstream `/models` calls. - **vLLM** — set `VLLM_API_KEY` only if the upstream server was started with `--api-key`. +- **Multiple API keys for one provider** — set `OPENAI_API_KEY_2`, + `OPENAI_API_KEY_3`, and so on to spread requests across keys round robin and + lift per-key rate limits. This costs you provider prompt caching; see + [API key rotation](/providers/key-rotation). - **Multiple instances of one provider type** — without `config.yaml`, use suffixed env vars such as `OPENAI_EAST_API_KEY` and `OPENAI_EAST_BASE_URL`; add `OPENAI_EAST_MODELS` to configure that instance's model list. This diff --git a/internal/embedding/embedding.go b/internal/embedding/embedding.go index fa5f4d32..45b289d2 100644 --- a/internal/embedding/embedding.go +++ b/internal/embedding/embedding.go @@ -13,6 +13,7 @@ import ( "github.com/goccy/go-json" "gomodel/config" + "gomodel/internal/providers" ) const defaultTimeout = 120 * time.Second @@ -55,9 +56,12 @@ func NewEmbedder(cfg config.EmbedderConfig, resolvedProviders map[string]config. } else if typ == "gemini" { model = normalizeGeminiEmbeddingModel(model) } + // APIKey leads APIKeys and is de-duplicated away when it repeats there, so + // this works whether raw came from provider resolution (which normalizes + // both fields) or was built by hand with only APIKey set. return &apiEmbedder{ endpointURL: endpointURL, - apiKey: raw.APIKey, + keys: providers.NewKeyring(append([]string{raw.APIKey}, raw.APIKeys...)...), model: model, httpClient: &http.Client{Timeout: defaultTimeout}, }, nil @@ -93,9 +97,12 @@ func openAIEmbeddingsEndpointURL(base string) (string, error) { } // apiEmbedder calls POST …/v1/embeddings on any OpenAI-compatible endpoint. +// It shares the provider's configured key set but keeps its own rotation +// counter, since it does not route through the provider's HTTP client. +// Embeddings have no prompt cache to lose, so rotating here is free. type apiEmbedder struct { endpointURL string - apiKey string + keys *providers.Keyring model string httpClient *http.Client } @@ -124,8 +131,8 @@ func (e *apiEmbedder) Embed(ctx context.Context, text string) ([]float32, error) return nil, fmt.Errorf("embedding: build request: %w", err) } req.Header.Set("Content-Type", "application/json") - if e.apiKey != "" { - req.Header.Set("Authorization", "Bearer "+e.apiKey) + if apiKey := e.keys.Next(); apiKey != "" { + req.Header.Set("Authorization", "Bearer "+apiKey) } resp, err := e.httpClient.Do(req) if err != nil { diff --git a/internal/embedding/embedding_test.go b/internal/embedding/embedding_test.go index d75f7ae3..3fbfefce 100644 --- a/internal/embedding/embedding_test.go +++ b/internal/embedding/embedding_test.go @@ -137,8 +137,8 @@ func TestAPIEmbedder_UsesProviderCredentials(t *testing.T) { if !ok { t.Fatalf("expected *apiEmbedder, got %T", emb) } - if a.apiKey != "gsk-abc" { - t.Errorf("expected apiKey gsk-abc, got %q", a.apiKey) + if got := a.keys.Primary(); got != "gsk-abc" { + t.Errorf("expected primary key gsk-abc, got %q", got) } if want := "https://api.groq.com/openai/v1/embeddings"; a.endpointURL != want { t.Errorf("endpointURL = %q, want %q", a.endpointURL, want) diff --git a/internal/providers/anthropic/anthropic.go b/internal/providers/anthropic/anthropic.go index 9a6a7ed4..10222ea4 100644 --- a/internal/providers/anthropic/anthropic.go +++ b/internal/providers/anthropic/anthropic.go @@ -45,7 +45,7 @@ var allowedAnthropicImageMediaTypes = map[string]struct{}{ // Provider implements the core.Provider interface for Anthropic type Provider struct { client *llmclient.Client - apiKey string + keys *providers.Keyring batchEndpointsMu sync.RWMutex // batchResultEndpoints keeps endpoint hints by provider batch id and custom_id. @@ -56,7 +56,7 @@ type Provider struct { // New creates a new Anthropic provider. func New(providerCfg providers.ProviderConfig, opts providers.ProviderOptions) core.Provider { p := &Provider{ - apiKey: providerCfg.APIKey, + keys: opts.Keyring(providerCfg.APIKey), batchResultEndpoints: make(map[string]map[string]string), } clientCfg := llmclient.Config{ @@ -77,7 +77,7 @@ func NewWithHTTPClient(apiKey string, httpClient *http.Client, hooks llmclient.H httpClient = http.DefaultClient } p := &Provider{ - apiKey: apiKey, + keys: providers.NewKeyring(apiKey), batchResultEndpoints: make(map[string]map[string]string), } cfg := llmclient.DefaultConfig("anthropic", defaultBaseURL) @@ -155,9 +155,10 @@ func (p *Provider) getBatchResultEndpoints(batchID string) map[string]string { return cloned } -// setHeaders sets the required headers for Anthropic API requests +// setHeaders sets the required headers for Anthropic API requests. It runs once +// per outbound request, so p.keys.Next advances the rotation per call. func (p *Provider) setHeaders(req *http.Request) { - req.Header.Set("x-api-key", p.apiKey) + req.Header.Set("x-api-key", p.keys.Next()) req.Header.Set("anthropic-version", anthropicAPIVersion) // Forward request ID if present in context diff --git a/internal/providers/anthropic/anthropic_test.go b/internal/providers/anthropic/anthropic_test.go index 33f424d1..d02a0b00 100644 --- a/internal/providers/anthropic/anthropic_test.go +++ b/internal/providers/anthropic/anthropic_test.go @@ -22,8 +22,8 @@ func TestNew(t *testing.T) { // Use NewWithHTTPClient to get concrete type for internal testing provider := NewWithHTTPClient(apiKey, nil, llmclient.Hooks{}) - if provider.apiKey != apiKey { - t.Errorf("apiKey = %q, want %q", provider.apiKey, apiKey) + if got := provider.keys.Primary(); got != apiKey { + t.Errorf("primary key = %q, want %q", got, apiKey) } if provider.client == nil { t.Error("client should not be nil") diff --git a/internal/providers/azure/azure.go b/internal/providers/azure/azure.go index d4ef7553..8559e0fb 100644 --- a/internal/providers/azure/azure.go +++ b/internal/providers/azure/azure.go @@ -29,13 +29,14 @@ type Provider struct { resourceProvider *openai.CompatibleProvider openAIResourceProvider *openai.CompatibleProvider apiVersion string - apiKey string // retained to inject the api-key header on the realtime target + keys *providers.Keyring // retained to inject the api-key header on the realtime target } func New(providerCfg providers.ProviderConfig, opts providers.ProviderOptions) core.Provider { baseURL := providers.ResolveBaseURL(providerCfg.BaseURL, "https://example.invalid") apiVersion := providers.ResolveAPIVersion(providerCfg.APIVersion, defaultAPIVersion) - p := &Provider{apiVersion: apiVersion, apiKey: providerCfg.APIKey} + // All three clients share opts.Keys, so the rotation is even across them. + p := &Provider{apiVersion: apiVersion, keys: opts.Keyring(providerCfg.APIKey)} clientCfg := openai.CompatibleProviderConfig{ ProviderName: "azure", BaseURL: baseURL, @@ -52,7 +53,7 @@ func New(providerCfg providers.ProviderConfig, opts providers.ProviderOptions) c } func NewWithHTTPClient(apiKey string, httpClient *http.Client, hooks llmclient.Hooks) *Provider { - p := &Provider{apiVersion: defaultAPIVersion, apiKey: apiKey} + p := &Provider{apiVersion: defaultAPIVersion, keys: providers.NewKeyring(apiKey)} cfg := openai.CompatibleProviderConfig{ ProviderName: "azure", BaseURL: "https://example.invalid", diff --git a/internal/providers/azure/realtime.go b/internal/providers/azure/realtime.go index d5b61c62..d64493ec 100644 --- a/internal/providers/azure/realtime.go +++ b/internal/providers/azure/realtime.go @@ -122,8 +122,8 @@ func (p *Provider) realtimeRoot(insecureScheme, secureScheme string) (*url.URL, func (p *Provider) realtimeAuthHeaders() http.Header { headers := http.Header{} - if p.apiKey != "" { - headers.Set("api-key", p.apiKey) + if apiKey := p.keys.Next(); apiKey != "" { + headers.Set("api-key", apiKey) } return headers } diff --git a/internal/providers/bailian/bailian.go b/internal/providers/bailian/bailian.go index 9eddead8..c81387cd 100644 --- a/internal/providers/bailian/bailian.go +++ b/internal/providers/bailian/bailian.go @@ -43,28 +43,28 @@ type Provider struct { *openai.BatchSurface *openai.FileSurface compatible *openai.CompatibleProvider - apiKey string // retained to inject auth on the realtime websocket target + keys *providers.Keyring // retained to inject auth on the realtime websocket target } // New creates a new Bailian provider from a resolved ProviderConfig. func New(cfg providers.ProviderConfig, opts providers.ProviderOptions) core.Provider { compat := openai.NewCompatibleProvider(cfg.APIKey, opts, compatibleConfig(providers.ResolveBaseURL(cfg.BaseURL, defaultBaseURL))) - return newProvider(compat, cfg.APIKey) + return newProvider(compat, opts.Keyring(cfg.APIKey)) } // NewWithHTTPClient creates a new Bailian provider with a custom HTTP client. // If httpClient is nil, http.DefaultClient is used. func NewWithHTTPClient(apiKey string, httpClient *http.Client, hooks llmclient.Hooks) *Provider { compat := openai.NewCompatibleProviderWithHTTPClient(apiKey, httpClient, hooks, compatibleConfig(defaultBaseURL)) - return newProvider(compat, apiKey) + return newProvider(compat, providers.NewKeyring(apiKey)) } -func newProvider(compat *openai.CompatibleProvider, apiKey string) *Provider { +func newProvider(compat *openai.CompatibleProvider, keys *providers.Keyring) *Provider { return &Provider{ BatchSurface: openai.NewBatchSurface(compat), FileSurface: openai.NewFileSurface(compat), compatible: compat, - apiKey: apiKey, + keys: keys, } } diff --git a/internal/providers/bailian/realtime.go b/internal/providers/bailian/realtime.go index 1814a530..c1f20ac5 100644 --- a/internal/providers/bailian/realtime.go +++ b/internal/providers/bailian/realtime.go @@ -25,8 +25,8 @@ func (p *Provider) RealtimeTarget(_ context.Context, req *core.RealtimeRequest) } headers := http.Header{} - if p.apiKey != "" { - headers.Set("Authorization", "Bearer "+p.apiKey) + if apiKey := p.keys.Next(); apiKey != "" { + headers.Set("Authorization", "Bearer "+apiKey) } return &core.RealtimeTarget{URL: endpoint, Headers: headers}, nil diff --git a/internal/providers/config.go b/internal/providers/config.go index df3ed77a..b15bc279 100644 --- a/internal/providers/config.go +++ b/internal/providers/config.go @@ -4,6 +4,7 @@ import ( "maps" "os" "sort" + "strconv" "strings" "unicode" @@ -14,8 +15,15 @@ import ( // ProviderConfig holds the fully resolved provider configuration after merging // global defaults with per-provider overrides. type ProviderConfig struct { - Type string - APIKey string + Type string + // APIKey is the provider's primary credential: the first entry of APIKeys, + // or "" for keyless providers. Prefer APIKeys for anything that + // authenticates a request, so rotation is honoured. + APIKey string + // APIKeys is the provider's full, ordered, de-duplicated key set. Requests + // rotate across it round robin when it holds more than one key. It is nil + // for keyless providers and holds exactly one entry in the common case. + APIKeys []string BaseURL string APIVersion string Backend string @@ -42,11 +50,53 @@ type ProviderConfig struct { // (same keys as the first); use it for auxiliary clients that need the same // API keys and base URLs as the live router (e.g. semantic-cache embeddings). func resolveProviders(raw map[string]config.RawProviderConfig, global config.ResilienceConfig, discovery map[string]DiscoveryConfig) (map[string]ProviderConfig, map[string]config.RawProviderConfig) { - merged := applyProviderEnvVars(raw, discovery) + merged := normalizeProviderAPIKeys(applyProviderEnvVars(raw, discovery)) filtered := filterEmptyProviders(merged, discovery) return buildProviderConfigs(filtered, global), filtered } +// normalizeProviderAPIKeys collapses each provider's `api_key` and `api_keys` +// into one canonical ordered set: APIKeys holds every usable key and APIKey +// holds the first. Unresolved `${VAR}` placeholders are dropped here rather +// than forwarded as literal credentials, so a provider whose only key failed +// to resolve ends up keyless and is then dropped by filterEmptyProviders -- +// the same outcome as before rotation existed. +func normalizeProviderAPIKeys(raw map[string]config.RawProviderConfig) map[string]config.RawProviderConfig { + result := make(map[string]config.RawProviderConfig, len(raw)) + for name, p := range raw { + keys := resolvedAPIKeys(append([]string{p.APIKey}, p.APIKeys...)) + p.APIKeys = keys + p.APIKey = "" + if len(keys) > 0 { + p.APIKey = keys[0] + } + result[name] = p + } + return result +} + +// resolvedAPIKeys trims, drops unresolved and empty entries, and de-duplicates +// while preserving order. +func resolvedAPIKeys(keys []string) []string { + resolved := make([]string, 0, len(keys)) + seen := make(map[string]struct{}, len(keys)) + for _, key := range keys { + key = strings.TrimSpace(key) + if !HasResolvedProviderValue(key) { + continue + } + if _, dup := seen[key]; dup { + continue + } + seen[key] = struct{}{} + resolved = append(resolved, key) + } + if len(resolved) == 0 { + return nil + } + return resolved +} + // applyProviderEnvVars overlays well-known provider env vars onto the raw YAML map. // Env var values always win over YAML values for the same provider name. func applyProviderEnvVars(raw map[string]config.RawProviderConfig, discovery map[string]DiscoveryConfig) map[string]config.RawProviderConfig { @@ -101,7 +151,11 @@ type providerEnvSource struct { } type providerEnvValues struct { - APIKey string + APIKey string + // APIKeysByIndex holds `_API_KEY_` values keyed by n. The + // unsuffixed `_API_KEY` is kept apart in APIKey so it can claim + // slot 1 regardless of the order os.Environ happens to return. + APIKeysByIndex map[int]string BaseURL string APIVersion string Backend string @@ -116,8 +170,58 @@ type providerEnvValues struct { Models []string } +// apiKeys returns the ordered key set this env group declares: the unsuffixed +// key leads, then the numbered keys in ascending index order. Gaps are ignored, +// so setting only `_API_KEY` and `_API_KEY_3` yields two keys, and a key +// repeated across `_API_KEY` and `_API_KEY_1` is de-duplicated to one. +func (v providerEnvValues) apiKeys() []string { + if strings.TrimSpace(v.APIKey) == "" && len(v.APIKeysByIndex) == 0 { + return nil + } + + // The unsuffixed key sorts ahead of every numbered slot, which are 1-based. + byIndex := make(map[int]string, len(v.APIKeysByIndex)+1) + maps.Copy(byIndex, v.APIKeysByIndex) + if strings.TrimSpace(v.APIKey) != "" { + byIndex[0] = v.APIKey + } + + indexes := make([]int, 0, len(byIndex)) + for index := range byIndex { + indexes = append(indexes, index) + } + sort.Ints(indexes) + + keys := make([]string, 0, len(indexes)) + for _, index := range indexes { + keys = append(keys, byIndex[index]) + } + return resolvedAPIKeys(keys) +} + +// hasAPIKey reports whether this env group carries any credential, numbered or +// not. Base-URL defaulting keys off it, so a provider configured only through +// `_API_KEY_2` still resolves its default endpoint. +// +// It probes the fields directly rather than calling apiKeys: empty() asks this +// question for every env group, and ordering the keys to then discard them +// costs a map, a sort, and two slices. Both spellings agree because a key that +// fails HasResolvedProviderValue -- blank, whitespace, or an unresolved +// `${VAR}` -- is one that apiKeys would drop. +func (v providerEnvValues) hasAPIKey() bool { + if HasResolvedProviderValue(v.APIKey) { + return true + } + for _, key := range v.APIKeysByIndex { + if HasResolvedProviderValue(key) { + return true + } + } + return false +} + func (v providerEnvValues) empty() bool { - return strings.TrimSpace(v.APIKey) == "" && + return !v.hasAPIKey() && strings.TrimSpace(v.BaseURL) == "" && strings.TrimSpace(v.APIVersion) == "" && strings.TrimSpace(v.Backend) == "" && @@ -155,7 +259,7 @@ func collectProviderEnvValues(prefix string, spec DiscoveryConfig, environ []str continue } - suffix, field, ok := parseProviderEnvKey(prefix, key, spec) + suffix, field, index, ok := parseProviderEnvKey(prefix, key, spec) if !ok { continue } @@ -163,7 +267,14 @@ func collectProviderEnvValues(prefix string, spec DiscoveryConfig, environ []str values := groups[suffix] switch field { case providerEnvFieldAPIKey: - values.APIKey = value + if index == 0 { + values.APIKey = value + break + } + if values.APIKeysByIndex == nil { + values.APIKeysByIndex = make(map[int]string) + } + values.APIKeysByIndex[index] = value case providerEnvFieldBaseURL: values.BaseURL = normalizeResolvedBaseURL(value) case providerEnvFieldAPIVersion: @@ -201,10 +312,29 @@ func collectProviderEnvValues(prefix string, spec DiscoveryConfig, environ []str return groups } -func parseProviderEnvKey(prefix, key string, spec DiscoveryConfig) (string, providerEnvField, bool) { +// parseProviderEnvKey splits a provider env var into the provider-name suffix, +// the field it sets, and (for API keys) the 1-based rotation index. An index of +// 0 means the unsuffixed `_API_KEY`. +func parseProviderEnvKey(prefix, key string, spec DiscoveryConfig) (string, providerEnvField, int, bool) { rest, ok := strings.CutPrefix(key, prefix+"_") if !ok { - return "", 0, false + return "", 0, 0, false + } + + // A trailing `_` on an API key names a rotation slot, so check it before + // the field table: `OPENAI_API_KEY_2` is key 2 of provider `openai`, and + // `OPENAI_EU_API_KEY_2` is key 2 of provider `openai-eu`. A suffix that + // merely ends in a number is unambiguous the other way -- in + // `OPENAI_REGION_2_API_KEY` the digits do not trail the key, so it stays + // provider `openai-region-2`. + if base, index, isIndexed := cutAPIKeyIndex(rest); isIndexed { + if base == "API_KEY" { + return "", providerEnvFieldAPIKey, index, true + } + if suffix, found := strings.CutSuffix(base, "_API_KEY"); found && validProviderEnvSuffix(suffix) { + return suffix, providerEnvFieldAPIKey, index, true + } + return "", 0, 0, false } // Match field names from the right so suffixes can contain underscores. @@ -243,15 +373,40 @@ func parseProviderEnvKey(prefix, key string, spec DiscoveryConfig) (string, prov continue } if rest == candidate.name { - return "", candidate.field, true + return "", candidate.field, 0, true } suffix, found := strings.CutSuffix(rest, "_"+candidate.name) if found && validProviderEnvSuffix(suffix) { - return suffix, candidate.field, true + return suffix, candidate.field, 0, true } } - return "", 0, false + return "", 0, 0, false +} + +// cutAPIKeyIndex splits a trailing rotation index off an API-key env var, +// reporting the remaining base and the 1-based index. `_1` is accepted as well +// as `_2` and up: operators who spell every slot out (`_1`, `_2`, `_3`) get +// the keys they configured rather than a silently dropped first one. +func cutAPIKeyIndex(rest string) (string, int, bool) { + base, digits, found := lastCut(rest, "_") + if !found || !strings.HasSuffix(base, "API_KEY") { + return "", 0, false + } + index, err := strconv.Atoi(digits) + if err != nil || index < 1 { + return "", 0, false + } + return base, index, true +} + +// lastCut is strings.Cut anchored at the final separator. +func lastCut(s, sep string) (string, string, bool) { + i := strings.LastIndex(s, sep) + if i < 0 { + return s, "", false + } + return s[:i], s[i+len(sep):], true } func validProviderEnvSuffix(suffix string) bool { @@ -332,9 +487,15 @@ func applySuffixedProviderEnvVars(result map[string]config.RawProviderConfig, pr func (v providerEnvValues) rawConfig(providerType string, spec DiscoveryConfig) config.RawProviderConfig { backend := v.Backend + keys := v.apiKeys() + primary := "" + if len(keys) > 0 { + primary = keys[0] + } return config.RawProviderConfig{ Type: providerType, - APIKey: v.APIKey, + APIKey: primary, + APIKeys: keys, BaseURL: v.resolvedBaseURL(spec), APIVersion: v.APIVersion, Backend: backend, @@ -352,19 +513,23 @@ func (v providerEnvValues) rawConfig(providerType string, spec DiscoveryConfig) func (v providerEnvValues) resolvedBaseURL(spec DiscoveryConfig) string { baseURL := strings.TrimSpace(v.BaseURL) - if baseURL == "" && strings.TrimSpace(v.APIKey) != "" && spec.DefaultBaseURL != "" { + if baseURL == "" && v.hasAPIKey() && spec.DefaultBaseURL != "" { return spec.DefaultBaseURL } return baseURL } func overlayProviderEnvValues(existing config.RawProviderConfig, values providerEnvValues, spec DiscoveryConfig) config.RawProviderConfig { - if values.APIKey != "" { - existing.APIKey = values.APIKey + // Env replaces the provider's whole key set rather than merging into it, so + // dropping `OPENAI_API_KEY_2` from the environment removes that key instead + // of leaving a stale YAML entry rotating behind it. + if keys := values.apiKeys(); len(keys) > 0 { + existing.APIKey = keys[0] + existing.APIKeys = keys } if values.BaseURL != "" { existing.BaseURL = values.BaseURL - } else if normalizeResolvedBaseURL(existing.BaseURL) == "" && values.APIKey != "" && spec.DefaultBaseURL != "" { + } else if normalizeResolvedBaseURL(existing.BaseURL) == "" && values.hasAPIKey() && spec.DefaultBaseURL != "" { existing.BaseURL = spec.DefaultBaseURL } if values.APIVersion != "" { @@ -639,6 +804,7 @@ func buildProviderConfig(raw config.RawProviderConfig, global config.ResilienceC resolved := ProviderConfig{ Type: normalizeProviderType(raw), APIKey: raw.APIKey, + APIKeys: raw.APIKeys, BaseURL: raw.BaseURL, APIVersion: raw.APIVersion, Backend: raw.Backend, diff --git a/internal/providers/factory.go b/internal/providers/factory.go index bbfa8931..c5c9c16a 100644 --- a/internal/providers/factory.go +++ b/internal/providers/factory.go @@ -16,6 +16,23 @@ type ProviderOptions struct { Hooks llmclient.Hooks Models []string Resilience config.ResilienceConfig + // Keys carries every API key configured for this provider instance. It is + // nil for keyless providers and for constructors invoked outside the + // factory; use the Keyring method rather than reading it directly. + Keys *Keyring +} + +// Keyring returns the key source a provider should authenticate with, falling +// back to a single-key ring over apiKey when the factory supplied none. Every +// provider constructor takes an API key and ProviderOptions, so this one call +// gives a provider rotation support without changing its signature, and keeps +// constructors invoked outside the factory (tests, the NewWithHTTPClient +// variants) working unchanged. +func (o ProviderOptions) Keyring(apiKey string) *Keyring { + if o.Keys != nil { + return o.Keys + } + return NewKeyring(apiKey) } // ProviderConstructor is the constructor signature for providers. @@ -96,10 +113,13 @@ func (f *ProviderFactory) Create(cfg ProviderConfig) (core.Provider, error) { return nil, fmt.Errorf("unknown provider type: %s", cfg.Type) } + // One Keyring per provider instance: every client this provider builds + // shares the rotation, so keys are used evenly across all its endpoints. opts := ProviderOptions{ Hooks: hooks, Models: cfg.Models, Resilience: cfg.Resilience, + Keys: NewKeyring(cfg.APIKeys...), } return builder(cfg, opts), nil diff --git a/internal/providers/gemini/gemini.go b/internal/providers/gemini/gemini.go index ac00cff8..313476a3 100644 --- a/internal/providers/gemini/gemini.go +++ b/internal/providers/gemini/gemini.go @@ -48,7 +48,7 @@ type Provider struct { client *llmclient.Client nativeClient *llmclient.Client modelsClient *llmclient.Client - apiKey string + keys *providers.Keyring backend string authType string useNativeAPI bool @@ -75,7 +75,7 @@ func newProvider(providerCfg providers.ProviderConfig, opts providers.ProviderOp baseURL, nativeBaseURL := geminiBaseURLs(providerCfg, backend) modelsURL := geminiModelsBaseURL(backend, nativeBaseURL) p := &Provider{ - apiKey: providerCfg.APIKey, + keys: opts.Keyring(providerCfg.APIKey), backend: backend, authType: authType, useNativeAPI: useNativeAPI(providerCfg.APIMode), @@ -123,7 +123,7 @@ func NewWithHTTPClient(apiKey string, httpClient *http.Client, hooks llmclient.H baseURL, nativeBaseURL := geminiBaseURLs(providerCfg, geminiBackendAIStudio) modelsURL := geminiModelsBaseURL(geminiBackendAIStudio, nativeBaseURL) p := &Provider{ - apiKey: apiKey, + keys: providers.NewKeyring(apiKey), backend: geminiBackendAIStudio, authType: geminiAuthTypeAPIKey, useNativeAPI: useNativeAPIFromEnv(), @@ -225,10 +225,12 @@ func (p *Provider) responseProviderName() string { return "gemini" } -// setHeaders sets the required headers for Gemini API requests +// setHeaders sets the required headers for Gemini API requests. +// Vertex backends authenticate through a token source on the HTTP client +// instead, so only the API-key path consumes the rotation. func (p *Provider) setHeaders(req *http.Request) { if p.authType == geminiAuthTypeAPIKey { - req.Header.Set("Authorization", "Bearer "+p.apiKey) + req.Header.Set("Authorization", "Bearer "+p.keys.Next()) } // Forward request ID if present in context for request tracing @@ -240,7 +242,7 @@ func (p *Provider) setHeaders(req *http.Request) { // setNativeHeaders sets the required headers for Gemini native API requests. func (p *Provider) setNativeHeaders(req *http.Request) { if p.authType == geminiAuthTypeAPIKey { - req.Header.Set("x-goog-api-key", p.apiKey) + req.Header.Set("x-goog-api-key", p.keys.Next()) } if requestID := core.GetRequestID(req.Context()); requestID != "" { diff --git a/internal/providers/gemini/gemini_test.go b/internal/providers/gemini/gemini_test.go index 93c96778..6bdbc861 100644 --- a/internal/providers/gemini/gemini_test.go +++ b/internal/providers/gemini/gemini_test.go @@ -23,8 +23,8 @@ func TestNew(t *testing.T) { // Use NewWithHTTPClient to get concrete type for internal testing provider := NewWithHTTPClient(apiKey, nil, llmclient.Hooks{}) - if provider.apiKey != apiKey { - t.Errorf("apiKey = %q, want %q", provider.apiKey, apiKey) + if got := provider.keys.Primary(); got != apiKey { + t.Errorf("primary key = %q, want %q", got, apiKey) } if provider.modelsURL != defaultModelsBaseURL { t.Errorf("modelsURL = %q, want %q", provider.modelsURL, defaultModelsBaseURL) diff --git a/internal/providers/keyring.go b/internal/providers/keyring.go new file mode 100644 index 00000000..b17e19bf --- /dev/null +++ b/internal/providers/keyring.go @@ -0,0 +1,88 @@ +package providers + +import "sync/atomic" + +// Keyring holds the API keys configured for a single provider instance and +// hands them out one at a time, round robin. +// +// A provider is built once and serves every request, so the credential can no +// longer be a string captured at construction: it is resolved per outbound +// request by calling Next from the provider's header hook. One Keyring is +// shared by all of a provider's HTTP clients, so the rotation is even across +// every endpoint that provider serves. +// +// The zero value is not useful; build one with NewKeyring. A nil *Keyring is +// safe to call and behaves as an empty ring, which lets keyless providers +// (Ollama, vLLM) and direct test constructors skip it entirely. +type Keyring struct { + keys []string + next atomic.Uint64 +} + +// NewKeyring returns a Keyring over keys, preserving order while dropping +// empty and duplicate entries. Duplicates are dropped so that a key repeated +// across `OPENAI_API_KEY` and `OPENAI_API_KEY_1` does not take a double share +// of the rotation. It returns nil when no usable key remains, so callers can +// treat "no credentials" and "no keyring" identically. +func NewKeyring(keys ...string) *Keyring { + unique := make([]string, 0, len(keys)) + seen := make(map[string]struct{}, len(keys)) + for _, key := range keys { + if key == "" { + continue + } + if _, dup := seen[key]; dup { + continue + } + seen[key] = struct{}{} + unique = append(unique, key) + } + if len(unique) == 0 { + return nil + } + return &Keyring{keys: unique} +} + +// Next returns the key to authenticate the next outbound request, advancing +// the rotation. It is safe for concurrent use. Next returns "" for an empty +// ring, matching the unconfigured-credential behaviour providers already +// handle (see AuthHeaderConfig.OptionalAPIKey). +// +// Rotation advances per outbound HTTP request, which includes retries: a +// request retried after a 429 is re-sent under the next key rather than +// hammering the one that was just throttled. +func (k *Keyring) Next() string { + if k == nil || len(k.keys) == 0 { + return "" + } + if len(k.keys) == 1 { + return k.keys[0] + } + i := k.next.Add(1) - 1 + return k.keys[i%uint64(len(k.keys))] +} + +// Primary returns the first configured key without advancing the rotation. +// It is the key to use where a stable identity matters more than spreading +// load, and where an empty ring must stay empty. +func (k *Keyring) Primary() string { + if k == nil || len(k.keys) == 0 { + return "" + } + return k.keys[0] +} + +// Len reports how many distinct keys back the rotation. +func (k *Keyring) Len() int { + if k == nil { + return 0 + } + return len(k.keys) +} + +// Rotates reports whether more than one key is configured, and therefore +// whether successive requests will present different credentials. Callers use +// it to warn about the prompt-caching cost of rotation. +func (k *Keyring) Rotates() bool { + return k.Len() > 1 +} diff --git a/internal/providers/keyring_config_test.go b/internal/providers/keyring_config_test.go new file mode 100644 index 00000000..25d68e45 --- /dev/null +++ b/internal/providers/keyring_config_test.go @@ -0,0 +1,248 @@ +package providers + +import ( + "testing" + + "gomodel/config" + "gomodel/internal/core" +) + +// resolveKeys is a shorthand for the API key set the given provider ends up with. +func resolveKeys(t *testing.T, raw map[string]config.RawProviderConfig, provider string) ProviderConfig { + t.Helper() + got, _ := resolveProviders(raw, globalResilience, testDiscoveryConfigs) + cfg, ok := got[provider] + if !ok { + t.Fatalf("provider %q not resolved; got %v", provider, got) + } + return cfg +} + +func TestResolveProviders_NumberedAPIKeyEnvVars(t *testing.T) { + tests := []struct { + name string + env map[string]string + want []string + }{{ + name: "unsuffixed key plus a numbered one", + env: map[string]string{"OPENAI_API_KEY": "a", "OPENAI_API_KEY_2": "b"}, + want: []string{"a", "b"}, + }, { + name: "only a numbered key still configures the provider", + env: map[string]string{"OPENAI_API_KEY_2": "b"}, + want: []string{"b"}, + }, { + name: "gaps in the numbering are ignored", + env: map[string]string{"OPENAI_API_KEY": "a", "OPENAI_API_KEY_3": "c"}, + want: []string{"a", "c"}, + }, { + name: "every slot spelled out, no unsuffixed key", + env: map[string]string{"OPENAI_API_KEY_1": "a", "OPENAI_API_KEY_2": "b"}, + want: []string{"a", "b"}, + }, { + name: "the unsuffixed key leads the numbered ones", + env: map[string]string{"OPENAI_API_KEY_1": "b", "OPENAI_API_KEY": "a"}, + want: []string{"a", "b"}, + }, { + name: "a key repeated across slots is used once", + env: map[string]string{"OPENAI_API_KEY": "a", "OPENAI_API_KEY_1": "a"}, + want: []string{"a"}, + }, { + // Guards against sorting the indexes as strings, which would order + // slot 10 ahead of slot 2. + name: "slots are ordered numerically, not lexically", + env: map[string]string{"OPENAI_API_KEY": "a", "OPENAI_API_KEY_2": "b", "OPENAI_API_KEY_10": "j"}, + want: []string{"a", "b", "j"}, + }, { + name: "slot zero is not a rotation slot", + env: map[string]string{"OPENAI_API_KEY": "a", "OPENAI_API_KEY_0": "ignored"}, + want: []string{"a"}, + }} + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Setenv("OPENAI_API_KEY", "") + for name, value := range tt.env { + t.Setenv(name, value) + } + + cfg := resolveKeys(t, map[string]config.RawProviderConfig{}, "openai") + + if !equalStrings(cfg.APIKeys, tt.want) { + t.Errorf("APIKeys = %v, want %v", cfg.APIKeys, tt.want) + } + if cfg.APIKey != tt.want[0] { + t.Errorf("APIKey = %q, want the first key %q", cfg.APIKey, tt.want[0]) + } + // A provider configured only through a numbered key still needs its + // default endpoint. + if cfg.BaseURL != "https://api.openai.com/v1" { + t.Errorf("BaseURL = %q, want the OpenAI default", cfg.BaseURL) + } + }) + } +} + +// `OPENAI_EU_API_KEY_2` is key 2 of provider `openai-eu`, while +// `OPENAI_REGION_2_API_KEY` is the only key of provider `openai-region-2`. +// The trailing digits mean different things and must not be conflated. +func TestResolveProviders_NumberedKeysOnSuffixedProviders(t *testing.T) { + t.Run("numbered key on a suffixed provider", func(t *testing.T) { + t.Setenv("OPENAI_EU_API_KEY", "a") + t.Setenv("OPENAI_EU_API_KEY_2", "b") + + cfg := resolveKeys(t, map[string]config.RawProviderConfig{}, "openai-eu") + + if want := []string{"a", "b"}; !equalStrings(cfg.APIKeys, want) { + t.Errorf("APIKeys = %v, want %v", cfg.APIKeys, want) + } + }) + + t.Run("suffix ending in a digit is not a rotation slot", func(t *testing.T) { + t.Setenv("OPENAI_REGION_2_API_KEY", "a") + + cfg := resolveKeys(t, map[string]config.RawProviderConfig{}, "openai-region-2") + + if want := []string{"a"}; !equalStrings(cfg.APIKeys, want) { + t.Errorf("APIKeys = %v, want %v", cfg.APIKeys, want) + } + }) +} + +func TestResolveProviders_APIKeysFromYAML(t *testing.T) { + tests := []struct { + name string + raw config.RawProviderConfig + want []string + }{{ + name: "api_keys list alone", + raw: config.RawProviderConfig{Type: "openai", APIKeys: []string{"a", "b"}}, + want: []string{"a", "b"}, + }, { + name: "api_key leads the api_keys list", + raw: config.RawProviderConfig{Type: "openai", APIKey: "a", APIKeys: []string{"b"}}, + want: []string{"a", "b"}, + }, { + name: "api_key repeated inside api_keys is used once", + raw: config.RawProviderConfig{Type: "openai", APIKey: "a", APIKeys: []string{"a", "b"}}, + want: []string{"a", "b"}, + }, { + name: "unresolved placeholders are dropped, not sent as credentials", + raw: config.RawProviderConfig{Type: "openai", APIKey: "a", APIKeys: []string{"${OPENAI_API_KEY_2}"}}, + want: []string{"a"}, + }} + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Setenv("OPENAI_API_KEY", "") + + cfg := resolveKeys(t, map[string]config.RawProviderConfig{"openai": tt.raw}, "openai") + + if !equalStrings(cfg.APIKeys, tt.want) { + t.Errorf("APIKeys = %v, want %v", cfg.APIKeys, tt.want) + } + if cfg.APIKey != tt.want[0] { + t.Errorf("APIKey = %q, want %q", cfg.APIKey, tt.want[0]) + } + }) + } +} + +// hasAPIKey probes the env fields directly instead of building the ordered key +// list. A value that cannot serve as a credential -- whitespace, or an +// unresolved `${VAR}` -- must read as "no key", not merely "field is set". +func TestProviderEnvValues_HasAPIKeyMatchesAPIKeys(t *testing.T) { + tests := []struct { + name string + v providerEnvValues + want bool + }{ + {name: "no fields set"}, + {name: "blank unsuffixed key", v: providerEnvValues{APIKey: " "}}, + {name: "unresolved unsuffixed key", v: providerEnvValues{APIKey: "${MISSING}"}}, + {name: "blank numbered key", v: providerEnvValues{APIKeysByIndex: map[int]string{2: " "}}}, + {name: "unresolved numbered key", v: providerEnvValues{APIKeysByIndex: map[int]string{2: "${MISSING}"}}}, + {name: "usable unsuffixed key", v: providerEnvValues{APIKey: "a"}, want: true}, + {name: "usable numbered key only", v: providerEnvValues{APIKeysByIndex: map[int]string{2: "b"}}, want: true}, + { + name: "unresolved unsuffixed key beside a usable numbered one", + v: providerEnvValues{APIKey: "${MISSING}", APIKeysByIndex: map[int]string{2: "b"}}, + want: true, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + if got := tt.v.hasAPIKey(); got != tt.want { + t.Errorf("hasAPIKey() = %v, want %v", got, tt.want) + } + // The cheap probe must never disagree with the full key list. + if got, want := tt.v.hasAPIKey(), len(tt.v.apiKeys()) > 0; got != want { + t.Errorf("hasAPIKey() = %v, but len(apiKeys()) > 0 = %v", got, want) + } + }) + } +} + +// A provider whose only key is an unresolved placeholder has no credentials at +// all and must be dropped, exactly as before rotation existed. +func TestResolveProviders_ProviderWithOnlyUnresolvedKeysIsDropped(t *testing.T) { + t.Setenv("OPENAI_API_KEY", "") + + raw := map[string]config.RawProviderConfig{ + "openai": {Type: "openai", APIKey: "${OPENAI_API_KEY}"}, + } + got, _ := resolveProviders(raw, globalResilience, testDiscoveryConfigs) + + if _, ok := got["openai"]; ok { + t.Error("provider with no resolvable key should be dropped") + } +} + +// Env replaces the provider's whole key set rather than merging into it, so a +// key removed from the environment stops being used. +func TestResolveProviders_EnvKeysReplaceYAMLKeySet(t *testing.T) { + t.Setenv("OPENAI_API_KEY", "env-a") + t.Setenv("OPENAI_API_KEY_2", "env-b") + + raw := map[string]config.RawProviderConfig{ + "openai": {Type: "openai", APIKeys: []string{"yaml-a", "yaml-b", "yaml-c"}}, + } + cfg := resolveKeys(t, raw, "openai") + + if want := []string{"env-a", "env-b"}; !equalStrings(cfg.APIKeys, want) { + t.Errorf("APIKeys = %v, want %v", cfg.APIKeys, want) + } +} + +// The factory hands every provider one shared ring, so all the clients a +// provider builds rotate together. +func TestProviderFactory_CreateBuildsKeyring(t *testing.T) { + var got ProviderOptions + factory := NewProviderFactory() + factory.Add(Registration{ + Type: "openai", + New: func(_ ProviderConfig, opts ProviderOptions) core.Provider { + got = opts + return nil + }, + }) + + if _, err := factory.Create(ProviderConfig{Type: "openai", APIKey: "a", APIKeys: []string{"a", "b"}}); err != nil { + t.Fatalf("Create() error = %v", err) + } + + if got.Keys.Len() != 2 { + t.Fatalf("opts.Keys.Len() = %d, want 2", got.Keys.Len()) + } + if !got.Keys.Rotates() { + t.Error("opts.Keys.Rotates() = false, want true") + } + // The provider's own constructor key must not override the factory ring. + if key := got.Keyring("a").Next(); key != "a" { + t.Errorf("first key = %q, want a", key) + } + if key := got.Keyring("a").Next(); key != "b" { + t.Errorf("second key = %q, want b: the ring must be shared, not rebuilt", key) + } +} diff --git a/internal/providers/keyring_test.go b/internal/providers/keyring_test.go new file mode 100644 index 00000000..13f6cfe0 --- /dev/null +++ b/internal/providers/keyring_test.go @@ -0,0 +1,165 @@ +package providers + +import ( + "sync" + "testing" +) + +func TestNewKeyring(t *testing.T) { + tests := []struct { + name string + keys []string + want []string + }{ + {name: "no keys", keys: nil}, + {name: "only empty keys", keys: []string{"", ""}}, + {name: "single key", keys: []string{"k1"}, want: []string{"k1"}}, + {name: "preserves order", keys: []string{"k1", "k2", "k3"}, want: []string{"k1", "k2", "k3"}}, + {name: "drops empty keys", keys: []string{"k1", "", "k2"}, want: []string{"k1", "k2"}}, + {name: "drops duplicates", keys: []string{"k1", "k2", "k1"}, want: []string{"k1", "k2"}}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + ring := NewKeyring(tt.keys...) + + if len(tt.want) == 0 { + if ring != nil { + t.Fatalf("NewKeyring(%q) = %v, want nil for an unusable key set", tt.keys, ring.keys) + } + return + } + if ring.Len() != len(tt.want) { + t.Fatalf("Len() = %d, want %d", ring.Len(), len(tt.want)) + } + // One full cycle reproduces the configured order. + for i, want := range tt.want { + if got := ring.Next(); got != want { + t.Errorf("Next() call %d = %q, want %q", i+1, got, want) + } + } + }) + } +} + +func TestKeyringNextCyclesRoundRobin(t *testing.T) { + ring := NewKeyring("k1", "k2", "k3") + + // Two full cycles: the ring must wrap, not run dry. + want := []string{"k1", "k2", "k3", "k1", "k2", "k3"} + for i, expected := range want { + if got := ring.Next(); got != expected { + t.Errorf("Next() call %d = %q, want %q", i+1, got, expected) + } + } +} + +// A single key must behave exactly as it did before rotation existed: every +// request presents the same credential, so provider prompt caching still works. +func TestKeyringSingleKeyNeverRotates(t *testing.T) { + ring := NewKeyring("only") + + if ring.Rotates() { + t.Error("Rotates() = true, want false for one key") + } + for i := range 3 { + if got := ring.Next(); got != "only" { + t.Errorf("Next() call %d = %q, want %q", i+1, got, "only") + } + } +} + +// Keyless providers (Ollama, vLLM) and constructors invoked outside the factory +// hold a nil ring; every method must stay safe. +func TestKeyringNilIsEmpty(t *testing.T) { + var ring *Keyring + + if got := ring.Next(); got != "" { + t.Errorf("Next() = %q, want empty", got) + } + if got := ring.Primary(); got != "" { + t.Errorf("Primary() = %q, want empty", got) + } + if got := ring.Len(); got != 0 { + t.Errorf("Len() = %d, want 0", got) + } + if ring.Rotates() { + t.Error("Rotates() = true, want false") + } +} + +func TestKeyringPrimaryDoesNotAdvance(t *testing.T) { + ring := NewKeyring("k1", "k2") + + for range 3 { + if got := ring.Primary(); got != "k1" { + t.Fatalf("Primary() = %q, want k1", got) + } + } + if got := ring.Next(); got != "k1" { + t.Errorf("Next() = %q, want k1: Primary must not consume a slot", got) + } +} + +// Providers are shared across concurrent requests, so the rotation must both be +// race-free and hand out each key an equal number of times. +func TestKeyringNextIsConcurrentAndEven(t *testing.T) { + ring := NewKeyring("k1", "k2", "k3") + + const perKey = 200 + total := perKey * ring.Len() + + var mu sync.Mutex + counts := make(map[string]int, ring.Len()) + + var wg sync.WaitGroup + for range total { + wg.Add(1) + go func() { + defer wg.Done() + key := ring.Next() + mu.Lock() + counts[key]++ + mu.Unlock() + }() + } + wg.Wait() + + for _, key := range []string{"k1", "k2", "k3"} { + if counts[key] != perKey { + t.Errorf("key %q used %d times, want %d", key, counts[key], perKey) + } + } +} + +func TestProviderOptionsKeyringFallsBackToStaticKey(t *testing.T) { + // Constructed outside the factory: no ring supplied, so the single + // constructor key is used. + opts := ProviderOptions{} + if got := opts.Keyring("sk-static").Next(); got != "sk-static" { + t.Errorf("Next() = %q, want sk-static", got) + } + + // Built by the factory: the configured ring wins over the primary key that + // the provider constructor happens to pass along. + opts = ProviderOptions{Keys: NewKeyring("k1", "k2")} + ring := opts.Keyring("k1") + if !ring.Rotates() { + t.Fatal("Rotates() = false, want true") + } + if got, want := []string{ring.Next(), ring.Next(), ring.Next()}, []string{"k1", "k2", "k1"}; !equalStrings(got, want) { + t.Errorf("keys = %v, want %v", got, want) + } +} + +func equalStrings(a, b []string) bool { + if len(a) != len(b) { + return false + } + for i := range a { + if a[i] != b[i] { + return false + } + } + return true +} diff --git a/internal/providers/ollama/ollama.go b/internal/providers/ollama/ollama.go index e84d1381..ec63f78f 100644 --- a/internal/providers/ollama/ollama.go +++ b/internal/providers/ollama/ollama.go @@ -43,12 +43,12 @@ const ( type Provider struct { compat *openai.CompatibleProvider nativeClient *llmclient.Client - apiKey string // Accepted but ignored by Ollama + keys *providers.Keyring // Optional; Ollama accepts a bearer token but does not require one } // New creates a new Ollama provider. func New(providerCfg providers.ProviderConfig, opts providers.ProviderOptions) core.Provider { - p := &Provider{apiKey: providerCfg.APIKey} + p := &Provider{keys: opts.Keyring(providerCfg.APIKey)} p.compat = openai.NewCompatibleProvider(providerCfg.APIKey, opts, compatibleConfig(defaultBaseURL)) nativeCfg := llmclient.Config{ @@ -69,7 +69,7 @@ func NewWithHTTPClient(apiKey string, httpClient *http.Client, hooks llmclient.H if httpClient == nil { httpClient = http.DefaultClient } - p := &Provider{apiKey: apiKey} + p := &Provider{keys: providers.NewKeyring(apiKey)} p.compat = openai.NewCompatibleProviderWithHTTPClient(apiKey, httpClient, hooks, compatibleConfig(defaultBaseURL)) nativeCfg := llmclient.DefaultConfig("ollama", defaultNativeBaseURL) @@ -117,7 +117,7 @@ func setHeaders(req *http.Request, apiKey string) { // setNativeHeaders applies the same header policy on the native /api client. func (p *Provider) setNativeHeaders(req *http.Request) { - setHeaders(req, p.apiKey) + setHeaders(req, p.keys.Next()) } // ChatCompletion sends a chat completion request to Ollama diff --git a/internal/providers/ollama/ollama_test.go b/internal/providers/ollama/ollama_test.go index 90b4cdae..3a457092 100644 --- a/internal/providers/ollama/ollama_test.go +++ b/internal/providers/ollama/ollama_test.go @@ -20,8 +20,8 @@ func TestNew(t *testing.T) { // Use NewWithHTTPClient to get concrete type for internal testing provider := NewWithHTTPClient(apiKey, nil, llmclient.Hooks{}) - if provider.apiKey != apiKey { - t.Errorf("apiKey = %q, want %q", provider.apiKey, apiKey) + if got := provider.keys.Primary(); got != apiKey { + t.Errorf("primary key = %q, want %q", got, apiKey) } if provider.compat == nil { t.Error("compat should not be nil") @@ -43,8 +43,8 @@ func TestNew_WithoutAPIKey(t *testing.T) { // Ollama doesn't require an API key provider := NewWithHTTPClient("", nil, llmclient.Hooks{}) - if provider.apiKey != "" { - t.Errorf("apiKey = %q, want empty", provider.apiKey) + if got := provider.keys.Primary(); got != "" { + t.Errorf("primary key = %q, want empty", got) } if provider.compat == nil { t.Error("compat should not be nil") @@ -722,8 +722,8 @@ func TestNewWithHTTPClient(t *testing.T) { provider := NewWithHTTPClient(apiKey, customClient, llmclient.Hooks{}) - if provider.apiKey != apiKey { - t.Errorf("apiKey = %q, want %q", provider.apiKey, apiKey) + if got := provider.keys.Primary(); got != apiKey { + t.Errorf("primary key = %q, want %q", got, apiKey) } if provider.compat == nil { t.Error("compat should not be nil") diff --git a/internal/providers/openai/compatible_provider.go b/internal/providers/openai/compatible_provider.go index 00562c75..ecf2ffeb 100644 --- a/internal/providers/openai/compatible_provider.go +++ b/internal/providers/openai/compatible_provider.go @@ -58,8 +58,11 @@ type CompatibleProviderConfig struct { // AdaptChatRequest, ChatRequestHeaders, RequestMutator), not in copies of // the transport methods. type CompatibleProvider struct { - client *llmclient.Client - apiKey string + client *llmclient.Client + // keys resolves the credential for each outbound request. Providers in this + // package read it directly (see realtime.go) so a websocket dial picks up + // the same rotation as the HTTP endpoints. + keys *providers.Keyring providerName string requestMutator RequestMutator adaptChatRequest func(*core.ChatRequest) (*core.ChatRequest, error) @@ -68,7 +71,7 @@ type CompatibleProvider struct { func NewCompatibleProvider(apiKey string, opts providers.ProviderOptions, cfg CompatibleProviderConfig) *CompatibleProvider { p := &CompatibleProvider{ - apiKey: apiKey, + keys: opts.Keyring(apiKey), providerName: cfg.ProviderName, requestMutator: cfg.RequestMutator, adaptChatRequest: cfg.AdaptChatRequest, @@ -81,9 +84,11 @@ func NewCompatibleProvider(apiKey string, opts providers.ProviderOptions, cfg Co Hooks: opts.Hooks, CircuitBreaker: opts.Resilience.CircuitBreaker, } + // Resolved per request, not captured: with several keys configured this is + // what spreads successive calls across them. p.client = llmclient.New(clientCfg, func(req *http.Request) { if cfg.SetHeaders != nil { - cfg.SetHeaders(req, apiKey) + cfg.SetHeaders(req, p.keys.Next()) } }) return p @@ -94,7 +99,7 @@ func NewCompatibleProviderWithHTTPClient(apiKey string, httpClient *http.Client, httpClient = http.DefaultClient } p := &CompatibleProvider{ - apiKey: apiKey, + keys: providers.NewKeyring(apiKey), providerName: cfg.ProviderName, requestMutator: cfg.RequestMutator, adaptChatRequest: cfg.AdaptChatRequest, @@ -104,7 +109,7 @@ func NewCompatibleProviderWithHTTPClient(apiKey string, httpClient *http.Client, clientCfg.Hooks = hooks p.client = llmclient.NewWithHTTPClient(httpClient, clientCfg, func(req *http.Request) { if cfg.SetHeaders != nil { - cfg.SetHeaders(req, apiKey) + cfg.SetHeaders(req, p.keys.Next()) } }) return p diff --git a/internal/providers/openai/keyrotation_test.go b/internal/providers/openai/keyrotation_test.go new file mode 100644 index 00000000..19e35098 --- /dev/null +++ b/internal/providers/openai/keyrotation_test.go @@ -0,0 +1,157 @@ +package openai + +import ( + "context" + "net/http" + "net/http/httptest" + "sync" + "testing" + "time" + + "gomodel/config" + "gomodel/internal/providers" +) + +// recordAuthServer serves /models and records the Authorization header of every +// request it receives. status, when non-empty, is replayed one status code per +// request so retry behaviour can be exercised. +func recordAuthServer(t *testing.T, statuses ...int) (*httptest.Server, func() []string) { + t.Helper() + + var mu sync.Mutex + var seen []string + + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + mu.Lock() + attempt := len(seen) + seen = append(seen, r.Header.Get("Authorization")) + mu.Unlock() + + if attempt < len(statuses) && statuses[attempt] != http.StatusOK { + w.WriteHeader(statuses[attempt]) + _, _ = w.Write([]byte(`{"error":{"message":"slow down"}}`)) + return + } + w.Header().Set("Content-Type", "application/json") + _, _ = w.Write([]byte(`{"object":"list","data":[]}`)) + })) + t.Cleanup(server.Close) + + return server, func() []string { + mu.Lock() + defer mu.Unlock() + return append([]string(nil), seen...) + } +} + +func rotatingProvider(t *testing.T, baseURL string, retry config.RetryConfig, keys ...string) *CompatibleProvider { + t.Helper() + opts := providers.ProviderOptions{ + Keys: providers.NewKeyring(keys...), + Resilience: config.ResilienceConfig{Retry: retry}, + } + return NewCompatibleProvider(keys[0], opts, CompatibleProviderConfig{ + ProviderName: "rotating", + BaseURL: baseURL, + SetHeaders: bearerHeaders, + }) +} + +// The core promise: with several keys configured, successive calls authenticate +// with different keys, cycling in the configured order. +func TestCompatibleProvider_RotatesKeysAcrossRequests(t *testing.T) { + server, seen := recordAuthServer(t) + provider := rotatingProvider(t, server.URL, config.RetryConfig{}, "k1", "k2", "k3") + + for range 6 { + if _, err := provider.ListModels(context.Background()); err != nil { + t.Fatalf("ListModels() error = %v", err) + } + } + + want := []string{ + "Bearer k1", "Bearer k2", "Bearer k3", + "Bearer k1", "Bearer k2", "Bearer k3", + } + got := seen() + if len(got) != len(want) { + t.Fatalf("got %d requests, want %d", len(got), len(want)) + } + for i := range want { + if got[i] != want[i] { + t.Errorf("request %d Authorization = %q, want %q", i+1, got[i], want[i]) + } + } +} + +// One key must behave exactly as before rotation existed: the same credential +// every time, so upstream prompt caching keeps hitting. +func TestCompatibleProvider_SingleKeyIsStableAcrossRequests(t *testing.T) { + server, seen := recordAuthServer(t) + provider := rotatingProvider(t, server.URL, config.RetryConfig{}, "only") + + for range 3 { + if _, err := provider.ListModels(context.Background()); err != nil { + t.Fatalf("ListModels() error = %v", err) + } + } + + for i, auth := range seen() { + if auth != "Bearer only" { + t.Errorf("request %d Authorization = %q, want %q", i+1, auth, "Bearer only") + } + } +} + +// The header hook runs per HTTP attempt, so a request retried after a 429 is +// re-sent under the next key rather than hammering the throttled one. +func TestCompatibleProvider_RetryUsesNextKey(t *testing.T) { + server, seen := recordAuthServer(t, http.StatusTooManyRequests, http.StatusOK) + retry := config.RetryConfig{ + MaxRetries: 2, + InitialBackoff: time.Millisecond, + MaxBackoff: 2 * time.Millisecond, + BackoffFactor: 1, + } + provider := rotatingProvider(t, server.URL, retry, "k1", "k2") + + if _, err := provider.ListModels(context.Background()); err != nil { + t.Fatalf("ListModels() error = %v", err) + } + + got := seen() + if len(got) != 2 { + t.Fatalf("got %d attempts, want 2 (one 429 then one retry)", len(got)) + } + if got[0] != "Bearer k1" { + t.Errorf("first attempt Authorization = %q, want %q", got[0], "Bearer k1") + } + if got[1] != "Bearer k2" { + t.Errorf("retry Authorization = %q, want %q: the throttled key must not be reused", got[1], "Bearer k2") + } +} + +// Keyless providers must not grow an Authorization header just because the +// rotation machinery is in place. +func TestCompatibleProvider_NoKeysSendsNoCredential(t *testing.T) { + server, seen := recordAuthServer(t) + opts := providers.ProviderOptions{} + provider := NewCompatibleProvider("", opts, CompatibleProviderConfig{ + ProviderName: "keyless", + BaseURL: server.URL, + SetHeaders: func(req *http.Request, apiKey string) { + providers.SetAuthHeaders(req, apiKey, providers.AuthHeaderConfig{ + AuthScheme: "Bearer ", + OptionalAPIKey: true, + }) + }, + }) + + if _, err := provider.ListModels(context.Background()); err != nil { + t.Fatalf("ListModels() error = %v", err) + } + + if auth := seen()[0]; auth != "" { + t.Errorf("Authorization = %q, want no header", auth) + } +} diff --git a/internal/providers/openai/openai.go b/internal/providers/openai/openai.go index c9e1986f..14159fd7 100644 --- a/internal/providers/openai/openai.go +++ b/internal/providers/openai/openai.go @@ -28,12 +28,11 @@ const ( ) // Provider implements the core.Provider interface for OpenAI. -// apiKey is retained so the provider can inject auth on the realtime websocket -// dial target (see realtime.go); the realtime base URL is read live from the -// embedded CompatibleProvider so SetBaseURL overrides are honored. +// Credentials and the realtime base URL are both read live from the embedded +// CompatibleProvider, so SetBaseURL overrides and key rotation are honored on +// the realtime websocket dial target too (see realtime.go). type Provider struct { *CompatibleProvider - apiKey string } // New creates a new OpenAI provider. @@ -45,7 +44,6 @@ func New(cfg providers.ProviderConfig, opts providers.ProviderOptions) core.Prov BaseURL: baseURL, SetHeaders: setHeaders, }), - apiKey: cfg.APIKey, } } @@ -58,7 +56,6 @@ func NewWithHTTPClient(apiKey string, httpClient *http.Client, hooks llmclient.H BaseURL: defaultBaseURL, SetHeaders: setHeaders, }), - apiKey: apiKey, } } diff --git a/internal/providers/openai/openai_test.go b/internal/providers/openai/openai_test.go index 7e59ad35..44a18865 100644 --- a/internal/providers/openai/openai_test.go +++ b/internal/providers/openai/openai_test.go @@ -20,8 +20,8 @@ func TestNew(t *testing.T) { // Use NewWithHTTPClient to get concrete type for internal testing provider := NewWithHTTPClient(apiKey, nil, llmclient.Hooks{}) - if provider.apiKey != apiKey { - t.Errorf("apiKey = %q, want %q", provider.apiKey, apiKey) + if got := provider.keys.Primary(); got != apiKey { + t.Errorf("primary key = %q, want %q", got, apiKey) } if provider.client == nil { t.Error("client should not be nil") diff --git a/internal/providers/openai/realtime.go b/internal/providers/openai/realtime.go index 470dc1b5..628a566e 100644 --- a/internal/providers/openai/realtime.go +++ b/internal/providers/openai/realtime.go @@ -63,10 +63,12 @@ func (p *Provider) realtimeHTTPTarget(req *core.RealtimeRequest, endpoint string return &core.RealtimeHTTPTarget{URL: target, Headers: p.realtimeAuthHeaders()}, nil } +// realtimeAuthHeaders picks the next key in the rotation. A realtime session is +// long-lived, so the key is chosen once per session rather than per event. func (p *Provider) realtimeAuthHeaders() http.Header { headers := http.Header{} - if p.apiKey != "" { - headers.Set("Authorization", "Bearer "+p.apiKey) + if apiKey := p.keys.Next(); apiKey != "" { + headers.Set("Authorization", "Bearer "+apiKey) } return headers } diff --git a/internal/providers/opencodego/opencodego.go b/internal/providers/opencodego/opencodego.go index d00cfac0..f6e98d7b 100644 --- a/internal/providers/opencodego/opencodego.go +++ b/internal/providers/opencodego/opencodego.go @@ -76,7 +76,9 @@ func New(cfg providers.ProviderConfig, opts providers.ProviderOptions) core.Prov ProviderName: "opencode_go", BaseURL: baseURL, }) - messages := anthropic.New(providers.ProviderConfig{APIKey: cfg.APIKey, BaseURL: baseURL}, opts) + // opts carries the shared keyring, so the /messages client rotates in step + // with the chat client above rather than pinning the primary key. + messages := anthropic.New(providers.ProviderConfig{APIKey: cfg.APIKey, APIKeys: cfg.APIKeys, BaseURL: baseURL}, opts) return &Provider{ ChatCompatible: chat, messages: messages, diff --git a/internal/providers/xai/realtime.go b/internal/providers/xai/realtime.go index 27751332..384cf5b4 100644 --- a/internal/providers/xai/realtime.go +++ b/internal/providers/xai/realtime.go @@ -63,8 +63,8 @@ func (p *Provider) realtimeHTTPTarget(req *core.RealtimeRequest, endpoint string func (p *Provider) realtimeAuthHeaders() http.Header { headers := http.Header{} - if p.apiKey != "" { - headers.Set("Authorization", "Bearer "+p.apiKey) + if apiKey := p.keys.Next(); apiKey != "" { + headers.Set("Authorization", "Bearer "+apiKey) } return headers } diff --git a/internal/providers/xai/xai.go b/internal/providers/xai/xai.go index 93acce70..eacea479 100644 --- a/internal/providers/xai/xai.go +++ b/internal/providers/xai/xai.go @@ -43,26 +43,28 @@ type Provider struct { *openai.BatchSurface *openai.FileSurface compat *openai.CompatibleProvider - apiKey string // retained to inject auth on the realtime websocket target + keys *providers.Keyring // retained to inject auth on the realtime websocket target } // New creates a new xAI provider. func New(providerCfg providers.ProviderConfig, opts providers.ProviderOptions) core.Provider { - return newProvider(openai.NewCompatibleProvider(providerCfg.APIKey, opts, compatibleConfig(providers.ResolveBaseURL(providerCfg.BaseURL, defaultBaseURL))), providerCfg.APIKey) + compat := openai.NewCompatibleProvider(providerCfg.APIKey, opts, compatibleConfig(providers.ResolveBaseURL(providerCfg.BaseURL, defaultBaseURL))) + return newProvider(compat, opts.Keyring(providerCfg.APIKey)) } // NewWithHTTPClient creates a new xAI provider with a custom HTTP client. // If httpClient is nil, http.DefaultClient is used. func NewWithHTTPClient(apiKey string, httpClient *http.Client, hooks llmclient.Hooks) *Provider { - return newProvider(openai.NewCompatibleProviderWithHTTPClient(apiKey, httpClient, hooks, compatibleConfig(defaultBaseURL)), apiKey) + compat := openai.NewCompatibleProviderWithHTTPClient(apiKey, httpClient, hooks, compatibleConfig(defaultBaseURL)) + return newProvider(compat, providers.NewKeyring(apiKey)) } -func newProvider(compat *openai.CompatibleProvider, apiKey string) *Provider { +func newProvider(compat *openai.CompatibleProvider, keys *providers.Keyring) *Provider { return &Provider{ BatchSurface: openai.NewBatchSurface(compat), FileSurface: openai.NewFileSurface(compat), compat: compat, - apiKey: apiKey, + keys: keys, } } diff --git a/internal/providers/xai/xai_test.go b/internal/providers/xai/xai_test.go index 2d89c597..05abbcb5 100644 --- a/internal/providers/xai/xai_test.go +++ b/internal/providers/xai/xai_test.go @@ -19,8 +19,8 @@ func TestNew(t *testing.T) { // Use NewWithHTTPClient to get concrete type for internal testing provider := NewWithHTTPClient(apiKey, nil, llmclient.Hooks{}) - if provider.apiKey != apiKey { - t.Errorf("apiKey = %q, want %q", provider.apiKey, apiKey) + if got := provider.keys.Primary(); got != apiKey { + t.Errorf("primary key = %q, want %q", got, apiKey) } if provider.compat == nil { t.Error("compat should not be nil") @@ -93,8 +93,8 @@ func TestNewWithHTTPClient(t *testing.T) { if provider.compat == nil { t.Fatal("provider.compat should not be nil") } - if provider.apiKey != "test-api-key" { - t.Errorf("apiKey = %q, want %q", provider.apiKey, "test-api-key") + if got := provider.keys.Primary(); got != "test-api-key" { + t.Errorf("primary key = %q, want %q", got, "test-api-key") } // Set base URL to our test server diff --git a/internal/providers/zai/realtime.go b/internal/providers/zai/realtime.go index b37b3131..7e727440 100644 --- a/internal/providers/zai/realtime.go +++ b/internal/providers/zai/realtime.go @@ -34,8 +34,8 @@ func (p *Provider) RealtimeTarget(_ context.Context, req *core.RealtimeRequest) } headers := http.Header{} - if p.apiKey != "" { - headers.Set("Authorization", "Bearer "+p.apiKey) + if apiKey := p.keys.Next(); apiKey != "" { + headers.Set("Authorization", "Bearer "+apiKey) } return &core.RealtimeTarget{URL: endpoint, Headers: headers}, nil diff --git a/internal/providers/zai/zai.go b/internal/providers/zai/zai.go index a64bde1d..b33c67e7 100644 --- a/internal/providers/zai/zai.go +++ b/internal/providers/zai/zai.go @@ -22,10 +22,10 @@ var Registration = providers.Registration{ } // Provider implements the core.Provider interface for Z.ai. -// apiKey is retained to inject auth on the GLM-Realtime websocket target. +// keys is retained to inject auth on the GLM-Realtime websocket target. type Provider struct { *openai.ChatCompatible - apiKey string + keys *providers.Keyring } var _ core.Provider = (*Provider)(nil) @@ -37,7 +37,7 @@ func New(cfg providers.ProviderConfig, opts providers.ProviderOptions) core.Prov ProviderName: "zai", BaseURL: providers.ResolveBaseURL(cfg.BaseURL, defaultBaseURL), }), - apiKey: cfg.APIKey, + keys: opts.Keyring(cfg.APIKey), } } @@ -49,6 +49,6 @@ func NewWithHTTPClient(apiKey string, baseURL string, httpClient *http.Client, h ProviderName: "zai", BaseURL: providers.ResolveBaseURL(baseURL, defaultBaseURL), }), - apiKey: apiKey, + keys: providers.NewKeyring(apiKey), } }