Skip to main content

lean_ctx/proxy/
mod.rs

1//! LLM reverse proxy — the core of the **Gateway pillar**.
2//!
3//! Intercepts Anthropic, OpenAI, Gemini and ChatGPT traffic, compresses
4//! prompts, meters usage, and optionally translates request shapes.
5//!
6//! # TLS Requirements
7//!
8//! The proxy does NOT terminate TLS. For production/multi-tenant deployments,
9//! place behind a TLS-terminating reverse proxy (nginx, caddy, traefik).
10//! For mTLS, configure the reverse proxy with client certificates.
11//! See: <https://github.com/yvgude/lean-ctx/issues/1371>
12//!
13//! # Org gateway coupling (ADR-023)
14//!
15//! The self-hosted org gateway (lean-ctx-enterprise) extends this proxy with
16//! admin/usage-store routes. In the OSS build the proxy runs standalone.
17
18pub mod anthropic;
19#[cfg(test)]
20mod auth_tests;
21pub mod bedrock;
22pub mod break_even;
23pub mod cache_aligner;
24pub mod cache_attribution;
25pub mod cache_breakpoint;
26pub mod cache_policy;
27pub mod cache_safety;
28pub mod ccr;
29#[cfg(test)]
30mod ccr_robustness_tests;
31pub mod chatgpt;
32pub mod chatgpt_cookies;
33pub mod chatgpt_ws;
34mod codec;
35pub mod cold_prefix;
36pub mod compress;
37pub mod compress_api;
38pub mod compress_shared;
39mod connector;
40pub(crate) mod conversation;
41pub mod cost;
42pub mod counterfactual;
43#[allow(dead_code)]
44pub mod dedup;
45pub mod effort;
46pub mod effort_routing;
47pub mod forward;
48pub mod gateway_identity;
49pub mod google;
50pub mod history_prune;
51pub mod holdout;
52pub mod image_compression;
53mod intent;
54pub mod introspect;
55pub mod latency_guard;
56mod lineage;
57pub mod metrics;
58pub mod model_router;
59pub mod models_api;
60pub mod ocla_cache_bridge;
61pub mod openai;
62pub mod openai_responses;
63pub mod openai_responses_ws;
64pub mod output_savings;
65pub mod pii;
66pub mod policy_gate;
67pub mod prefix_cache_stats;
68pub mod prefix_replay;
69pub mod prose;
70pub mod prose_ranker;
71pub mod providers;
72pub mod quality_lab_api;
73pub mod response_optimizer;
74#[allow(dead_code)]
75pub(crate) mod response_shaper;
76pub mod routing;
77pub mod routing_feedback;
78#[cfg(feature = "shape-xlat")]
79pub mod shape_xlat;
80#[allow(dead_code)]
81pub(crate) mod shaping_hook;
82pub mod sse_keepalive;
83#[cfg(test)]
84mod stats_tests;
85pub mod sticky_tools;
86pub mod tool_kind;
87#[allow(dead_code)]
88pub mod tool_output;
89#[cfg(test)]
90mod upstream_tests;
91pub mod usage;
92pub mod usage_accounting;
93pub mod usage_azure;
94pub mod usage_meter;
95pub mod usage_parity;
96pub mod usage_sink;
97pub mod verbosity;
98pub mod web_app;
99#[allow(dead_code)]
100pub(crate) mod web_app_middleware;
101
102use std::net::SocketAddr;
103use std::sync::Arc;
104use std::sync::atomic::{AtomicU64, Ordering};
105
106use crate::core::config::Upstreams;
107
108use axum::{
109    Router,
110    body::Body,
111    extract::State,
112    http::{Request, StatusCode},
113    response::{IntoResponse, Response},
114    routing::{any, get, post},
115};
116
117#[derive(Clone)]
118pub struct ProxyState {
119    pub client: reqwest::Client,
120    pub port: u16,
121    pub stats: Arc<ProxyStats>,
122    pub(crate) break_even: Arc<break_even::BreakEvenCalculator>,
123    pub introspect: Arc<introspect::IntrospectState>,
124    pub ocla_cache: Option<Arc<ocla_cache_bridge::OclaCacheBridge>>,
125    /// Live provider upstreams, refreshed from config.toml without a proxy
126    /// restart (#449). Read per request via [`ProxyState::openai_upstream`] etc.
127    pub upstreams: tokio::sync::watch::Receiver<Arc<Upstreams>>,
128    /// Shared Cloudflare cookie jar (also wired into `client`), so the Codex
129    /// ChatGPT WebSocket passthrough replays the same clearance to chatgpt.com
130    /// that the reqwest rail accumulated (#597).
131    pub(crate) chatgpt_cookies: Arc<chatgpt_cookies::ChatGptCloudflareCookieStore>,
132    /// Resolved `[[gateway_server.mcp_servers]]` registry snapshot (GL#100).
133    /// Empty when none are registered; the `/mcp/{server}` routes are only
134    /// mounted with the `gateway-server` feature. Startup snapshot, restart
135    /// to reload — the same lifecycle as `gateway-keys.toml`.
136    pub mcp_servers: Arc<Vec<crate::core::config::ResolvedMcpServer>>,
137    pub(crate) web_app_tracker:
138        Arc<std::sync::Mutex<web_app::conversation_tracker::ConversationTracker>>,
139}
140
141impl ProxyState {
142    /// Consistent snapshot of all upstreams for the current request/response.
143    pub fn upstream_snapshot(&self) -> Arc<Upstreams> {
144        self.upstreams.borrow().clone()
145    }
146
147    /// Snapshot MCP schema-cost break-even status for `/status`.
148    pub(crate) fn break_even_summary(&self) -> break_even::BreakEvenSummary {
149        self.break_even.summary()
150    }
151
152    /// Current Anthropic upstream (live).
153    pub fn anthropic_upstream(&self) -> String {
154        self.upstreams.borrow().anthropic.clone()
155    }
156
157    /// Current OpenAI upstream (live).
158    pub fn openai_upstream(&self) -> String {
159        self.upstreams.borrow().openai.clone()
160    }
161
162    /// Current ChatGPT upstream (live).
163    pub fn chatgpt_upstream(&self) -> String {
164        self.upstreams.borrow().chatgpt.clone()
165    }
166
167    /// Current Gemini upstream (live).
168    pub fn gemini_upstream(&self) -> String {
169        self.upstreams.borrow().gemini.clone()
170    }
171
172    /// Cloudflare `Cookie` header for the current ChatGPT upstream, used by the
173    /// WebSocket passthrough handshake (#597). `None` until a request on the
174    /// reqwest rail has seen Cloudflare clearance.
175    pub fn chatgpt_cookie_header(&self) -> Option<String> {
176        let url = reqwest::Url::parse(&self.chatgpt_upstream()).ok()?;
177        self.chatgpt_cookies
178            .cookie_header(&url)
179            .and_then(|v| v.to_str().ok().map(str::to_owned))
180    }
181}
182
183pub struct ProxyStats {
184    pub requests_total: AtomicU64,
185    pub requests_compressed: AtomicU64,
186    pub tokens_saved: AtomicU64,
187    pub bytes_original: AtomicU64,
188    pub bytes_compressed: AtomicU64,
189    pub anthropic: ProviderStats,
190    pub openai: ProviderStats,
191    pub chatgpt: ProviderStats,
192    pub gemini: ProviderStats,
193    /// Grok dual-rail (`/providers/grok-chat`, `/providers/xai`) — OpenAI wire
194    /// shape, separate identity so `proxy status` does not fold xAI traffic
195    /// into the OpenAI bucket.
196    pub grok: ProviderStats,
197    /// Command Code gateway rail (`/providers/commandcode`) — separate identity
198    /// so `proxy status` does not fold Command Code traffic into the wire-shape
199    /// bucket.
200    pub commandcode: ProviderStats,
201}
202
203#[derive(Default)]
204pub struct ProviderStats {
205    pub requests_total: AtomicU64,
206    pub requests_compressed: AtomicU64,
207    pub tokens_saved: AtomicU64,
208    pub bytes_original: AtomicU64,
209    pub bytes_compressed: AtomicU64,
210}
211
212impl Default for ProxyStats {
213    fn default() -> Self {
214        Self {
215            requests_total: AtomicU64::new(0),
216            requests_compressed: AtomicU64::new(0),
217            tokens_saved: AtomicU64::new(0),
218            bytes_original: AtomicU64::new(0),
219            bytes_compressed: AtomicU64::new(0),
220            anthropic: ProviderStats::default(),
221            openai: ProviderStats::default(),
222            chatgpt: ProviderStats::default(),
223            gemini: ProviderStats::default(),
224            grok: ProviderStats::default(),
225            commandcode: ProviderStats::default(),
226        }
227    }
228}
229
230impl ProxyStats {
231    pub fn record_request(&self, original: usize, compressed: usize) {
232        self.record_totals(original, compressed);
233    }
234
235    pub fn record_provider_request(
236        &self,
237        provider_label: &str,
238        original: usize,
239        compressed: usize,
240    ) {
241        let (effective_compressed, saved_tokens, compressed_request) =
242            self.record_totals(original, compressed);
243
244        if let Some(provider) = self.provider(provider_label) {
245            provider.record(
246                original,
247                effective_compressed,
248                compressed_request,
249                saved_tokens,
250            );
251        }
252    }
253
254    fn record_totals(&self, original: usize, compressed: usize) -> (usize, u64, bool) {
255        self.requests_total.fetch_add(1, Ordering::Relaxed);
256        self.bytes_original
257            .fetch_add(original as u64, Ordering::Relaxed);
258        let effective_compressed = compressed.min(original);
259        self.bytes_compressed
260            .fetch_add(effective_compressed as u64, Ordering::Relaxed);
261        if compressed < original {
262            self.requests_compressed.fetch_add(1, Ordering::Relaxed);
263        }
264        let saved_tokens = (original.saturating_sub(effective_compressed) / 4) as u64;
265        self.tokens_saved.fetch_add(saved_tokens, Ordering::Relaxed);
266        (effective_compressed, saved_tokens, compressed < original)
267    }
268
269    pub fn compression_ratio(&self) -> f64 {
270        let original = self.bytes_original.load(Ordering::Relaxed);
271        if original == 0 {
272            return 0.0;
273        }
274        let compressed = self.bytes_compressed.load(Ordering::Relaxed);
275        (1.0 - compressed as f64 / original as f64) * 100.0
276    }
277
278    /// Maps a proxy `provider_label` to its per-upstream bucket. Unknown labels
279    /// return `None` (still counted in the totals, never misattributed to a bucket);
280    /// every real upstream — Gemini included — passes an explicit label.
281    /// `"Grok"` is the identity label for registry routes `grok-chat` / `xai`
282    /// (OpenAI wire shape; see `providers::stats_label`). `"CommandCode"` is the
283    /// identity label for the `commandcode` registry route.
284    fn provider(&self, provider_label: &str) -> Option<&ProviderStats> {
285        match provider_label {
286            "Anthropic" => Some(&self.anthropic),
287            "OpenAI" => Some(&self.openai),
288            "ChatGPT" => Some(&self.chatgpt),
289            "Gemini" => Some(&self.gemini),
290            "Grok" => Some(&self.grok),
291            "CommandCode" => Some(&self.commandcode),
292            _ => None,
293        }
294    }
295
296    pub fn provider_summary(&self) -> serde_json::Value {
297        serde_json::json!({
298            "anthropic": self.anthropic.summary(),
299            "openai": self.openai.summary(),
300            "chatgpt": self.chatgpt.summary(),
301            "gemini": self.gemini.summary(),
302            "grok": self.grok.summary(),
303            "commandcode": self.commandcode.summary(),
304        })
305    }
306}
307
308impl ProviderStats {
309    fn record(
310        &self,
311        original: usize,
312        effective_compressed: usize,
313        compressed_request: bool,
314        saved_tokens: u64,
315    ) {
316        self.requests_total.fetch_add(1, Ordering::Relaxed);
317        if compressed_request {
318            self.requests_compressed.fetch_add(1, Ordering::Relaxed);
319        }
320        self.tokens_saved.fetch_add(saved_tokens, Ordering::Relaxed);
321        self.bytes_original
322            .fetch_add(original as u64, Ordering::Relaxed);
323        self.bytes_compressed
324            .fetch_add(effective_compressed as u64, Ordering::Relaxed);
325    }
326
327    fn compression_ratio(&self) -> f64 {
328        let original = self.bytes_original.load(Ordering::Relaxed);
329        if original == 0 {
330            return 0.0;
331        }
332        let compressed = self.bytes_compressed.load(Ordering::Relaxed);
333        (1.0 - compressed as f64 / original as f64) * 100.0
334    }
335
336    fn summary(&self) -> serde_json::Value {
337        serde_json::json!({
338            "requests_total": self.requests_total.load(Ordering::Relaxed),
339            "requests_compressed": self.requests_compressed.load(Ordering::Relaxed),
340            "tokens_saved": self.tokens_saved.load(Ordering::Relaxed),
341            "bytes_original": self.bytes_original.load(Ordering::Relaxed),
342            "bytes_compressed": self.bytes_compressed.load(Ordering::Relaxed),
343            "compression_ratio_pct": format!("{:.1}", self.compression_ratio()),
344        })
345    }
346}
347
348/// TCP connect timeout (seconds). Configurable via `LEAN_CTX_PROXY_CONNECT_TIMEOUT_SECS`.
349fn connect_timeout_secs() -> u64 {
350    std::env::var("LEAN_CTX_PROXY_CONNECT_TIMEOUT_SECS")
351        .ok()
352        .and_then(|v| v.trim().parse::<u64>().ok())
353        .filter(|s| *s > 0)
354        .unwrap_or(15)
355}
356
357/// Idle read timeout (seconds) between bytes from upstream. Generous by default
358/// so long extended-thinking phases (which still emit SSE keepalives) are never
359/// cut, while a truly dead connection eventually fails. Configurable via
360/// `LEAN_CTX_PROXY_READ_TIMEOUT_SECS`.
361fn read_idle_timeout_secs() -> u64 {
362    std::env::var("LEAN_CTX_PROXY_READ_TIMEOUT_SECS")
363        .ok()
364        .and_then(|v| v.trim().parse::<u64>().ok())
365        .filter(|s| *s > 0)
366        .unwrap_or(300)
367}
368
369/// How often (seconds) a running proxy re-reads config.toml for upstream
370/// changes. `LEAN_CTX_PROXY_RELOAD_SECS` overrides; default 5s.
371fn upstream_reload_secs() -> u64 {
372    std::env::var("LEAN_CTX_PROXY_RELOAD_SECS")
373        .ok()
374        .and_then(|v| v.trim().parse::<u64>().ok())
375        .filter(|s| *s > 0)
376        .unwrap_or(5)
377}
378
379/// Background task: re-resolves the provider upstreams from config.toml on an
380/// interval and publishes any change to the live request handlers (#449). Ends
381/// once every receiver (the proxy itself) has been dropped.
382///
383/// `Config::load()` already keeps an internal content-hash cache, so re-reading
384/// an unchanged `config.toml` skips the TOML parse + merge and costs only a small
385/// file read; combined with the relaxed default interval (#453) the idle steady
386/// state is negligible without needing a separate stat pre-check.
387fn spawn_upstream_refresh(tx: tokio::sync::watch::Sender<Arc<Upstreams>>, initial: Upstreams) {
388    let interval = std::time::Duration::from_secs(upstream_reload_secs());
389    tokio::spawn(async move {
390        let mut last = initial;
391        loop {
392            tokio::time::sleep(interval).await;
393            let next = crate::core::config::Config::load()
394                .proxy
395                .refresh_upstreams(&last);
396            if next != last {
397                log_upstream_change(&last, &next);
398                last = next.clone();
399                if tx.send(Arc::new(next)).is_err() {
400                    break;
401                }
402            }
403        }
404    });
405}
406
407/// One stdout line per changed provider, matching the startup banner style so a
408/// running proxy's log shows when (and to what) an upstream switched.
409fn log_upstream_change(old: &Upstreams, new: &Upstreams) {
410    if old.anthropic != new.anthropic {
411        println!("  ↻ Anthropic upstream → {}", new.anthropic);
412    }
413    if old.openai != new.openai {
414        println!("  ↻ OpenAI upstream → {}", new.openai);
415    }
416    if old.chatgpt != new.chatgpt {
417        println!("  ↻ ChatGPT upstream → {}", new.chatgpt);
418    }
419    if old.gemini != new.gemini {
420        println!("  ↻ Gemini upstream → {}", new.gemini);
421    }
422    if old.providers != new.providers {
423        let ids: Vec<&str> = new.providers.iter().map(|p| p.id.as_str()).collect();
424        println!("  ↻ provider registry → [{}]", ids.join(", "));
425    }
426}
427
428pub async fn start_proxy(port: u16) -> anyhow::Result<()> {
429    let token = crate::core::session_token::resolve_proxy_token("LEAN_CTX_PROXY_TOKEN");
430    start_proxy_with_token(port, Some(token)).await
431}
432
433/// Security invariant: the proxy NEVER runs unauthenticated. `None` does not
434/// mean "no auth" — it means "resolve the session token for me". Provider
435/// routes additionally accept provider API keys (see `proxy_auth_guard`), so
436/// IDE clients keep working without any setup.
437fn effective_auth_token(auth_token: Option<String>) -> String {
438    auth_token
439        .filter(|t| !t.trim().is_empty())
440        .unwrap_or_else(|| crate::core::session_token::resolve_proxy_token("LEAN_CTX_PROXY_TOKEN"))
441}
442
443/// Install the process-default rustls `CryptoProvider` for the Codex ChatGPT
444/// WebSocket passthrough (#597).
445///
446/// `tokio-tungstenite`'s rustls connector builds its `ClientConfig` from the
447/// process-default provider. Our tree pulls *both* aws-lc-rs (reqwest) and ring
448/// (lettre/ureq), so rustls cannot auto-pick one and the `wss://chatgpt.com`
449/// handshake aborts with *"Could not automatically determine the process-level
450/// CryptoProvider"*. reqwest is unaffected (it configures aws-lc-rs explicitly),
451/// so we match it here. Idempotent: a prior install just returns the provider
452/// back as `Err`, which we ignore.
453fn install_default_crypto_provider() {
454    let _ = rustls::crypto::aws_lc_rs::default_provider().install_default();
455}
456
457pub async fn start_proxy_with_token(port: u16, auth_token: Option<String>) -> anyhow::Result<()> {
458    use crate::core::config::{Config, is_local_proxy_url};
459
460    // Must run before any WebSocket passthrough opens a wss:// upstream (#597).
461    install_default_crypto_provider();
462
463    let auth_token = effective_auth_token(auth_token);
464
465    // A single total timeout aborts long streaming generations (e.g. Opus doing
466    // a big refactor) mid-response. Use a connect timeout plus a read (idle)
467    // timeout instead: a genuinely hung upstream still fails, but a slow-but-
468    // alive stream is never cut off. Both are configurable for edge networks.
469    let chatgpt_cookies = chatgpt_cookies::shared_chatgpt_cloudflare_cookie_store();
470    let client = chatgpt_cookies::with_chatgpt_cloudflare_cookie_store(
471        reqwest::Client::builder(),
472        chatgpt_cookies.clone(),
473    )
474    .connect_timeout(std::time::Duration::from_secs(connect_timeout_secs()))
475    .read_timeout(std::time::Duration::from_secs(read_idle_timeout_secs()))
476    .build()?;
477
478    // Seed the measured-spend meter from disk so a proxy restart never zeroes
479    // the user's cumulative real provider bill.
480    usage_meter::resume_from_disk();
481    // Live model prices (#1179): load the cached provider price list and keep
482    // it fresh in the background, so new models are billed at their real
483    // market rates instead of family heuristics. Fail-open, kill switch:
484    // LEAN_CTX_LIVE_PRICING=off.
485    crate::core::gain::live_pricing::spawn_background_refresh();
486    // Seed the cold-prefix baselines too so a long idle gap that straddles a
487    // proxy restart is still detected and the repack can fire (#499).
488    cold_prefix::resume_from_disk();
489
490    let cfg = Config::load();
491    let break_even_schema_overhead = match cfg.tool_profile_effective() {
492        crate::core::tool_profiles::ToolProfile::Minimal
493        | crate::core::tool_profiles::ToolProfile::Auto => 1500,
494        crate::core::tool_profiles::ToolProfile::Standard => 3000,
495        crate::core::tool_profiles::ToolProfile::Power
496        | crate::core::tool_profiles::ToolProfile::Custom(_) => 5000,
497    };
498    // Read once at startup — avoids a Config::load() on every proxied request.
499    let bind_host = cfg.resolved_proxy_bind_host();
500    let loopback_bind = bind_host.is_loopback();
501    // Gateway mode (non-loopback bind, enterprise#8) hard-requires the Bearer
502    // token: the provider-key fallback's whole justification is "loopback only",
503    // so it is disabled by construction once the listener is reachable from the
504    // network — regardless of the config flag.
505    let require_token = cfg.proxy_require_token || !loopback_bind;
506    let loopback_open = cfg.proxy_loopback_open && loopback_bind;
507    let allowed_hosts: Arc<Vec<String>> = Arc::new(
508        cfg.proxy_allowed_hosts
509            .iter()
510            .map(|h| h.trim().trim_end_matches('.').to_ascii_lowercase())
511            .filter(|h| !h.is_empty())
512            .collect(),
513    );
514    // Rate limit (enterprise#37): explicit config wins; gateway mode ships a
515    // sane default floor; loopback stays unlimited unless configured. `0`
516    // disables the limiter explicitly.
517    let rate_limiter = match (cfg.proxy_max_rps, loopback_bind) {
518        (Some(rps), _) if rps > 0 => Some(Arc::new(RateLimiter::new(rps, rps.saturating_mul(2)))),
519        (None, false) => Some(Arc::new(RateLimiter::new(50, 100))),
520        _ => None,
521    };
522    let initial = cfg.proxy.resolve_all();
523
524    // The proxy reads its upstreams live from a watch channel: a background task
525    // re-resolves them from config.toml on an interval and publishes any change,
526    // so `lean-ctx config set proxy.*_upstream` (or any config.toml edit) takes
527    // effect on the running proxy within seconds, without a restart (#449).
528    let (upstream_tx, upstream_rx) = tokio::sync::watch::channel(Arc::new(initial.clone()));
529    // Clone before moving into ProxyState — auth guard needs live registry
530    // snapshots to decide client-auth provider fallback under require_token.
531    let auth_upstreams = upstream_rx.clone();
532    spawn_upstream_refresh(upstream_tx, initial.clone());
533
534    let Upstreams {
535        anthropic: anthropic_upstream,
536        openai: openai_upstream,
537        chatgpt: chatgpt_upstream,
538        gemini: gemini_upstream,
539        providers: initial_providers,
540    } = initial;
541
542    // Governed MCP reverse proxy registry (GL#91/#100): resolved once at
543    // startup (restart to reload — the gateway-keys.toml lifecycle). The
544    // `/mcp/{server}` routes are mounted below with the gateway-server
545    // feature; the snapshot lives in ProxyState either way.
546    let mcp_servers = Arc::new(
547        cfg.gateway_server
548            .resolve_mcp_servers(cfg.proxy.allows_insecure_http_upstream()),
549    );
550
551    let state = ProxyState {
552        client,
553        port,
554        stats: Arc::new(ProxyStats::default()),
555        break_even: Arc::new(break_even::BreakEvenCalculator::new(
556            break_even_schema_overhead,
557        )),
558        introspect: Arc::new(introspect::IntrospectState::default()),
559        ocla_cache: None,
560        upstreams: upstream_rx,
561        chatgpt_cookies,
562        mcp_servers: mcp_servers.clone(),
563        web_app_tracker: Arc::new(std::sync::Mutex::new(
564            web_app::conversation_tracker::ConversationTracker::default(),
565        )),
566    };
567
568    let app = Router::new()
569        .route("/health", get(health))
570        .route("/status", get(status_handler))
571        .route("/v1/messages", any(anthropic::handler))
572        .route("/v1/messages/{*rest}", any(anthropic::handler))
573        .route("/v1/chat/completions", any(openai::handler))
574        // POST → HTTP/SSE forwarder; GET → Codex/OpenAI WebSocket bridge (#440).
575        .route(
576            "/v1/responses",
577            post(openai_responses::handler).get(openai_responses::ws_handler),
578        )
579        .route("/v1/responses/{*rest}", any(openai_responses::handler))
580        // Bare provider endpoints (no `/v1` prefix). Clients whose base URL points
581        // at the proxy root — notably OpenCode via `@ai-sdk/openai`, whose
582        // Responses-API requests hit `/responses` — dispatch here. The
583        // `normalize_provider_path` layer rewrites the URI to its canonical
584        // `/v1/...` form before the handler forwards upstream (#353).
585        .route("/messages", any(anthropic::handler))
586        .route("/messages/{*rest}", any(anthropic::handler))
587        .route("/chat/completions", any(openai::handler))
588        .route(
589            "/responses",
590            post(openai_responses::handler).get(openai_responses::ws_handler),
591        )
592        .route("/responses/{*rest}", any(openai_responses::handler))
593        .route(
594            "/backend-api/codex/responses",
595            post(chatgpt::codex_responses_handler).get(chatgpt::codex_responses_ws_handler),
596        )
597        .route(
598            "/backend-api/codex/responses/{*rest}",
599            any(chatgpt::codex_responses_handler),
600        )
601        // Non-model ChatGPT backend calls (including codex_apps MCP) are not
602        // prompt JSON. Keep them as credential-preserving passthrough traffic.
603        .route("/backend-api", any(chatgpt::backend_api_handler))
604        .route("/backend-api/{*rest}", any(chatgpt::backend_api_handler))
605        .route("/v1/references/{id}", get(v1_resolve_reference))
606        // LiteLLM headroom-guardrail CCR retrieval (#702): resolves the 24-hex
607        // `hash=` marker `/v1/compress` emits back to the verbatim original.
608        .route("/v1/retrieve/{hash}", get(v1_retrieve_ccr))
609        // Org model catalog (enterprise#63): IDE clients discover the curated
610        // alias namespace (`zuehlke/fast` → provider:model) and verify their
611        // key. Exact-match only — `/v1/models/{...}` subpaths stay Gemini
612        // passthrough in the fallback router.
613        .route("/v1/models", get(models_api::handler))
614        .route("/models", get(models_api::handler))
615        // Drop-in `compress(messages, model)` contract (#739): deterministic
616        // messages-in / messages-out compression for SDK clients.
617        .route("/v1/compress", post(compress_api::handler))
618        .route(
619            "/v1/quality-lab",
620            post(quality_lab_api::handler).get(quality_lab_api::handler_get),
621        )
622        // Universal provider registry (`[[proxy.providers]]`, enterprise#7):
623        // `/providers/{id}/...` forwards to the registry entry with that id,
624        // speaking its declared wire shape. New provider = config, not code.
625        .route("/providers/{id}/{*rest}", any(providers::handler))
626        .fallback(fallback_router);
627
628    let mut app = app
629        .layer(axum::middleware::from_fn(
630            web_app_middleware::detect_web_app,
631        ))
632        .layer(axum::middleware::from_fn(move |req, next| {
633            let allowed = allowed_hosts.clone();
634            host_guard(req, next, allowed)
635        }))
636        .with_state(state);
637
638    // Per-person gateway keys (enterprise#11): sha256(bearer) → person/team/
639    // default_project. Loaded once at startup; rotation = restart (the standard
640    // secret-mount flow). A malformed file fails the start loudly.
641    let gateway_keys = match gateway_identity::GatewayKeys::load_default() {
642        Ok(keys) => {
643            if !keys.is_empty() {
644                println!(
645                    "  Identity:  {} gateway key(s) loaded ({})",
646                    keys.len(),
647                    gateway_identity::GatewayKeys::default_path().display()
648                );
649            }
650            Arc::new(keys)
651        }
652        Err(e) => anyhow::bail!("gateway-keys.toml: {e}"),
653    };
654
655    {
656        let expected = auth_token.clone();
657        let keys = gateway_keys.clone();
658        let upstreams = auth_upstreams;
659        app = app.layer(axum::middleware::from_fn(move |req, next| {
660            let expected = expected.clone();
661            let keys = keys.clone();
662            let upstreams = upstreams.clone();
663            proxy_auth_guard(
664                req,
665                next,
666                expected,
667                require_token,
668                loopback_open,
669                keys,
670                upstreams,
671            )
672        }));
673    }
674
675    if let Some(limiter) = rate_limiter {
676        app = app.layer(axum::middleware::from_fn(move |req, next| {
677            let limiter = limiter.clone();
678            rate_limit_guard(req, next, limiter)
679        }));
680    }
681
682    // Outermost layer (runs first): normalize bare provider endpoints to their
683    // canonical `/v1/...` form so auth, routing and upstream forwarding all agree,
684    // regardless of whether the client's base URL includes `/v1` (#353).
685    app = app.layer(axum::middleware::from_fn(normalize_provider_path));
686
687    let addr = SocketAddr::from((bind_host, port));
688    if loopback_open {
689        println!("lean-ctx proxy listening on http://{addr} (loopback-open: auth disabled)");
690    } else {
691        println!("lean-ctx proxy listening on http://{addr} (token auth enabled)");
692    }
693    if !loopback_bind {
694        println!(
695            "  ⚠ gateway mode: non-loopback bind — Bearer token REQUIRED for gateway-held \
696             keys and non-registry routes; client-auth /providers/* (api_key_env unset) \
697             still accept provider credentials. Host allowlist + rate limit active"
698        );
699    }
700    println!("  Anthropic: POST /v1/messages → {anthropic_upstream}");
701    println!("  OpenAI:    POST /v1/chat/completions → {openai_upstream}");
702    println!(
703        "  OpenAI:    POST /v1/responses → {openai_upstream}  (bare /responses also accepted)"
704    );
705    println!("  ChatGPT:   POST /backend-api/codex/responses → {chatgpt_upstream}");
706    println!("  ChatGPT:   any  /backend-api/* → {chatgpt_upstream}");
707    println!("  Gemini:    POST /v1beta/models/... → {gemini_upstream}");
708    println!("  Compress:  POST /v1/compress (deterministic messages-in/out, local)");
709    // Codex defaults to a WebSocket Responses transport (ws://…/responses). The
710    // proxy now bridges it to the HTTP/SSE upstream (#440), so Codex works as a
711    // drop-in without a `supports_websockets = false` workaround.
712    println!(
713        "  Codex:     WS  ws://{addr}/responses → bridged to {openai_upstream} (HTTP/SSE, #440)"
714    );
715    for p in &initial_providers {
716        println!(
717            "  Provider:  any  /providers/{}/... → {} ({} shape{})",
718            p.id,
719            p.base_url,
720            p.shape.as_str(),
721            if p.api_key_env.is_some() {
722                ", gateway-held key"
723            } else {
724                ""
725            }
726        );
727    }
728    for s in mcp_servers.iter() {
729        println!(
730            "  MCP:       any  /mcp/{} → {} (observed{})",
731            s.id,
732            s.url,
733            if s.auth_env.is_some() {
734                ", gateway-held credential"
735            } else {
736                ""
737            }
738        );
739    }
740    if openai_upstream.starts_with("http://") && !is_local_proxy_url(&openai_upstream) {
741        println!(
742            "  ⚠ OpenAI upstream is plaintext HTTP to a non-loopback host \
743             (allow_insecure_http_upstream) — use only on a trusted local network"
744        );
745    }
746
747    let listener = tokio::net::TcpListener::bind(addr).await?;
748    axum::serve(listener, app)
749        .with_graceful_shutdown(shutdown_signal())
750        .await?;
751
752    println!("lean-ctx proxy shut down cleanly.");
753    Ok(())
754}
755
756async fn shutdown_signal() {
757    let ctrl_c = tokio::signal::ctrl_c();
758
759    #[cfg(unix)]
760    {
761        // Fall back to Ctrl-C only if the SIGTERM handler cannot be installed,
762        // rather than panicking the proxy on startup.
763        match tokio::signal::unix::signal(tokio::signal::unix::SignalKind::terminate()) {
764            Ok(mut sigterm) => {
765                tokio::select! {
766                    _ = ctrl_c => {},
767                    _ = sigterm.recv() => {},
768                }
769            }
770            Err(e) => {
771                tracing::warn!("lean-ctx proxy: SIGTERM handler unavailable ({e}); Ctrl-C only");
772                ctrl_c.await.ok();
773            }
774        }
775    }
776
777    #[cfg(not(unix))]
778    {
779        ctrl_c.await.ok();
780    }
781
782    println!("lean-ctx proxy: received shutdown signal, draining…");
783}
784
785async fn health() -> impl IntoResponse {
786    let body = serde_json::json!({
787        "status": "ok",
788        "service": "lean-ctx-proxy",
789        "pid": std::process::id(),
790    });
791    (StatusCode::OK, axum::Json(body))
792}
793
794async fn v1_resolve_reference(
795    axum::extract::Path(id): axum::extract::Path<String>,
796) -> impl IntoResponse {
797    match crate::server::reference_store::resolve(&id) {
798        Some(content) => (StatusCode::OK, content),
799        None => (
800            StatusCode::NOT_FOUND,
801            "Reference expired or not found".to_string(),
802        ),
803    }
804}
805
806/// `GET /v1/retrieve/{hash}` (#702) — LiteLLM headroom-guardrail CCR contract
807/// (BerriAI/litellm#31681): resolve a 24-hex `hash=` marker emitted by
808/// `/v1/compress` back to the verbatim original from the tee store. The reply
809/// carries `original_content` (the field LiteLLM's `_call_retrieve` reads
810/// first). The optional `?query=` the guardrail forwards is accepted but the
811/// full original is always returned — a superset of any ranked slice, and the
812/// stored blobs are single tool outputs, not corpora worth ranking. Auth: the
813/// standard proxy bearer guard wraps this route (LiteLLM sends the configured
814/// `api_key` as a Bearer token); the tee store is loopback-scoped local state.
815async fn v1_retrieve_ccr(
816    axum::extract::Path(hash): axum::extract::Path<String>,
817) -> impl IntoResponse {
818    match ccr::retrieve_litellm(&hash) {
819        Some(content) => (
820            StatusCode::OK,
821            axum::Json(serde_json::json!({ "original_content": content })),
822        ),
823        None => (
824            StatusCode::NOT_FOUND,
825            axum::Json(serde_json::json!({
826                "error": "hash not found or expired",
827                "hash": hash,
828            })),
829        ),
830    }
831}
832
833async fn status_handler(State(state): State<ProxyState>) -> impl IntoResponse {
834    use std::sync::atomic::Ordering::Relaxed;
835    let s = &state.stats;
836    let i = &state.introspect;
837
838    let last_breakdown = i
839        .last_breakdown
840        .lock()
841        .ok()
842        .and_then(|guard| guard.as_ref().map(|b| serde_json::to_value(b).ok()))
843        .flatten();
844
845    let spend = usage_meter::snapshot();
846    let spend_total: f64 = spend.iter().map(|m| m.cost_usd).sum();
847
848    // Live upstreams the proxy is forwarding to right now (#449). This is the
849    // single source of truth for "where is my traffic actually going" — it
850    // reflects config.toml hot-reloads and any start-time env override.
851    let up = state.upstream_snapshot();
852
853    // Resolve the effort level fresh so /status reflects config.toml hot-reloads
854    // and env overrides, matching the upstream snapshot above (#834).
855    let active_effort = crate::core::config::Config::load().proxy.resolved_effort();
856
857    let introspect_total_input = i.total_input_tokens.load(Relaxed);
858    let introspect_bulk_candidates = i.total_bulk_candidate_tokens.load(Relaxed);
859    let introspect_bulk_share =
860        introspect::token_share_basis_points(introspect_bulk_candidates, introspect_total_input);
861    let break_even = state.break_even_summary();
862
863    let body = serde_json::json!({
864        "status": "running",
865        "proxy_mode": format!("{:?}", crate::core::config::Config::load().proxy.resolved_proxy_mode()),
866        "port": state.port,
867        "upstreams": {
868            "anthropic": up.anthropic.clone(),
869            "openai": up.openai.clone(),
870            "chatgpt": up.chatgpt.clone(),
871            "gemini": up.gemini.clone(),
872        },
873        // Universal registry (`[[proxy.providers]]`, enterprise#7). Key names
874        // only — never the key material.
875        "providers": up.providers.iter().map(|p| serde_json::json!({
876            "id": p.id,
877            "shape": p.shape.as_str(),
878            "base_url": p.base_url,
879            "gateway_key": p.api_key_env.is_some(),
880        })).collect::<Vec<_>>(),
881        "requests_total": s.requests_total.load(Relaxed),
882        "requests_compressed": s.requests_compressed.load(Relaxed),
883        "tokens_saved": s.tokens_saved.load(Relaxed),
884        "tokens_saved_estimated": true,
885        // Provider-verified savings (#701, opt-in counterfactual metering):
886        // both sides counted by Anthropic on the same request. `null` until
887        // the first probe-covered request lands.
888        "verified_savings": usage_meter::verified_savings(),
889        "bytes_original": s.bytes_original.load(Relaxed),
890        "bytes_compressed": s.bytes_compressed.load(Relaxed),
891        "compression_ratio_pct": format!("{:.1}", s.compression_ratio()),
892        "break_even": {
893            "proxy_savings": break_even.proxy_savings,
894            "estimated_mcp_overhead": break_even.estimated_mcp_overhead,
895            "net": break_even.net,
896            "mcp_recommended": break_even.mcp_recommended,
897            "reason": break_even.reason,
898        },
899        "per_upstream": s.provider_summary(),
900        "prefix_cache": prefix_cache_stats::snapshot(),
901        "cache_safety": cache_safety::snapshot(),
902        "cache_attribution": cache_attribution::snapshot(),
903        "effort": effort::snapshot(active_effort),
904        "per_model": cost::snapshot(),
905        "spend": {
906            "source": "measured",
907            "total_usd": spend_total,
908            "per_model": spend,
909            "note": "Actual provider bill: real model + billed tokens (incl. cache reads/writes & reasoning) read from upstream responses for proxy-routed clients."
910        },
911        "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.",
912        "introspect": {
913            "total_requests_analyzed": i.total_requests.load(Relaxed),
914            "total_input_tokens": introspect_total_input,
915            "total_system_prompt_tokens": i.total_system_prompt_tokens.load(Relaxed),
916            "total_bulk_candidate_tokens": introspect_bulk_candidates,
917            "bulk_candidate_share_basis_points": introspect_bulk_share,
918            "vision_encoding_decision_gate_met": introspect_bulk_share.is_some_and(|share| share >= 2_000),
919            "last_breakdown": last_breakdown,
920        }
921    });
922    (StatusCode::OK, axum::Json(body))
923}
924
925#[allow(clippy::result_large_err)]
926async fn proxy_auth_guard(
927    mut req: axum::extract::Request,
928    next: axum::middleware::Next,
929    expected_token: String,
930    require_token: bool,
931    loopback_open: bool,
932    gateway_keys: Arc<gateway_identity::GatewayKeys>,
933    upstreams: tokio::sync::watch::Receiver<Arc<Upstreams>>,
934) -> Result<Response, Response> {
935    let path = req.uri().path();
936    if path == "/health" {
937        return Ok(next.run(req).await);
938    }
939
940    // #755: loopback-open mode skips all auth — every local process is trusted.
941    if loopback_open {
942        req.extensions_mut()
943            .insert(gateway_identity::TrustedGatewayRequest);
944        attach_gateway_tags(&mut req, gateway_identity::GatewayTags::default());
945        return Ok(next.run(req).await);
946    }
947
948    let bearer = req
949        .headers()
950        .get("authorization")
951        .and_then(|v| v.to_str().ok())
952        .and_then(|auth| auth.strip_prefix("Bearer "))
953        .map(str::to_string);
954
955    if let Some(token) = bearer.as_deref()
956        && !proxy_token_expired(token, std::time::SystemTime::now())
957        && constant_time_eq(token.as_bytes(), expected_token.as_bytes())
958    {
959        req.extensions_mut()
960            .insert(gateway_identity::TrustedGatewayRequest);
961        attach_gateway_tags(&mut req, gateway_identity::GatewayTags::default());
962        return Ok(next.run(req).await);
963    }
964
965    // Per-person gateway keys (enterprise#11): a bearer key whose SHA-256 is in
966    // gateway-keys.toml authenticates AND identifies — its person/team/project
967    // tags travel with the request and end up on the usage record.
968    if let Some(token) = bearer.as_deref()
969        && let Some(tags) = gateway_keys.lookup(token)
970    {
971        req.extensions_mut()
972            .insert(gateway_identity::TrustedGatewayRequest);
973        attach_gateway_tags(&mut req, tags);
974        return Ok(next.run(req).await);
975    }
976
977    // Accept provider API keys on provider routes (host_guard runs first).
978    // Client-auth registry rails (`/providers/{id}` with api_key_env=None,
979    // e.g. Grok dual-rail) stay open even when require_token is set — the
980    // Authorization / x-xai-token-auth header must carry the upstream session
981    // and cannot also hold the lean-ctx Bearer.
982    let client_auth_registry = is_client_auth_registry_route(path, &upstreams.borrow());
983    if provider_key_fallback_allowed(
984        require_token,
985        has_provider_api_key(&req),
986        is_provider_route(path),
987        client_auth_registry,
988    ) {
989        attach_gateway_tags(&mut req, gateway_identity::GatewayTags::default());
990        return Ok(next.run(req).await);
991    }
992
993    Err(auth_error_response(path))
994}
995
996/// Returns true only for `TOKEN.TIMESTAMP` bearer tokens more than 24 hours
997/// old. Tokens without a valid timestamp suffix remain valid for backwards
998/// compatibility.
999fn proxy_token_expired(token: &str, now: std::time::SystemTime) -> bool {
1000    const TOKEN_MAX_AGE: std::time::Duration = std::time::Duration::from_hours(24);
1001
1002    let Some((value, timestamp)) = token.rsplit_once('.') else {
1003        return false;
1004    };
1005    if value.is_empty() {
1006        return false;
1007    }
1008    let Ok(timestamp) = timestamp.parse::<u64>() else {
1009        return false;
1010    };
1011
1012    let issued_at = std::time::UNIX_EPOCH + std::time::Duration::from_secs(timestamp);
1013    now.duration_since(issued_at)
1014        .is_ok_and(|age| age > TOKEN_MAX_AGE)
1015}
1016
1017fn auth_error_response(path: &str) -> Response {
1018    let is_mcp = path.starts_with("/mcp");
1019    let hint = if is_mcp {
1020        "MCP Streamable HTTP requires a Bearer token. Get it with: lean-ctx proxy token"
1021    } else if is_provider_route(path) {
1022        "Provider route requires authentication. Set your API key header or use: lean-ctx proxy token"
1023    } else {
1024        "This endpoint requires a lean-ctx Bearer token. Get it with: lean-ctx proxy token  \
1025         — or set proxy_loopback_open = true to disable auth on localhost"
1026    };
1027
1028    let body = serde_json::json!({
1029        "type": "error",
1030        "error": {
1031            "type": "authentication_error",
1032            "message": format!("401 Unauthorized — {hint}")
1033        }
1034    });
1035
1036    (StatusCode::UNAUTHORIZED, axum::Json(body)).into_response()
1037}
1038
1039/// Resolves the final identity tags for an authenticated request and inserts
1040/// them as a request extension (read by `forward.rs::wire_context`).
1041///
1042/// Project resolution (enterprise#11): the `x-leanctx-project` header wins over
1043/// the key's `default_project` — one person books work onto different projects
1044/// per request. The header also works without a gateway key (solo/local mode:
1045/// project tagging without identity). It is an internal gateway header, not on
1046/// `ALLOWED_REQUEST_HEADERS`, so it never reaches the upstream.
1047fn attach_gateway_tags(req: &mut axum::extract::Request, mut tags: gateway_identity::GatewayTags) {
1048    if let Some(project) = req
1049        .headers()
1050        .get("x-leanctx-project")
1051        .and_then(|v| v.to_str().ok())
1052        .map(str::trim)
1053        .filter(|p| !p.is_empty() && p.len() <= 128 && !p.chars().any(char::is_control))
1054    {
1055        tags.project = Some(project.to_string());
1056    }
1057    // GDPR pseudonymization (enterprise#39): applied at this single
1058    // choke-point, so budgets, usage rows, dashboards and logs only ever see
1059    // the pseudonym. No-op unless [gateway_server].pseudonymize_persons.
1060    if let Some(person) = tags.person.as_deref()
1061        && pii::enabled()
1062    {
1063        tags.person = Some(pii::pseudonymize(person));
1064    }
1065    if !tags.is_empty() {
1066        req.extensions_mut().insert(tags);
1067    }
1068}
1069
1070fn has_provider_api_key(req: &axum::extract::Request) -> bool {
1071    let headers = req.headers();
1072    // Provider-specific key headers: Anthropic `x-api-key`, Google
1073    // `x-goog-api-key`, Azure `api-key`, Grok subscription `x-xai-token-auth`.
1074    // Any non-empty value authenticates. Grok dual-rail may send only the
1075    // session header (no Authorization) against `/providers/grok-chat`.
1076    for key in ["x-api-key", "x-goog-api-key", "api-key", "x-xai-token-auth"] {
1077        if headers
1078            .get(key)
1079            .and_then(|v| v.to_str().ok())
1080            .is_some_and(|v| !v.trim().is_empty())
1081        {
1082            return true;
1083        }
1084    }
1085    // OpenAI-style `Authorization` auth. Accept ANY non-empty credential, not
1086    // just `Bearer sk-`/`gsk_`: OpenAI-*compatible* providers driven through
1087    // OpenCode/Codex (Azure, OpenRouter, Groq, vLLM/Ollama gateways, project &
1088    // service-account keys) issue keys that don't carry those prefixes. The proxy
1089    // binds to loopback only and never injects upstream credentials — it forwards
1090    // this header verbatim, so an invalid key is rejected by the real upstream,
1091    // never silently honoured. Gating provider routes on key *shape* only ever
1092    // produced false 401s for those clients (#362).
1093    if let Some(auth) = headers.get("authorization").and_then(|v| v.to_str().ok()) {
1094        let auth = auth.trim();
1095        let credential = auth
1096            .strip_prefix("Bearer ")
1097            .or_else(|| auth.strip_prefix("bearer "))
1098            .unwrap_or(auth)
1099            .trim();
1100        // Reject an empty value or a bare scheme keyword carrying no token.
1101        return !credential.is_empty() && !credential.eq_ignore_ascii_case("bearer");
1102    }
1103    false
1104}
1105
1106/// `/providers/{id}/...` where the registered provider has no gateway-held
1107/// key (`api_key_env = None`). These dual-rail / client-auth rails put the
1108/// upstream session in Authorization or provider-specific headers.
1109fn is_client_auth_registry_route(path: &str, upstreams: &Upstreams) -> bool {
1110    let Some(rest) = path.strip_prefix("/providers/") else {
1111        return false;
1112    };
1113    let id = rest.split('/').next().unwrap_or("");
1114    if id.is_empty() {
1115        return false;
1116    }
1117    upstreams
1118        .providers
1119        .iter()
1120        .any(|p| p.id == id && !p.injects_gateway_credential())
1121}
1122
1123fn is_provider_route(path: &str) -> bool {
1124    path.starts_with("/v1/")
1125        || path.starts_with("/v1beta/")
1126        || path.starts_with("/chat/completions")
1127        || path.starts_with("/responses")
1128        || path.starts_with("/messages")
1129        || path.starts_with("/backend-api")
1130        // Bare model-catalog discovery (enterprise#63): clients whose base URL
1131        // omits `/v1` send `GET /models` with their provider key.
1132        || path == "/models"
1133        // Universal provider registry (enterprise#7): clients point their base
1134        // URL at `/providers/{id}/v1` (Grok subscription:
1135        // `GROK_CLI_CHAT_PROXY_BASE_URL`, Azure Foundry, OpenRouter, …) and
1136        // authenticate with their own provider/session Bearer. Without this,
1137        // model-list + chat land as lean-ctx 401s — Grok cannot select a mode.
1138        || path.starts_with("/providers/")
1139}
1140
1141/// Decides whether a request authenticates via a provider API key alone, without
1142/// the lean-ctx Bearer token.
1143///
1144/// * Built-in provider routes (`/v1/*`, bare endpoints, …): only when
1145///   `require_token` is false (loopback-friendly default). Gateway mode
1146///   (non-loopback bind) forces `require_token` so these require the lean-ctx
1147///   Bearer (enterprise#8).
1148/// * Client-auth registry routes (`/providers/{id}` with `api_key_env = None`,
1149///   e.g. Grok dual-rail): always allowed when a provider credential is present.
1150///   The Authorization / `x-xai-token-auth` header must carry the upstream
1151///   session and cannot also hold the lean-ctx Bearer. Gateway-held keys
1152///   (`api_key_env` set) never take this path — those still need the lean-ctx
1153///   token so inject cannot be abused with a forged provider header.
1154///
1155/// Pure (except the precomputed `client_auth_registry` flag), so the policy is
1156/// unit-testable without axum middleware plumbing.
1157#[allow(clippy::fn_params_excessive_bools)] // deliberate pure predicate for unit tests
1158fn provider_key_fallback_allowed(
1159    require_token: bool,
1160    has_provider_key: bool,
1161    is_provider_route: bool,
1162    client_auth_registry: bool,
1163) -> bool {
1164    if !has_provider_key || !is_provider_route {
1165        return false;
1166    }
1167    if client_auth_registry {
1168        return true;
1169    }
1170    !require_token
1171}
1172
1173/// Maps a bare provider endpoint to its canonical `/v1/...` form, preserving any
1174/// sub-path. Returns `None` when the path is already canonical or not a known
1175/// provider endpoint.
1176///
1177/// Some OpenAI-compatible clients treat the configured base URL as the API root
1178/// and append the bare endpoint, so they send `POST /responses` or
1179/// `/chat/completions` instead of `/v1/responses` — notably OpenCode via
1180/// `@ai-sdk/openai`, whose Responses-API requests land on `/responses`. The proxy
1181/// and every upstream only know the `/v1/...` paths, so an un-prefixed request
1182/// would 401 (not a provider route) and then 404 (no handler). (#353)
1183fn canonical_provider_path(path: &str) -> Option<String> {
1184    // Inverse case of the bare-endpoint rewrite below: the advertised
1185    // OPENAI_BASE_URL includes `/v1` (#366), so a client that treats the base URL
1186    // as an origin and appends `/v1/...` itself produces `/v1/v1/...`.
1187    if let Some(rest) = path.strip_prefix("/v1/v1/") {
1188        return Some(format!("/v1/{rest}"));
1189    }
1190    const BARE_TO_CANONICAL: &[(&str, &str, &str)] = &[
1191        ("/responses", "/v1/responses", "/responses/"),
1192        (
1193            "/chat/completions",
1194            "/v1/chat/completions",
1195            "/chat/completions/",
1196        ),
1197        ("/messages", "/v1/messages", "/messages/"),
1198    ];
1199    for (bare, canonical, bare_with_slash) in BARE_TO_CANONICAL {
1200        if path == *bare {
1201            return Some((*canonical).to_string());
1202        }
1203        if let Some(rest) = path.strip_prefix(bare_with_slash) {
1204            return Some(format!("{canonical}/{rest}"));
1205        }
1206    }
1207    None
1208}
1209
1210/// Returns the canonicalized URI for a bare provider endpoint (query preserved),
1211/// or `None` when no rewrite is needed. Pure, so the rewrite is unit-testable
1212/// without constructing axum middleware plumbing.
1213fn normalized_provider_uri(uri: &axum::http::Uri) -> Option<axum::http::Uri> {
1214    let canonical = canonical_provider_path(uri.path())?;
1215    let new_path_and_query = match uri.query() {
1216        Some(q) => format!("{canonical}?{q}"),
1217        None => canonical,
1218    };
1219    new_path_and_query.parse::<axum::http::Uri>().ok()
1220}
1221
1222/// Rewrites the request URI in place when it targets a bare provider endpoint, so
1223/// downstream auth (`is_provider_route`), routing and upstream forwarding all see
1224/// the canonical `/v1/...` path. (#353)
1225async fn normalize_provider_path(
1226    mut req: axum::extract::Request,
1227    next: axum::middleware::Next,
1228) -> Response {
1229    if let Some(uri) = normalized_provider_uri(req.uri()) {
1230        *req.uri_mut() = uri;
1231    }
1232    next.run(req).await
1233}
1234
1235fn constant_time_eq(a: &[u8], b: &[u8]) -> bool {
1236    use subtle::ConstantTimeEq;
1237    if a.len() != b.len() {
1238        return false;
1239    }
1240    bool::from(a.ct_eq(b))
1241}
1242
1243async fn host_guard(
1244    req: axum::extract::Request,
1245    next: axum::middleware::Next,
1246    allowed_hosts: Arc<Vec<String>>,
1247) -> Result<Response, StatusCode> {
1248    if let Some(host) = req.headers().get("host").and_then(|v| v.to_str().ok())
1249        && host_allowed(host, &allowed_hosts)
1250    {
1251        return Ok(next.run(req).await);
1252    }
1253    Err(StatusCode::FORBIDDEN)
1254}
1255
1256/// DNS-rebinding guard: loopback Host headers always pass (today's local
1257/// behavior); in gateway mode the operator additionally allowlists the names
1258/// the gateway is reachable under (`proxy_allowed_hosts`, enterprise#8).
1259/// Matching is case-insensitive on the host with the port stripped.
1260fn host_allowed(host_header: &str, allowed: &[String]) -> bool {
1261    // `[::1]:8080` carries the port after the bracket; plain hosts after `:`.
1262    let host = host_header.trim();
1263    let h = if let Some(bracketed) = host.strip_prefix('[') {
1264        bracketed
1265            .split(']')
1266            .next()
1267            .map(|inner| format!("[{inner}]"))
1268    } else {
1269        host.split(':').next().map(str::to_string)
1270    };
1271    let Some(h) = h else {
1272        return false;
1273    };
1274    let h = h.trim_end_matches('.').to_ascii_lowercase();
1275    matches!(h.as_str(), "127.0.0.1" | "localhost" | "[::1]") || allowed.contains(&h)
1276}
1277
1278/// Proxy-wide token-bucket rate limit (enterprise#37). `/health` is exempt so
1279/// orchestrator liveness probes never get throttled into a false restart.
1280async fn rate_limit_guard(
1281    req: axum::extract::Request,
1282    next: axum::middleware::Next,
1283    limiter: Arc<RateLimiter>,
1284) -> Result<Response, StatusCode> {
1285    if req.uri().path() != "/health" && !limiter.allow().await {
1286        return Err(StatusCode::TOO_MANY_REQUESTS);
1287    }
1288    Ok(next.run(req).await)
1289}
1290
1291/// Token bucket: `max_rps` sustained, `burst` peak. Mirrors the team server's
1292/// limiter; lives here so the proxy stays independent of `http_server`
1293/// internals.
1294pub(crate) struct RateLimiter {
1295    max_rps: f64,
1296    burst: f64,
1297    state: tokio::sync::Mutex<RateLimiterState>,
1298}
1299
1300struct RateLimiterState {
1301    tokens: f64,
1302    last: std::time::Instant,
1303}
1304
1305impl RateLimiter {
1306    pub(crate) fn new(max_rps: u32, burst: u32) -> Self {
1307        Self {
1308            max_rps: f64::from(max_rps.max(1)),
1309            burst: f64::from(burst.max(1)),
1310            state: tokio::sync::Mutex::new(RateLimiterState {
1311                tokens: f64::from(burst.max(1)),
1312                last: std::time::Instant::now(),
1313            }),
1314        }
1315    }
1316
1317    pub(crate) async fn allow(&self) -> bool {
1318        let mut s = self.state.lock().await;
1319        let now = std::time::Instant::now();
1320        let refill = now.saturating_duration_since(s.last).as_secs_f64() * self.max_rps;
1321        s.tokens = (s.tokens + refill).min(self.burst);
1322        s.last = now;
1323        if s.tokens >= 1.0 {
1324            s.tokens -= 1.0;
1325            true
1326        } else {
1327            false
1328        }
1329    }
1330}
1331
1332async fn fallback_router(State(state): State<ProxyState>, req: Request<Body>) -> Response {
1333    let path = req.uri().path().to_string();
1334
1335    if path.starts_with("/v1beta/models/") || path.starts_with("/v1/models/") {
1336        match google::handler(State(state), req).await {
1337            Ok(resp) => resp,
1338            Err(status) => Response::builder()
1339                .status(status)
1340                .body(Body::from("proxy error"))
1341                .expect("BUG: building error response with valid status should never fail"),
1342        }
1343    } else {
1344        let method = req.method().to_string();
1345        eprintln!("lean-ctx proxy: unmatched {method} {path}");
1346        Response::builder()
1347            .status(StatusCode::NOT_FOUND)
1348            .body(Body::from(format!(
1349                "lean-ctx proxy: no handler for {method} {path}"
1350            )))
1351            .expect("BUG: building 404 response should never fail")
1352    }
1353}