diff --git a/Cargo.lock b/Cargo.lock index 67c596e..38a429f 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1341,7 +1341,7 @@ checksum = "b6d2cec3eae94f9f509c767b45932f1ada8350c4bdb85af2fcab4a3c14807981" [[package]] name = "link-assistant-router" -version = "0.20.0" +version = "0.21.0" dependencies = [ "aes-gcm", "async-trait", diff --git a/README.md b/README.md index 9f9771d..65d0bae 100644 --- a/README.md +++ b/README.md @@ -20,7 +20,7 @@ Link.Assistant.Router is a transparent proxy that sits between API clients (such - **Optional Gonka upstream** — `UPSTREAM_PROVIDER=gonka` forwards OpenAI-compatible routes to Gonka instead of translating them to Anthropic - **Optional Crater ForgeFed upstream** — `UPSTREAM_PROVIDER=crater` turns OpenAI chat requests into ForgeFed `Offer{Ticket}` tasks and waits for resolved task results - **Optional LiteLLM/OpenAI-compatible upstream** — `UPSTREAM_PROVIDER=openai-compatible` routes OpenAI SDK traffic to a stored provider such as LiteLLM -- **Multi-account routing** — pool any number of Claude MAX accounts; round-robin / priority / least-used; automatic cooldowns on 429 +- **Multi-account routing** — pool any number of Claude, Codex, Gemini, or Qwen subscriptions; session affinity, strict token pins, round-robin / fill-first / least-used selection, request caps, and `Retry-After`-aware cooldowns - **Issues custom `la_sk_...` JWT tokens** with expiration and revocation for multi-tenant access - **Persistent token store** — text (Lino) **and** binary backends, both on by default; tokens survive restarts - **Live observability** — Prometheus `/metrics`, JSON `/v1/usage`, per-account health at `/v1/accounts` @@ -83,6 +83,15 @@ Gemini uses the Code Assist envelope with synthesized SSE for streaming; Qwen is OpenAI-compatible). Run `router doctor` to verify each credential file is present and its token valid. +To pool subscriptions, set `ADDITIONAL_ACCOUNT_DIRS` to vendor-specific +credential homes. The active `UPSTREAM_PROVIDER` determines how every directory +is parsed. New sessions use `ACCOUNT_ROUTING_STRATEGY`; a session remains on its +chosen account for `SESSION_AFFINITY_TTL_SECS`, and an `la_sk_...` token issued +with an `account` claim is a strict pin. Automatic selection skips accounts +whose configured `ACCOUNT_REQUEST_LIMITS` cap is spent or whose upstream +returned HTTP 429. Pinned and session-bound requests fail instead of silently +changing identity. + ## Quick Start ### Prerequisites @@ -281,6 +290,8 @@ Claude Code will work exactly as normal, with all requests transparently proxied |---|---|---| | `/v1/messages` | POST | Anthropic Messages — preserves SSE streaming | | `/v1/messages/count_tokens` | POST | Token-count helper | +| `/api/anthropic/v1/messages` | POST | Namespaced Anthropic Messages alias | +| `/api/anthropic/v1/messages/count_tokens` | POST | Namespaced token-count alias | | `/invoke` | POST | Bedrock-format invoke | | `/invoke-with-response-stream` | POST | Bedrock streaming invoke | | `/api/latest/anthropic/*` | ANY | Legacy prefix; stripped and forwarded | @@ -293,6 +304,19 @@ Claude Code will work exactly as normal, with all requests transparently proxied | `/v1/chat/completions` | POST | Chat Completions, translated to Anthropic Messages, forwarded to the selected OpenAI-compatible provider, or delivered as a Crater ForgeFed task | | `/v1/responses` | POST | Responses API, translated to Anthropic Messages or forwarded to the selected OpenAI-compatible provider | | `/v1/models` | GET | OpenAI-shaped model list | +| `/api/openai/v1/*` | GET/POST | Namespaced aliases for models, Chat Completions, and Responses | +| `/api/codex/v1/*` | GET/POST | Codex namespace; Responses is the subscription's native protocol | +| `/api/qwen/v1/*` | GET/POST | Qwen namespace; forwards its native OpenAI-compatible protocol | +| `/api/gemini/v1beta/models` | GET | Native Gemini model list | +| `/api/gemini/v1beta/models/{model}` | GET | Native Gemini model metadata | +| `/api/gemini/v1beta/models/{model}:generateContent` | POST | Native Gemini generation | +| `/api/gemini/v1beta/models/{model}:streamGenerateContent` | POST | Native Gemini SSE response | +| `/api/vertex/v1/projects/.../models/{model}:generateContent` | POST | Native Vertex-style generation through Gemini Code Assist | + +Provider-specific namespaces use the matching active subscription configured +by `UPSTREAM_PROVIDER`; for example, native Gemini and Vertex generation +requires `UPSTREAM_PROVIDER=gemini`. They are additive client-facing protocol +aliases, not cross-provider fallback rules. `gpt-4o`, `gpt-4o-mini`, `gpt-4`, and the `o*` reasoning families auto-map to the Claude Sonnet / Haiku / Opus tiers respectively. Native `claude-*` IDs pass through unchanged. @@ -337,7 +361,7 @@ method-specific verifier is configured. |---|---|---| | `/metrics` | GET | Prometheus text-exposition counters | | `/v1/usage` | GET | JSON snapshot of all counters | -| `/v1/accounts` | GET | Multi-account health: cooldowns, last error, used-count | +| `/v1/accounts` | GET | Multi-account health: cooldowns, last error, used count, configured limit, and remaining requests | ### POST /api/tokens @@ -422,7 +446,7 @@ Every flag listed in `--help` has an env-var alias and can be configured from | `--port` / `ROUTER_PORT` | `8080` | No | Port to listen on | | `--host` / `ROUTER_HOST` | `0.0.0.0` | No | Host/IP to bind to | | `--claude-code-home` / `CLAUDE_CODE_HOME` | `~/.claude` | No | Primary Claude Code credentials directory | -| `--upstream-provider` / `UPSTREAM_PROVIDER` | `anthropic` | No | Upstream provider: `anthropic`, `gonka`, `crater`, or `openai-compatible` | +| `--upstream-provider` / `UPSTREAM_PROVIDER` | `anthropic` | No | Upstream provider: `anthropic`, `codex`, `gemini`, `qwen`, `gonka`, `crater`, or `openai-compatible` | | `--upstream-base-url` / `UPSTREAM_BASE_URL` | `https://api.anthropic.com` | No | Upstream Anthropic API URL | | `--api-format` / `UPSTREAM_API_FORMAT` | (auto) | No | Restrict the proxy to `anthropic` / `bedrock` / `vertex` | | `--verbose` / `VERBOSE` | `false` | No | Verbose tracing | @@ -550,7 +574,11 @@ The HTTP API accepts the same shape at `POST /api/providers`: | `--storage-policy` / `STORAGE_POLICY` | `both` | Persistent token store: `memory`, `text` (Lino), `binary`, or `both` | | `--data-dir` / `DATA_DIR` | platform-specific | Where `tokens.lino` / `tokens.bin` live | | `--claude-cli-bin` / `CLAUDE_CLI_BIN` | `claude` | Local Claude CLI binary used by the `cli` backend | -| `--additional-account-dirs` / `ADDITIONAL_ACCOUNT_DIRS` | (empty) | Comma-separated extra credential dirs for multi-account routing | +| `--additional-account-dirs` / `ADDITIONAL_ACCOUNT_DIRS` | (empty) | Comma-separated extra credential homes for the active subscription provider | +| `--account-routing-strategy` / `ACCOUNT_ROUTING_STRATEGY` | `round-robin` | New-session policy: `round-robin`, `priority`/`fill-first`, or `least-used`/`quota-first` | +| `--account-cooldown-secs` / `ACCOUNT_COOLDOWN_SECS` | `60` | Minimum cooldown after a quota response; a longer upstream `Retry-After` wins | +| `--session-affinity-ttl-secs` / `SESSION_AFFINITY_TTL_SECS` | `3600` | Inactive seconds before a conversation can be assigned again; `0` disables affinity | +| `--account-request-limits` / `ACCOUNT_REQUEST_LIMITS` | (unknown) | Comma-separated request caps, primary first then extras; must match pool size, and `0` means unknown/unlimited | ### Feature toggles @@ -777,9 +805,9 @@ curl -s -X POST http://localhost:8080/api/tokens \ cargo test ``` -This runs: -- **Unit tests** in every module under `src/` (44 tests covering config, oauth, token, storage, accounts, openai, metrics, cli) -- **Integration tests** in `tests/integration_test.rs` cover API path routing, OpenAI translation, metrics rendering, and CLI parsing +This runs the unit, integration, release-workflow, and documentation tests, +including account affinity/caps, provider-scoped token caching, request-routing +metadata, protocol translation, metrics, and configuration validation. ### Run specific test suites @@ -891,9 +919,11 @@ This demonstrates token issuance, validation, and revocation programmatically. │ ├── crater.rs # Crater ForgeFed task provider │ ├── oauth.rs # Claude Code OAuth credential reader │ ├── accounts.rs # Multi-account router (round-robin/priority/least-used + cooldowns) +│ ├── app_state.rs # Shared HTTP handler state │ ├── storage.rs # Persistent token store (text Lino + binary backends) │ ├── providers.rs # OpenAI-compatible provider store + encrypted secrets │ ├── proxy.rs # Transparent API proxy with token swap, OpenAI shim, ops endpoints +│ ├── request_routing.rs # Session/account routing signal extraction │ ├── openai.rs # OpenAI <-> Anthropic translation helpers │ ├── metrics.rs # Atomic counters, Prometheus rendering, JSON snapshots │ └── token.rs # Custom JWT token management (la_sk_...) diff --git a/changelog.d/20260722_070000_issue_42_multi_subscription_routing.md b/changelog.d/20260722_070000_issue_42_multi_subscription_routing.md new file mode 100644 index 0000000..ab08b73 --- /dev/null +++ b/changelog.d/20260722_070000_issue_42_multi_subscription_routing.md @@ -0,0 +1,20 @@ +--- +bump: minor +--- + +### Added + +- Provider-neutral multi-subscription pools for Claude, Codex, Gemini, and Qwen + with strict token pins, session affinity, round-robin/fill-first/least-used + selection, configurable per-account request caps, and `Retry-After`-aware + quota cooldowns. +- Formal AI-style namespaced protocol routes for Anthropic, OpenAI, Codex, + Qwen, native Gemini `generateContent`, and Vertex publisher-model requests. +- Issue #42 research and requirement trace under + `docs/case-studies/issue-42`. + +### Fixed + +- Subscription requests now enforce router-token request budgets, keep + refreshed credentials isolated per account, and preserve original request + metadata when selecting an account before protocol translation. diff --git a/docs/case-studies/issue-42/README.md b/docs/case-studies/issue-42/README.md new file mode 100644 index 0000000..10b9b4c --- /dev/null +++ b/docs/case-studies/issue-42/README.md @@ -0,0 +1,126 @@ +# Issue 42: Multi-subscription routing + +## Scope + +Issue [#42](https://github.com/link-assistant/router/issues/42) asks the router +to adopt Claudexor's multi-subscription practices, support both single- and +multi-account deployments for every existing vendor subscription, minimize +spent capacity when selecting new work, preserve sessions, and expose the +provider-native API namespaces used by Formal AI. + +The implementation extends the existing `SubscriptionProvider` abstraction; +it does not create a second vendor registry. Claude, Codex, Gemini, and Qwen +therefore share one account-selection policy while retaining their existing +credential readers, refresh flows, request translations, and upstreams. + +## Root cause + +PR #38 added four subscription readers, but the account pool was still built +from Claude-only `OAuthProvider` instances. Codex, Gemini, and Qwen bypassed the +pool, refreshed tokens shared a cache slot per provider instead of per account, +and routing happened before the proxy copied stable session metadata. The +result was four usable single subscriptions but only one usable multi-account +provider, with no identity-stability contract. + +The public HTTP surface also exposed generic Anthropic/OpenAI paths only. A +Gemini client could reach Code Assist only after translating through the OpenAI +shape, even though Formal AI demonstrates that native protocol namespaces are +useful integration boundaries. + +## Adopted design + +### One provider-neutral pool + +`AccountRouter::new_for_provider` builds an ordered pool of +`SubscriptionReader`s for the active provider. A pool is optional: the existing +single-account path remains the default when no extra directories or account +caps are configured. + +New work supports three policies: + +- `round-robin` distributes requests across available accounts; +- `priority` / `fill-first` keeps using the first available account; +- `least-used` / `quota-first` compares normalized configured consumption + (`used / limit`) and chooses the lowest spent account. Unknown limits remain + eligible fallback capacity instead of being presented as zero usage. + +`ACCOUNT_REQUEST_LIMITS` supplies the router's internal per-account limits, +ordered primary then additional accounts. A value of zero explicitly means +unknown/unlimited. Counters are atomic, and the cap check plus increment is a +compare-and-swap operation so concurrent requests cannot oversubscribe a cap. + +### Stable and strict identity + +The router copies routing signals before request translation: + +1. a router-issued token's `account` claim is an explicit strict pin; +2. session headers (`x-claude-code-session-id`, `x-codex-session-id`, + `x-session-id`, `session-id`) win over body metadata; +3. JSON session/conversation metadata and OpenAI `prompt_cache_key` are + accepted as fallback session keys; +4. requests without either signal enter the configured automatic policy. + +Session bindings expire after `SESSION_AFFINITY_TTL_SECS` of inactivity. A +pinned or bound account that becomes spent/unavailable fails selection; the +router does not move an in-flight conversation to a different identity. + +### Quota failure behavior + +Only an upstream HTTP 429 starts account cooldown. `Retry-After` is parsed in +both delta-seconds and HTTP-date forms, and the longer of that value and +`ACCOUNT_COOLDOWN_SECS` wins. Concurrent failures can extend but never shorten +an existing cooldown. Network, parsing, model, and request-shape errors do not +rotate an account. + +The account-scoped refresh cache key is `(provider, account name)`. This avoids +returning account A's refreshed credential after account B has been selected. +Per-token request budgets are now enforced on subscription paths as well as the +Anthropic proxy path. + +### Native protocol namespaces + +The unprefixed compatibility routes remain stable. The following additive +namespaces mirror Formal AI's protocol organization: + +| Namespace | Routes | Native role | +| --- | --- | --- | +| `/api/anthropic/v1` | Messages, count tokens | Anthropic subscription protocol | +| `/api/openai/v1` | models, Chat Completions, Responses | Generic OpenAI compatibility | +| `/api/codex/v1` | models, Chat Completions, Responses | Codex's native Responses upstream plus compatibility projection | +| `/api/qwen/v1` | models, Chat Completions, Responses | Qwen/DashScope's native OpenAI-compatible upstream | +| `/api/gemini/v1beta` | models, model metadata, `generateContent`, `streamGenerateContent` | Native Gemini request/response bodies | +| `/api/vertex/v1` | publisher-model generation paths | Native Vertex-style request paths over Gemini Code Assist | + +Each provider-specific namespace uses the matching configured +`UPSTREAM_PROVIDER`; namespaces do not introduce cross-provider fallback. + +Native Gemini streaming currently returns the complete Code Assist response as +one valid SSE data event. This preserves the native response shape and stream +transport without claiming token-by-token upstream streaming. + +## Verification + +The regression suite covers: + +- same-session account affinity and strict no-fallback behavior; +- strict explicit account claims and unknown pins; +- all three automatic policies and normalized least-used selection; +- atomic configured limits, spent-account exclusion, cooldown extension, and + `Retry-After` parsing; +- Codex pool selection and per-account refresh-cache isolation; +- token account lookup, request metadata extraction, and subscription budget + enforcement; +- Gemini/Vertex native path parsing and response-envelope handling; +- configuration validation, including exact cap-to-pool cardinality. + +Before implementation, compiling the new regression tests produced 26 missing +API/type errors. After implementation, the complete `cargo test --all-features` +suite passes, and a real `cargo run -- serve` startup confirms that the Axum +route table has no conflicting paths. + +## Files in this case study + +- [requirements.md](requirements.md) maps every issue requirement to code and + verification. +- [online-research.md](online-research.md) records repositories, immutable + snapshots, inspected files, findings, and design decisions. diff --git a/docs/case-studies/issue-42/online-research.md b/docs/case-studies/issue-42/online-research.md new file mode 100644 index 0000000..5728d5e --- /dev/null +++ b/docs/case-studies/issue-42/online-research.md @@ -0,0 +1,73 @@ +# Online research and source inventory + +Research was performed on 2026-07-22. Repository analyses use immutable commit +links so later upstream changes do not rewrite the evidence. + +## Primary requested references + +### Claudexor + +- Repository: [razzant/claudexor](https://github.com/razzant/claudexor) +- Snapshot: [`2b0a5e17c6bc1298c43f360da32630042170c28e`](https://github.com/razzant/claudexor/tree/2b0a5e17c6bc1298c43f360da32630042170c28e) +- Inspected sources: + [architecture](https://github.com/razzant/claudexor/blob/2b0a5e17c6bc1298c43f360da32630042170c28e/docs/ARCHITECTURE.md), + [credential profile policy](https://github.com/razzant/claudexor/blob/2b0a5e17c6bc1298c43f360da32630042170c28e/packages/orchestrator/src/credential-profiles.ts), + [orchestrator integration](https://github.com/razzant/claudexor/blob/2b0a5e17c6bc1298c43f360da32630042170c28e/packages/orchestrator/src/orchestrator.ts), + and their adjacent tests/schema. + +Findings adopted: + +- route precedence is explicit pin, then sticky conversation identity, then + automatic pool selection; +- explicit/sticky identities are strict and do not silently fall back; +- rotation is driven by typed quota evidence rather than arbitrary failures; +- eligible targets are ordered, enabled, unspent, and the same credential kind; +- unknown quota stays unknown and eligible instead of appearing unused; +- the credential identity follows the request/session through observability. + +Claudexor has richer live vendor quota snapshots and preflight policies than +this proxy currently has. The portable subset implemented here is explicit +caps plus normalized request consumption, 429 cooldowns, and strict identity. + +### Formal AI + +- Repository: [link-assistant/formal-ai](https://github.com/link-assistant/formal-ai) +- Snapshot: [`20ed7700656e32f7d6285b98c999ba7ad0c5342f`](https://github.com/link-assistant/formal-ai/tree/20ed7700656e32f7d6285b98c999ba7ad0c5342f) +- Inspected sources: + [server protocol dispatcher](https://github.com/link-assistant/formal-ai/blob/20ed7700656e32f7d6285b98c999ba7ad0c5342f/src/server.rs), + [multi-protocol HTTP test](https://github.com/link-assistant/formal-ai/blob/20ed7700656e32f7d6285b98c999ba7ad0c5342f/tests/integration/multi_protocol_api.rs), + and [server API documentation](https://github.com/link-assistant/formal-ai/blob/20ed7700656e32f7d6285b98c999ba7ad0c5342f/docs/configuration/server-api.md). + +Findings adopted: + +- protocol-specific namespaces make client configuration explicit and avoid + collisions: `/api/openai/v1`, `/api/anthropic/v1`, `/api/gemini/v1beta`, and + `/api/vertex/v1`; +- native model discovery and native generation should be tested alongside + compatibility APIs; +- Gemini `generateContent` and `streamGenerateContent` keep Gemini request and + response shapes instead of forcing an OpenAI projection. + +## Related implementations evaluated + +| Project | Relevant behavior | Decision | +| --- | --- | --- | +| [hjanuschka/pi-multi-pass](https://github.com/hjanuschka/pi-multi-pass) | Quota-first identity selection for Codex/Gemini and runtime failover | Adopt the least-spent policy name/intent; retain provider-local pools | +| [Neurolink Claude proxy](https://github.com/juspay/neurolink/blob/release/docs/features/claude-proxy.md) | Fill-first identity stability, cooldowns using `Retry-After`, and no failover for request/model errors | Adopt fill-first alias, typed 429 cooldown, and strict non-quota behavior | +| [raine/claude-code-proxy](https://github.com/raine/claude-code-proxy) | Small transparent Claude OAuth proxy with compatibility headers | Preserve transparent token substitution and existing Anthropic behavior | +| [kittors/CliRelay](https://github.com/kittors/CliRelay) | Multi-provider CLI credential relay and OpenAI-compatible surface | Confirms shared provider abstraction; avoid importing its broader relay scope | + +## Decision synthesis + +The references converge on four rules that fit this repository: + +1. select identity once from stable request/session data; +2. distinguish explicit identity from automatic selection; +3. use quota/rate-limit evidence for account availability, not generic errors; +4. keep each wire protocol accessible under an unambiguous namespace. + +The implementation intentionally leaves live vendor quota polling for a future +change. That feature needs provider-specific freshness/provenance contracts; +guessing it from nonstandard headers would violate the research's central rule +that unknown quota must remain unknown. + diff --git a/docs/case-studies/issue-42/requirements.md b/docs/case-studies/issue-42/requirements.md new file mode 100644 index 0000000..6a60e81 --- /dev/null +++ b/docs/case-studies/issue-42/requirements.md @@ -0,0 +1,33 @@ +# Issue 42 requirement trace + +| ID | Requirement | Implementation | Verification | +| --- | --- | --- | --- | +| R1 | Research existing implementations online | Pinned Claudexor/Formal AI source review plus four related routing projects in `online-research.md` | Snapshot SHAs and source links are recorded | +| R2 | Collect research under `docs/case-studies/issue-42` | This case-study package | Documentation is committed with the implementation | +| R3 | Support one subscription | Existing no-pool path remains unchanged; a one-account router is synthesized only for inspection | Full existing unit/integration suite | +| R4 | Support multiple subscriptions for every supported vendor | Provider-neutral `AccountRouter` backed by `SubscriptionReader` | Claude and Codex pool tests; shared reader/parser tests cover all providers | +| R5 | Switch automatically when an account is unavailable or spent | Automatic strategies skip capped/cooling/unreadable accounts | cap, cooldown, failover, and no-healthy-account tests | +| R6 | Identify internal limits and prefer minimum spent capacity | `ACCOUNT_REQUEST_LIMITS`; atomic usage; normalized `least-used` / `quota-first` policy | uneven-cap least-used and cap-exhaustion tests | +| R7 | Keep sessions stable | copied routing context plus TTL affinity map | repeated session selects one account | +| R8 | Do not silently change explicit/session identity | strict pin and strict bound-session selection modes | unavailable pin/session tests | +| R9 | Copy request data used for routing before translation | headers plus original JSON metadata passed separately through every translated proxy | header/body precedence and prompt-cache-key tests | +| R10 | React only to typed limit failures | 429-only cooldown, standard `Retry-After`, extend-only timers | delta/date parser and cooldown tests | +| R11 | Isolate account credentials and refresh state | cache key includes provider and stable account name | provider/account cache-isolation test | +| R12 | Enforce caller-token budgets on subscription routes | common budget enforcement before vendor selection | token budget regression coverage | +| R13 | Add Formal AI-style API namespaces | additive Anthropic, OpenAI, Codex, Qwen, Gemini, and Vertex routes | path/helper tests plus real server startup | +| R14 | Support native Gemini and Qwen APIs | native Gemini/Vertex request shapes; Qwen's native OpenAI-compatible body under its own namespace | Gemini parser/envelope tests and existing Qwen forwarding tests | +| R15 | Preserve observability | selected account on request metrics; health includes limit and remaining requests | health snapshot and metrics tests | +| R16 | Expose configuration consistently | CLI/env fields, validation, doctor output, README | config and CLI conversion tests | +| R17 | Keep code maintainable | state and request-routing extraction keep all Rust files under 1,000 lines | repository file-size check | + +## Deliberate boundaries + +- The router does not guess undocumented vendor quota from error text. Operators + can configure known request caps; unknown capacity remains explicitly + unknown. HTTP 429 and `Retry-After` are the runtime authority. +- Cross-provider fallback is not introduced. A configured pool contains one + credential kind/provider, avoiding an invisible billing or protocol change. +- Credentials never change during an upstream request. A later unpinned request + may select a recovered account; pinned/session work remains strict. +- Gemini native streaming is transport-compatible synthesized SSE over the + Code Assist response, not a claim of token-level upstream streaming. diff --git a/src/accounts.rs b/src/accounts.rs index 78640eb..8810113 100644 --- a/src/accounts.rs +++ b/src/accounts.rs @@ -1,22 +1,19 @@ //! Multi-account OAuth credential routing. //! -//! Issue #7 R7 / R10 require the router to support multiple Claude MAX -//! accounts: one primary + an optional list of additional credential -//! directories. Inbound requests are dispatched to a healthy account using -//! a configurable selection strategy (round-robin by default), with -//! cooldowns and quota windows that automatically remove an account from -//! the rotation when it returns `429/insufficient_quota`. -//! -//! The single-account `OAuthProvider` is reused as the building block for -//! each account so existing tests and behaviour are preserved when only a -//! primary directory is configured. +//! A pool contains one primary subscription plus optional additional +//! credential directories for Claude, Codex, Gemini, or Qwen. New sessions +//! use a configurable selection strategy; existing sessions stay on their +//! selected account. Typed quota failures and configured request caps remove +//! accounts from automatic selection without silently moving pinned work. +use std::cmp::Ordering as CmpOrdering; +use std::collections::HashMap; use std::path::PathBuf; use std::sync::atomic::{AtomicUsize, Ordering}; use std::sync::{Arc, Mutex}; use std::time::{Duration, Instant}; -use crate::oauth::{OAuthError, OAuthProvider}; +use crate::subscription::{SubscriptionProvider, SubscriptionReader, SubscriptionToken}; /// Strategy used to pick the next account on each request. #[derive(Debug, Clone, Copy, PartialEq, Eq, Default)] @@ -33,21 +30,78 @@ pub enum SelectionStrategy { impl SelectionStrategy { #[must_use] pub fn from_str_opt(s: &str) -> Option { - match s.to_lowercase().as_str() { + match s.trim().to_lowercase().as_str() { "round-robin" | "roundrobin" | "rr" => Some(Self::RoundRobin), - "priority" | "prio" => Some(Self::Priority), - "least-used" | "leastused" | "lru" => Some(Self::LeastUsed), + "priority" | "prio" | "fill-first" | "fillfirst" => Some(Self::Priority), + "least-used" | "leastused" | "least-utilized" | "quota-first" | "lru" => { + Some(Self::LeastUsed) + } _ => None, } } } +/// Tunable multi-account behavior. +#[derive(Debug, Clone)] +pub struct AccountRouterOptions { + /// Account selection policy for new sessions. + pub strategy: SelectionStrategy, + /// Default cooldown after an upstream quota failure. + pub cooldown: Duration, + /// How long an inactive session remains bound to its account. Zero disables + /// session affinity. + pub session_affinity_ttl: Duration, + /// Optional request cap for each account, ordered primary then additional. + pub request_limits: Vec>, +} + +impl Default for AccountRouterOptions { + fn default() -> Self { + Self { + strategy: SelectionStrategy::default(), + cooldown: Duration::from_secs(60), + session_affinity_ttl: Duration::from_secs(60 * 60), + request_limits: Vec::new(), + } + } +} + +/// Stable routing signals copied from an inbound request. +#[derive(Debug, Clone, Default, PartialEq, Eq)] +pub struct RoutingContext { + /// Conversation/session identifier detected from headers or JSON metadata. + pub session_key: Option, + /// Explicit account selected by the router-issued caller token. + pub pinned_account: Option, +} + +impl RoutingContext { + /// Build a context containing only a session binding. + #[must_use] + pub fn for_session(session: impl Into) -> Self { + Self { + session_key: Some(session.into()), + pinned_account: None, + } + } + + /// Build a context containing an explicit, strict account pin. + #[must_use] + pub fn pinned(account: impl Into) -> Self { + Self { + session_key: None, + pinned_account: Some(account.into()), + } + } +} + /// Per-account runtime state (cooldowns, request counts, last error). struct AccountState { name: String, - provider: OAuthProvider, + reader: SubscriptionReader, home: PathBuf, used: AtomicUsize, + request_limit: Option, cooldown_until: Mutex>, last_error: Mutex>, } @@ -60,12 +114,43 @@ impl AccountState { .unwrap_or_else(std::sync::PoisonError::into_inner); !matches!(*guard, Some(t) if t > Instant::now()) } + + fn is_available(&self) -> bool { + self.is_healthy() + && self + .request_limit + .is_none_or(|limit| self.used.load(Ordering::Relaxed) < limit) + } + + fn try_record_use(&self) -> bool { + let mut used = self.used.load(Ordering::Relaxed); + loop { + if self.request_limit.is_some_and(|limit| used >= limit) { + return false; + } + match self.used.compare_exchange_weak( + used, + used.saturating_add(1), + Ordering::Relaxed, + Ordering::Relaxed, + ) { + Ok(_) => return true, + Err(actual) => used = actual, + } + } + } +} + +#[derive(Debug, Clone)] +struct AffinityBinding { + account_index: usize, + expires_at: Instant, } /// Multi-account router. /// -/// Holds an ordered list of [`OAuthProvider`]s and dispatches requests -/// using the configured selection strategy. Cheap to clone (Arc-wrapped). +/// Holds an ordered list of vendor subscription readers and dispatches +/// requests using the configured selection strategy. Cheap to clone. #[derive(Clone)] pub struct AccountRouter { inner: Arc, @@ -74,8 +159,11 @@ pub struct AccountRouter { struct AccountRouterInner { accounts: Vec, cursor: AtomicUsize, + provider: SubscriptionProvider, strategy: SelectionStrategy, cooldown: Duration, + session_affinity_ttl: Duration, + affinities: Mutex>, } /// Information returned to the caller for use in upstream calls. @@ -85,6 +173,20 @@ pub struct SelectedAccount { pub token: String, } +/// A normalized vendor subscription token and its selected account. +#[derive(Debug, Clone)] +pub struct SelectedSubscriptionAccount { + pub name: String, + pub token: SubscriptionToken, +} + +#[derive(Debug, Clone, Copy)] +enum SelectionMode { + Automatic, + Pinned, + Session, +} + impl AccountRouter { /// Build a new router with one primary account and any additional /// account directories. @@ -95,21 +197,50 @@ impl AccountRouter { strategy: SelectionStrategy, cooldown: Duration, ) -> Self { + Self::new_for_provider( + primary, + additional, + SubscriptionProvider::Claude, + AccountRouterOptions { + strategy, + cooldown, + ..AccountRouterOptions::default() + }, + ) + } + + /// Build a router for any supported vendor subscription. + #[must_use] + pub fn new_for_provider( + primary: PathBuf, + additional: &[PathBuf], + provider: SubscriptionProvider, + options: AccountRouterOptions, + ) -> Self { + let AccountRouterOptions { + strategy, + cooldown, + session_affinity_ttl, + request_limits, + } = options; let mut accounts = Vec::with_capacity(1 + additional.len()); + let request_limit = |index: usize| request_limits.get(index).copied().flatten(); accounts.push(AccountState { name: "primary".to_string(), - provider: OAuthProvider::new(primary.to_string_lossy().as_ref()), + reader: SubscriptionReader::new(provider, &primary), home: primary, used: AtomicUsize::new(0), + request_limit: request_limit(0), cooldown_until: Mutex::new(None), last_error: Mutex::new(None), }); for (i, p) in additional.iter().enumerate() { accounts.push(AccountState { name: format!("account-{}", i + 1), - provider: OAuthProvider::new(p.to_string_lossy().as_ref()), + reader: SubscriptionReader::new(provider, p), home: p.clone(), used: AtomicUsize::new(0), + request_limit: request_limit(i + 1), cooldown_until: Mutex::new(None), last_error: Mutex::new(None), }); @@ -118,12 +249,21 @@ impl AccountRouter { inner: Arc::new(AccountRouterInner { accounts, cursor: AtomicUsize::new(0), + provider, strategy, cooldown, + session_affinity_ttl, + affinities: Mutex::new(HashMap::new()), }), } } + /// Provider whose credential layout is used by every account in the pool. + #[must_use] + pub fn provider(&self) -> SubscriptionProvider { + self.inner.provider + } + /// Number of configured accounts (incl. primary). #[must_use] pub fn len(&self) -> usize { @@ -144,8 +284,12 @@ impl AccountRouter { .map(|a| AccountHealth { name: a.name.clone(), home: a.home.clone(), - healthy: a.is_healthy(), + healthy: a.is_available(), used: a.used.load(Ordering::Relaxed), + request_limit: a.request_limit, + remaining_requests: a + .request_limit + .map(|limit| limit.saturating_sub(a.used.load(Ordering::Relaxed))), last_error: a .last_error .lock() @@ -167,56 +311,165 @@ impl AccountRouter { /// case; the legacy single-account path treats this as a fatal config /// error today. pub fn select(&self) -> Result { - if self.inner.accounts.is_empty() { - return Err(AccountError::NoAccountsConfigured); - } - let mut tried = 0usize; - let total = self.inner.accounts.len(); - let start_idx = match self.inner.strategy { - SelectionStrategy::RoundRobin => self.inner.cursor.fetch_add(1, Ordering::Relaxed), - SelectionStrategy::Priority => 0, - SelectionStrategy::LeastUsed => self.least_used_index(), - }; - while tried < total { - let idx = (start_idx + tried) % total; - let acc = &self.inner.accounts[idx]; - if !acc.is_healthy() { - tried += 1; + self.select_with_context(&RoutingContext::default()) + } + + /// Select a Claude-compatible access token using explicit/session routing. + pub fn select_with_context( + &self, + context: &RoutingContext, + ) -> Result { + let selected = self.select_subscription(context)?; + Ok(SelectedAccount { + name: selected.name, + token: selected.token.access_token, + }) + } + + /// Select and normalize a credential for the pool's vendor provider. + pub fn select_subscription( + &self, + context: &RoutingContext, + ) -> Result { + let (indices, mode) = self.selection_plan(context)?; + for idx in indices { + let account = &self.inner.accounts[idx]; + if !account.is_available() { + if !matches!(mode, SelectionMode::Automatic) { + return Err(Self::unavailable_error(mode, &account.name)); + } continue; } - match acc.provider.get_token() { - Ok(tok) => { - acc.used.fetch_add(1, Ordering::Relaxed); - return Ok(SelectedAccount { - name: acc.name.clone(), - token: tok, + match account.reader.read_token() { + Ok(token) if account.try_record_use() => { + self.bind_session(context, idx); + return Ok(SelectedSubscriptionAccount { + name: account.name.clone(), + token, }); } - Err(e) => { - self.record_error(idx, &e.to_string()); - self.start_cooldown(idx); + Ok(_) => { + if !matches!(mode, SelectionMode::Automatic) { + return Err(Self::unavailable_error(mode, &account.name)); + } + } + Err(error) => { + self.record_error(idx, &error.to_string()); + self.start_cooldown(idx, self.inner.cooldown); + if !matches!(mode, SelectionMode::Automatic) { + return Err(Self::unavailable_error(mode, &account.name)); + } } } - tried += 1; } Err(AccountError::NoHealthyAccounts) } - fn least_used_index(&self) -> usize { - let mut best = 0usize; - let mut best_count = usize::MAX; - for (i, a) in self.inner.accounts.iter().enumerate() { - let c = a.used.load(Ordering::Relaxed); - if a.is_healthy() && c < best_count { - best_count = c; - best = i; + fn selection_plan( + &self, + context: &RoutingContext, + ) -> Result<(Vec, SelectionMode), AccountError> { + if self.inner.accounts.is_empty() { + return Err(AccountError::NoAccountsConfigured); + } + if let Some(pin) = context.pinned_account.as_deref() { + let Some(index) = self.inner.accounts.iter().position(|a| a.name == pin) else { + return Err(AccountError::UnknownPinnedAccount(pin.to_string())); + }; + return Ok((vec![index], SelectionMode::Pinned)); + } + if let Some(session) = context.session_key.as_deref() { + if let Some(index) = self.bound_account(session) { + return Ok((vec![index], SelectionMode::Session)); } } - best + let mut indices: Vec = (0..self.inner.accounts.len()).collect(); + match self.inner.strategy { + SelectionStrategy::RoundRobin => { + let start = self.inner.cursor.fetch_add(1, Ordering::Relaxed) % indices.len(); + indices.rotate_left(start); + } + SelectionStrategy::Priority => {} + SelectionStrategy::LeastUsed => indices.sort_by(|left, right| { + Self::compare_usage(&self.inner.accounts[*left], &self.inner.accounts[*right]) + }), + } + Ok((indices, SelectionMode::Automatic)) + } + + fn compare_usage(left: &AccountState, right: &AccountState) -> CmpOrdering { + let left_used = left.used.load(Ordering::Relaxed); + let right_used = right.used.load(Ordering::Relaxed); + match (left.request_limit, right.request_limit) { + (Some(left_limit), Some(right_limit)) => left_used + .saturating_mul(right_limit) + .cmp(&right_used.saturating_mul(left_limit)) + .then_with(|| left_used.cmp(&right_used)), + // Prefer measurable quota headroom; unknown quotas remain eligible + // as a fallback instead of being treated as unlimited. + (Some(_), None) => CmpOrdering::Less, + (None, Some(_)) => CmpOrdering::Greater, + (None, None) => left_used.cmp(&right_used), + } + } + + fn bound_account(&self, session: &str) -> Option { + if self.inner.session_affinity_ttl.is_zero() { + return None; + } + let now = Instant::now(); + let mut affinities = self + .inner + .affinities + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + affinities.retain(|_, binding| binding.expires_at > now); + affinities.get(session).map(|binding| binding.account_index) + } + + fn bind_session(&self, context: &RoutingContext, account_index: usize) { + let Some(session) = context.session_key.as_ref() else { + return; + }; + if self.inner.session_affinity_ttl.is_zero() { + return; + } + let mut affinities = self + .inner + .affinities + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + affinities.insert( + session.clone(), + AffinityBinding { + account_index, + expires_at: Instant::now() + self.inner.session_affinity_ttl, + }, + ); + } + + fn unavailable_error(mode: SelectionMode, account: &str) -> AccountError { + match mode { + SelectionMode::Pinned => AccountError::PinnedAccountUnavailable(account.to_string()), + SelectionMode::Session => AccountError::SessionAccountUnavailable(account.to_string()), + SelectionMode::Automatic => AccountError::NoHealthyAccounts, + } } /// Mark the named account as having failed (e.g., upstream returned 429). pub fn report_failure(&self, account_name: &str, err: &str) { + self.report_failure_with_retry_after(account_name, err, None); + } + + /// Cool an account after a typed quota failure. A vendor `Retry-After` + /// duration overrides a shorter configured default; concurrent failures + /// never shorten an existing cooldown. + pub fn report_failure_with_retry_after( + &self, + account_name: &str, + err: &str, + retry_after: Option, + ) { if let Some(idx) = self .inner .accounts @@ -224,7 +477,10 @@ impl AccountRouter { .position(|a| a.name == account_name) { self.record_error(idx, err); - self.start_cooldown(idx); + self.start_cooldown( + idx, + retry_after.map_or(self.inner.cooldown, |retry| retry.max(self.inner.cooldown)), + ); } } @@ -236,12 +492,15 @@ impl AccountRouter { *guard = Some(err.to_string()); } - fn start_cooldown(&self, idx: usize) { + fn start_cooldown(&self, idx: usize, duration: Duration) { let mut guard = self.inner.accounts[idx] .cooldown_until .lock() .unwrap_or_else(std::sync::PoisonError::into_inner); - *guard = Some(Instant::now() + self.inner.cooldown); + let proposed = Instant::now() + duration; + if guard.is_none_or(|current| current < proposed) { + *guard = Some(proposed); + } } } @@ -252,6 +511,8 @@ pub struct AccountHealth { pub home: PathBuf, pub healthy: bool, pub used: usize, + pub request_limit: Option, + pub remaining_requests: Option, pub last_error: Option, pub cooldown_remaining: Option, } @@ -263,8 +524,12 @@ pub enum AccountError { NoAccountsConfigured, /// Every configured account is currently on cooldown or failing. NoHealthyAccounts, - /// An underlying [`OAuthError`] (rare — usually wrapped into a cooldown). - OAuth(OAuthError), + /// An explicit token pin named no configured account. + UnknownPinnedAccount(String), + /// A strict token-pinned account is cooling down or spent. + PinnedAccountUnavailable(String), + /// A session's bound account is cooling down or spent. + SessionAccountUnavailable(String), } impl std::fmt::Display for AccountError { @@ -272,22 +537,25 @@ impl std::fmt::Display for AccountError { match self { Self::NoAccountsConfigured => write!(f, "no accounts configured"), Self::NoHealthyAccounts => write!(f, "no healthy accounts available"), - Self::OAuth(e) => write!(f, "oauth error: {e}"), + Self::UnknownPinnedAccount(account) => { + write!(f, "token is pinned to unknown account {account}") + } + Self::PinnedAccountUnavailable(account) => { + write!(f, "pinned account {account} is unavailable") + } + Self::SessionAccountUnavailable(account) => { + write!(f, "session account {account} is unavailable") + } } } } impl std::error::Error for AccountError {} -impl From for AccountError { - fn from(e: OAuthError) -> Self { - Self::OAuth(e) - } -} - #[cfg(test)] mod tests { use super::*; + use crate::subscription::SubscriptionProvider; use std::fs; fn tempdir(slug: &str) -> PathBuf { @@ -380,4 +648,228 @@ mod tests { // both accounts should be exercised (LeastUsed prefers the unused one) assert!(snap.iter().any(|s| s.used >= 1)); } + + #[test] + fn strategy_aliases_ignore_surrounding_whitespace() { + assert_eq!( + SelectionStrategy::from_str_opt(" quota-first "), + Some(SelectionStrategy::LeastUsed) + ); + } + + #[test] + fn least_used_compares_normalized_spend_for_uneven_limits() { + let a = tempdir("normalized-a"); + let b = tempdir("normalized-b"); + write_creds(&a, "tok-a"); + write_creds(&b, "tok-b"); + let router = AccountRouter::new_for_provider( + a, + &[b], + SubscriptionProvider::Claude, + AccountRouterOptions { + strategy: SelectionStrategy::LeastUsed, + request_limits: vec![Some(2), Some(100)], + ..AccountRouterOptions::default() + }, + ); + + assert_eq!(router.select().unwrap().name, "primary"); + assert_eq!(router.select().unwrap().name, "account-1"); + assert_eq!(router.select().unwrap().name, "account-1"); + } + + #[test] + fn session_affinity_keeps_a_conversation_on_one_account() { + let a = tempdir("session-a"); + let b = tempdir("session-b"); + write_creds(&a, "tok-a"); + write_creds(&b, "tok-b"); + let router = AccountRouter::new_for_provider( + a, + &[b], + SubscriptionProvider::Claude, + AccountRouterOptions::default(), + ); + + let first = router + .select_with_context(&RoutingContext::for_session("conversation-1")) + .unwrap(); + let again = router + .select_with_context(&RoutingContext::for_session("conversation-1")) + .unwrap(); + let other = router + .select_with_context(&RoutingContext::for_session("conversation-2")) + .unwrap(); + + assert_eq!(first.name, again.name); + assert_ne!(first.name, other.name); + } + + #[test] + fn session_activity_renews_the_affinity_timeout() { + let a = tempdir("session-renew-a"); + let b = tempdir("session-renew-b"); + write_creds(&a, "tok-a"); + write_creds(&b, "tok-b"); + let router = AccountRouter::new_for_provider( + a, + &[b], + SubscriptionProvider::Claude, + AccountRouterOptions::default(), + ); + let context = RoutingContext::for_session("active-conversation"); + router.select_with_context(&context).unwrap(); + + let shortened_expiry = Instant::now() + Duration::from_secs(1); + router + .inner + .affinities + .lock() + .unwrap() + .get_mut("active-conversation") + .unwrap() + .expires_at = shortened_expiry; + + router.select_with_context(&context).unwrap(); + let renewed_expiry = router + .inner + .affinities + .lock() + .unwrap() + .get("active-conversation") + .unwrap() + .expires_at; + assert!(renewed_expiry > shortened_expiry); + } + + #[test] + fn an_unavailable_session_account_is_not_silently_changed() { + let a = tempdir("strict-session-a"); + let b = tempdir("strict-session-b"); + write_creds(&a, "tok-a"); + write_creds(&b, "tok-b"); + let router = AccountRouter::new_for_provider( + a, + &[b], + SubscriptionProvider::Claude, + AccountRouterOptions::default(), + ); + let context = RoutingContext::for_session("strict-conversation"); + let selected = router.select_with_context(&context).unwrap(); + router.report_failure(&selected.name, "quota exhausted"); + + assert!(matches!( + router.select_with_context(&context), + Err(AccountError::SessionAccountUnavailable(_)) + )); + } + + #[test] + fn explicit_account_pins_are_strict() { + let a = tempdir("pin-a"); + let b = tempdir("pin-b"); + write_creds(&a, "tok-a"); + write_creds(&b, "tok-b"); + let router = AccountRouter::new_for_provider( + a, + &[b], + SubscriptionProvider::Claude, + AccountRouterOptions::default(), + ); + + let selected = router + .select_with_context(&RoutingContext::pinned("account-1")) + .unwrap(); + assert_eq!(selected.name, "account-1"); + router.report_failure("account-1", "quota exhausted"); + assert!(matches!( + router.select_with_context(&RoutingContext::pinned("account-1")), + Err(AccountError::PinnedAccountUnavailable(_)) + )); + assert!(matches!( + router.select_with_context(&RoutingContext::pinned("missing")), + Err(AccountError::UnknownPinnedAccount(_)) + )); + } + + #[test] + fn configured_request_limits_remove_spent_accounts() { + let a = tempdir("limits-a"); + let b = tempdir("limits-b"); + write_creds(&a, "tok-a"); + write_creds(&b, "tok-b"); + let options = AccountRouterOptions { + request_limits: vec![Some(1), Some(2)], + ..AccountRouterOptions::default() + }; + let router = + AccountRouter::new_for_provider(a, &[b], SubscriptionProvider::Claude, options); + + assert_eq!(router.select().unwrap().name, "primary"); + assert_eq!(router.select().unwrap().name, "account-1"); + assert_eq!(router.select().unwrap().name, "account-1"); + assert!(matches!( + router.select(), + Err(AccountError::NoHealthyAccounts) + )); + let health = router.health_snapshot(); + assert_eq!(health[0].remaining_requests, Some(0)); + assert_eq!(health[1].remaining_requests, Some(0)); + } + + #[test] + fn concurrent_selection_cannot_oversubscribe_an_account_cap() { + let a = tempdir("atomic-limit"); + write_creds(&a, "tok-a"); + let router = AccountRouter::new_for_provider( + a, + &[], + SubscriptionProvider::Claude, + AccountRouterOptions { + request_limits: vec![Some(1)], + ..AccountRouterOptions::default() + }, + ); + let successful = (0..16) + .map(|_| { + let router = router.clone(); + std::thread::spawn(move || router.select().is_ok()) + }) + .map(|worker| worker.join().unwrap()) + .filter(|successful| *successful) + .count(); + + assert_eq!(successful, 1); + assert_eq!(router.health_snapshot()[0].used, 1); + } + + #[test] + fn vendor_subscription_accounts_use_the_same_pool() { + let a = tempdir("codex-a"); + let b = tempdir("codex-b"); + fs::write( + a.join("auth.json"), + r#"{"tokens":{"access_token":"codex-a","account_id":"acct-a"}}"#, + ) + .unwrap(); + fs::write( + b.join("auth.json"), + r#"{"tokens":{"access_token":"codex-b","account_id":"acct-b"}}"#, + ) + .unwrap(); + let router = AccountRouter::new_for_provider( + a, + &[b], + SubscriptionProvider::Codex, + AccountRouterOptions::default(), + ); + + let selected = router + .select_subscription(&RoutingContext::pinned("account-1")) + .unwrap(); + assert_eq!(selected.name, "account-1"); + assert_eq!(selected.token.access_token, "codex-b"); + assert_eq!(selected.token.account_id.as_deref(), Some("acct-b")); + } } diff --git a/src/app_state.rs b/src/app_state.rs new file mode 100644 index 0000000..d4fbc1d --- /dev/null +++ b/src/app_state.rs @@ -0,0 +1,56 @@ +//! Shared state used by HTTP route handlers. + +use std::sync::Arc; + +use log_lazy::LogLazy; +use reqwest::Client; + +use crate::accounts::AccountRouter; +use crate::config::UpstreamProvider; +use crate::gonka::GonkaConfig; +use crate::oauth::OAuthProvider; +use crate::providers::{OpenAICompatibleConfig, ProviderStore}; +use crate::token::TokenManager; + +/// Shared application state accessible by all route handlers. +#[derive(Clone)] +pub struct AppState { + /// HTTP client for upstream requests. + pub client: Client, + /// Token manager for validating custom tokens. + pub token_manager: TokenManager, + /// OAuth provider for obtaining upstream credentials (legacy single-account). + pub oauth_provider: OAuthProvider, + /// Multi-account router (when configured). When `None`, the legacy + /// `oauth_provider` is used directly. + pub account_router: Option, + /// Subscription credential reader for vendor OAuth providers + /// (Codex/Gemini/Qwen). `None` for non-subscription upstreams. + pub subscription_reader: Option, + /// In-memory cache of refreshed subscription tokens (Codex/Gemini/Qwen). + pub subscription_cache: Arc, + /// Base URL for the upstream Anthropic API. + pub upstream_base_url: String, + /// Selected upstream inference provider. + pub upstream_provider: UpstreamProvider, + /// Gonka provider configuration when selected. + pub gonka: Option, + /// Crater `ForgeFed` task provider when selected. + pub crater: Option>, + /// Boot-time generic OpenAI-compatible provider config. + pub openai_compatible: OpenAICompatibleConfig, + /// Persisted provider records with encrypted upstream secrets. + pub provider_store: ProviderStore, + /// Lazy logger for verbose output. + pub logger: LogLazy, + /// Optional admin key (Bearer) required for `/api/tokens` issuance. + pub admin_key: Option, + /// Live metrics counter handle. + pub metrics: Arc, + /// Public base URL for `ActivityPub` actor documents. + pub activitypub_actor_base_url: String, + /// Public key PEM advertised by the `ActivityPub` actor. + pub activitypub_public_key_pem: String, + /// Optional MPP charge settings for OpenAI-compatible endpoints. + pub mpp: crate::mpp::MppConfig, +} diff --git a/src/cli.rs b/src/cli.rs index cf41317..a89dacb 100644 --- a/src/cli.rs +++ b/src/cli.rs @@ -221,6 +221,42 @@ pub struct Cli { )] pub additional_account_dirs: Vec, + /// New-session account policy: round-robin, fill-first, or least-used. + #[arg( + long, + env = "ACCOUNT_ROUTING_STRATEGY", + default_value = "round-robin", + global = true + )] + pub account_routing_strategy: String, + + /// Default seconds to cool an account after a quota response. + #[arg( + long, + env = "ACCOUNT_COOLDOWN_SECS", + default_value_t = 60, + global = true + )] + pub account_cooldown_secs: u64, + + /// Seconds an inactive conversation remains on its selected account. + #[arg( + long, + env = "SESSION_AFFINITY_TTL_SECS", + default_value_t = 3600, + global = true + )] + pub session_affinity_ttl_secs: u64, + + /// Per-account request caps (primary first); zero means unknown/unlimited. + #[arg( + long, + env = "ACCOUNT_REQUEST_LIMITS", + value_delimiter = ',', + global = true + )] + pub account_request_limits: Vec, + /// Enable experimental compatibility shims (XML history, spoofing, …). #[arg(long, env = "EXPERIMENTAL_COMPATIBILITY", global = true)] pub experimental_compatibility: bool, @@ -351,6 +387,9 @@ impl Cli { let upstream_provider = UpstreamProvider::from_str_opt(&self.upstream_provider) .unwrap_or(UpstreamProvider::Anthropic); let storage_policy = StoragePolicy::from_str_opt(&self.storage_policy).unwrap_or_default(); + let account_routing_strategy = + crate::accounts::SelectionStrategy::from_str_opt(&self.account_routing_strategy) + .ok_or(ConfigError::InvalidAccountRoutingStrategy)?; let data_dir = self.data_dir.clone().unwrap_or_else(default_data_dir); let activitypub_actor_base_url = self .activitypub_actor_base_url @@ -422,6 +461,10 @@ impl Cli { enable_anthropic_api: !self.disable_anthropic_api, enable_metrics: !self.disable_metrics, additional_account_dirs: self.additional_account_dirs.clone(), + account_routing_strategy, + account_cooldown_secs: self.account_cooldown_secs, + session_affinity_ttl_secs: self.session_affinity_ttl_secs, + account_request_limits: self.account_request_limits.clone(), experimental_compatibility: self.experimental_compatibility, admin_key: self.admin_key.clone().filter(|s| !s.is_empty()), mpp: crate::mpp::MppConfig { @@ -478,6 +521,10 @@ mod tests { disable_anthropic_api: false, disable_metrics: false, additional_account_dirs: vec![], + account_routing_strategy: "round-robin".into(), + account_cooldown_secs: 60, + session_affinity_ttl_secs: 3600, + account_request_limits: vec![], experimental_compatibility: false, admin_key: None, mpp_enable: false, @@ -531,6 +578,10 @@ mod tests { disable_anthropic_api: false, disable_metrics: false, additional_account_dirs: vec![], + account_routing_strategy: "round-robin".into(), + account_cooldown_secs: 60, + session_affinity_ttl_secs: 3600, + account_request_limits: vec![], experimental_compatibility: false, admin_key: None, mpp_enable: false, diff --git a/src/config.rs b/src/config.rs index 4a2514c..f456617 100644 --- a/src/config.rs +++ b/src/config.rs @@ -15,6 +15,8 @@ use std::path::PathBuf; use std::str::FromStr; use std::time::Duration; +use crate::accounts::SelectionStrategy; + /// Supported upstream inference providers. #[derive(Debug, Clone, Copy, PartialEq, Eq, Default)] pub enum UpstreamProvider { @@ -200,9 +202,18 @@ pub struct Config { pub enable_anthropic_api: bool, /// Whether to expose `/metrics` and other operational endpoints. pub enable_metrics: bool, - /// Optional comma-separated list of additional Claude account credential - /// directories — used by the multi-account router. + /// Optional comma-separated list of additional credential directories for + /// the active vendor-subscription provider. pub additional_account_dirs: Vec, + /// Selection policy applied to new sessions in a multi-account pool. + pub account_routing_strategy: SelectionStrategy, + /// Default cooldown after an account returns a typed quota failure. + pub account_cooldown_secs: u64, + /// Inactive session-affinity lifetime. Zero disables affinity. + pub session_affinity_ttl_secs: u64, + /// Per-account request caps, ordered primary then additional. Zero means + /// unknown/unlimited. + pub account_request_limits: Vec, /// Whether to enable experimental compatibility features (spoofing, /// XML history reconstruction, etc.). Off by default. pub experimental_compatibility: bool, @@ -313,6 +324,18 @@ impl Config { .collect() }) .unwrap_or_default(); + let account_routing_strategy = match env::var("ACCOUNT_ROUTING_STRATEGY") { + Ok(value) => SelectionStrategy::from_str_opt(&value) + .ok_or(ConfigError::InvalidAccountRoutingStrategy)?, + Err(_) => SelectionStrategy::default(), + }; + let account_cooldown_secs = parse_u64_env("ACCOUNT_COOLDOWN_SECS", 60); + let session_affinity_ttl_secs = parse_u64_env("SESSION_AFFINITY_TTL_SECS", 3600); + let account_request_limits = env::var("ACCOUNT_REQUEST_LIMITS") + .ok() + .map(|raw| parse_usize_csv(&raw)) + .transpose()? + .unwrap_or_default(); let experimental_compatibility = env::var("EXPERIMENTAL_COMPATIBILITY") .is_ok_and(|v| v == "1" || v.eq_ignore_ascii_case("true")); let admin_key = env::var("TOKEN_ADMIN_KEY").ok().filter(|s| !s.is_empty()); @@ -349,6 +372,10 @@ impl Config { enable_anthropic_api, enable_metrics, additional_account_dirs, + account_routing_strategy, + account_cooldown_secs, + session_affinity_ttl_secs, + account_request_limits, experimental_compatibility, admin_key, mpp, @@ -377,6 +404,11 @@ impl Config { if args.upstream_provider == UpstreamProvider::Crater && args.crater.inbox.is_none() { return Err(ConfigError::MissingCraterForgeFedInbox); } + if !args.account_request_limits.is_empty() + && args.account_request_limits.len() != args.additional_account_dirs.len() + 1 + { + return Err(ConfigError::MismatchedAccountRequestLimits); + } Ok(Self { listen_addr, @@ -404,6 +436,10 @@ impl Config { enable_anthropic_api: args.enable_anthropic_api, enable_metrics: args.enable_metrics, additional_account_dirs: args.additional_account_dirs, + account_routing_strategy: args.account_routing_strategy, + account_cooldown_secs: args.account_cooldown_secs, + session_affinity_ttl_secs: args.session_affinity_ttl_secs, + account_request_limits: args.account_request_limits, experimental_compatibility: args.experimental_compatibility, admin_key: args.admin_key, mpp: args.mpp, @@ -436,6 +472,10 @@ pub struct BuildArgs<'a> { pub enable_anthropic_api: bool, pub enable_metrics: bool, pub additional_account_dirs: Vec, + pub account_routing_strategy: SelectionStrategy, + pub account_cooldown_secs: u64, + pub session_affinity_ttl_secs: u64, + pub account_request_limits: Vec, pub experimental_compatibility: bool, pub admin_key: Option, pub mpp: crate::mpp::MppConfig, @@ -529,6 +569,18 @@ fn parse_u64_env(name: &str, default: u64) -> u64 { .unwrap_or(default) } +fn parse_usize_csv(raw: &str) -> Result, ConfigError> { + raw.split(',') + .map(str::trim) + .filter(|value| !value.is_empty()) + .map(|value| { + value + .parse::() + .map_err(|_| ConfigError::InvalidAccountRequestLimits) + }) + .collect() +} + /// Errors that can occur during configuration loading. #[derive(Debug)] pub enum ConfigError { @@ -540,6 +592,12 @@ pub enum ConfigError { MissingTokenSecret, /// Routing mode was not recognised. InvalidRoutingMode, + /// The multi-account strategy was not recognised. + InvalidAccountRoutingStrategy, + /// An account request cap was not a non-negative integer. + InvalidAccountRequestLimits, + /// Request caps did not align with primary plus additional accounts. + MismatchedAccountRequestLimits, /// Gonka was selected without `GONKA_PRIVATE_KEY`. MissingGonkaPrivateKey, /// Crater was selected without `CRATER_FORGEFED_INBOX`. @@ -557,6 +615,18 @@ impl std::fmt::Display for ConfigError { Self::InvalidRoutingMode => { write!(f, "ROUTING_MODE must be one of: direct, cli, hybrid") } + Self::InvalidAccountRoutingStrategy => write!( + f, + "ACCOUNT_ROUTING_STRATEGY must be one of: round-robin, fill-first, least-used" + ), + Self::InvalidAccountRequestLimits => write!( + f, + "ACCOUNT_REQUEST_LIMITS must be comma-separated non-negative integers" + ), + Self::MismatchedAccountRequestLimits => write!( + f, + "ACCOUNT_REQUEST_LIMITS must contain one entry for primary and each additional account" + ), Self::MissingGonkaPrivateKey => write!( f, "Gonka provider requires GONKA_PRIVATE_KEY. Make sure your Gonka account is activated for inference, funded, and has a published on-chain public key." @@ -599,6 +669,10 @@ mod tests { enable_anthropic_api: true, enable_metrics: true, additional_account_dirs: vec![], + account_routing_strategy: SelectionStrategy::default(), + account_cooldown_secs: 60, + session_affinity_ttl_secs: 3600, + account_request_limits: vec![], experimental_compatibility: false, admin_key: None, mpp: default_mpp_config(), @@ -635,6 +709,28 @@ mod tests { assert_eq!(config.upstream_provider, UpstreamProvider::Anthropic); } + #[test] + fn account_limits_must_align_with_the_configured_pool() { + let mut args = gonka_args(None); + args.upstream_provider = UpstreamProvider::Anthropic; + args.additional_account_dirs = vec![PathBuf::from("/tmp/second")]; + args.account_request_limits = vec![100]; + + assert!(matches!( + Config::build(args), + Err(ConfigError::MismatchedAccountRequestLimits) + )); + } + + #[test] + fn account_limit_parser_accepts_zero_as_unlimited() { + assert_eq!(parse_usize_csv("100, 0,250").unwrap(), vec![100, 0, 250]); + assert!(matches!( + parse_usize_csv("100,nope"), + Err(ConfigError::InvalidAccountRequestLimits) + )); + } + #[test] fn gonka_provider_requires_private_key() { let result = Config::build(gonka_args(None)); @@ -709,6 +805,10 @@ mod tests { enable_anthropic_api: true, enable_metrics: true, additional_account_dirs: vec![], + account_routing_strategy: SelectionStrategy::default(), + account_cooldown_secs: 60, + session_affinity_ttl_secs: 3600, + account_request_limits: vec![], experimental_compatibility: false, admin_key: None, mpp: default_mpp_config(), @@ -741,6 +841,10 @@ mod tests { enable_anthropic_api: true, enable_metrics: true, additional_account_dirs: vec![], + account_routing_strategy: SelectionStrategy::default(), + account_cooldown_secs: 60, + session_affinity_ttl_secs: 3600, + account_request_limits: vec![], experimental_compatibility: false, admin_key: None, mpp: default_mpp_config(), diff --git a/src/gemini.rs b/src/gemini.rs index db4c33d..ec9ac07 100644 --- a/src/gemini.rs +++ b/src/gemini.rs @@ -15,12 +15,16 @@ #![allow(clippy::unused_async)] use axum::body::Body; +use axum::extract::{Path, State}; use axum::http::{HeaderMap, StatusCode}; -use axum::response::Response; +use axum::response::{IntoResponse, Response}; use serde_json::{Value, json}; use crate::metrics::Surface; -use crate::proxy::{AppState, error_response, extract_client_token, maybe_mpp_challenge}; +use crate::proxy::{ + AppState, error_response, extract_client_token, maybe_mpp_challenge, request_routing_context, + retry_after_duration, +}; /// Environment variable carrying the Google Cloud project id for Code Assist. pub const PROJECT_ENV: &str = "GEMINI_PROJECT"; @@ -230,60 +234,115 @@ enum ShapeIn { Responses, } -async fn forward( +struct RoutedGeminiToken { + token: crate::subscription::SubscriptionToken, + account: String, +} + +async fn route_gemini_token( state: &AppState, headers: &HeaderMap, - body: Value, - surface: Surface, - shape: ShapeIn, -) -> Response { - if let Some(resp) = maybe_mpp_challenge(state, headers, "/v1/chat/completions") { - return resp; - } + body: &Value, +) -> Result { let Some(token) = extract_client_token(headers) else { - return error_response( + return Err(error_response( StatusCode::UNAUTHORIZED, "authentication_error", "Missing Authorization Bearer token or x-api-key", - ); + )); }; - if let Err(e) = state.token_manager.validate_token(token) { - let status = match &e { - crate::token::TokenError::Revoked => StatusCode::FORBIDDEN, - _ => StatusCode::UNAUTHORIZED, + let claims = state.token_manager.validate_token(token).map_err(|error| { + let status = if matches!(error, crate::token::TokenError::Revoked) { + StatusCode::FORBIDDEN + } else { + StatusCode::UNAUTHORIZED }; - return error_response(status, "authentication_error", &format!("{e}")); - } - - let Some(reader) = state.subscription_reader.as_ref() else { - return error_response( - StatusCode::INTERNAL_SERVER_ERROR, - "api_error", - "subscription credentials reader is not configured", - ); - }; - let disk_token = match reader.read_token() { - Ok(token) => token, - Err(e) => { - return error_response( + error_response(status, "authentication_error", &error.to_string()) + })?; + state + .token_manager + .enforce_request_budget(&claims.sub) + .map_err(|error| { + error_response( + StatusCode::TOO_MANY_REQUESTS, + "rate_limit_error", + &error.to_string(), + ) + })?; + let pinned_account = state + .token_manager + .account_for(&claims.sub) + .map_err(|error| { + error_response( + StatusCode::INTERNAL_SERVER_ERROR, + "api_error", + &format!("failed to resolve token account binding: {error}"), + ) + })?; + let routing_context = request_routing_context(headers, body, pinned_account); + let selected = if let Some(router) = state.account_router.as_ref() { + router + .select_subscription(&routing_context) + .map_err(|error| { + error_response( + StatusCode::SERVICE_UNAVAILABLE, + "account_unavailable", + &error.to_string(), + ) + })? + } else { + let reader = state.subscription_reader.as_ref().ok_or_else(|| { + error_response( + StatusCode::INTERNAL_SERVER_ERROR, + "api_error", + "subscription credentials reader is not configured", + ) + })?; + let token = reader.read_token().map_err(|error| { + error_response( StatusCode::BAD_GATEWAY, "authentication_error", - &format!("failed to read Gemini subscription credentials: {e}"), - ); + &format!("failed to read Gemini subscription credentials: {error}"), + ) + })?; + crate::accounts::SelectedSubscriptionAccount { + name: "primary".to_string(), + token, } }; - // Refresh in memory if the on-disk token has expired; vendor files stay - // read-only. let now_ms = chrono::Utc::now().timestamp_millis(); - let sub_token = state + let token = state .subscription_cache - .get_fresh( + .get_fresh_for( &state.client, crate::subscription::SubscriptionProvider::Gemini, - disk_token, + &selected.name, + selected.token, now_ms, ) .await; + Ok(RoutedGeminiToken { + token, + account: selected.name, + }) +} + +async fn forward( + state: &AppState, + headers: &HeaderMap, + body: Value, + surface: Surface, + shape: ShapeIn, +) -> Response { + if let Some(resp) = maybe_mpp_challenge(state, headers, "/v1/chat/completions") { + return resp; + } + let routed = match route_gemini_token(state, headers, &body).await { + Ok(routed) => routed, + Err(response) => return response, + }; + let sub_token = routed.token; + let selected_account = Some(routed.account); // Normalize Responses input into the Chat `messages` shape so a single // translator handles both surfaces. @@ -337,7 +396,9 @@ async fn forward( { Ok(resp) => resp, Err(e) => { - state.metrics.record_request(surface, 502, None); + state + .metrics + .record_request(surface, 502, selected_account.as_deref()); return error_response( StatusCode::BAD_GATEWAY, "api_error", @@ -347,12 +408,28 @@ async fn forward( }; let status = StatusCode::from_u16(upstream_resp.status().as_u16()) .unwrap_or(StatusCode::INTERNAL_SERVER_ERROR); - state.metrics.record_request(surface, status.as_u16(), None); + state + .metrics + .record_request(surface, status.as_u16(), selected_account.as_deref()); + let retry_after = retry_after_duration(upstream_resp.headers()); + if status == StatusCode::TOO_MANY_REQUESTS { + if let (Some(router), Some(account)) = + (state.account_router.as_ref(), selected_account.as_deref()) + { + router.report_failure_with_retry_after( + account, + "Gemini subscription upstream returned 429", + retry_after, + ); + } + } let upstream_body = match upstream_resp.bytes().await { Ok(bytes) => bytes, Err(e) => { - state.metrics.record_request(surface, 502, None); + state + .metrics + .record_request(surface, 502, selected_account.as_deref()); return error_response( StatusCode::BAD_GATEWAY, "api_error", @@ -442,6 +519,203 @@ fn sse_from_chat_completion(chat: &Value, model: &str) -> Response { response } +fn native_model_document(model: &str) -> Value { + let model = model.trim_start_matches("models/"); + json!({ + "name": format!("models/{model}"), + "displayName": model, + "description": "Gemini Code Assist subscription model routed by Link.Assistant.Router", + "inputTokenLimit": 1_048_576, + "outputTokenLimit": 65_536, + "supportedGenerationMethods": ["generateContent", "streamGenerateContent"] + }) +} + +/// Native Gemini `ListModels` response. +pub async fn native_models() -> impl IntoResponse { + let models = [ + "gemini-2.5-pro", + "gemini-2.5-flash", + "gemini-2.0-flash", + "gemini-2.0-flash-lite", + ] + .iter() + .map(|model| native_model_document(model)) + .collect::>(); + (StatusCode::OK, axum::Json(json!({"models": models}))) +} + +/// Native Gemini `GetModel` response. +pub async fn native_model(Path(model): Path) -> impl IntoResponse { + (StatusCode::OK, axum::Json(native_model_document(&model))) +} + +fn parse_native_target(path: &str) -> Option<(String, bool)> { + let (resource, action) = path.rsplit_once(':')?; + let streaming = match action { + "generateContent" => false, + "streamGenerateContent" => true, + _ => return None, + }; + let model = resource + .rsplit_once("/models/") + .map_or(resource, |(_, model)| model) + .trim_start_matches("models/") + .trim_matches('/'); + (!model.is_empty()).then(|| (model.to_string(), streaming)) +} + +/// Forward a native Gemini `GenerateContentRequest` without translating it to +/// an `OpenAI` shape. +pub async fn forward_native_gemini( + State(state): State, + Path(path): Path, + headers: HeaderMap, + axum::Json(body): axum::Json, +) -> Response { + forward_native(&state, &headers, &path, body).await +} + +/// Forward a Vertex publisher-model request using the same native Gemini body. +pub async fn forward_native_vertex( + State(state): State, + Path(path): Path, + headers: HeaderMap, + axum::Json(body): axum::Json, +) -> Response { + forward_native(&state, &headers, &path, body).await +} + +async fn forward_native( + state: &AppState, + headers: &HeaderMap, + path: &str, + body: Value, +) -> Response { + if state.upstream_provider != crate::config::UpstreamProvider::Gemini { + return error_response( + StatusCode::BAD_REQUEST, + "invalid_request_error", + "native Gemini and Vertex routes require UPSTREAM_PROVIDER=gemini", + ); + } + let Some((model, streaming)) = parse_native_target(path) else { + return error_response( + StatusCode::NOT_FOUND, + "not_found_error", + "expected a model :generateContent or :streamGenerateContent action", + ); + }; + let routed = match route_gemini_token(state, headers, &body).await { + Ok(routed) => routed, + Err(response) => return response, + }; + let envelope = code_assist_envelope(&model, &body); + let serialized = match serde_json::to_vec(&envelope) { + Ok(serialized) => serialized, + Err(error) => { + return error_response( + StatusCode::INTERNAL_SERVER_ERROR, + "api_error", + &format!("failed to serialize Gemini request: {error}"), + ); + } + }; + let base = routed + .token + .base_url(crate::subscription::SubscriptionProvider::Gemini); + // Code Assist accepts the same envelope for both client modes. A + // non-streaming upstream call lets us unwrap the native response reliably; + // stream callers receive that response as a valid Gemini SSE data event. + let upstream_url = format!("{}/v1internal:generateContent", base.trim_end_matches('/')); + let upstream = match state + .client + .post(upstream_url) + .header("content-type", "application/json") + .header( + "authorization", + format!("Bearer {}", routed.token.access_token), + ) + .body(serialized.clone()) + .send() + .await + { + Ok(response) => response, + Err(error) => { + state + .metrics + .record_request(Surface::OpenAIChat, 502, Some(&routed.account)); + return error_response( + StatusCode::BAD_GATEWAY, + "api_error", + &format!("Gemini subscription upstream request failed: {error}"), + ); + } + }; + let status = StatusCode::from_u16(upstream.status().as_u16()) + .unwrap_or(StatusCode::INTERNAL_SERVER_ERROR); + let retry_after = retry_after_duration(upstream.headers()); + let retry_header = upstream.headers().get("retry-after").cloned(); + let response_body = match upstream.bytes().await { + Ok(bytes) => bytes, + Err(error) => { + return error_response( + StatusCode::BAD_GATEWAY, + "api_error", + &format!("Gemini subscription upstream body read failed: {error}"), + ); + } + }; + state + .metrics + .record_bytes(serialized.len() as u64, response_body.len() as u64); + state + .metrics + .record_request(Surface::OpenAIChat, status.as_u16(), Some(&routed.account)); + if status == StatusCode::TOO_MANY_REQUESTS { + if let Some(router) = state.account_router.as_ref() { + router.report_failure_with_retry_after( + &routed.account, + "Gemini subscription upstream returned 429", + retry_after, + ); + } + } + if !status.is_success() { + let mut response = Response::new(Body::from(response_body)); + *response.status_mut() = status; + response.headers_mut().insert( + "content-type", + axum::http::HeaderValue::from_static("application/json"), + ); + if let Some(value) = retry_header { + response.headers_mut().insert("retry-after", value); + } + return response; + } + let parsed: Value = match serde_json::from_slice(&response_body) { + Ok(parsed) => parsed, + Err(error) => { + return error_response( + StatusCode::BAD_GATEWAY, + "api_error", + &format!("failed to parse Gemini response: {error}"), + ); + } + }; + let native = parsed.get("response").cloned().unwrap_or(parsed); + if streaming { + let mut response = Response::new(Body::from(format!("data: {native}\n\n"))); + *response.status_mut() = StatusCode::OK; + response.headers_mut().insert( + "content-type", + axum::http::HeaderValue::from_static("text/event-stream"), + ); + return response; + } + (StatusCode::OK, axum::Json(native)).into_response() +} + /// Project an `OpenAI` Responses request onto the Chat Completions shape. fn responses_to_chat(body: &Value) -> Value { let mut messages: Vec = Vec::new(); @@ -574,4 +848,19 @@ mod tests { assert_eq!(map_model("gemini-2.5-flash"), "gemini-2.5-flash"); assert_eq!(map_model("gpt-4o"), DEFAULT_MODEL); } + + #[test] + fn parses_gemini_and_vertex_native_actions() { + assert_eq!( + parse_native_target("models/gemini-2.5-pro:generateContent"), + Some(("gemini-2.5-pro".into(), false)) + ); + assert_eq!( + parse_native_target( + "projects/p/locations/us/publishers/google/models/gemini-2.5-flash:streamGenerateContent" + ), + Some(("gemini-2.5-flash".into(), true)) + ); + assert!(parse_native_target("models/gemini-2.5-pro:countTokens").is_none()); + } } diff --git a/src/lib.rs b/src/lib.rs index 4574077..4603a23 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -6,6 +6,7 @@ pub mod accounts; pub mod activitypub; +pub mod app_state; pub mod cli; pub mod config; pub mod crater; @@ -19,6 +20,7 @@ pub mod provider_proxy; pub mod providers; pub mod proxy; pub mod refresh; +mod request_routing; pub mod responses; pub mod storage; pub mod subscription; diff --git a/src/main.rs b/src/main.rs index c701556..f8dd22b 100644 --- a/src/main.rs +++ b/src/main.rs @@ -16,7 +16,7 @@ use std::time::Duration; use axum::Router; use axum::routing::{get, post}; -use link_assistant_router::accounts::{AccountRouter, SelectionStrategy}; +use link_assistant_router::accounts::{AccountRouter, AccountRouterOptions}; use link_assistant_router::activitypub; use link_assistant_router::cli::{AccountOp, Cli, Command, ProviderOp, TokenOp}; use link_assistant_router::config::{Config, RoutingMode, StoragePolicy}; @@ -27,6 +27,7 @@ use link_assistant_router::provider_proxy; use link_assistant_router::providers::{ProviderStore, ProviderUpsert}; use link_assistant_router::proxy::{self, AppState}; use link_assistant_router::storage::{TokenStore, build_token_store}; +use link_assistant_router::subscription::{SubscriptionProvider, SubscriptionReader}; use link_assistant_router::token::TokenManager; use link_assistant_router::token_admin; use log_lazy::{LogLazy, levels}; @@ -98,6 +99,20 @@ fn build_logger(verbose: bool) -> LogLazy { }) } +fn subscription_pool(config: &Config) -> (SubscriptionProvider, std::path::PathBuf) { + let provider = config + .upstream_provider + .subscription_provider() + .unwrap_or(SubscriptionProvider::Claude); + let primary = if provider == SubscriptionProvider::Claude { + std::path::PathBuf::from(&config.claude_code_home) + } else { + let user_home = std::env::var("HOME").unwrap_or_else(|_| ".".to_string()); + provider.resolve_home(&user_home) + }; + (provider, primary) +} + /// Construct the persistent token store and the optional multi-account router /// for the given [`Config`]. Both are needed by both the server and the CLI /// subcommands. @@ -106,23 +121,39 @@ fn build_shared_state(config: &Config) -> Result { std::fs::create_dir_all(&config.data_dir)?; } let store = build_token_store(config.storage_policy, &config.data_dir)?; - let account_router = if config.additional_account_dirs.is_empty() { - None - } else { - Some(AccountRouter::new( - std::path::PathBuf::from(&config.claude_code_home), - &config.additional_account_dirs, - SelectionStrategy::default(), - Duration::from_secs(60), - )) - }; + let account_router = + if config.additional_account_dirs.is_empty() && config.account_request_limits.is_empty() { + None + } else { + let (provider, primary) = subscription_pool(config); + let options = AccountRouterOptions { + strategy: config.account_routing_strategy, + cooldown: Duration::from_secs(config.account_cooldown_secs), + session_affinity_ttl: Duration::from_secs(config.session_affinity_ttl_secs), + request_limits: config + .account_request_limits + .iter() + .map(|limit| (*limit != 0).then_some(*limit)) + .collect(), + }; + Some(AccountRouter::new_for_provider( + primary, + &config.additional_account_dirs, + provider, + options, + )) + }; Ok((store, account_router)) } async fn run_server(config: Config, logger: LogLazy) -> Result<(), Box> { tracing::info!("Upstream: {}", config.upstream_base_url); tracing::info!("Upstream provider: {:?}", config.upstream_provider); - tracing::info!("Claude Code home: {}", config.claude_code_home); + let (subscription_provider, subscription_home) = subscription_pool(&config); + tracing::info!( + "Subscription home ({subscription_provider}): {}", + subscription_home.display() + ); tracing::info!("Routing mode: {:?}", config.routing_mode); tracing::info!("Storage policy: {:?}", config.storage_policy); if config.routing_mode == RoutingMode::Cli || config.routing_mode == RoutingMode::Hybrid { @@ -225,6 +256,11 @@ async fn run_server(config: Config, logger: LogLazy) -> Result<(), Box Result<(), Box ExitCode { Ok((_, Some(r))) => r, Ok((_, None)) => { // Single-account mode: synthesise a one-account router for inspection. - AccountRouter::new( - std::path::PathBuf::from(&config.claude_code_home), + let (provider, primary) = subscription_pool(config); + AccountRouter::new_for_provider( + primary, &[], - SelectionStrategy::default(), - Duration::from_secs(60), + provider, + AccountRouterOptions { + strategy: config.account_routing_strategy, + cooldown: Duration::from_secs(config.account_cooldown_secs), + session_affinity_ttl: Duration::from_secs(config.session_affinity_ttl_secs), + request_limits: config + .account_request_limits + .iter() + .map(|limit| (*limit != 0).then_some(*limit)) + .collect(), + }, ) } Err(e) => { @@ -354,13 +431,24 @@ fn run_accounts(config: &Config, op: &AccountOp) -> ExitCode { match op { AccountOp::List => { let snap = router.health_snapshot(); - println!("{:<16} {:<8} {:<6} home", "name", "healthy", "used"); + println!( + "{:<16} {:<8} {:<6} {:<9} {:<9} home", + "name", "healthy", "used", "limit", "remaining" + ); for h in snap { + let limit = h + .request_limit + .map_or_else(|| "-".to_string(), |value| value.to_string()); + let remaining = h + .remaining_requests + .map_or_else(|| "-".to_string(), |value| value.to_string()); println!( - "{:<16} {:<8} {:<6} {}", + "{:<16} {:<8} {:<6} {:<9} {:<9} {}", h.name, h.healthy, h.used, + limit, + remaining, h.home.display() ); } @@ -505,6 +593,19 @@ fn run_doctor(config: &Config) -> ExitCode { "additional_account_dirs: {} configured", config.additional_account_dirs.len() ); + println!( + "account_routing_strategy: {:?}", + config.account_routing_strategy + ); + println!("account_cooldown_secs : {}", config.account_cooldown_secs); + println!( + "account_request_limits : {:?}", + config.account_request_limits + ); + println!( + "session_affinity_ttl : {}", + config.session_affinity_ttl_secs + ); println!( "admin_key : {}", if config.admin_key.is_some() { @@ -522,13 +623,12 @@ fn run_doctor(config: &Config) -> ExitCode { } ); - // Probe credentials. Use the OAuth provider so the report covers every - // candidate file the router actually reads, including the dotfile - // `.credentials.json` and the nested `claudeAiOauth` layout. - let primary = OAuthProvider::new(&config.claude_code_home); + // Probe the active pool with its vendor-specific credential layout. + let (active_provider, primary_home) = subscription_pool(config); + let primary = SubscriptionReader::new(active_provider, &primary_home); match primary.discover_credential_path() { Some(path) => { - let readable = primary.get_token().is_ok(); + let readable = primary.read_token().is_ok(); println!( "primary credentials : {} ({})", path.display(), @@ -541,14 +641,12 @@ fn run_doctor(config: &Config) -> ExitCode { } None => println!( "primary credentials : {} (MISSING)", - std::path::Path::new(&config.claude_code_home) - .join(".credentials.json") - .display() + primary_home.display() ), } for (i, dir) in config.additional_account_dirs.iter().enumerate() { - let provider = OAuthProvider::new(&dir.to_string_lossy()); - match provider.discover_credential_path() { + let reader = SubscriptionReader::new(active_provider, dir); + match reader.discover_credential_path() { Some(path) => println!( "extra account {} : {} (found)", i + 1, @@ -557,7 +655,7 @@ fn run_doctor(config: &Config) -> ExitCode { None => println!( "extra account {} : {} (MISSING)", i + 1, - dir.join(".credentials.json").display() + dir.display() ), } } @@ -567,9 +665,8 @@ fn run_doctor(config: &Config) -> ExitCode { // when the matching upstream provider is active. let user_home = std::env::var("HOME").unwrap_or_else(|_| ".".to_string()); for provider in link_assistant_router::subscription::SubscriptionProvider::ALL { - use link_assistant_router::subscription::SubscriptionProvider; - if provider == SubscriptionProvider::Claude { - continue; // covered by the primary Claude probe above + if provider == active_provider { + continue; // covered by the active primary probe above } let reader = link_assistant_router::subscription::SubscriptionReader::from_user_home( provider, &user_home, diff --git a/src/proxy.rs b/src/proxy.rs index 7dc221b..6e04e6f 100644 --- a/src/proxy.rs +++ b/src/proxy.rs @@ -19,61 +19,14 @@ use axum::http::{HeaderMap, HeaderValue, StatusCode}; use axum::response::{IntoResponse, Response}; use futures_util::StreamExt; use log_lazy::LogLazy; -use reqwest::Client; use std::collections::BTreeMap; -use std::sync::Arc; -use crate::accounts::AccountRouter; +use crate::accounts::RoutingContext; +pub use crate::app_state::AppState; use crate::config::UpstreamProvider; -use crate::gonka::GonkaConfig; -use crate::oauth::OAuthProvider; use crate::openai; -use crate::providers::{OpenAICompatibleConfig, ProviderStore}; +pub(crate) use crate::request_routing::{request_routing_context, retry_after_duration}; use crate::responses; -use crate::token::TokenManager; - -/// Shared application state accessible by all route handlers. -#[derive(Clone)] -pub struct AppState { - /// HTTP client for upstream requests. - pub client: Client, - /// Token manager for validating custom tokens. - pub token_manager: TokenManager, - /// OAuth provider for obtaining upstream credentials (legacy single-account). - pub oauth_provider: OAuthProvider, - /// Multi-account router (when configured). When `None`, the legacy - /// `oauth_provider` is used directly. - pub account_router: Option, - /// Subscription credential reader for vendor OAuth providers - /// (Codex/Gemini/Qwen). `None` for non-subscription upstreams. - pub subscription_reader: Option, - /// In-memory cache of refreshed subscription tokens (Codex/Gemini/Qwen). - pub subscription_cache: std::sync::Arc, - /// Base URL for the upstream Anthropic API. - pub upstream_base_url: String, - /// Selected upstream inference provider. - pub upstream_provider: UpstreamProvider, - /// Gonka provider configuration when selected. - pub gonka: Option, - /// Crater `ForgeFed` task provider when selected. - pub crater: Option>, - /// Boot-time generic OpenAI-compatible provider config. - pub openai_compatible: OpenAICompatibleConfig, - /// Persisted provider records with encrypted upstream secrets. - pub provider_store: ProviderStore, - /// Lazy logger for verbose output. - pub logger: LogLazy, - /// Optional admin key (Bearer) required for `/api/tokens` issuance. - pub admin_key: Option, - /// Live metrics counter handle. - pub metrics: Arc, - /// Public base URL for `ActivityPub` actor documents. - pub activitypub_actor_base_url: String, - /// Public key PEM advertised by the `ActivityPub` actor. - pub activitypub_public_key_pem: String, - /// Optional MPP charge settings for OpenAI-compatible endpoints. - pub mpp: crate::mpp::MppConfig, -} /// The legacy API path prefix used to route requests through the proxy. pub const API_PREFIX: &str = "/api/latest/anthropic/"; @@ -154,6 +107,7 @@ pub(crate) fn extract_client_token(headers: &HeaderMap) -> Option<&str> { pub async fn proxy_handler(State(state): State, req: Request) -> impl IntoResponse { let path = req.uri().path().to_string(); let method = req.method().clone(); + let incoming_headers = req.headers().clone(); state.logger.verbose(|| format!("Incoming {method} {path}")); @@ -178,14 +132,14 @@ pub async fn proxy_handler(State(state): State, req: Request) -> impl }; // Log session tracking header if present - if let Some(session_id) = req.headers().get("x-claude-code-session-id") { + if let Some(session_id) = incoming_headers.get("x-claude-code-session-id") { state .logger .verbose(|| format!("Session: {}", session_id.to_str().unwrap_or(""))); } // Extract and validate the bearer token from the Authorization header - let Some(token) = extract_client_token(req.headers()) else { + let Some(token) = extract_client_token(&incoming_headers) else { state.logger.debug(|| "Missing Authorization header"); return error_response( StatusCode::UNAUTHORIZED, @@ -224,23 +178,8 @@ pub async fn proxy_handler(State(state): State, req: Request) -> impl ); } - // Get the real OAuth token (multi-account aware). - let (oauth_token, selected_account) = match resolve_upstream_credentials(&state) { - Ok(pair) => pair, - Err(e) => { - tracing::error!("Failed to resolve upstream credentials: {e}"); - return error_response( - StatusCode::BAD_GATEWAY, - "api_error", - "Upstream authentication unavailable", - ); - } - }; - - // Build upstream headers - let upstream_headers = build_upstream_headers(req.headers(), &oauth_token, &state.logger); - - // Read the request body + // Read the body before account selection so the router gets a copy of + // stable request metadata and can preserve conversation affinity. let body_bytes = match axum::body::to_bytes(req.into_body(), 10 * 1024 * 1024).await { Ok(bytes) => bytes, Err(e) => { @@ -251,6 +190,35 @@ pub async fn proxy_handler(State(state): State, req: Request) -> impl ); } }; + let routing_body = serde_json::from_slice(&body_bytes).unwrap_or(serde_json::Value::Null); + let pinned_account = match state.token_manager.account_for(&claims.sub) { + Ok(account) => account, + Err(error) => { + return error_response( + StatusCode::INTERNAL_SERVER_ERROR, + "api_error", + &format!("failed to resolve token account binding: {error}"), + ); + } + }; + let routing_context = request_routing_context(&incoming_headers, &routing_body, pinned_account); + + // Get the real OAuth token (multi-account aware). + let (oauth_token, selected_account) = + match resolve_upstream_credentials(&state, &routing_context) { + Ok(pair) => pair, + Err(e) => { + tracing::error!("Failed to resolve upstream credentials: {e}"); + return error_response( + StatusCode::BAD_GATEWAY, + "api_error", + "Upstream authentication unavailable", + ); + } + }; + + // Build upstream headers + let upstream_headers = build_upstream_headers(&incoming_headers, &oauth_token, &state.logger); state.logger.verbose(|| { format!( @@ -280,6 +248,7 @@ pub async fn proxy_handler(State(state): State, req: Request) -> impl let status = StatusCode::from_u16(upstream_resp.status().as_u16()) .unwrap_or(StatusCode::INTERNAL_SERVER_ERROR); + let retry_after = retry_after_duration(upstream_resp.headers()); state .logger @@ -295,7 +264,7 @@ pub async fn proxy_handler(State(state): State, req: Request) -> impl if let (Some(router), Some(name)) = (state.account_router.as_ref(), selected_account.as_deref()) { - router.report_failure(name, "upstream returned 429"); + router.report_failure_with_retry_after(name, "upstream returned 429", retry_after); } } @@ -334,6 +303,9 @@ pub async fn proxy_handler(State(state): State, req: Request) -> impl /// - `/api/latest/anthropic/*` -> `/*` (legacy) #[must_use] pub fn resolve_upstream_path(path: &str) -> String { + if let Some(rest) = path.strip_prefix("/api/anthropic") { + return rest.to_string(); + } // Legacy prefix: strip and forward if let Some(rest) = path.strip_prefix("/api/latest/anthropic") { return rest.to_string(); @@ -410,9 +382,10 @@ pub(crate) fn build_upstream_headers( /// router; otherwise we fall back to the single-account legacy provider. fn resolve_upstream_credentials( state: &AppState, + context: &RoutingContext, ) -> Result<(String, Option), Box> { if let Some(router) = state.account_router.as_ref() { - let sel = router.select()?; + let sel = router.select_with_context(context)?; return Ok((sel.token, Some(sel.name))); } let token = state.oauth_provider.get_token()?; @@ -483,10 +456,12 @@ pub async fn openai_chat_completions( .await; } if state.upstream_provider == UpstreamProvider::Qwen { + let routing_body = body.clone(); return crate::subscription_proxy::forward_subscription_openai( &state, &headers, body, + &routing_body, "/v1/chat/completions", crate::metrics::Surface::OpenAIChat, ) @@ -503,11 +478,13 @@ pub async fn openai_chat_completions( &state, &headers, responses_body, + &body, "/v1/responses", crate::metrics::Surface::OpenAIChat, ) .await; } + let routing_body = body.clone(); let req = match serde_json::from_value::(body) { Ok(req) => req, Err(e) => { @@ -518,15 +495,14 @@ pub async fn openai_chat_completions( ); } }; - let requested_model = req.model.clone(); let stream_requested = req.stream.unwrap_or(false) || stream_from_query; let body = openai::chat_completion_to_anthropic(&req); forward_openai( &state, &headers, body, + &routing_body, crate::metrics::Surface::OpenAIChat, - &requested_model, stream_requested, OpenAIShape::Chat, ) @@ -569,10 +545,12 @@ pub async fn openai_responses( state.upstream_provider, UpstreamProvider::Codex | UpstreamProvider::Qwen ) { + let routing_body = body.clone(); return crate::subscription_proxy::forward_subscription_openai( &state, &headers, body, + &routing_body, "/v1/responses", crate::metrics::Surface::OpenAIResponses, ) @@ -581,6 +559,7 @@ pub async fn openai_responses( if state.upstream_provider == UpstreamProvider::Gemini { return crate::gemini::forward_responses(&state, &headers, body).await; } + let routing_body = body.clone(); let req = match serde_json::from_value::(body) { Ok(req) => req, Err(e) => { @@ -591,15 +570,14 @@ pub async fn openai_responses( ); } }; - let requested_model = req.model.clone(); let stream_requested = req.stream.unwrap_or(false); let body = responses::response_to_anthropic(&req); forward_openai( &state, &headers, body, + &routing_body, crate::metrics::Surface::OpenAIResponses, - &requested_model, stream_requested, OpenAIShape::Response, ) @@ -736,11 +714,15 @@ async fn forward_openai( state: &AppState, headers: &HeaderMap, body: serde_json::Value, + routing_body: &serde_json::Value, surface: crate::metrics::Surface, - requested_model: &str, stream_requested: bool, shape: OpenAIShape, ) -> Response { + let requested_model = routing_body + .get("model") + .and_then(serde_json::Value::as_str) + .unwrap_or_default(); let path = match shape { OpenAIShape::Chat => "/v1/chat/completions", OpenAIShape::Response => "/v1/responses", @@ -775,18 +757,31 @@ async fn forward_openai( ); } - // Resolve OAuth credentials. - let (oauth_token, selected_account) = match resolve_upstream_credentials(state) { - Ok(p) => p, - Err(e) => { - tracing::error!("openai: upstream credentials unavailable: {e}"); + let pinned_account = match state.token_manager.account_for(&claims.sub) { + Ok(account) => account, + Err(error) => { return error_response( - StatusCode::BAD_GATEWAY, + StatusCode::INTERNAL_SERVER_ERROR, "api_error", - "Upstream authentication unavailable", + &format!("failed to resolve token account binding: {error}"), ); } }; + let routing_context = request_routing_context(headers, routing_body, pinned_account); + + // Resolve OAuth credentials. + let (oauth_token, selected_account) = + match resolve_upstream_credentials(state, &routing_context) { + Ok(p) => p, + Err(e) => { + tracing::error!("openai: upstream credentials unavailable: {e}"); + return error_response( + StatusCode::BAD_GATEWAY, + "api_error", + "Upstream authentication unavailable", + ); + } + }; let upstream_url = format!( "{}/v1/messages", @@ -829,6 +824,7 @@ async fn forward_openai( } }; let upstream_status = upstream_resp.status(); + let retry_after = retry_after_duration(upstream_resp.headers()); if stream_requested && upstream_status.is_success() { state .metrics @@ -873,7 +869,7 @@ async fn forward_openai( if let (Some(router), Some(name)) = (state.account_router.as_ref(), selected_account.as_deref()) { - router.report_failure(name, "upstream returned 429"); + router.report_failure_with_retry_after(name, "upstream returned 429", retry_after); } } state.metrics.record_request( @@ -971,6 +967,8 @@ pub async fn accounts_endpoint(State(state): State) -> impl IntoRespon "home": h.home.display().to_string(), "healthy": h.healthy, "used": h.used, + "request_limit": h.request_limit, + "remaining_requests": h.remaining_requests, "last_error": h.last_error, "cooldown_remaining_seconds": h.cooldown_remaining.map(|d| d.as_secs()), }) diff --git a/src/proxy_tests.rs b/src/proxy_tests.rs index e1363f0..0cc4606 100644 --- a/src/proxy_tests.rs +++ b/src/proxy_tests.rs @@ -3,6 +3,7 @@ use log_lazy::{LogLazy, levels}; use crate::proxy::{ OAUTH_BETA_FLAG, build_upstream_headers, extract_client_token, merge_oauth_beta, + request_routing_context, retry_after_duration, }; #[test] @@ -95,3 +96,53 @@ fn merge_oauth_beta_is_idempotent_and_dedups() { let multi = format!("foo,{OAUTH_BETA_FLAG},bar"); assert_eq!(merge_oauth_beta(Some(&multi)), multi); } + +#[test] +fn routing_context_prefers_token_pin_and_detects_sessions() { + let mut headers = HeaderMap::new(); + headers.insert( + "x-claude-code-session-id", + HeaderValue::from_static("header-session"), + ); + let body = serde_json::json!({"metadata": {"session_id": "body-session"}}); + + let context = request_routing_context(&headers, &body, Some("account-3".into())); + + assert_eq!(context.pinned_account.as_deref(), Some("account-3")); + assert_eq!(context.session_key.as_deref(), Some("header-session")); +} + +#[test] +fn routing_context_falls_back_to_standard_body_session_fields() { + let headers = HeaderMap::new(); + let body = serde_json::json!({"metadata": {"session_id": "body-session"}}); + + let context = request_routing_context(&headers, &body, None); + + assert_eq!(context.session_key.as_deref(), Some("body-session")); +} + +#[test] +fn retry_after_delta_seconds_is_used_for_account_cooldown() { + let mut headers = HeaderMap::new(); + headers.insert("retry-after", HeaderValue::from_static("120")); + + assert_eq!( + retry_after_duration(&headers), + Some(std::time::Duration::from_secs(120)) + ); +} + +#[test] +fn retry_after_http_date_is_used_for_account_cooldown() { + let retry_at = chrono::Utc::now() + chrono::Duration::seconds(120); + let mut headers = HeaderMap::new(); + headers.insert( + "retry-after", + HeaderValue::from_str(&retry_at.to_rfc2822()).unwrap(), + ); + + let parsed = retry_after_duration(&headers).unwrap(); + assert!(parsed >= std::time::Duration::from_secs(118)); + assert!(parsed <= std::time::Duration::from_secs(120)); +} diff --git a/src/refresh.rs b/src/refresh.rs index b761d37..883e049 100644 --- a/src/refresh.rs +++ b/src/refresh.rs @@ -241,13 +241,15 @@ pub async fn refresh( .ok_or_else(|| RefreshError::Parse("response contained no access_token".to_string())) } -/// Process-wide cache of refreshed subscription tokens, keyed by provider. +/// Process-wide cache of refreshed subscription tokens, keyed by provider and +/// account. Two subscriptions for the same vendor must never reuse each +/// other's bearer token. /// /// Holds only in-memory copies obtained via OAuth refresh; vendor credential /// files on disk are never modified. #[derive(Debug, Default)] pub struct TokenCache { - inner: Mutex>, + inner: Mutex>, } impl TokenCache { @@ -273,16 +275,29 @@ impl TokenCache { provider: SubscriptionProvider, disk_token: SubscriptionToken, now_ms: i64, + ) -> SubscriptionToken { + self.get_fresh_for(client, provider, "primary", disk_token, now_ms) + .await + } + + /// Account-scoped variant of [`Self::get_fresh`]. + pub async fn get_fresh_for( + &self, + client: &reqwest::Client, + provider: SubscriptionProvider, + account: &str, + disk_token: SubscriptionToken, + now_ms: i64, ) -> SubscriptionToken { if !disk_token.is_expired(now_ms) { return disk_token; } - if let Some(cached) = self.cached_valid(provider, now_ms) { + if let Some(cached) = self.cached_valid_for(provider, account, now_ms) { return cached; } match refresh(client, provider, &disk_token, now_ms).await { Ok(fresh) => { - self.store(provider, fresh.clone()); + self.store_for(provider, account, fresh.clone()); fresh } Err(e) => { @@ -292,21 +307,22 @@ impl TokenCache { } } - fn cached_valid( + fn cached_valid_for( &self, provider: SubscriptionProvider, + account: &str, now_ms: i64, ) -> Option { let guard = self.inner.lock().ok()?; guard - .get(&provider) + .get(&(provider, account.to_string())) .filter(|token| !token.is_expired(now_ms)) .cloned() } - fn store(&self, provider: SubscriptionProvider, token: SubscriptionToken) { + fn store_for(&self, provider: SubscriptionProvider, account: &str, token: SubscriptionToken) { if let Ok(mut guard) = self.inner.lock() { - guard.insert(provider, token); + guard.insert((provider, account.to_string()), token); } } } @@ -401,11 +417,38 @@ mod tests { account_id: None, resource_url: None, }; - cache.store(SubscriptionProvider::Qwen, cached); + cache.store_for(SubscriptionProvider::Qwen, "primary", cached); let expired_disk = token(Some("r1"), Some(0)); let out = cache .get_fresh(&client, SubscriptionProvider::Qwen, expired_disk, 1_000) .await; assert_eq!(out.access_token, "cached-access"); } + + #[test] + fn refreshed_tokens_are_isolated_by_account() { + let cache = TokenCache::new(); + let mut first = token(Some("refresh-a"), Some(10_000)); + first.access_token = "access-a".into(); + let mut second = token(Some("refresh-b"), Some(10_000)); + second.access_token = "access-b".into(); + + cache.store_for(SubscriptionProvider::Qwen, "primary", first); + cache.store_for(SubscriptionProvider::Qwen, "account-1", second); + + assert_eq!( + cache + .cached_valid_for(SubscriptionProvider::Qwen, "primary", 1_000) + .unwrap() + .access_token, + "access-a" + ); + assert_eq!( + cache + .cached_valid_for(SubscriptionProvider::Qwen, "account-1", 1_000) + .unwrap() + .access_token, + "access-b" + ); + } } diff --git a/src/request_routing.rs b/src/request_routing.rs new file mode 100644 index 0000000..ec47e81 --- /dev/null +++ b/src/request_routing.rs @@ -0,0 +1,62 @@ +//! Provider-neutral extraction of stable account-routing signals. + +use std::time::Duration; + +use axum::http::HeaderMap; + +use crate::accounts::RoutingContext; + +/// Copy only stable routing signals from a request. Header signals take +/// precedence over JSON metadata and the caller token's account binding is +/// carried separately as a strict pin. +pub fn request_routing_context( + headers: &HeaderMap, + body: &serde_json::Value, + pinned_account: Option, +) -> RoutingContext { + const SESSION_HEADERS: [&str; 4] = [ + "x-claude-code-session-id", + "x-codex-session-id", + "x-session-id", + "session-id", + ]; + let header_session = SESSION_HEADERS.iter().find_map(|name| { + headers + .get(*name) + .and_then(|value| value.to_str().ok()) + .filter(|value| !value.is_empty()) + .map(ToString::to_string) + }); + let body_session = [ + "/metadata/session_id", + "/metadata/conversation_id", + "/session_id", + "/conversation_id", + "/prompt_cache_key", + ] + .iter() + .find_map(|pointer| { + body.pointer(pointer) + .and_then(serde_json::Value::as_str) + .filter(|value| !value.is_empty()) + .map(ToString::to_string) + }); + RoutingContext { + session_key: header_session.or(body_session), + pinned_account, + } +} + +/// Parse the standard `Retry-After` delta-seconds or HTTP-date forms. +pub fn retry_after_duration(headers: &HeaderMap) -> Option { + let value = headers.get("retry-after")?.to_str().ok()?.trim(); + if let Ok(seconds) = value.parse::() { + return Some(Duration::from_secs(seconds)); + } + let retry_at = chrono::DateTime::parse_from_rfc2822(value).ok()?; + let seconds = retry_at + .signed_duration_since(chrono::Utc::now()) + .num_seconds() + .max(0); + Some(Duration::from_secs(u64::try_from(seconds).ok()?)) +} diff --git a/src/subscription.rs b/src/subscription.rs index dd9b10d..3e13ab2 100644 --- a/src/subscription.rs +++ b/src/subscription.rs @@ -83,7 +83,15 @@ impl SubscriptionProvider { #[must_use] pub const fn credential_filenames(self) -> &'static [&'static str] { match self { - Self::Claude => &[".credentials.json", "credentials.json", "auth.json"], + // Keep parity with the legacy OAuthProvider search order so + // enabling a pool does not make an existing Claude login vanish. + Self::Claude => &[ + "credentials.json", + ".credentials.json", + "auth.json", + "oauth.json", + "config.json", + ], Self::Codex => &["auth.json"], Self::Gemini | Self::Qwen => &["oauth_creds.json"], } @@ -556,6 +564,17 @@ mod tests { assert_eq!(token.expires_at_ms, Some(9_999_999_999_999)); } + #[test] + fn claude_pool_reader_preserves_legacy_credential_candidates() { + let dir = tempdir(); + fs::write(dir.join("oauth.json"), r#"{"accessToken":"legacy"}"#).unwrap(); + let token = SubscriptionReader::new(SubscriptionProvider::Claude, &dir) + .read_token() + .unwrap(); + + assert_eq!(token.access_token, "legacy"); + } + #[test] fn missing_credentials_errors() { let reader = SubscriptionReader::new( diff --git a/src/subscription_proxy.rs b/src/subscription_proxy.rs index 0379570..1c2897c 100644 --- a/src/subscription_proxy.rs +++ b/src/subscription_proxy.rs @@ -20,7 +20,10 @@ use futures_util::StreamExt; use crate::config::UpstreamProvider; use crate::metrics::Surface; -use crate::proxy::{AppState, error_response, extract_client_token, maybe_mpp_challenge}; +use crate::proxy::{ + AppState, error_response, extract_client_token, maybe_mpp_challenge, request_routing_context, + retry_after_duration, +}; use crate::subscription::{SubscriptionProvider, SubscriptionToken}; /// Forward one `OpenAI`-shaped request to the active subscription upstream. @@ -31,6 +34,7 @@ pub async fn forward_subscription_openai( state: &AppState, headers: &HeaderMap, mut body: serde_json::Value, + routing_body: &serde_json::Value, path: &str, surface: Surface, ) -> Response { @@ -45,12 +49,22 @@ pub async fn forward_subscription_openai( "Missing Authorization Bearer token or x-api-key", ); }; - if let Err(e) = state.token_manager.validate_token(token) { - let status = match &e { - crate::token::TokenError::Revoked => StatusCode::FORBIDDEN, - _ => StatusCode::UNAUTHORIZED, - }; - return error_response(status, "authentication_error", &format!("{e}")); + let claims = match state.token_manager.validate_token(token) { + Ok(claims) => claims, + Err(e) => { + let status = match &e { + crate::token::TokenError::Revoked => StatusCode::FORBIDDEN, + _ => StatusCode::UNAUTHORIZED, + }; + return error_response(status, "authentication_error", &format!("{e}")); + } + }; + if let Err(e) = state.token_manager.enforce_request_budget(&claims.sub) { + return error_response( + StatusCode::TOO_MANY_REQUESTS, + "rate_limit_error", + &format!("{e}"), + ); } let Some(provider) = state.upstream_provider.subscription_provider() else { @@ -60,21 +74,49 @@ pub async fn forward_subscription_openai( "active upstream is not a subscription provider", ); }; - let Some(reader) = state.subscription_reader.as_ref() else { - return error_response( - StatusCode::INTERNAL_SERVER_ERROR, - "api_error", - "subscription credentials reader is not configured", - ); + let pinned_account = match state.token_manager.account_for(&claims.sub) { + Ok(account) => account, + Err(error) => { + return error_response( + StatusCode::INTERNAL_SERVER_ERROR, + "api_error", + &format!("failed to resolve token account binding: {error}"), + ); + } }; - let disk_token = match reader.read_token() { - Ok(token) => token, - Err(e) => { + let routing_context = request_routing_context(headers, routing_body, pinned_account); + let selected = if let Some(router) = state.account_router.as_ref() { + match router.select_subscription(&routing_context) { + Ok(selected) => selected, + Err(error) => { + return error_response( + StatusCode::SERVICE_UNAVAILABLE, + "account_unavailable", + &error.to_string(), + ); + } + } + } else { + let Some(reader) = state.subscription_reader.as_ref() else { return error_response( - StatusCode::BAD_GATEWAY, - "authentication_error", - &format!("failed to read {provider} subscription credentials: {e}"), + StatusCode::INTERNAL_SERVER_ERROR, + "api_error", + "subscription credentials reader is not configured", ); + }; + let disk_token = match reader.read_token() { + Ok(token) => token, + Err(e) => { + return error_response( + StatusCode::BAD_GATEWAY, + "authentication_error", + &format!("failed to read {provider} subscription credentials: {e}"), + ); + } + }; + crate::accounts::SelectedSubscriptionAccount { + name: "primary".to_string(), + token: disk_token, } }; // Refresh in memory if the on-disk token has expired; vendor files stay @@ -82,8 +124,15 @@ pub async fn forward_subscription_openai( let now_ms = chrono::Utc::now().timestamp_millis(); let sub_token = state .subscription_cache - .get_fresh(&state.client, provider, disk_token, now_ms) + .get_fresh_for( + &state.client, + provider, + &selected.name, + selected.token, + now_ms, + ) .await; + let selected_account = Some(selected.name); let stream_requested = body .get("stream") @@ -127,7 +176,9 @@ pub async fn forward_subscription_openai( let upstream_resp = match upstream_req.send().await { Ok(resp) => resp, Err(e) => { - state.metrics.record_request(surface, 502, None); + state + .metrics + .record_request(surface, 502, selected_account.as_deref()); return error_response( StatusCode::BAD_GATEWAY, "api_error", @@ -137,7 +188,21 @@ pub async fn forward_subscription_openai( }; let status = StatusCode::from_u16(upstream_resp.status().as_u16()) .unwrap_or(StatusCode::INTERNAL_SERVER_ERROR); - state.metrics.record_request(surface, status.as_u16(), None); + state + .metrics + .record_request(surface, status.as_u16(), selected_account.as_deref()); + let retry_after = retry_after_duration(upstream_resp.headers()); + if status == StatusCode::TOO_MANY_REQUESTS { + if let (Some(router), Some(account)) = + (state.account_router.as_ref(), selected_account.as_deref()) + { + router.report_failure_with_retry_after( + account, + "subscription upstream returned 429", + retry_after, + ); + } + } let content_type = upstream_resp .headers() @@ -172,7 +237,9 @@ pub async fn forward_subscription_openai( let upstream_body = match upstream_resp.bytes().await { Ok(bytes) => bytes, Err(e) => { - state.metrics.record_request(surface, 502, None); + state + .metrics + .record_request(surface, 502, selected_account.as_deref()); return error_response( StatusCode::BAD_GATEWAY, "api_error", diff --git a/src/token.rs b/src/token.rs index 1a9a80d..b42271e 100644 --- a/src/token.rs +++ b/src/token.rs @@ -144,6 +144,14 @@ impl TokenManager { } } + /// Return the strict account binding stored for a router-issued token. + pub fn account_for(&self, token_id: &str) -> Result, TokenError> { + self.store + .get(token_id) + .map(|record| record.and_then(|record| record.account)) + .map_err(|error| TokenError::Storage(error.to_string())) + } + /// Validate a custom token string. /// /// Strips the `la_sk_` prefix, decodes the JWT, checks expiration and @@ -304,6 +312,18 @@ mod tests { assert!(labels.contains(&"two")); } + #[test] + fn account_binding_is_available_during_request_routing() { + let mgr = test_manager(); + let token = mgr.issue_token_for(1, "bound", Some("account-2")).unwrap(); + let claims = mgr.validate_token(&token).unwrap(); + + assert_eq!( + mgr.account_for(&claims.sub).unwrap().as_deref(), + Some("account-2") + ); + } + #[test] fn test_unlimited_token_never_hits_budget() { let mgr = test_manager(); diff --git a/tests/integration_test.rs b/tests/integration_test.rs index 486752e..7e0c274 100644 --- a/tests/integration_test.rs +++ b/tests/integration_test.rs @@ -162,6 +162,14 @@ mod path_routing_tests { assert_eq!(resolve_upstream_path("/v1/messages"), "/v1/messages"); } + #[test] + fn test_namespaced_anthropic_messages_path() { + assert_eq!( + resolve_upstream_path("/api/anthropic/v1/messages"), + "/v1/messages" + ); + } + #[test] fn test_anthropic_count_tokens_path() { assert_eq!( @@ -359,6 +367,10 @@ mod config_verbose_tests { enable_anthropic_api: true, enable_metrics: true, additional_account_dirs: vec![], + account_routing_strategy: link_assistant_router::accounts::SelectionStrategy::default(), + account_cooldown_secs: 60, + session_affinity_ttl_secs: 3600, + account_request_limits: vec![], experimental_compatibility: false, admin_key: None, mpp: link_assistant_router::config::default_mpp_config(),