From f86490123878c6c9ebf6db3465399b904beb38fd Mon Sep 17 00:00:00 2001 From: "Jakub A. W" Date: Thu, 9 Jul 2026 13:02:01 +0200 Subject: [PATCH 1/2] feat(providers): rotate requests across multiple API keys per provider Configure several API keys for one provider and GoModel spreads outbound requests across them round robin, raising the per-key rate limit before a provider starts returning 429. Keys come from `[_SUFFIX]_API_KEY_` env vars (numbered from 2; `_1` is accepted as a synonym for the unsuffixed key) or `api_keys:` in config.yaml. The trailing number names a key, not a provider: OPENAI_API_KEY_2 is key 2 of `openai`, while OPENAI_REGION_2_API_KEY remains the sole key of provider `openai-region-2`. Every provider already resolved its credential through a per-request header hook, so rotation hooks in there rather than at construction: providers.Keyring replaces the captured API key string and the factory hands each provider instance one shared ring. No provider constructor signature changes, and constructors invoked outside the factory fall back to a single-key ring. Rotation advances per outbound HTTP request, including retries, so a request throttled with 429 retries under the next key. Realtime websocket sessions pick a key per session. Keyless providers (Ollama, vLLM) and providers that authenticate another way (Vertex, Bedrock) are unaffected. User-visible caveat, documented prominently: rotation defeats provider prompt caching, since providers scope the cache to the key that filled it. One configured key behaves exactly as before, so nothing regresses by default. Co-Authored-By: Claude Opus 4.8 (1M context) --- .env.template | 7 + CLAUDE.md | 1 + config/config.example.yaml | 7 + config/providers.go | 8 +- docs/docs.json | 1 + docs/providers/key-rotation.mdx | 130 +++++++++++ docs/providers/overview.mdx | 4 + internal/embedding/embedding.go | 15 +- internal/embedding/embedding_test.go | 4 +- internal/providers/anthropic/anthropic.go | 11 +- .../providers/anthropic/anthropic_test.go | 4 +- internal/providers/azure/azure.go | 7 +- internal/providers/azure/realtime.go | 4 +- internal/providers/bailian/bailian.go | 10 +- internal/providers/bailian/realtime.go | 4 +- internal/providers/config.go | 186 +++++++++++++-- internal/providers/factory.go | 20 ++ internal/providers/gemini/gemini.go | 14 +- internal/providers/gemini/gemini_test.go | 4 +- internal/providers/keyring.go | 88 ++++++++ internal/providers/keyring_config_test.go | 212 ++++++++++++++++++ internal/providers/keyring_test.go | 165 ++++++++++++++ internal/providers/ollama/ollama.go | 8 +- internal/providers/ollama/ollama_test.go | 12 +- .../providers/openai/compatible_provider.go | 17 +- internal/providers/openai/keyrotation_test.go | 157 +++++++++++++ internal/providers/openai/openai.go | 9 +- internal/providers/openai/openai_test.go | 4 +- internal/providers/openai/realtime.go | 6 +- internal/providers/opencodego/opencodego.go | 4 +- internal/providers/xai/realtime.go | 4 +- internal/providers/xai/xai.go | 12 +- internal/providers/xai/xai_test.go | 8 +- internal/providers/zai/realtime.go | 4 +- internal/providers/zai/zai.go | 8 +- 35 files changed, 1063 insertions(+), 96 deletions(-) create mode 100644 docs/providers/key-rotation.mdx create mode 100644 internal/providers/keyring.go create mode 100644 internal/providers/keyring_config_test.go create mode 100644 internal/providers/keyring_test.go create mode 100644 internal/providers/openai/keyrotation_test.go 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..40343d2d 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,44 @@ 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. +func (v providerEnvValues) hasAPIKey() bool { + return len(v.apiKeys()) > 0 +} + 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 +245,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 +253,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 +298,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 +359,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 +473,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 +499,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 +790,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..d8e6f5e0 --- /dev/null +++ b/internal/providers/keyring_config_test.go @@ -0,0 +1,212 @@ +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]) + } + }) + } +} + +// 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), } } From e41e906d3c54868070cd3ec46eee540675afcd2c Mon Sep 17 00:00:00 2001 From: "Jakub A. W" Date: Thu, 9 Jul 2026 13:33:14 +0200 Subject: [PATCH 2/2] perf(providers): make the env-group API key probe allocation-free hasAPIKey answered a yes/no question by building the full ordered key list -- a map, a sort, and two slices -- and then discarding it. empty() asks it for every provider env group. Probe the fields directly instead. The check stays equivalent because it reuses HasResolvedProviderValue, so a blank, whitespace, or unresolved `${VAR}` value still reads as "no key" rather than merely "field is set"; a plain non-empty test on APIKey/APIKeysByIndex would have disagreed on exactly those inputs. A table test pins the two spellings together. Co-Authored-By: Claude Opus 4.8 (1M context) --- internal/providers/config.go | 16 +++++++++- internal/providers/keyring_config_test.go | 36 +++++++++++++++++++++++ 2 files changed, 51 insertions(+), 1 deletion(-) diff --git a/internal/providers/config.go b/internal/providers/config.go index 40343d2d..b15bc279 100644 --- a/internal/providers/config.go +++ b/internal/providers/config.go @@ -202,8 +202,22 @@ func (v providerEnvValues) apiKeys() []string { // 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 { - return len(v.apiKeys()) > 0 + if HasResolvedProviderValue(v.APIKey) { + return true + } + for _, key := range v.APIKeysByIndex { + if HasResolvedProviderValue(key) { + return true + } + } + return false } func (v providerEnvValues) empty() bool { diff --git a/internal/providers/keyring_config_test.go b/internal/providers/keyring_config_test.go index d8e6f5e0..25d68e45 100644 --- a/internal/providers/keyring_config_test.go +++ b/internal/providers/keyring_config_test.go @@ -148,6 +148,42 @@ func TestResolveProviders_APIKeysFromYAML(t *testing.T) { } } +// 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) {