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