kaizen/proxy/
server.rs

1// SPDX-License-Identifier: AGPL-3.0-or-later
2//! Local HTTP server: any path and method, forward to Anthropic base URL.
3
4use crate::proxy::forward::run_forward_inner;
5use crate::proxy::opts::ProxyRunOptions;
6use crate::proxy::state::ProxyState;
7use axum::Router;
8use axum::http::{Request, StatusCode};
9use axum::response::IntoResponse;
10use axum::response::Response;
11use axum::routing::any;
12use std::net::SocketAddr;
13use std::path::PathBuf;
14use std::sync::Arc;
15
16use axum::extract::DefaultBodyLimit;
17
18/// Forward any method+path+query+body, buffer upstream response, append one store row.
19async fn handle(
20    axum::extract::State(st): axum::extract::State<Arc<ProxyState>>,
21    request: Request<axum::body::Body>,
22) -> axum::response::Response {
23    match do_forward(&st, request).await {
24        Ok(r) => r,
25        Err(e) => (StatusCode::BAD_GATEWAY, e.to_string()).into_response(),
26    }
27}
28
29async fn do_forward(
30    st: &Arc<ProxyState>,
31    request: Request<axum::body::Body>,
32) -> Result<Response, anyhow::Error> {
33    let (parts, body) = request.into_parts();
34    let method = parts.method;
35    let path = parts.uri.path().trim_start_matches('/').to_string();
36    let q = parts.uri.query().unwrap_or("").to_string();
37    let headers = &parts.headers;
38    let body = axum::body::to_bytes(body, st.options.max_request_bytes as usize).await?;
39    let path_ref = if path.is_empty() { "" } else { &path };
40    run_forward_inner(st, method, path_ref, &q, headers, &body).await
41}
42
43/// Build `Client`, bind, run until the process is killed.
44pub async fn run(
45    options: Arc<ProxyRunOptions>,
46    workspace: PathBuf,
47    config: crate::core::config::Config,
48) -> Result<(), anyhow::Error> {
49    let store_path = crate::core::workspace::db_path(&workspace)?;
50    let client = build_client(&options)?;
51    let st = Arc::new(ProxyState {
52        options: options.clone(),
53        store_path,
54        workspace: workspace.clone(),
55        config: Arc::new(config),
56        client,
57    });
58    let limit = usize::try_from(st.options.max_request_bytes).unwrap_or(usize::MAX);
59    let app = Router::new()
60        .route("/{*path}", any(handle))
61        .layer(DefaultBodyLimit::max(limit))
62        .with_state(st);
63    let addr: SocketAddr = options
64        .listen
65        .parse()
66        .map_err(|e: std::net::AddrParseError| {
67            anyhow::anyhow!(r#"bad --listen (expected e.g. "127.0.0.1:3847"): {e}"#)
68        })?;
69    let listener = tokio::net::TcpListener::bind(addr).await?;
70    let seen = listener.local_addr()?;
71    tracing::info!(addr = %seen, "kaizen LLM proxy listening (set ANTHROPIC_BASE_URL to this /)");
72    axum::serve(listener, app).await?;
73    Ok(())
74}
75
76fn build_client(o: &ProxyRunOptions) -> Result<reqwest::Client, reqwest::Error> {
77    use std::time::Duration;
78    let mut b = reqwest::Client::builder()
79        .connect_timeout(Duration::from_secs(30))
80        .timeout(Duration::from_secs(300));
81    if !o.compress_transport {
82        b = b.no_gzip();
83    }
84    b.build()
85}
kaizen/proxy/server.rs

kaizen/proxy/
server.rs
