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, 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 openapi::ApiDoc;
23pub use router::make_router;
24pub use usage::{UsageRecorder, UsageStats};
25
26use arc_swap::ArcSwap;
27use byokey_auth::AuthManager;
28use byokey_provider::DeviceProfileCache;
29use byokey_types::{RateLimitStore, UsageStore};
30use std::sync::Arc;
31
32/// Shared application state passed to all route handlers.
33pub struct AppState {
34    /// Server configuration (providers, listen address, etc.).
35    /// Atomically swappable for hot-reloading.
36    pub config: Arc<ArcSwap<byokey_config::Config>>,
37    /// Token manager for OAuth-based providers.
38    pub auth: Arc<AuthManager>,
39    /// HTTP client for upstream requests.
40    pub http: wreq::Client,
41    /// In-memory usage statistics with optional persistent backing.
42    pub usage: Arc<UsageRecorder>,
43    /// Per-provider, per-account rate limit snapshots from upstream responses.
44    pub ratelimits: Arc<RateLimitStore>,
45    /// Per-auth device fingerprint cache for Claude API headers.
46    pub device_profiles: Arc<DeviceProfileCache>,
47    /// Remote version/fingerprint info fetched from assets.byokey.io at startup.
48    pub versions: VersionStore,
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        versions: VersionStore,
61    ) -> Arc<Self> {
62        let snapshot = config.load();
63        let http = build_http_client(snapshot.proxy_url.as_deref());
64        Arc::new(Self {
65            config,
66            auth,
67            http,
68            usage: Arc::new(UsageRecorder::new(usage_store)),
69            ratelimits: Arc::new(RateLimitStore::new()),
70            device_profiles: Arc::new(DeviceProfileCache::new()),
71            versions,
72        })
73    }
74}
75
76/// Build an HTTP client, optionally configured with a proxy URL.
77fn build_http_client(proxy_url: Option<&str>) -> wreq::Client {
78    if let Some(url) = proxy_url {
79        match wreq::Proxy::all(url) {
80            Ok(proxy) => {
81                return wreq::Client::builder()
82                    .proxy(proxy)
83                    .build()
84                    .unwrap_or_else(|_| wreq::Client::new());
85            }
86            Err(e) => {
87                tracing::warn!(url = url, error = %e, "invalid proxy_url, using direct connection");
88            }
89        }
90    }
91    wreq::Client::new()
92}
byokey_proxy/lib.rs

byokey_proxy/
lib.rs
