lean_ctx/proxy/
mod.rs

1pub mod anthropic;
2pub mod compress;
3pub mod forward;
4pub mod google;
5pub mod openai;
6
7use std::net::SocketAddr;
8use std::sync::atomic::{AtomicU64, Ordering};
9use std::sync::Arc;
10
11use axum::{
12    body::Body,
13    extract::State,
14    http::{Request, StatusCode},
15    response::{IntoResponse, Response},
16    routing::{any, get},
17    Router,
18};
19
20#[derive(Clone)]
21pub struct ProxyState {
22    pub client: reqwest::Client,
23    pub port: u16,
24    pub stats: Arc<ProxyStats>,
25    pub anthropic_upstream: String,
26    pub openai_upstream: String,
27    pub gemini_upstream: String,
28}
29
30pub struct ProxyStats {
31    pub requests_total: AtomicU64,
32    pub requests_compressed: AtomicU64,
33    pub tokens_saved: AtomicU64,
34    pub bytes_original: AtomicU64,
35    pub bytes_compressed: AtomicU64,
36}
37
38impl Default for ProxyStats {
39    fn default() -> Self {
40        Self {
41            requests_total: AtomicU64::new(0),
42            requests_compressed: AtomicU64::new(0),
43            tokens_saved: AtomicU64::new(0),
44            bytes_original: AtomicU64::new(0),
45            bytes_compressed: AtomicU64::new(0),
46        }
47    }
48}
49
50impl ProxyStats {
51    pub fn record_request(&self) {
52        self.requests_total.fetch_add(1, Ordering::Relaxed);
53    }
54
55    pub fn record_compression(&self, original: usize, compressed: usize) {
56        self.requests_compressed.fetch_add(1, Ordering::Relaxed);
57        self.bytes_original
58            .fetch_add(original as u64, Ordering::Relaxed);
59        self.bytes_compressed
60            .fetch_add(compressed as u64, Ordering::Relaxed);
61        let saved_tokens = (original.saturating_sub(compressed) / 4) as u64;
62        self.tokens_saved.fetch_add(saved_tokens, Ordering::Relaxed);
63    }
64
65    pub fn compression_ratio(&self) -> f64 {
66        let original = self.bytes_original.load(Ordering::Relaxed);
67        if original == 0 {
68            return 0.0;
69        }
70        let compressed = self.bytes_compressed.load(Ordering::Relaxed);
71        (1.0 - compressed as f64 / original as f64) * 100.0
72    }
73}
74
75pub async fn start_proxy(port: u16) -> anyhow::Result<()> {
76    use crate::core::config::{Config, ProxyProvider};
77
78    let client = reqwest::Client::builder()
79        .timeout(std::time::Duration::from_mins(2))
80        .build()?;
81
82    let cfg = Config::load();
83    let anthropic_upstream = cfg.proxy.resolve_upstream(ProxyProvider::Anthropic);
84    let openai_upstream = cfg.proxy.resolve_upstream(ProxyProvider::OpenAi);
85    let gemini_upstream = cfg.proxy.resolve_upstream(ProxyProvider::Gemini);
86
87    let state = ProxyState {
88        client,
89        port,
90        stats: Arc::new(ProxyStats::default()),
91        anthropic_upstream: anthropic_upstream.clone(),
92        openai_upstream: openai_upstream.clone(),
93        gemini_upstream: gemini_upstream.clone(),
94    };
95
96    let app = Router::new()
97        .route("/health", get(health))
98        .route("/status", get(status_handler))
99        .route("/v1/messages", any(anthropic::handler))
100        .route("/v1/chat/completions", any(openai::handler))
101        .fallback(fallback_router)
102        .with_state(state);
103
104    let addr = SocketAddr::from(([127, 0, 0, 1], port));
105    println!("lean-ctx proxy listening on http://{addr}");
106    println!("  Anthropic: POST /v1/messages → {anthropic_upstream}");
107    println!("  OpenAI:    POST /v1/chat/completions → {openai_upstream}");
108    println!("  Gemini:    POST /v1beta/models/... → {gemini_upstream}");
109
110    let listener = tokio::net::TcpListener::bind(addr).await?;
111    axum::serve(listener, app).await?;
112
113    Ok(())
114}
115
116async fn health() -> impl IntoResponse {
117    (StatusCode::OK, "ok")
118}
119
120async fn status_handler(State(state): State<ProxyState>) -> impl IntoResponse {
121    use std::sync::atomic::Ordering::Relaxed;
122    let s = &state.stats;
123    let body = serde_json::json!({
124        "status": "running",
125        "port": state.port,
126        "requests_total": s.requests_total.load(Relaxed),
127        "requests_compressed": s.requests_compressed.load(Relaxed),
128        "tokens_saved": s.tokens_saved.load(Relaxed),
129        "bytes_original": s.bytes_original.load(Relaxed),
130        "bytes_compressed": s.bytes_compressed.load(Relaxed),
131        "compression_ratio_pct": format!("{:.1}", s.compression_ratio()),
132    });
133    (StatusCode::OK, axum::Json(body))
134}
135
136async fn fallback_router(State(state): State<ProxyState>, req: Request<Body>) -> Response {
137    let path = req.uri().path().to_string();
138
139    if path.starts_with("/v1beta/models/") || path.starts_with("/v1/models/") {
140        match google::handler(State(state), req).await {
141            Ok(resp) => resp,
142            Err(status) => Response::builder()
143                .status(status)
144                .body(Body::from("proxy error"))
145                .expect("BUG: building error response with valid status should never fail"),
146        }
147    } else {
148        let method = req.method().to_string();
149        eprintln!("lean-ctx proxy: unmatched {method} {path}");
150        Response::builder()
151            .status(StatusCode::NOT_FOUND)
152            .body(Body::from(format!(
153                "lean-ctx proxy: no handler for {method} {path}"
154            )))
155            .expect("BUG: building 404 response should never fail")
156    }
157}
lean_ctx/proxy/mod.rs

lean_ctx/proxy/
mod.rs
