byokey_proxy/
lib.rs

1//! HTTP proxy layer — axum router, route handlers, and error mapping.
2//!
3//! ## Module layout
4//!
5//! - [`handler`]  — HTTP route handlers (API, Amp, management).
6//! - [`router`]   — Axum router construction and route registration.
7//! - [`error`]    — [`ApiError`] type for OpenAI-compatible error responses.
8//! - [`openapi`]  — `OpenAPI` specification generation.
9//! - [`usage`]    — In-memory request/token usage tracking.
10
11pub mod error;
12pub mod handler;
13#[allow(clippy::needless_for_each)]
14pub mod openapi;
15pub mod router;
16pub mod usage;
17
18pub use error::ApiError;
19pub use handler::amp_threads::AmpThreadIndex;
20pub use openapi::ApiDoc;
21pub use router::make_router;
22pub use usage::{UsageRecorder, UsageStats};
23
24use arc_swap::ArcSwap;
25use byokey_auth::AuthManager;
26use byokey_provider::DeviceProfileCache;
27use byokey_types::{AmpQuotaStore, RateLimitStore, UsageStore};
28use std::sync::Arc;
29
30/// Shared application state passed to all route handlers.
31pub struct AppState {
32    /// Server configuration (providers, listen address, etc.).
33    /// Atomically swappable for hot-reloading.
34    pub config: Arc<ArcSwap<byokey_config::Config>>,
35    /// Token manager for OAuth-based providers.
36    pub auth: Arc<AuthManager>,
37    /// HTTP client for upstream requests.
38    pub http: rquest::Client,
39    /// In-memory usage statistics with optional persistent backing.
40    pub usage: Arc<UsageRecorder>,
41    /// Per-provider, per-account rate limit snapshots from upstream responses.
42    pub ratelimits: Arc<RateLimitStore>,
43    /// Per-auth device fingerprint cache for Claude API headers.
44    pub device_profiles: Arc<DeviceProfileCache>,
45    /// Cached `AmpCode` quota data (free-tier status + balance).
46    pub amp_quota: Arc<AmpQuotaStore>,
47    /// Pre-built, file-watched index of local Amp CLI thread summaries.
48    pub amp_threads: Arc<AmpThreadIndex>,
49}
50
51impl AppState {
52    /// Creates a new shared application state wrapped in an `Arc`.
53    ///
54    /// If the config specifies a `proxy_url`, the HTTP client is built with that proxy.
55    /// An optional [`UsageStore`] enables persistent usage tracking.
56    pub fn new(
57        config: Arc<ArcSwap<byokey_config::Config>>,
58        auth: Arc<AuthManager>,
59        usage_store: Option<Arc<dyn UsageStore>>,
60    ) -> Arc<Self> {
61        Self::with_thread_index(config, auth, usage_store, {
62            let idx = Arc::new(AmpThreadIndex::build());
63            idx.watch();
64            idx
65        })
66    }
67
68    /// Create state with a pre-built thread index (avoids filesystem scan in tests).
69    pub fn with_thread_index(
70        config: Arc<ArcSwap<byokey_config::Config>>,
71        auth: Arc<AuthManager>,
72        usage_store: Option<Arc<dyn UsageStore>>,
73        amp_threads: Arc<AmpThreadIndex>,
74    ) -> Arc<Self> {
75        let snapshot = config.load();
76        let http = build_http_client(snapshot.proxy_url.as_deref());
77        Arc::new(Self {
78            config,
79            auth,
80            http,
81            usage: Arc::new(UsageRecorder::new(usage_store)),
82            ratelimits: Arc::new(RateLimitStore::new()),
83            device_profiles: Arc::new(DeviceProfileCache::new()),
84            amp_quota: Arc::new(AmpQuotaStore::new()),
85            amp_threads,
86        })
87    }
88}
89
90/// Build an HTTP client, optionally configured with a proxy URL.
91fn build_http_client(proxy_url: Option<&str>) -> rquest::Client {
92    if let Some(url) = proxy_url {
93        match rquest::Proxy::all(url) {
94            Ok(proxy) => {
95                return rquest::Client::builder()
96                    .proxy(proxy)
97                    .build()
98                    .unwrap_or_else(|_| rquest::Client::new());
99            }
100            Err(e) => {
101                tracing::warn!(url = url, error = %e, "invalid proxy_url, using direct connection");
102            }
103        }
104    }
105    rquest::Client::new()
106}
byokey_proxy/lib.rs

byokey_proxy/
lib.rs
