forge_gateway/
lib.rs

1//! `forge-gateway` - the account pool: one backend per
2//! `forge.toml` `provider` token, plus account selection, health,
3//! probe scheduling, backoff, and the inference listener spawned
4//! sessions route through.
5//!
6//! Each [`ProviderBackend`] owns credential resolution, the probe
7//! request and its payload mapping, the billing shape, and what repair
8//! a failed probe allows for one provider. The [`ProviderHost`] port,
9//! implemented by forge-agent, is the only process plumbing a backend
10//! may reach, so this crate stays HTTP + mapping and is testable
11//! offline.
12//!
13//! [`AccountStateMap`] holds each account's health, usage snapshot and
14//! probe schedule; [`Gateway::select_for`] is the one place the account
15//! serving a session's model is decided.
16
17pub mod account;
18mod anthropic;
19pub mod binding;
20mod codex;
21pub mod forward;
22pub mod helpers;
23pub mod listener;
24mod openrouter;
25pub mod pool;
26pub mod rotation;
27pub mod selection;
28pub mod splice;
29mod zai;
30
31use std::collections::HashMap;
32use std::time::Duration;
33
34use async_trait::async_trait;
35pub use forge_primitives::account::Provider;
36pub use forge_primitives::usage::AccountBudget;
37pub use forge_primitives::usage::UsageSnapshot;
38pub use forge_primitives::usage::UsageSourceKind;
39pub use forge_primitives::usage::oauth::OauthUsageError;
40
41pub use crate::account::{
42    AccountKey, AccountState, AccountStateMap, LoadingState, Unusable, UsageFetchStatus,
43};
44pub use crate::anthropic::{Anthropic, CLAUDE_CODE_OAUTH_TOKEN_ENV, token_bearer};
45pub use crate::codex::Codex;
46pub use crate::forward::{Gateway, SelectFailure};
47pub use crate::openrouter::Openrouter;
48pub use crate::pool::AccountPool;
49pub use crate::zai::Zai;
50
51/// Everything a backend may read about one account. `env` is the
52/// merged global `[env]` + `[accounts.env]` block; the merge happens
53/// in forge-workspace config load and stays there.
54pub struct AccountEnv<'a> {
55    pub env: &'a HashMap<String, String>,
56}
57
58/// Billing shape: plan windows that reset (a percentage of an
59/// allowance) or per-key spend over a period.
60#[derive(Clone, Copy, Debug, PartialEq, Eq)]
61pub enum BillingModel {
62    Windows,
63    Spend,
64}
65
66/// Why a probe did not produce a snapshot: a credential miss, a wire
67/// error class from the round-trip, or a 200 whose body maps to
68/// nothing.
69#[derive(Debug, thiserror::Error)]
70pub enum ProbeError {
71    #[error("no CLAUDE_CODE_OAUTH_TOKEN in the account env")]
72    NoCredentials,
73    #[error(transparent)]
74    Fetch(#[from] OauthUsageError),
75    #[error("probe returned 200 but the body maps to nothing: {0}")]
76    Unmappable(String),
77}
78
79impl ProbeError {
80    /// The plan windows a rate-limited probe still carried, if any: a 429
81    /// is a reading, and it is the one that says which window is spent.
82    pub fn rate_limited_snapshot(&self) -> Option<&UsageSnapshot> {
83        match self {
84            Self::Fetch(OauthUsageError::RateLimited { snapshot, .. }) => snapshot.as_deref(),
85            Self::NoCredentials | Self::Fetch(_) | Self::Unmappable(_) => None,
86        }
87    }
88}
89
90/// What a backend allows after one of its probes failed, decided by
91/// [`ProviderBackend::repair`]. The boot loader executes the verdict
92/// against its state machine.
93#[derive(Debug, PartialEq, Eq)]
94pub enum RepairAction {
95    /// Auth failure forge cannot repair: the account bails and the
96    /// poller re-probes it later.
97    Terminal,
98    /// Transient failure: record it, back off, and retry. Carries the
99    /// server `Retry-After` for a 429.
100    Retry { retry_after: Option<Duration> },
101}
102
103/// The host port, implemented by forge-agent. The only process
104/// plumbing a backend may reach, so the crate stays HTTP + mapping and
105/// is testable offline.
106#[async_trait]
107pub trait ProviderHost: Send + Sync {
108    /// A reqwest client with the NODE_EXTRA_CA_CERTS roots applied
109    /// and the caller's timeout baked in.
110    fn http_client(&self, timeout: Duration) -> Result<reqwest::Client, String>;
111    /// The same client for a streamed response: the NODE_EXTRA_CA_CERTS
112    /// roots applied, `connect_timeout` bounding the connect and the TLS
113    /// handshake, and `idle_timeout` bounding each read. The idle bound
114    /// covers the wait for the response headers and then resets on every
115    /// chunk; a streamed inference response has no known length, so the
116    /// total deadline [`ProviderHost::http_client`] bakes in would cut a
117    /// healthy long turn, while an idle bound still ends one whose
118    /// upstream has stopped producing bytes.
119    fn streaming_http_client(
120        &self,
121        connect_timeout: Duration,
122        idle_timeout: Duration,
123    ) -> Result<reqwest::Client, String>;
124    /// `claude-code/<version>`, resolved by one `claude --version`
125    /// shell-out per process and cached, off the async runtime via
126    /// spawn_blocking. Err preserves the probe's UaProbe failure
127    /// class.
128    async fn user_agent(&self) -> Result<String, String>;
129}
130
131/// One backend per `forge.toml` `provider` token. The impls are
132/// stateless unit structs; per-account state arrives through
133/// [`AccountEnv`], never through `self`.
134#[async_trait]
135pub trait ProviderBackend: Send + Sync {
136    /// The token this backend serves; the registry keys on it.
137    fn token(&self) -> Provider;
138
139    /// One probe round-trip: credential resolution, URL derivation,
140    /// headers, status classification, payload-to-snapshot mapping.
141    /// Strict-vs-lenient mapping lives in the impls, not in a flag.
142    async fn probe(
143        &self,
144        account: &AccountEnv<'_>,
145        host: &dyn ProviderHost,
146    ) -> Result<UsageSnapshot, ProbeError>;
147
148    /// Billing shape.
149    fn billing(&self) -> BillingModel;
150
151    /// What this backend allows after a failed probe. The shared
152    /// default: the auth classes are terminal (an env credential is
153    /// fixed when the session spawns and there is no local credential
154    /// repair to fire) and
155    /// everything else retries. Backends whose provider repairs
156    /// differently override this.
157    fn repair(&self, _account: &AccountEnv<'_>, err: &ProbeError) -> RepairAction {
158        match err {
159            ProbeError::NoCredentials
160            | ProbeError::Fetch(
161                OauthUsageError::NoCredentials
162                | OauthUsageError::Expired
163                | OauthUsageError::Unauthorized(_),
164            ) => RepairAction::Terminal,
165            ProbeError::Fetch(OauthUsageError::RateLimited { retry_after, .. }) => {
166                RepairAction::Retry { retry_after: *retry_after }
167            }
168            ProbeError::Fetch(_) | ProbeError::Unmappable(_) => {
169                RepairAction::Retry { retry_after: None }
170            }
171        }
172    }
173
174    /// The source kind this backend's probe stamps on the snapshots
175    /// it returns; `budget` refuses a cached snapshot of any other
176    /// kind.
177    fn source(&self) -> UsageSourceKind;
178
179    /// The budget shape this backend renders for one account's cached
180    /// usage snapshot.
181    ///
182    /// A snapshot whose source is not this backend's is treated as
183    /// absent: the redb cache survives a `forge.toml` provider edit
184    /// and is re-seeded at every boot, so a stale row is normal, and
185    /// rendering windows as money (or the reverse) is worse than
186    /// saying nothing.
187    fn budget(&self, account: &str, snapshot: Option<&UsageSnapshot>) -> AccountBudget {
188        let unknown =
189            AccountBudget::Unknown { spend_billed: self.billing() == BillingModel::Spend };
190        let Some(snapshot) = snapshot else {
191            return unknown;
192        };
193        if snapshot.source != self.source() {
194            warn_unusable_snapshot(account, self.token(), snapshot.source);
195            return unknown;
196        }
197        match self.billing() {
198            BillingModel::Windows => AccountBudget::Subscription {
199                five_hour_util: snapshot.five_hour_util(),
200                seven_day_util: snapshot.seven_day_util(),
201                resets_at: snapshot.binding_reset_at(),
202            },
203            BillingModel::Spend => {
204                if let Some(spend) = snapshot.spend.as_ref() {
205                    AccountBudget::Api {
206                        daily: spend.daily,
207                        weekly: spend.weekly,
208                        monthly: spend.monthly,
209                    }
210                } else {
211                    // Unreachable from today's mapper, which refuses a
212                    // body with no figures - same warn as a source
213                    // mismatch rather than a second silent path.
214                    warn_unusable_snapshot(account, self.token(), snapshot.source);
215                    unknown
216                }
217            }
218        }
219    }
220}
221
222/// A cached snapshot that cannot be used under this backend.
223///
224/// Warns rather than falling back quietly: the redb row is rewritten
225/// only after a successful poll, so a stale one outlives a `forge.toml`
226/// edit and is re-seeded every boot. An account whose provider changed
227/// and whose new endpoint is failing would otherwise carry no figures
228/// indefinitely with nothing anywhere saying why.
229fn warn_unusable_snapshot(account: &str, provider: Provider, source: UsageSourceKind) {
230    tracing::warn!(
231        target: "forge_gateway",
232        account = %account,
233        provider = ?provider,
234        source = source.label(),
235        "cached usage snapshot does not fit the account's provider; showing no figures until a \
236         fresh probe lands",
237    );
238}
239
240static ANTHROPIC: Anthropic = Anthropic;
241static CODEX: Codex = Codex;
242static OPENROUTER: Openrouter = Openrouter;
243static ZAI: Zai = Zai;
244static BACKENDS: &[&dyn ProviderBackend] = &[&ANTHROPIC, &CODEX, &OPENROUTER, &ZAI];
245
246/// The backend registered for `token`. Every `Provider` variant
247/// resolves; the Option keeps the trait open for a token whose
248/// backend is not worth a probe.
249pub fn backend(token: Provider) -> Option<&'static dyn ProviderBackend> {
250    match token {
251        Provider::Anthropic => Some(&ANTHROPIC),
252        Provider::Codex => Some(&CODEX),
253        Provider::Openrouter => Some(&OPENROUTER),
254        Provider::Zai => Some(&ZAI),
255    }
256}
257
258/// Every registered backend.
259pub fn all() -> &'static [&'static dyn ProviderBackend] {
260    BACKENDS
261}
262
263#[cfg(test)]
264mod tests {
265    use super::*;
266
267    #[test]
268    fn registry_round_trips_every_registered_backend() {
269        for entry in all() {
270            let resolved = backend(entry.token());
271            assert_eq!(
272                resolved.map(ProviderBackend::token),
273                Some(entry.token()),
274                "backend() must resolve every backend all() lists",
275            );
276        }
277    }
278
279    /// Every token a user may write in `forge.toml` resolves to a
280    /// backend. A missing registration is invisible at the callers:
281    /// their fallbacks render Unknown forever, and the debug_assert
282    /// guarding them compiles out in release.
283    #[test]
284    fn backend_resolves_every_accepted_token() {
285        for token in Provider::ACCEPTED.split(',').map(str::trim) {
286            let provider: Provider =
287                serde_json::from_str(token).expect("an ACCEPTED token parses to a Provider");
288            assert!(
289                backend(provider).is_some(),
290                "{provider:?} is accepted by config but has no backend registered",
291            );
292        }
293    }
294
295    #[test]
296    fn anthropic_backend_is_windowed() {
297        let backend = backend(Provider::Anthropic).expect("registered");
298        assert_eq!(backend.token(), Provider::Anthropic);
299        assert_eq!(backend.billing(), BillingModel::Windows);
300        assert_eq!(backend.source(), UsageSourceKind::Oauth);
301    }
302
303    #[test]
304    fn codex_backend_is_windowed() {
305        let backend = backend(Provider::Codex).expect("registered");
306        assert_eq!(backend.token(), Provider::Codex);
307        assert_eq!(backend.billing(), BillingModel::Windows);
308        assert_eq!(backend.source(), UsageSourceKind::Oauth);
309    }
310
311    #[test]
312    fn openrouter_backend_is_spend() {
313        let backend = backend(Provider::Openrouter).expect("registered");
314        assert_eq!(backend.token(), Provider::Openrouter);
315        assert_eq!(backend.billing(), BillingModel::Spend);
316        assert_eq!(backend.source(), UsageSourceKind::OpenRouterKey);
317    }
318
319    #[test]
320    fn zai_backend_is_windowed() {
321        let backend = backend(Provider::Zai).expect("registered");
322        assert_eq!(backend.token(), Provider::Zai);
323        assert_eq!(backend.billing(), BillingModel::Windows);
324        assert_eq!(backend.source(), UsageSourceKind::ZaiMonitor);
325    }
326
327    fn repair_account(env: &HashMap<String, String>) -> AccountEnv<'_> {
328        AccountEnv { env }
329    }
330
331    /// The repair pin: every auth failure is terminal. forge has no
332    /// credential repair of its own to fire - an account whose token is
333    /// dead bails and waits for the pollers, and nothing spawns a
334    /// `claude` subprocess to rotate anything.
335    #[test]
336    fn repair_terminals_every_auth_failure() {
337        let empty = HashMap::new();
338        let account = repair_account(&empty);
339        let auth_classes = [
340            ProbeError::NoCredentials,
341            ProbeError::Fetch(OauthUsageError::Expired),
342            ProbeError::Fetch(OauthUsageError::Unauthorized(401)),
343        ];
344        for provider in [Provider::Anthropic, Provider::Codex, Provider::Openrouter, Provider::Zai]
345        {
346            let resolved = backend(provider).expect("registered");
347            for err in &auth_classes {
348                assert_eq!(
349                    resolved.repair(&account, err),
350                    RepairAction::Terminal,
351                    "{provider:?}; {err:?} must not fire a credential repair",
352                );
353            }
354        }
355    }
356
357    /// Every non-auth class retries: the network class,
358    /// a 429 carrying the server Retry-After, the scope refusal (the
359    /// token arms convert it before classification) and a 200 that
360    /// maps to nothing - mapping drift is not an auth failure.
361    #[test]
362    fn repair_retries_every_transient_class() {
363        let empty = HashMap::new();
364        let account = repair_account(&empty);
365        for provider in [Provider::Anthropic, Provider::Codex] {
366            let resolved = backend(provider).expect("registered");
367            assert_eq!(
368                resolved.repair(
369                    &account,
370                    &ProbeError::Fetch(OauthUsageError::Network("dns".to_owned()))
371                ),
372                RepairAction::Retry { retry_after: None },
373                "{provider:?}",
374            );
375        }
376
377        let anthropic = backend(Provider::Anthropic).expect("registered");
378        let retry_after = Some(Duration::from_secs(60));
379        assert_eq!(
380            anthropic.repair(
381                &account,
382                &ProbeError::Fetch(OauthUsageError::RateLimited { retry_after, snapshot: None }),
383            ),
384            RepairAction::Retry { retry_after },
385        );
386        assert_eq!(
387            anthropic.repair(&account, &ProbeError::Fetch(OauthUsageError::ScopeInsufficient)),
388            RepairAction::Retry { retry_after: None },
389        );
390        assert_eq!(
391            anthropic.repair(&account, &ProbeError::Unmappable("no window".to_owned())),
392            RepairAction::Retry { retry_after: None },
393        );
394    }
395
396    fn budget_snapshot(
397        future: std::time::SystemTime,
398        source: UsageSourceKind,
399        spend: Option<forge_primitives::usage::ApiSpend>,
400    ) -> UsageSnapshot {
401        let window = |utilization| forge_primitives::usage::UsageWindow {
402            utilization,
403            resets_at: Some(future),
404            reset_description: None,
405            status: None,
406        };
407        UsageSnapshot {
408            source,
409            fetched_at: std::time::SystemTime::UNIX_EPOCH,
410            five_hour: Some(window(100.0)),
411            seven_day: Some(window(20.0)),
412            seven_day_opus: None,
413            seven_day_sonnet: None,
414            extra_usage: None,
415            balance: None,
416            spend,
417        }
418    }
419
420    /// Every `(backend, source)` pair `budget` can be handed. The
421    /// mismatch pairs are the reason the stale-cache refusal exists - a
422    /// stale cached row survives a `forge.toml` provider change and is
423    /// re-seeded at every boot - and a `budget` that keyed on billing
424    /// alone would render windows as money with every other test still
425    /// green.
426    #[test]
427    fn budget_covers_every_backend_and_source_pair() {
428        use forge_primitives::usage::ApiSpend;
429
430        let future = std::time::SystemTime::now() + std::time::Duration::from_secs(60);
431        let spend = ApiSpend {
432            daily: 0.5,
433            weekly: 1.0,
434            monthly: 2.0,
435            limit: None,
436            limit_remaining: None,
437            limit_reset: None,
438            expires_at: None,
439        };
440
441        for backend in all() {
442            let spend_billed = backend.billing() == BillingModel::Spend;
443            let unknown = AccountBudget::Unknown { spend_billed };
444
445            // No snapshot still carries the billing model, so the empty
446            // row sits under the labels the account would really have.
447            assert_eq!(backend.budget("Acct", None), unknown, "{:?}", backend.token());
448
449            for source in [
450                UsageSourceKind::Oauth,
451                UsageSourceKind::OpenRouterKey,
452                UsageSourceKind::ZaiMonitor,
453            ] {
454                let snapshot = budget_snapshot(future, source, Some(spend.clone()));
455                let expected = if source == backend.source() {
456                    match backend.billing() {
457                        BillingModel::Windows => AccountBudget::Subscription {
458                            five_hour_util: Some(100.0),
459                            seven_day_util: Some(20.0),
460                            resets_at: Some(future),
461                        },
462                        BillingModel::Spend => {
463                            AccountBudget::Api { daily: 0.5, weekly: 1.0, monthly: 2.0 }
464                        }
465                    }
466                } else {
467                    unknown.clone()
468                };
469                assert_eq!(
470                    backend.budget("Acct", Some(&snapshot)),
471                    expected,
472                    "{:?} with a {:?} snapshot",
473                    backend.token(),
474                    source,
475                );
476            }
477
478            // A spend backend whose matching-source snapshot carries no
479            // figures maps to nothing - today's openrouter mapper
480            // refuses such a body, so this takes the same warn path a
481            // stale row takes.
482            if spend_billed {
483                let bare = budget_snapshot(future, backend.source(), None);
484                assert_eq!(backend.budget("Acct", Some(&bare)), unknown);
485            }
486        }
487    }
488}
forge_gateway/lib.rs

forge_gateway/
lib.rs
