Skip to main content

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}