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

byokey_proxy/
lib.rs
