crabllm_proxy/
lib.rs

1use axum::{
2    Json, Router,
3    extract::Request,
4    middleware,
5    response::Response,
6    routing::{get, post},
7};
8use crabllm_core::{Prefix, Provider, Storage};
9
10pub use auth::KeyName;
11pub use state::{AppState, UsageEvent};
12
13// Storage table prefixes. Each 4-byte prefix namespaces a logical table
14// in the key-value storage backend.
15pub const PREFIX_KEYS: Prefix = *b"keys";
16pub const PREFIX_RATE_LIMIT: Prefix = *b"rlim";
17pub const PREFIX_USAGE: Prefix = *b"usge";
18pub const PREFIX_CACHE: Prefix = *b"cach";
19pub const PREFIX_BUDGET: Prefix = *b"bdgt";
20pub const PREFIX_AUDIT: Prefix = *b"alog";
21pub const PREFIX_PROVIDERS: Prefix = *b"prvd";
22
23pub mod admin;
24pub mod admin_providers;
25pub mod anthropic;
26pub mod auth;
27pub mod ext;
28pub(crate) mod handlers;
29#[cfg(feature = "openapi")]
30pub mod openapi;
31mod state;
32pub mod storage;
33
34/// Middleware that tracks the number of in-flight API requests.
35/// For SSE streams, the gauge decrements when the response starts (not when the
36/// stream ends), so it undercounts long-lived streaming connections.
37async fn track_active_connections(request: Request, next: middleware::Next) -> Response {
38    metrics::gauge!("crabllm_active_connections").increment(1.0);
39    let response = next.run(request).await;
40    metrics::gauge!("crabllm_active_connections").decrement(1.0);
41    response
42}
43
44/// Middleware that logs every incoming HTTP request with method, path,
45/// status, and latency. 2xx/3xx log at `info`, 4xx/5xx at `warn`.
46/// `/health` and `/metrics` log at `debug` so probes don't flood the
47/// default output.
48pub async fn log_request(request: Request, next: middleware::Next) -> Response {
49    let method = request.method().clone();
50    let path = request.uri().path().to_string();
51    let start = std::time::Instant::now();
52
53    let response = next.run(request).await;
54    let status = response.status();
55    let latency_ms = start.elapsed().as_millis() as u64;
56    let is_probe = path == "/health" || path == "/metrics";
57
58    if is_probe {
59        tracing::debug!(%method, path, status = status.as_u16(), latency_ms, "request");
60    } else if status.is_client_error() || status.is_server_error() {
61        tracing::warn!(%method, path, status = status.as_u16(), latency_ms, "request");
62    } else {
63        tracing::info!(%method, path, status = status.as_u16(), latency_ms, "request");
64    }
65
66    response
67}
68
69/// Build the Axum router with all API routes and admin routes.
70pub fn router<S, P>(state: AppState<S, P>, admin_routes: Vec<Router>) -> Router
71where
72    S: Storage + 'static,
73    P: Provider + 'static,
74{
75    let mut app = Router::<AppState<S, P>>::new()
76        .route(
77            "/v1/chat/completions",
78            post(handlers::chat_completions::<S, P>),
79        )
80        .route("/v1/messages", post(anthropic::messages::<S, P>))
81        .route("/v1/embeddings", post(handlers::embeddings::<S, P>))
82        .route(
83            "/v1/images/generations",
84            post(handlers::image_generations::<S, P>),
85        )
86        .route("/v1/audio/speech", post(handlers::audio_speech::<S, P>))
87        .route(
88            "/v1/audio/transcriptions",
89            post(handlers::audio_transcriptions::<S, P>),
90        )
91        .route("/v1/models", get(handlers::models::<S, P>))
92        .route("/v1/usage", get(handlers::usage::<S, P>))
93        .layer(middleware::from_fn_with_state(
94            state.clone(),
95            auth::auth::<S, P>,
96        ))
97        .layer(middleware::from_fn(track_active_connections))
98        .with_state(state);
99
100    // Health check — outside auth middleware so load balancers can probe it.
101    app = app.route(
102        "/health",
103        get(|| async { Json(serde_json::json!({"status": "ok"})) }),
104    );
105
106    // Merge extension-provided admin routes (stateless — extensions
107    // capture their own state via closures in the Router<()>).
108    for admin_router in admin_routes {
109        app = app.merge(admin_router);
110    }
111
112    app
113}
crabllm_proxy/lib.rs

crabllm_proxy/
lib.rs
