1pub mod anthropic;
2pub mod compress;
3pub mod cost;
4pub mod forward;
5pub mod google;
6pub mod history_prune;
7pub mod introspect;
8pub mod metrics;
9pub mod openai;
10pub mod openai_responses;
11pub mod tool_kind;
12
13use std::net::SocketAddr;
14use std::sync::atomic::{AtomicU64, Ordering};
15use std::sync::Arc;
16
17use axum::{
18 body::Body,
19 extract::State,
20 http::{Request, StatusCode},
21 response::{IntoResponse, Response},
22 routing::{any, get},
23 Router,
24};
25
26#[derive(Clone)]
27pub struct ProxyState {
28 pub client: reqwest::Client,
29 pub port: u16,
30 pub stats: Arc<ProxyStats>,
31 pub introspect: Arc<introspect::IntrospectState>,
32 pub anthropic_upstream: String,
33 pub openai_upstream: String,
34 pub gemini_upstream: String,
35}
36
37pub struct ProxyStats {
38 pub requests_total: AtomicU64,
39 pub requests_compressed: AtomicU64,
40 pub tokens_saved: AtomicU64,
41 pub bytes_original: AtomicU64,
42 pub bytes_compressed: AtomicU64,
43}
44
45impl Default for ProxyStats {
46 fn default() -> Self {
47 Self {
48 requests_total: AtomicU64::new(0),
49 requests_compressed: AtomicU64::new(0),
50 tokens_saved: AtomicU64::new(0),
51 bytes_original: AtomicU64::new(0),
52 bytes_compressed: AtomicU64::new(0),
53 }
54 }
55}
56
57impl ProxyStats {
58 pub fn record_request(&self) {
59 self.requests_total.fetch_add(1, Ordering::Relaxed);
60 }
61
62 pub fn record_compression(&self, original: usize, compressed: usize) {
63 self.requests_compressed.fetch_add(1, Ordering::Relaxed);
64 self.bytes_original
65 .fetch_add(original as u64, Ordering::Relaxed);
66 self.bytes_compressed
67 .fetch_add(compressed as u64, Ordering::Relaxed);
68 let saved_tokens = (original.saturating_sub(compressed) / 4) as u64;
69 self.tokens_saved.fetch_add(saved_tokens, Ordering::Relaxed);
70 }
71
72 pub fn compression_ratio(&self) -> f64 {
73 let original = self.bytes_original.load(Ordering::Relaxed);
74 if original == 0 {
75 return 0.0;
76 }
77 let compressed = self.bytes_compressed.load(Ordering::Relaxed);
78 (1.0 - compressed as f64 / original as f64) * 100.0
79 }
80}
81
82fn connect_timeout_secs() -> u64 {
84 std::env::var("LEAN_CTX_PROXY_CONNECT_TIMEOUT_SECS")
85 .ok()
86 .and_then(|v| v.trim().parse::<u64>().ok())
87 .filter(|s| *s > 0)
88 .unwrap_or(15)
89}
90
91fn read_idle_timeout_secs() -> u64 {
96 std::env::var("LEAN_CTX_PROXY_READ_TIMEOUT_SECS")
97 .ok()
98 .and_then(|v| v.trim().parse::<u64>().ok())
99 .filter(|s| *s > 0)
100 .unwrap_or(300)
101}
102
103pub async fn start_proxy(port: u16) -> anyhow::Result<()> {
104 let token = crate::core::session_token::resolve_proxy_token("LEAN_CTX_PROXY_TOKEN");
105 start_proxy_with_token(port, Some(token)).await
106}
107
108fn effective_auth_token(auth_token: Option<String>) -> String {
113 auth_token
114 .filter(|t| !t.trim().is_empty())
115 .unwrap_or_else(|| crate::core::session_token::resolve_proxy_token("LEAN_CTX_PROXY_TOKEN"))
116}
117
118pub async fn start_proxy_with_token(port: u16, auth_token: Option<String>) -> anyhow::Result<()> {
119 use crate::core::config::{Config, ProxyProvider};
120
121 let auth_token = effective_auth_token(auth_token);
122
123 let client = reqwest::Client::builder()
128 .connect_timeout(std::time::Duration::from_secs(connect_timeout_secs()))
129 .read_timeout(std::time::Duration::from_secs(read_idle_timeout_secs()))
130 .build()?;
131
132 let cfg = Config::load();
133 let anthropic_upstream = cfg.proxy.resolve_upstream(ProxyProvider::Anthropic);
134 let openai_upstream = cfg.proxy.resolve_upstream(ProxyProvider::OpenAi);
135 let gemini_upstream = cfg.proxy.resolve_upstream(ProxyProvider::Gemini);
136
137 let state = ProxyState {
138 client,
139 port,
140 stats: Arc::new(ProxyStats::default()),
141 introspect: Arc::new(introspect::IntrospectState::default()),
142 anthropic_upstream: anthropic_upstream.clone(),
143 openai_upstream: openai_upstream.clone(),
144 gemini_upstream: gemini_upstream.clone(),
145 };
146
147 let mut app = Router::new()
148 .route("/health", get(health))
149 .route("/status", get(status_handler))
150 .route("/v1/messages", any(anthropic::handler))
151 .route("/v1/messages/{*rest}", any(anthropic::handler))
152 .route("/v1/chat/completions", any(openai::handler))
153 .route("/v1/responses", any(openai_responses::handler))
154 .route("/v1/responses/{*rest}", any(openai_responses::handler))
155 .route("/messages", any(anthropic::handler))
161 .route("/messages/{*rest}", any(anthropic::handler))
162 .route("/chat/completions", any(openai::handler))
163 .route("/responses", any(openai_responses::handler))
164 .route("/responses/{*rest}", any(openai_responses::handler))
165 .route("/v1/references/{id}", get(v1_resolve_reference))
166 .fallback(fallback_router)
167 .layer(axum::middleware::from_fn(host_guard))
168 .with_state(state);
169
170 {
171 let expected = auth_token.clone();
172 app = app.layer(axum::middleware::from_fn(move |req, next| {
173 let expected = expected.clone();
174 proxy_auth_guard(req, next, expected)
175 }));
176 }
177
178 app = app.layer(axum::middleware::from_fn(normalize_provider_path));
182
183 let addr = SocketAddr::from(([127, 0, 0, 1], port));
184 println!("lean-ctx proxy listening on http://{addr} (token auth enabled)");
185 println!(" Anthropic: POST /v1/messages → {anthropic_upstream}");
186 println!(" OpenAI: POST /v1/chat/completions → {openai_upstream}");
187 println!(" OpenAI: POST /v1/responses → {openai_upstream}");
188 println!(" Gemini: POST /v1beta/models/... → {gemini_upstream}");
189
190 let listener = tokio::net::TcpListener::bind(addr).await?;
191 axum::serve(listener, app)
192 .with_graceful_shutdown(shutdown_signal())
193 .await?;
194
195 println!("lean-ctx proxy shut down cleanly.");
196 Ok(())
197}
198
199async fn shutdown_signal() {
200 let ctrl_c = tokio::signal::ctrl_c();
201
202 #[cfg(unix)]
203 {
204 match tokio::signal::unix::signal(tokio::signal::unix::SignalKind::terminate()) {
207 Ok(mut sigterm) => {
208 tokio::select! {
209 _ = ctrl_c => {},
210 _ = sigterm.recv() => {},
211 }
212 }
213 Err(e) => {
214 tracing::warn!("lean-ctx proxy: SIGTERM handler unavailable ({e}); Ctrl-C only");
215 ctrl_c.await.ok();
216 }
217 }
218 }
219
220 #[cfg(not(unix))]
221 {
222 ctrl_c.await.ok();
223 }
224
225 println!("lean-ctx proxy: received shutdown signal, draining…");
226}
227
228async fn health() -> impl IntoResponse {
229 let body = serde_json::json!({
230 "status": "ok",
231 "pid": std::process::id(),
232 });
233 (StatusCode::OK, axum::Json(body))
234}
235
236async fn v1_resolve_reference(
237 axum::extract::Path(id): axum::extract::Path<String>,
238) -> impl IntoResponse {
239 match crate::server::reference_store::resolve(&id) {
240 Some(content) => (StatusCode::OK, content),
241 None => (
242 StatusCode::NOT_FOUND,
243 "Reference expired or not found".to_string(),
244 ),
245 }
246}
247
248async fn status_handler(State(state): State<ProxyState>) -> impl IntoResponse {
249 use std::sync::atomic::Ordering::Relaxed;
250 let s = &state.stats;
251 let i = &state.introspect;
252
253 let last_breakdown = i
254 .last_breakdown
255 .lock()
256 .ok()
257 .and_then(|guard| guard.as_ref().map(|b| serde_json::to_value(b).ok()))
258 .flatten();
259
260 let body = serde_json::json!({
261 "status": "running",
262 "port": state.port,
263 "requests_total": s.requests_total.load(Relaxed),
264 "requests_compressed": s.requests_compressed.load(Relaxed),
265 "tokens_saved": s.tokens_saved.load(Relaxed),
266 "tokens_saved_estimated": true,
267 "bytes_original": s.bytes_original.load(Relaxed),
268 "bytes_compressed": s.bytes_compressed.load(Relaxed),
269 "compression_ratio_pct": format!("{:.1}", s.compression_ratio()),
270 "per_model": cost::snapshot(),
271 "note": "Savings are request-side (tokens removed before forwarding); they do not subtract any re-reads the agent performs. Token figures are estimates; USD uses the shared model price table.",
272 "introspect": {
273 "total_requests_analyzed": i.total_requests.load(Relaxed),
274 "total_system_prompt_tokens": i.total_system_prompt_tokens.load(Relaxed),
275 "last_breakdown": last_breakdown,
276 }
277 });
278 (StatusCode::OK, axum::Json(body))
279}
280
281async fn proxy_auth_guard(
282 req: axum::extract::Request,
283 next: axum::middleware::Next,
284 expected_token: String,
285) -> Result<Response, Response> {
286 let path = req.uri().path();
287 if path == "/health" {
288 return Ok(next.run(req).await);
289 }
290
291 if let Some(auth) = req
293 .headers()
294 .get("authorization")
295 .and_then(|v| v.to_str().ok())
296 {
297 if let Some(token) = auth.strip_prefix("Bearer ") {
298 if constant_time_eq(token.as_bytes(), expected_token.as_bytes()) {
299 return Ok(next.run(req).await);
300 }
301 }
302 }
303
304 if has_provider_api_key(&req) && is_provider_route(path) {
309 return Ok(next.run(req).await);
310 }
311
312 let cfg = crate::core::config::Config::load();
313 let hint = match cfg.proxy_enabled {
314 Some(true) => "lean-ctx proxy requires authentication. Use a Bearer token (LEAN_CTX_PROXY_TOKEN) or configure your AI tool's API key.",
315 Some(false) => "lean-ctx proxy is disabled but still running. Run: lean-ctx proxy cleanup",
316 None => "lean-ctx proxy is not configured. Your AI tool's ANTHROPIC_BASE_URL may be pointing here by mistake. Fix: lean-ctx proxy cleanup OR lean-ctx proxy enable",
317 };
318
319 let body = serde_json::json!({
320 "type": "error",
321 "error": {
322 "type": "authentication_error",
323 "message": format!("401 Unauthorized — {hint}")
324 }
325 });
326
327 Err((StatusCode::UNAUTHORIZED, axum::Json(body)).into_response())
328}
329
330fn has_provider_api_key(req: &axum::extract::Request) -> bool {
331 let headers = req.headers();
332 for key in ["x-api-key", "x-goog-api-key", "api-key"] {
335 if headers
336 .get(key)
337 .and_then(|v| v.to_str().ok())
338 .is_some_and(|v| !v.trim().is_empty())
339 {
340 return true;
341 }
342 }
343 if let Some(auth) = headers.get("authorization").and_then(|v| v.to_str().ok()) {
352 let auth = auth.trim();
353 let credential = auth
354 .strip_prefix("Bearer ")
355 .or_else(|| auth.strip_prefix("bearer "))
356 .unwrap_or(auth)
357 .trim();
358 return !credential.is_empty() && !credential.eq_ignore_ascii_case("bearer");
360 }
361 false
362}
363
364fn is_provider_route(path: &str) -> bool {
365 path.starts_with("/v1/")
366 || path.starts_with("/v1beta/")
367 || path.starts_with("/chat/completions")
368 || path.starts_with("/responses")
369 || path.starts_with("/messages")
370}
371
372fn canonical_provider_path(path: &str) -> Option<String> {
383 if let Some(rest) = path.strip_prefix("/v1/v1/") {
387 return Some(format!("/v1/{rest}"));
388 }
389 const BARE_TO_CANONICAL: &[(&str, &str)] = &[
390 ("/responses", "/v1/responses"),
391 ("/chat/completions", "/v1/chat/completions"),
392 ("/messages", "/v1/messages"),
393 ];
394 for (bare, canonical) in BARE_TO_CANONICAL {
395 if path == *bare {
396 return Some((*canonical).to_string());
397 }
398 if let Some(rest) = path.strip_prefix(&format!("{bare}/")) {
399 return Some(format!("{canonical}/{rest}"));
400 }
401 }
402 None
403}
404
405fn normalized_provider_uri(uri: &axum::http::Uri) -> Option<axum::http::Uri> {
409 let canonical = canonical_provider_path(uri.path())?;
410 let new_path_and_query = match uri.query() {
411 Some(q) => format!("{canonical}?{q}"),
412 None => canonical,
413 };
414 new_path_and_query.parse::<axum::http::Uri>().ok()
415}
416
417async fn normalize_provider_path(
421 mut req: axum::extract::Request,
422 next: axum::middleware::Next,
423) -> Response {
424 if let Some(uri) = normalized_provider_uri(req.uri()) {
425 *req.uri_mut() = uri;
426 }
427 next.run(req).await
428}
429
430fn constant_time_eq(a: &[u8], b: &[u8]) -> bool {
431 use subtle::ConstantTimeEq;
432 if a.len() != b.len() {
433 return false;
434 }
435 bool::from(a.ct_eq(b))
436}
437
438async fn host_guard(
439 req: axum::extract::Request,
440 next: axum::middleware::Next,
441) -> Result<Response, StatusCode> {
442 if let Some(host) = req.headers().get("host").and_then(|v| v.to_str().ok()) {
443 let h = host.split(':').next().unwrap_or(host);
444 if matches!(h, "127.0.0.1" | "localhost" | "[::1]") {
445 return Ok(next.run(req).await);
446 }
447 }
448 Err(StatusCode::FORBIDDEN)
449}
450
451async fn fallback_router(State(state): State<ProxyState>, req: Request<Body>) -> Response {
452 let path = req.uri().path().to_string();
453
454 if path.starts_with("/v1beta/models/") || path.starts_with("/v1/models/") {
455 match google::handler(State(state), req).await {
456 Ok(resp) => resp,
457 Err(status) => Response::builder()
458 .status(status)
459 .body(Body::from("proxy error"))
460 .expect("BUG: building error response with valid status should never fail"),
461 }
462 } else {
463 let method = req.method().to_string();
464 eprintln!("lean-ctx proxy: unmatched {method} {path}");
465 Response::builder()
466 .status(StatusCode::NOT_FOUND)
467 .body(Body::from(format!(
468 "lean-ctx proxy: no handler for {method} {path}"
469 )))
470 .expect("BUG: building 404 response should never fail")
471 }
472}
473
474#[cfg(test)]
475mod auth_tests {
476 use super::*;
477
478 #[test]
481 fn effective_auth_token_never_yields_empty() {
482 let _env = crate::core::data_dir::test_env_lock();
483 let tmp = tempfile::tempdir().unwrap();
484 std::env::set_var("LEAN_CTX_DATA_DIR", tmp.path());
485
486 assert_eq!(effective_auth_token(Some("tok".into())), "tok");
487 let auto = effective_auth_token(None);
488 assert!(!auto.trim().is_empty(), "None must auto-resolve a token");
489 let blank = effective_auth_token(Some(" ".into()));
490 assert!(!blank.trim().is_empty(), "blank tokens must be replaced");
491
492 std::env::remove_var("LEAN_CTX_DATA_DIR");
493 }
494
495 #[test]
496 fn is_provider_route_v1() {
497 assert!(is_provider_route("/v1/chat/completions"));
498 assert!(is_provider_route("/v1/messages"));
499 assert!(is_provider_route("/v1/completions"));
500 }
501
502 #[test]
503 fn is_provider_route_anthropic_subpaths() {
504 assert!(is_provider_route("/v1/messages/count_tokens"));
505 assert!(is_provider_route("/v1/messages/batches"));
506 assert!(is_provider_route("/v1/messages/batches/batch_123"));
507 }
508
509 #[test]
510 fn is_provider_route_v1beta() {
511 assert!(is_provider_route("/v1beta/models"));
512 }
513
514 #[test]
515 fn is_provider_route_chat() {
516 assert!(is_provider_route("/chat/completions"));
517 }
518
519 #[test]
520 fn is_provider_route_rejects_non_provider() {
521 assert!(!is_provider_route("/health"));
522 assert!(!is_provider_route("/api/v2/test"));
523 assert!(!is_provider_route("/"));
524 }
525
526 fn build_request(headers: &[(&str, &str)], path: &str) -> axum::extract::Request {
527 let mut builder = axum::http::Request::builder().uri(path);
528 for (k, v) in headers {
529 builder = builder.header(*k, *v);
530 }
531 builder.body(axum::body::Body::empty()).unwrap()
532 }
533
534 #[test]
535 fn has_provider_api_key_x_api_key() {
536 let req = build_request(&[("x-api-key", "sk-ant-abc123")], "/v1/messages");
537 assert!(has_provider_api_key(&req));
538 }
539
540 #[test]
541 fn has_provider_api_key_x_goog() {
542 let req = build_request(&[("x-goog-api-key", "AIzaSyAbc")], "/v1beta/models");
543 assert!(has_provider_api_key(&req));
544 }
545
546 #[test]
547 fn has_provider_api_key_azure() {
548 let req = build_request(&[("api-key", "deadbeef")], "/v1/completions");
549 assert!(has_provider_api_key(&req));
550 }
551
552 #[test]
553 fn has_provider_api_key_bearer_sk() {
554 let req = build_request(
555 &[("authorization", "Bearer sk-proj-abc123")],
556 "/v1/chat/completions",
557 );
558 assert!(has_provider_api_key(&req));
559 }
560
561 #[test]
562 fn has_provider_api_key_empty_rejected() {
563 let req = build_request(&[("x-api-key", " ")], "/v1/messages");
564 assert!(!has_provider_api_key(&req));
565 }
566
567 #[test]
568 fn has_provider_api_key_no_headers() {
569 let req = build_request(&[], "/v1/messages");
570 assert!(!has_provider_api_key(&req));
571 }
572
573 #[test]
574 fn has_provider_api_key_accepts_non_sk_bearer() {
575 for key in [
581 "Bearer or-v1-9f8e7d6c", "Bearer gsk_live_1234", "Bearer abc.def.ghi", "Bearer 0123456789", ] {
586 let req = build_request(&[("authorization", key)], "/v1/responses");
587 assert!(
588 has_provider_api_key(&req),
589 "non-sk Bearer must count as a provider credential: {key}"
590 );
591 }
592 }
593
594 #[test]
595 fn has_provider_api_key_empty_bearer_rejected() {
596 for bad in ["Bearer ", "", "Bearer", "bearer", " "] {
599 let req = build_request(&[("authorization", bad)], "/responses");
600 assert!(
601 !has_provider_api_key(&req),
602 "blank/scheme-only Authorization must not authenticate: {bad:?}"
603 );
604 }
605 }
606
607 #[test]
610 fn is_provider_route_bare_responses_and_messages() {
611 assert!(is_provider_route("/responses"));
614 assert!(is_provider_route("/responses/resp_123/input_items"));
615 assert!(is_provider_route("/messages"));
616 }
617
618 #[test]
619 fn canonical_provider_path_rewrites_bare_endpoints() {
620 assert_eq!(
621 canonical_provider_path("/responses").as_deref(),
622 Some("/v1/responses")
623 );
624 assert_eq!(
625 canonical_provider_path("/chat/completions").as_deref(),
626 Some("/v1/chat/completions")
627 );
628 assert_eq!(
629 canonical_provider_path("/messages").as_deref(),
630 Some("/v1/messages")
631 );
632 }
633
634 #[test]
635 fn canonical_provider_path_preserves_subpaths() {
636 assert_eq!(
637 canonical_provider_path("/responses/resp_abc/cancel").as_deref(),
638 Some("/v1/responses/resp_abc/cancel")
639 );
640 assert_eq!(
641 canonical_provider_path("/messages/batches/batch_1").as_deref(),
642 Some("/v1/messages/batches/batch_1")
643 );
644 }
645
646 #[test]
647 fn canonical_provider_path_ignores_already_canonical_and_unknown() {
648 assert_eq!(canonical_provider_path("/v1/responses"), None);
650 assert_eq!(canonical_provider_path("/v1/chat/completions"), None);
651 assert_eq!(canonical_provider_path("/health"), None);
653 assert_eq!(canonical_provider_path("/responsesx"), None);
654 assert_eq!(canonical_provider_path("/"), None);
655 }
656
657 #[test]
658 fn canonical_provider_path_collapses_double_v1_prefix() {
659 assert_eq!(
662 canonical_provider_path("/v1/v1/responses").as_deref(),
663 Some("/v1/responses")
664 );
665 assert_eq!(
666 canonical_provider_path("/v1/v1/chat/completions").as_deref(),
667 Some("/v1/chat/completions")
668 );
669 }
670
671 #[test]
672 fn normalized_provider_uri_rewrites_path_and_preserves_query() {
673 use axum::http::Uri;
674 let uri: Uri = "/responses?stream=true".parse().unwrap();
675 let rewritten = normalized_provider_uri(&uri).expect("bare /responses must rewrite");
676 assert_eq!(rewritten.path(), "/v1/responses");
677 assert_eq!(rewritten.query(), Some("stream=true"));
678 assert_eq!(
679 rewritten
680 .path_and_query()
681 .map(axum::http::uri::PathAndQuery::as_str),
682 Some("/v1/responses?stream=true")
683 );
684 }
685
686 #[test]
687 fn normalized_provider_uri_noop_for_canonical() {
688 use axum::http::Uri;
689 let uri: Uri = "/v1/responses".parse().unwrap();
690 assert!(normalized_provider_uri(&uri).is_none());
691 }
692}