1pub mod anthropic;
2pub mod cache_aligner;
3pub mod cache_breakpoint;
4pub mod cache_safety;
5pub mod ccr;
6pub mod chatgpt;
7pub mod chatgpt_cookies;
8pub mod cold_prefix;
9pub mod compress;
10pub mod compress_api;
11pub mod cost;
12pub mod effort;
13pub mod forward;
14pub mod google;
15pub mod history_prune;
16pub mod holdout;
17pub mod introspect;
18pub mod metrics;
19pub mod openai;
20pub mod openai_responses;
21pub mod openai_responses_ws;
22pub mod output_savings;
23pub mod prose;
24pub mod prose_ranker;
25pub mod tool_kind;
26pub mod usage;
27pub mod usage_meter;
28pub mod verbosity;
29
30use std::net::SocketAddr;
31use std::sync::Arc;
32use std::sync::atomic::{AtomicU64, Ordering};
33
34use crate::core::config::Upstreams;
35
36use axum::{
37 Router,
38 body::Body,
39 extract::State,
40 http::{Request, StatusCode},
41 response::{IntoResponse, Response},
42 routing::{any, get, post},
43};
44
45#[derive(Clone)]
46pub struct ProxyState {
47 pub client: reqwest::Client,
48 pub port: u16,
49 pub stats: Arc<ProxyStats>,
50 pub introspect: Arc<introspect::IntrospectState>,
51 pub upstreams: tokio::sync::watch::Receiver<Arc<Upstreams>>,
54}
55
56impl ProxyState {
57 pub fn upstream_snapshot(&self) -> Arc<Upstreams> {
59 self.upstreams.borrow().clone()
60 }
61
62 pub fn anthropic_upstream(&self) -> String {
64 self.upstreams.borrow().anthropic.clone()
65 }
66
67 pub fn openai_upstream(&self) -> String {
69 self.upstreams.borrow().openai.clone()
70 }
71
72 pub fn chatgpt_upstream(&self) -> String {
74 self.upstreams.borrow().chatgpt.clone()
75 }
76
77 pub fn gemini_upstream(&self) -> String {
79 self.upstreams.borrow().gemini.clone()
80 }
81}
82
83pub struct ProxyStats {
84 pub requests_total: AtomicU64,
85 pub requests_compressed: AtomicU64,
86 pub tokens_saved: AtomicU64,
87 pub bytes_original: AtomicU64,
88 pub bytes_compressed: AtomicU64,
89 pub anthropic: ProviderStats,
90 pub openai: ProviderStats,
91 pub chatgpt: ProviderStats,
92 pub gemini: ProviderStats,
93}
94
95#[derive(Default)]
96pub struct ProviderStats {
97 pub requests_total: AtomicU64,
98 pub requests_compressed: AtomicU64,
99 pub tokens_saved: AtomicU64,
100 pub bytes_original: AtomicU64,
101 pub bytes_compressed: AtomicU64,
102}
103
104impl Default for ProxyStats {
105 fn default() -> Self {
106 Self {
107 requests_total: AtomicU64::new(0),
108 requests_compressed: AtomicU64::new(0),
109 tokens_saved: AtomicU64::new(0),
110 bytes_original: AtomicU64::new(0),
111 bytes_compressed: AtomicU64::new(0),
112 anthropic: ProviderStats::default(),
113 openai: ProviderStats::default(),
114 chatgpt: ProviderStats::default(),
115 gemini: ProviderStats::default(),
116 }
117 }
118}
119
120impl ProxyStats {
121 pub fn record_request(&self, original: usize, compressed: usize) {
122 self.record_totals(original, compressed);
123 }
124
125 pub fn record_provider_request(
126 &self,
127 provider_label: &str,
128 original: usize,
129 compressed: usize,
130 ) {
131 let (effective_compressed, saved_tokens, compressed_request) =
132 self.record_totals(original, compressed);
133
134 if let Some(provider) = self.provider(provider_label) {
135 provider.record(
136 original,
137 effective_compressed,
138 compressed_request,
139 saved_tokens,
140 );
141 }
142 }
143
144 fn record_totals(&self, original: usize, compressed: usize) -> (usize, u64, bool) {
145 self.requests_total.fetch_add(1, Ordering::Relaxed);
146 self.bytes_original
147 .fetch_add(original as u64, Ordering::Relaxed);
148 let effective_compressed = compressed.min(original);
149 self.bytes_compressed
150 .fetch_add(effective_compressed as u64, Ordering::Relaxed);
151 if compressed < original {
152 self.requests_compressed.fetch_add(1, Ordering::Relaxed);
153 }
154 let saved_tokens = (original.saturating_sub(effective_compressed) / 4) as u64;
155 self.tokens_saved.fetch_add(saved_tokens, Ordering::Relaxed);
156 (effective_compressed, saved_tokens, compressed < original)
157 }
158
159 pub fn compression_ratio(&self) -> f64 {
160 let original = self.bytes_original.load(Ordering::Relaxed);
161 if original == 0 {
162 return 0.0;
163 }
164 let compressed = self.bytes_compressed.load(Ordering::Relaxed);
165 (1.0 - compressed as f64 / original as f64) * 100.0
166 }
167
168 fn provider(&self, provider_label: &str) -> Option<&ProviderStats> {
172 match provider_label {
173 "Anthropic" => Some(&self.anthropic),
174 "OpenAI" => Some(&self.openai),
175 "ChatGPT" => Some(&self.chatgpt),
176 "Gemini" => Some(&self.gemini),
177 _ => None,
178 }
179 }
180
181 pub fn provider_summary(&self) -> serde_json::Value {
182 serde_json::json!({
183 "anthropic": self.anthropic.summary(),
184 "openai": self.openai.summary(),
185 "chatgpt": self.chatgpt.summary(),
186 "gemini": self.gemini.summary(),
187 })
188 }
189}
190
191impl ProviderStats {
192 fn record(
193 &self,
194 original: usize,
195 effective_compressed: usize,
196 compressed_request: bool,
197 saved_tokens: u64,
198 ) {
199 self.requests_total.fetch_add(1, Ordering::Relaxed);
200 if compressed_request {
201 self.requests_compressed.fetch_add(1, Ordering::Relaxed);
202 }
203 self.tokens_saved.fetch_add(saved_tokens, Ordering::Relaxed);
204 self.bytes_original
205 .fetch_add(original as u64, Ordering::Relaxed);
206 self.bytes_compressed
207 .fetch_add(effective_compressed as u64, Ordering::Relaxed);
208 }
209
210 fn compression_ratio(&self) -> f64 {
211 let original = self.bytes_original.load(Ordering::Relaxed);
212 if original == 0 {
213 return 0.0;
214 }
215 let compressed = self.bytes_compressed.load(Ordering::Relaxed);
216 (1.0 - compressed as f64 / original as f64) * 100.0
217 }
218
219 fn summary(&self) -> serde_json::Value {
220 serde_json::json!({
221 "requests_total": self.requests_total.load(Ordering::Relaxed),
222 "requests_compressed": self.requests_compressed.load(Ordering::Relaxed),
223 "tokens_saved": self.tokens_saved.load(Ordering::Relaxed),
224 "bytes_original": self.bytes_original.load(Ordering::Relaxed),
225 "bytes_compressed": self.bytes_compressed.load(Ordering::Relaxed),
226 "compression_ratio_pct": format!("{:.1}", self.compression_ratio()),
227 })
228 }
229}
230
231#[cfg(test)]
232mod stats_tests {
233 use super::*;
234 use std::sync::atomic::Ordering;
235
236 #[test]
237 fn compression_ratio_includes_uncompressed_requests() {
238 let stats = ProxyStats::default();
239
240 stats.record_request(1_000, 500);
241 stats.record_request(1_000, 1_000);
242
243 assert_eq!(stats.requests_total.load(Ordering::Relaxed), 2);
244 assert_eq!(stats.requests_compressed.load(Ordering::Relaxed), 1);
245 assert_eq!(stats.tokens_saved.load(Ordering::Relaxed), 125);
246 assert_eq!(stats.compression_ratio(), 25.0);
247 }
248
249 #[test]
250 fn expanded_requests_count_as_zero_savings() {
251 let stats = ProxyStats::default();
252
253 stats.record_request(1_000, 1_500);
254
255 assert_eq!(stats.requests_total.load(Ordering::Relaxed), 1);
256 assert_eq!(stats.requests_compressed.load(Ordering::Relaxed), 0);
257 assert_eq!(stats.tokens_saved.load(Ordering::Relaxed), 0);
258 assert_eq!(stats.compression_ratio(), 0.0);
259 }
260
261 #[test]
262 fn provider_stats_are_separate() {
263 let stats = ProxyStats::default();
264
265 stats.record_provider_request("OpenAI", 1_000, 500);
266 stats.record_provider_request("ChatGPT", 2_000, 1_000);
267
268 assert_eq!(stats.requests_total.load(Ordering::Relaxed), 2);
269 assert_eq!(stats.openai.requests_total.load(Ordering::Relaxed), 1);
270 assert_eq!(stats.chatgpt.requests_total.load(Ordering::Relaxed), 1);
271 assert_eq!(stats.openai.tokens_saved.load(Ordering::Relaxed), 125);
272 assert_eq!(stats.chatgpt.tokens_saved.load(Ordering::Relaxed), 250);
273 assert_eq!(stats.openai.compression_ratio(), 50.0);
274 assert_eq!(stats.chatgpt.compression_ratio(), 50.0);
275 }
276
277 #[test]
278 fn unlabelled_requests_do_not_count_as_gemini() {
279 let stats = ProxyStats::default();
280
281 stats.record_request(1_000, 500);
282
283 assert_eq!(stats.requests_total.load(Ordering::Relaxed), 1);
284 assert_eq!(stats.gemini.requests_total.load(Ordering::Relaxed), 0);
285 }
286
287 #[test]
288 fn unknown_label_is_not_recorded_to_any_bucket() {
289 let stats = ProxyStats::default();
290
291 stats.record_provider_request("Mystery", 1_000, 500);
292
293 assert_eq!(stats.requests_total.load(Ordering::Relaxed), 1);
295 assert_eq!(stats.anthropic.requests_total.load(Ordering::Relaxed), 0);
296 assert_eq!(stats.openai.requests_total.load(Ordering::Relaxed), 0);
297 assert_eq!(stats.chatgpt.requests_total.load(Ordering::Relaxed), 0);
298 assert_eq!(stats.gemini.requests_total.load(Ordering::Relaxed), 0);
299 }
300}
301
302fn connect_timeout_secs() -> u64 {
304 std::env::var("LEAN_CTX_PROXY_CONNECT_TIMEOUT_SECS")
305 .ok()
306 .and_then(|v| v.trim().parse::<u64>().ok())
307 .filter(|s| *s > 0)
308 .unwrap_or(15)
309}
310
311fn read_idle_timeout_secs() -> u64 {
316 std::env::var("LEAN_CTX_PROXY_READ_TIMEOUT_SECS")
317 .ok()
318 .and_then(|v| v.trim().parse::<u64>().ok())
319 .filter(|s| *s > 0)
320 .unwrap_or(300)
321}
322
323fn upstream_reload_secs() -> u64 {
326 std::env::var("LEAN_CTX_PROXY_RELOAD_SECS")
327 .ok()
328 .and_then(|v| v.trim().parse::<u64>().ok())
329 .filter(|s| *s > 0)
330 .unwrap_or(5)
331}
332
333fn spawn_upstream_refresh(tx: tokio::sync::watch::Sender<Arc<Upstreams>>, initial: Upstreams) {
342 let interval = std::time::Duration::from_secs(upstream_reload_secs());
343 tokio::spawn(async move {
344 let mut last = initial;
345 loop {
346 tokio::time::sleep(interval).await;
347 let next = crate::core::config::Config::load()
348 .proxy
349 .refresh_upstreams(&last);
350 if next != last {
351 log_upstream_change(&last, &next);
352 last = next.clone();
353 if tx.send(Arc::new(next)).is_err() {
354 break;
355 }
356 }
357 }
358 });
359}
360
361fn log_upstream_change(old: &Upstreams, new: &Upstreams) {
364 if old.anthropic != new.anthropic {
365 println!(" ↻ Anthropic upstream → {}", new.anthropic);
366 }
367 if old.openai != new.openai {
368 println!(" ↻ OpenAI upstream → {}", new.openai);
369 }
370 if old.chatgpt != new.chatgpt {
371 println!(" ↻ ChatGPT upstream → {}", new.chatgpt);
372 }
373 if old.gemini != new.gemini {
374 println!(" ↻ Gemini upstream → {}", new.gemini);
375 }
376}
377
378pub async fn start_proxy(port: u16) -> anyhow::Result<()> {
379 let token = crate::core::session_token::resolve_proxy_token("LEAN_CTX_PROXY_TOKEN");
380 start_proxy_with_token(port, Some(token)).await
381}
382
383fn effective_auth_token(auth_token: Option<String>) -> String {
388 auth_token
389 .filter(|t| !t.trim().is_empty())
390 .unwrap_or_else(|| crate::core::session_token::resolve_proxy_token("LEAN_CTX_PROXY_TOKEN"))
391}
392
393pub async fn start_proxy_with_token(port: u16, auth_token: Option<String>) -> anyhow::Result<()> {
394 use crate::core::config::{Config, is_local_proxy_url};
395
396 let auth_token = effective_auth_token(auth_token);
397
398 let client = chatgpt_cookies::with_chatgpt_cloudflare_cookie_store(reqwest::Client::builder())
403 .connect_timeout(std::time::Duration::from_secs(connect_timeout_secs()))
404 .read_timeout(std::time::Duration::from_secs(read_idle_timeout_secs()))
405 .build()?;
406
407 usage_meter::resume_from_disk();
410 cold_prefix::resume_from_disk();
413
414 let cfg = Config::load();
415 let require_token = cfg.proxy_require_token;
417 let initial = cfg.proxy.resolve_all();
418
419 let (upstream_tx, upstream_rx) = tokio::sync::watch::channel(Arc::new(initial.clone()));
424 spawn_upstream_refresh(upstream_tx, initial.clone());
425
426 let Upstreams {
427 anthropic: anthropic_upstream,
428 openai: openai_upstream,
429 chatgpt: chatgpt_upstream,
430 gemini: gemini_upstream,
431 } = initial;
432
433 let state = ProxyState {
434 client,
435 port,
436 stats: Arc::new(ProxyStats::default()),
437 introspect: Arc::new(introspect::IntrospectState::default()),
438 upstreams: upstream_rx,
439 };
440
441 let mut app = Router::new()
442 .route("/health", get(health))
443 .route("/status", get(status_handler))
444 .route("/v1/messages", any(anthropic::handler))
445 .route("/v1/messages/{*rest}", any(anthropic::handler))
446 .route("/v1/chat/completions", any(openai::handler))
447 .route(
449 "/v1/responses",
450 post(openai_responses::handler).get(openai_responses::ws_handler),
451 )
452 .route("/v1/responses/{*rest}", any(openai_responses::handler))
453 .route("/messages", any(anthropic::handler))
459 .route("/messages/{*rest}", any(anthropic::handler))
460 .route("/chat/completions", any(openai::handler))
461 .route(
462 "/responses",
463 post(openai_responses::handler).get(openai_responses::ws_handler),
464 )
465 .route("/responses/{*rest}", any(openai_responses::handler))
466 .route(
467 "/backend-api/codex/responses",
468 post(chatgpt::codex_responses_handler).get(chatgpt::codex_responses_ws_handler),
469 )
470 .route(
471 "/backend-api/codex/responses/{*rest}",
472 any(chatgpt::codex_responses_handler),
473 )
474 .route("/backend-api", any(chatgpt::backend_api_handler))
477 .route("/backend-api/{*rest}", any(chatgpt::backend_api_handler))
478 .route("/v1/references/{id}", get(v1_resolve_reference))
479 .route("/v1/compress", post(compress_api::handler))
482 .fallback(fallback_router)
483 .layer(axum::middleware::from_fn(host_guard))
484 .with_state(state);
485
486 {
487 let expected = auth_token.clone();
488 app = app.layer(axum::middleware::from_fn(move |req, next| {
489 let expected = expected.clone();
490 proxy_auth_guard(req, next, expected, require_token)
491 }));
492 }
493
494 app = app.layer(axum::middleware::from_fn(normalize_provider_path));
498
499 let addr = SocketAddr::from(([127, 0, 0, 1], port));
500 println!("lean-ctx proxy listening on http://{addr} (token auth enabled)");
501 println!(" Anthropic: POST /v1/messages → {anthropic_upstream}");
502 println!(" OpenAI: POST /v1/chat/completions → {openai_upstream}");
503 println!(
504 " OpenAI: POST /v1/responses → {openai_upstream} (bare /responses also accepted)"
505 );
506 println!(" ChatGPT: POST /backend-api/codex/responses → {chatgpt_upstream}");
507 println!(" ChatGPT: any /backend-api/* → {chatgpt_upstream}");
508 println!(" Gemini: POST /v1beta/models/... → {gemini_upstream}");
509 println!(" Compress: POST /v1/compress (deterministic messages-in/out, local)");
510 println!(
514 " Codex: WS ws://{addr}/responses → bridged to {openai_upstream} (HTTP/SSE, #440)"
515 );
516 if openai_upstream.starts_with("http://") && !is_local_proxy_url(&openai_upstream) {
517 println!(
518 " ⚠ OpenAI upstream is plaintext HTTP to a non-loopback host \
519 (allow_insecure_http_upstream) — use only on a trusted local network"
520 );
521 }
522
523 let listener = tokio::net::TcpListener::bind(addr).await?;
524 axum::serve(listener, app)
525 .with_graceful_shutdown(shutdown_signal())
526 .await?;
527
528 println!("lean-ctx proxy shut down cleanly.");
529 Ok(())
530}
531
532async fn shutdown_signal() {
533 let ctrl_c = tokio::signal::ctrl_c();
534
535 #[cfg(unix)]
536 {
537 match tokio::signal::unix::signal(tokio::signal::unix::SignalKind::terminate()) {
540 Ok(mut sigterm) => {
541 tokio::select! {
542 _ = ctrl_c => {},
543 _ = sigterm.recv() => {},
544 }
545 }
546 Err(e) => {
547 tracing::warn!("lean-ctx proxy: SIGTERM handler unavailable ({e}); Ctrl-C only");
548 ctrl_c.await.ok();
549 }
550 }
551 }
552
553 #[cfg(not(unix))]
554 {
555 ctrl_c.await.ok();
556 }
557
558 println!("lean-ctx proxy: received shutdown signal, draining…");
559}
560
561async fn health() -> impl IntoResponse {
562 let body = serde_json::json!({
563 "status": "ok",
564 "pid": std::process::id(),
565 });
566 (StatusCode::OK, axum::Json(body))
567}
568
569async fn v1_resolve_reference(
570 axum::extract::Path(id): axum::extract::Path<String>,
571) -> impl IntoResponse {
572 match crate::server::reference_store::resolve(&id) {
573 Some(content) => (StatusCode::OK, content),
574 None => (
575 StatusCode::NOT_FOUND,
576 "Reference expired or not found".to_string(),
577 ),
578 }
579}
580
581async fn status_handler(State(state): State<ProxyState>) -> impl IntoResponse {
582 use std::sync::atomic::Ordering::Relaxed;
583 let s = &state.stats;
584 let i = &state.introspect;
585
586 let last_breakdown = i
587 .last_breakdown
588 .lock()
589 .ok()
590 .and_then(|guard| guard.as_ref().map(|b| serde_json::to_value(b).ok()))
591 .flatten();
592
593 let spend = usage_meter::snapshot();
594 let spend_total: f64 = spend.iter().map(|m| m.cost_usd).sum();
595
596 let up = state.upstream_snapshot();
600
601 let active_effort = crate::core::config::Config::load().proxy.resolved_effort();
604
605 let body = serde_json::json!({
606 "status": "running",
607 "port": state.port,
608 "upstreams": {
609 "anthropic": up.anthropic.clone(),
610 "openai": up.openai.clone(),
611 "chatgpt": up.chatgpt.clone(),
612 "gemini": up.gemini.clone(),
613 },
614 "requests_total": s.requests_total.load(Relaxed),
615 "requests_compressed": s.requests_compressed.load(Relaxed),
616 "tokens_saved": s.tokens_saved.load(Relaxed),
617 "tokens_saved_estimated": true,
618 "bytes_original": s.bytes_original.load(Relaxed),
619 "bytes_compressed": s.bytes_compressed.load(Relaxed),
620 "compression_ratio_pct": format!("{:.1}", s.compression_ratio()),
621 "per_upstream": s.provider_summary(),
622 "cache_safety": cache_safety::snapshot(),
623 "effort": effort::snapshot(active_effort),
624 "per_model": cost::snapshot(),
625 "spend": {
626 "source": "measured",
627 "total_usd": spend_total,
628 "per_model": spend,
629 "note": "Actual provider bill: real model + billed tokens (incl. cache reads/writes & reasoning) read from upstream responses for proxy-routed clients."
630 },
631 "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.",
632 "introspect": {
633 "total_requests_analyzed": i.total_requests.load(Relaxed),
634 "total_system_prompt_tokens": i.total_system_prompt_tokens.load(Relaxed),
635 "last_breakdown": last_breakdown,
636 }
637 });
638 (StatusCode::OK, axum::Json(body))
639}
640
641#[allow(clippy::result_large_err)]
642async fn proxy_auth_guard(
643 req: axum::extract::Request,
644 next: axum::middleware::Next,
645 expected_token: String,
646 require_token: bool,
647) -> Result<Response, Response> {
648 let path = req.uri().path();
649 if path == "/health" {
650 return Ok(next.run(req).await);
651 }
652
653 if let Some(auth) = req
654 .headers()
655 .get("authorization")
656 .and_then(|v| v.to_str().ok())
657 && let Some(token) = auth.strip_prefix("Bearer ")
658 && constant_time_eq(token.as_bytes(), expected_token.as_bytes())
659 {
660 return Ok(next.run(req).await);
661 }
662
663 if provider_key_fallback_allowed(
669 require_token,
670 has_provider_api_key(&req),
671 is_provider_route(path),
672 ) {
673 return Ok(next.run(req).await);
674 }
675
676 let cfg = crate::core::config::Config::load();
677 let hint = match cfg.proxy_enabled {
678 Some(true) => {
679 "lean-ctx proxy requires authentication. Use a Bearer token (LEAN_CTX_PROXY_TOKEN) or configure your AI tool's API key."
680 }
681 Some(false) => "lean-ctx proxy is disabled but still running. Run: lean-ctx proxy cleanup",
682 None => {
683 "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"
684 }
685 };
686
687 let body = serde_json::json!({
688 "type": "error",
689 "error": {
690 "type": "authentication_error",
691 "message": format!("401 Unauthorized — {hint}")
692 }
693 });
694
695 Err((StatusCode::UNAUTHORIZED, axum::Json(body)).into_response())
696}
697
698fn has_provider_api_key(req: &axum::extract::Request) -> bool {
699 let headers = req.headers();
700 for key in ["x-api-key", "x-goog-api-key", "api-key"] {
703 if headers
704 .get(key)
705 .and_then(|v| v.to_str().ok())
706 .is_some_and(|v| !v.trim().is_empty())
707 {
708 return true;
709 }
710 }
711 if let Some(auth) = headers.get("authorization").and_then(|v| v.to_str().ok()) {
720 let auth = auth.trim();
721 let credential = auth
722 .strip_prefix("Bearer ")
723 .or_else(|| auth.strip_prefix("bearer "))
724 .unwrap_or(auth)
725 .trim();
726 return !credential.is_empty() && !credential.eq_ignore_ascii_case("bearer");
728 }
729 false
730}
731
732fn is_provider_route(path: &str) -> bool {
733 path.starts_with("/v1/")
734 || path.starts_with("/v1beta/")
735 || path.starts_with("/chat/completions")
736 || path.starts_with("/responses")
737 || path.starts_with("/messages")
738 || path.starts_with("/backend-api")
739}
740
741fn provider_key_fallback_allowed(
748 require_token: bool,
749 has_provider_key: bool,
750 is_provider_route: bool,
751) -> bool {
752 !require_token && has_provider_key && is_provider_route
753}
754
755fn canonical_provider_path(path: &str) -> Option<String> {
766 if let Some(rest) = path.strip_prefix("/v1/v1/") {
770 return Some(format!("/v1/{rest}"));
771 }
772 const BARE_TO_CANONICAL: &[(&str, &str, &str)] = &[
773 ("/responses", "/v1/responses", "/responses/"),
774 (
775 "/chat/completions",
776 "/v1/chat/completions",
777 "/chat/completions/",
778 ),
779 ("/messages", "/v1/messages", "/messages/"),
780 ];
781 for (bare, canonical, bare_with_slash) in BARE_TO_CANONICAL {
782 if path == *bare {
783 return Some((*canonical).to_string());
784 }
785 if let Some(rest) = path.strip_prefix(bare_with_slash) {
786 return Some(format!("{canonical}/{rest}"));
787 }
788 }
789 None
790}
791
792fn normalized_provider_uri(uri: &axum::http::Uri) -> Option<axum::http::Uri> {
796 let canonical = canonical_provider_path(uri.path())?;
797 let new_path_and_query = match uri.query() {
798 Some(q) => format!("{canonical}?{q}"),
799 None => canonical,
800 };
801 new_path_and_query.parse::<axum::http::Uri>().ok()
802}
803
804async fn normalize_provider_path(
808 mut req: axum::extract::Request,
809 next: axum::middleware::Next,
810) -> Response {
811 if let Some(uri) = normalized_provider_uri(req.uri()) {
812 *req.uri_mut() = uri;
813 }
814 next.run(req).await
815}
816
817fn constant_time_eq(a: &[u8], b: &[u8]) -> bool {
818 use subtle::ConstantTimeEq;
819 if a.len() != b.len() {
820 return false;
821 }
822 bool::from(a.ct_eq(b))
823}
824
825async fn host_guard(
826 req: axum::extract::Request,
827 next: axum::middleware::Next,
828) -> Result<Response, StatusCode> {
829 if let Some(host) = req.headers().get("host").and_then(|v| v.to_str().ok()) {
830 let h = host.split(':').next().unwrap_or(host);
831 if matches!(h, "127.0.0.1" | "localhost" | "[::1]") {
832 return Ok(next.run(req).await);
833 }
834 }
835 Err(StatusCode::FORBIDDEN)
836}
837
838async fn fallback_router(State(state): State<ProxyState>, req: Request<Body>) -> Response {
839 let path = req.uri().path().to_string();
840
841 if path.starts_with("/v1beta/models/") || path.starts_with("/v1/models/") {
842 match google::handler(State(state), req).await {
843 Ok(resp) => resp,
844 Err(status) => Response::builder()
845 .status(status)
846 .body(Body::from("proxy error"))
847 .expect("BUG: building error response with valid status should never fail"),
848 }
849 } else {
850 let method = req.method().to_string();
851 eprintln!("lean-ctx proxy: unmatched {method} {path}");
852 Response::builder()
853 .status(StatusCode::NOT_FOUND)
854 .body(Body::from(format!(
855 "lean-ctx proxy: no handler for {method} {path}"
856 )))
857 .expect("BUG: building 404 response should never fail")
858 }
859}
860
861#[cfg(test)]
862mod auth_tests {
863 use super::*;
864
865 #[test]
868 fn effective_auth_token_never_yields_empty() {
869 let _env = crate::core::data_dir::test_env_lock();
870 let tmp = tempfile::tempdir().unwrap();
871 crate::test_env::set_var("LEAN_CTX_DATA_DIR", tmp.path());
872
873 assert_eq!(effective_auth_token(Some("tok".into())), "tok");
874 let auto = effective_auth_token(None);
875 assert!(!auto.trim().is_empty(), "None must auto-resolve a token");
876 let blank = effective_auth_token(Some(" ".into()));
877 assert!(!blank.trim().is_empty(), "blank tokens must be replaced");
878
879 crate::test_env::remove_var("LEAN_CTX_DATA_DIR");
880 }
881
882 #[test]
883 fn is_provider_route_v1() {
884 assert!(is_provider_route("/v1/chat/completions"));
885 assert!(is_provider_route("/v1/messages"));
886 assert!(is_provider_route("/v1/completions"));
887 }
888
889 #[test]
890 fn is_provider_route_anthropic_subpaths() {
891 assert!(is_provider_route("/v1/messages/count_tokens"));
892 assert!(is_provider_route("/v1/messages/batches"));
893 assert!(is_provider_route("/v1/messages/batches/batch_123"));
894 }
895
896 #[test]
897 fn is_provider_route_v1beta() {
898 assert!(is_provider_route("/v1beta/models"));
899 }
900
901 #[test]
902 fn is_provider_route_chat() {
903 assert!(is_provider_route("/chat/completions"));
904 }
905
906 #[test]
907 fn is_provider_route_chatgpt_backend_api() {
908 assert!(is_provider_route("/backend-api/codex/responses"));
909 assert!(is_provider_route("/backend-api/codex/responses/resp_123"));
910 assert!(is_provider_route("/backend-api/wham/session"));
911 assert!(is_provider_route("/backend-api/ps/mcp"));
912 assert!(is_provider_route("/backend-api/codex_apps"));
913 assert!(is_provider_route("/backend-api/codex_apps/mcp"));
914 assert!(is_provider_route("/backend-api/mcp/codex_apps"));
915 assert!(is_provider_route("/backend-api/apps/codex_apps/mcp"));
916 }
917
918 #[test]
919 fn is_provider_route_rejects_non_provider() {
920 assert!(!is_provider_route("/health"));
921 assert!(!is_provider_route("/api/v2/test"));
922 assert!(!is_provider_route("/"));
923 }
924
925 fn build_request(headers: &[(&str, &str)], path: &str) -> axum::extract::Request {
926 let mut builder = axum::http::Request::builder().uri(path);
927 for (k, v) in headers {
928 builder = builder.header(*k, *v);
929 }
930 builder.body(axum::body::Body::empty()).unwrap()
931 }
932
933 #[test]
934 fn has_provider_api_key_x_api_key() {
935 let req = build_request(&[("x-api-key", "sk-ant-abc123")], "/v1/messages");
936 assert!(has_provider_api_key(&req));
937 }
938
939 #[test]
940 fn has_provider_api_key_x_goog() {
941 let req = build_request(&[("x-goog-api-key", "AIzaSyAbc")], "/v1beta/models");
942 assert!(has_provider_api_key(&req));
943 }
944
945 #[test]
946 fn has_provider_api_key_azure() {
947 let req = build_request(&[("api-key", "deadbeef")], "/v1/completions");
948 assert!(has_provider_api_key(&req));
949 }
950
951 #[test]
952 fn has_provider_api_key_bearer_sk() {
953 let req = build_request(
954 &[("authorization", "Bearer sk-proj-abc123")],
955 "/v1/chat/completions",
956 );
957 assert!(has_provider_api_key(&req));
958 }
959
960 #[test]
961 fn has_provider_api_key_empty_rejected() {
962 let req = build_request(&[("x-api-key", " ")], "/v1/messages");
963 assert!(!has_provider_api_key(&req));
964 }
965
966 #[test]
967 fn has_provider_api_key_no_headers() {
968 let req = build_request(&[], "/v1/messages");
969 assert!(!has_provider_api_key(&req));
970 }
971
972 #[test]
973 fn has_provider_api_key_accepts_non_sk_bearer() {
974 for key in [
980 "Bearer or-v1-9f8e7d6c", "Bearer gsk_live_1234", "Bearer abc.def.ghi", "Bearer 0123456789", ] {
985 let req = build_request(&[("authorization", key)], "/v1/responses");
986 assert!(
987 has_provider_api_key(&req),
988 "non-sk Bearer must count as a provider credential: {key}"
989 );
990 }
991 }
992
993 #[test]
994 fn has_provider_api_key_empty_bearer_rejected() {
995 for bad in ["Bearer ", "", "Bearer", "bearer", " "] {
998 let req = build_request(&[("authorization", bad)], "/responses");
999 assert!(
1000 !has_provider_api_key(&req),
1001 "blank/scheme-only Authorization must not authenticate: {bad:?}"
1002 );
1003 }
1004 }
1005
1006 #[test]
1009 fn provider_key_fallback_allowed_in_default_mode() {
1010 assert!(provider_key_fallback_allowed(false, true, true));
1014 }
1015
1016 #[test]
1017 fn provider_key_fallback_denied_in_strict_mode() {
1018 assert!(!provider_key_fallback_allowed(true, true, true));
1022 }
1023
1024 #[test]
1025 fn provider_key_fallback_requires_key_and_provider_route() {
1026 assert!(!provider_key_fallback_allowed(false, false, true));
1029 assert!(!provider_key_fallback_allowed(false, true, false));
1030 assert!(!provider_key_fallback_allowed(true, false, true));
1031 }
1032
1033 #[test]
1034 fn proxy_require_token_defaults_off() {
1035 assert!(!crate::core::config::Config::default().proxy_require_token);
1039 }
1040
1041 #[test]
1044 fn is_provider_route_bare_responses_and_messages() {
1045 assert!(is_provider_route("/responses"));
1048 assert!(is_provider_route("/responses/resp_123/input_items"));
1049 assert!(is_provider_route("/messages"));
1050 }
1051
1052 #[test]
1053 fn canonical_provider_path_rewrites_bare_endpoints() {
1054 assert_eq!(
1055 canonical_provider_path("/responses").as_deref(),
1056 Some("/v1/responses")
1057 );
1058 assert_eq!(
1059 canonical_provider_path("/chat/completions").as_deref(),
1060 Some("/v1/chat/completions")
1061 );
1062 assert_eq!(
1063 canonical_provider_path("/messages").as_deref(),
1064 Some("/v1/messages")
1065 );
1066 }
1067
1068 #[test]
1069 fn canonical_provider_path_preserves_subpaths() {
1070 assert_eq!(
1071 canonical_provider_path("/responses/resp_abc/cancel").as_deref(),
1072 Some("/v1/responses/resp_abc/cancel")
1073 );
1074 assert_eq!(
1075 canonical_provider_path("/messages/batches/batch_1").as_deref(),
1076 Some("/v1/messages/batches/batch_1")
1077 );
1078 }
1079
1080 #[test]
1081 fn canonical_provider_path_ignores_already_canonical_and_unknown() {
1082 assert_eq!(canonical_provider_path("/v1/responses"), None);
1084 assert_eq!(canonical_provider_path("/v1/chat/completions"), None);
1085 assert_eq!(canonical_provider_path("/health"), None);
1087 assert_eq!(canonical_provider_path("/responsesx"), None);
1088 assert_eq!(canonical_provider_path("/"), None);
1089 }
1090
1091 #[test]
1092 fn canonical_provider_path_collapses_double_v1_prefix() {
1093 assert_eq!(
1096 canonical_provider_path("/v1/v1/responses").as_deref(),
1097 Some("/v1/responses")
1098 );
1099 assert_eq!(
1100 canonical_provider_path("/v1/v1/chat/completions").as_deref(),
1101 Some("/v1/chat/completions")
1102 );
1103 }
1104
1105 #[test]
1106 fn normalized_provider_uri_rewrites_path_and_preserves_query() {
1107 use axum::http::Uri;
1108 let uri: Uri = "/responses?stream=true".parse().unwrap();
1109 let rewritten = normalized_provider_uri(&uri).expect("bare /responses must rewrite");
1110 assert_eq!(rewritten.path(), "/v1/responses");
1111 assert_eq!(rewritten.query(), Some("stream=true"));
1112 assert_eq!(
1113 rewritten
1114 .path_and_query()
1115 .map(axum::http::uri::PathAndQuery::as_str),
1116 Some("/v1/responses?stream=true")
1117 );
1118 }
1119
1120 #[test]
1121 fn normalized_provider_uri_noop_for_canonical() {
1122 use axum::http::Uri;
1123 let uri: Uri = "/v1/responses".parse().unwrap();
1124 assert!(normalized_provider_uri(&uri).is_none());
1125 }
1126}
1127
1128#[cfg(test)]
1129mod upstream_tests {
1130 use super::*;
1131
1132 fn upstreams_with_openai(openai: &str) -> Upstreams {
1133 Upstreams {
1134 anthropic: "https://api.anthropic.com".into(),
1135 openai: openai.into(),
1136 chatgpt: "https://chatgpt.com".into(),
1137 gemini: "https://generativelanguage.googleapis.com".into(),
1138 }
1139 }
1140
1141 #[tokio::test]
1145 async fn proxy_state_reads_upstream_live_from_watch() {
1146 let (tx, rx) =
1147 tokio::sync::watch::channel(Arc::new(upstreams_with_openai("https://old.example")));
1148 let state = ProxyState {
1149 client: reqwest::Client::new(),
1150 port: 0,
1151 stats: Arc::new(ProxyStats::default()),
1152 introspect: Arc::new(introspect::IntrospectState::default()),
1153 upstreams: rx,
1154 };
1155 assert_eq!(state.openai_upstream(), "https://old.example");
1156
1157 tx.send(Arc::new(upstreams_with_openai("https://new.example")))
1158 .unwrap();
1159 assert_eq!(
1160 state.openai_upstream(),
1161 "https://new.example",
1162 "a live handler read must reflect the published change"
1163 );
1164 assert_eq!(state.upstream_snapshot().openai, "https://new.example");
1165 }
1166
1167 #[tokio::test]
1177 #[allow(clippy::await_holding_lock)]
1178 async fn config_change_is_picked_up_live_without_restart() {
1179 use crate::core::config::Config;
1180
1181 let _lock = crate::core::data_dir::test_env_lock();
1182 let tmp = tempfile::tempdir().unwrap();
1183 crate::test_env::set_var("LEAN_CTX_DATA_DIR", tmp.path());
1184 crate::test_env::remove_var("LEAN_CTX_OPENAI_UPSTREAM");
1187 crate::test_env::set_var("LEAN_CTX_PROXY_RELOAD_SECS", "1");
1188
1189 Config::update_global(|c| {
1191 c.proxy.openai_upstream = Some("http://127.0.0.1:19101".into());
1192 })
1193 .unwrap();
1194 let initial = Config::load().proxy.resolve_all();
1195 assert_eq!(initial.openai, "http://127.0.0.1:19101");
1196
1197 let (tx, rx) = tokio::sync::watch::channel(Arc::new(initial.clone()));
1198 spawn_upstream_refresh(tx, initial);
1199
1200 Config::update_global(|c| {
1202 c.proxy.openai_upstream = Some("http://127.0.0.1:19102".into());
1203 })
1204 .unwrap();
1205
1206 let mut live = rx.borrow().openai.clone();
1208 for _ in 0..80 {
1209 if live == "http://127.0.0.1:19102" {
1210 break;
1211 }
1212 tokio::time::sleep(std::time::Duration::from_millis(100)).await;
1213 live = rx.borrow().openai.clone();
1214 }
1215 assert_eq!(
1216 live, "http://127.0.0.1:19102",
1217 "running proxy must serve the new config.toml upstream without a restart"
1218 );
1219
1220 crate::test_env::remove_var("LEAN_CTX_PROXY_RELOAD_SECS");
1221 crate::test_env::remove_var("LEAN_CTX_DATA_DIR");
1222 }
1223}