byokey_proxy/
lib.rs

1//! HTTP proxy layer — axum router, route handlers, and error mapping.
2//!
3//! Exposes an OpenAI-compatible `/v1/chat/completions` endpoint, a `/v1/models`
4//! listing, and an Amp CLI compatibility layer under `/amp/*`.
5
6mod amp;
7mod amp_provider;
8mod chat;
9mod error;
10mod messages;
11mod models;
12
13pub use error::ApiError;
14
15use axum::{
16    Router,
17    routing::{any, get, post},
18};
19use byokey_auth::AuthManager;
20use byokey_config::Config;
21use std::sync::Arc;
22
23/// Shared application state passed to all route handlers.
24pub struct AppState {
25    /// Server configuration (providers, listen address, etc.).
26    pub config: Arc<Config>,
27    /// Token manager for OAuth-based providers.
28    pub auth: Arc<AuthManager>,
29    /// HTTP client for upstream requests.
30    pub http: rquest::Client,
31}
32
33impl AppState {
34    /// Creates a new shared application state wrapped in an `Arc`.
35    pub fn new(config: Config, auth: Arc<AuthManager>) -> Arc<Self> {
36        Arc::new(Self {
37            config: Arc::new(config),
38            auth,
39            http: rquest::Client::new(),
40        })
41    }
42}
43
44/// Build the full axum router.
45///
46/// Routes:
47/// - POST /v1/chat/completions                          OpenAI-compatible
48/// - POST /v1/messages                                  Anthropic native passthrough
49/// - GET  /v1/models
50/// - GET  /amp/v1/login
51/// - ANY  /amp/v0/management/{*path}
52/// - POST /amp/v1/chat/completions
53///
54/// `AmpCode` provider routes:
55/// - POST /api/provider/anthropic/v1/messages           Anthropic native (`AmpCode`)
56/// - POST /api/provider/openai/v1/chat/completions      `OpenAI`-compatible (`AmpCode`)
57/// - POST /api/provider/openai/v1/responses             Codex Responses API (`AmpCode`)
58/// - POST /api/provider/google/v1beta/models/{action}   Gemini native (`AmpCode`)
59/// - ANY  /api/{*path}                                  `ampcode.com` management proxy
60pub fn make_router(state: Arc<AppState>) -> Router {
61    Router::new()
62        // Standard routes
63        .route("/v1/chat/completions", post(chat::chat_completions))
64        .route("/v1/messages", post(messages::anthropic_messages))
65        .route("/v1/models", get(models::list_models))
66        // Legacy Amp CLI routes
67        .route("/amp/v1/login", get(amp::login_redirect))
68        .route("/amp/v0/management/{*path}", any(amp::management_proxy))
69        .route("/amp/v1/chat/completions", post(chat::chat_completions))
70        // AmpCode provider-specific routes (must be registered before the catch-all)
71        .route(
72            "/api/provider/anthropic/v1/messages",
73            post(messages::anthropic_messages),
74        )
75        .route(
76            "/api/provider/openai/v1/chat/completions",
77            post(chat::chat_completions),
78        )
79        .route(
80            "/api/provider/openai/v1/responses",
81            post(amp_provider::codex_responses_passthrough),
82        )
83        .route(
84            "/api/provider/google/v1beta/models/{action}",
85            post(amp_provider::gemini_native_passthrough),
86        )
87        // Catch-all: forward remaining /api/* routes to ampcode.com
88        .route("/api/{*path}", any(amp_provider::amp_management_proxy))
89        .with_state(state)
90}
91
92#[cfg(test)]
93mod tests {
94    use super::*;
95    use axum::{body::Body, http::Request};
96    use byokey_store::InMemoryTokenStore;
97    use http_body_util::BodyExt as _;
98    use serde_json::Value;
99    use tower::ServiceExt as _;
100
101    fn make_state() -> Arc<AppState> {
102        let store = Arc::new(InMemoryTokenStore::new());
103        let auth = Arc::new(AuthManager::new(store));
104        AppState::new(Config::default(), auth)
105    }
106
107    async fn body_json(resp: axum::response::Response) -> Value {
108        let bytes = resp.into_body().collect().await.unwrap().to_bytes();
109        serde_json::from_slice(&bytes).unwrap()
110    }
111
112    #[tokio::test]
113    async fn test_list_models_empty_config() {
114        let app = make_router(make_state());
115        let resp = app
116            .oneshot(
117                Request::builder()
118                    .uri("/v1/models")
119                    .body(Body::empty())
120                    .unwrap(),
121            )
122            .await
123            .unwrap();
124
125        assert_eq!(resp.status(), axum::http::StatusCode::OK);
126        let json = body_json(resp).await;
127        assert_eq!(json["object"], "list");
128        assert!(json["data"].is_array());
129        // All providers are enabled by default even without explicit config.
130        assert!(!json["data"].as_array().unwrap().is_empty());
131    }
132
133    #[tokio::test]
134    async fn test_amp_login_redirect() {
135        let app = make_router(make_state());
136        let resp = app
137            .oneshot(
138                Request::builder()
139                    .uri("/amp/v1/login")
140                    .body(Body::empty())
141                    .unwrap(),
142            )
143            .await
144            .unwrap();
145
146        assert_eq!(resp.status(), axum::http::StatusCode::FOUND);
147        assert_eq!(
148            resp.headers().get("location").and_then(|v| v.to_str().ok()),
149            Some("https://ampcode.com/login")
150        );
151    }
152
153    #[tokio::test]
154    async fn test_chat_unknown_model_returns_400() {
155        use serde_json::json;
156
157        let app = make_router(make_state());
158        let body = json!({"model": "nonexistent-model-xyz", "messages": []});
159        let resp = app
160            .oneshot(
161                Request::builder()
162                    .method("POST")
163                    .uri("/v1/chat/completions")
164                    .header("content-type", "application/json")
165                    .body(Body::from(serde_json::to_vec(&body).unwrap()))
166                    .unwrap(),
167            )
168            .await
169            .unwrap();
170
171        assert_eq!(resp.status(), axum::http::StatusCode::BAD_REQUEST);
172        let json = body_json(resp).await;
173        assert!(
174            json["error"]["message"]
175                .as_str()
176                .unwrap_or("")
177                .contains("nonexistent-model-xyz")
178        );
179    }
180
181    #[tokio::test]
182    async fn test_chat_missing_model_returns_400() {
183        use serde_json::json;
184
185        let app = make_router(make_state());
186        let body = json!({"messages": [{"role": "user", "content": "hi"}]});
187        let resp = app
188            .oneshot(
189                Request::builder()
190                    .method("POST")
191                    .uri("/v1/chat/completions")
192                    .header("content-type", "application/json")
193                    .body(Body::from(serde_json::to_vec(&body).unwrap()))
194                    .unwrap(),
195            )
196            .await
197            .unwrap();
198
199        // Empty model string → UnsupportedModel → 400
200        assert_eq!(resp.status(), axum::http::StatusCode::BAD_REQUEST);
201    }
202
203    #[tokio::test]
204    async fn test_amp_chat_route_exists() {
205        use serde_json::json;
206
207        let app = make_router(make_state());
208        let body = json!({"model": "nonexistent", "messages": []});
209        let resp = app
210            .oneshot(
211                Request::builder()
212                    .method("POST")
213                    .uri("/amp/v1/chat/completions")
214                    .header("content-type", "application/json")
215                    .body(Body::from(serde_json::to_vec(&body).unwrap()))
216                    .unwrap(),
217            )
218            .await
219            .unwrap();
220
221        // Route exists (not 404), even though model is invalid
222        assert_ne!(resp.status(), axum::http::StatusCode::NOT_FOUND);
223    }
224}
byokey_proxy/lib.rs

byokey_proxy/
lib.rs
