shift_proxy/routes/mod.rs
1//! Route handlers for the SHIFT proxy.
2
3pub mod anthropic;
4pub mod google;
5pub mod health;
6pub mod openai;
7pub mod passthrough;
8
9use crate::ProxyState;
10use axum::extract::DefaultBodyLimit;
11use axum::extract::State;
12use axum::http::{HeaderMap, Uri};
13use axum::response::Response;
14use axum::routing::{any, get, post};
15use axum::Router;
16
17/// Maximum request body size: 200 MB.
18/// AI payloads with base64 images can be large (50MB+). This limit prevents
19/// unbounded memory consumption from malicious clients while accommodating
20/// legitimate multi-image payloads.
21const MAX_BODY_SIZE: usize = 200 * 1024 * 1024;
22
23/// Fallback handler for `POST /messages` (without the `/v1` prefix).
24///
25/// Some clients (e.g. OpenCode with a misconfigured `baseURL` that omits `/v1`)
26/// send requests to `/messages` instead of `/v1/messages`. Rather than returning
27/// a 404 "Unknown route" error, we rewrite the URI to `/v1/messages` and delegate
28/// to the standard Anthropic handler.
29async fn messages_fallback_handler(
30 state: State<ProxyState>,
31 uri: Uri,
32 headers: HeaderMap,
33 body: String,
34) -> Response {
35 // Rewrite /messages → /v1/messages so the Anthropic handler builds the
36 // correct upstream URL (https://api.anthropic.com/v1/messages).
37 let query = uri.query().map(|q| format!("?{}", q)).unwrap_or_default();
38 let rewritten: Uri = format!("/v1/messages{}", query)
39 .parse()
40 .expect("/v1/messages is a valid URI");
41
42 anthropic::anthropic_handler(state, rewritten, headers, body).await
43}
44
45/// Build the complete proxy router with all routes.
46pub fn build_router(state: ProxyState) -> Router {
47 Router::new()
48 // Health and stats
49 .route("/health", get(health::health_handler))
50 .route("/stats", get(health::stats_handler))
51 // Provider-specific routes (with optimization)
52 .route("/v1/messages", post(anthropic::anthropic_handler))
53 .route("/v1/chat/completions", post(openai::openai_handler))
54 // Fallback: /messages → /v1/messages (resilience for misconfigured clients)
55 .route("/messages", post(messages_fallback_handler))
56 // Google routes (passthrough only)
57 .route("/v1beta/models/{*path}", post(google::google_handler))
58 .route("/v1/models/{*path}", post(google::google_handler))
59 // Catch-all passthrough for all HTTP methods (not just POST).
60 // Some provider APIs use GET (list models), PUT, DELETE, etc.
61 .fallback(any(passthrough::passthrough_handler))
62 // Explicit body size limit — prevents OOM from malicious payloads.
63 .layer(DefaultBodyLimit::max(MAX_BODY_SIZE))
64 .with_state(state)
65}