lean_ctx/proxy/
mod.rs

1pub mod anthropic;
2pub mod cache_aligner;
3pub mod cache_attribution;
4pub mod cache_breakpoint;
5pub mod cache_policy;
6pub mod cache_safety;
7pub mod ccr;
8#[cfg(test)]
9mod ccr_robustness_tests;
10pub mod chatgpt;
11pub mod chatgpt_cookies;
12pub mod chatgpt_ws;
13pub mod cold_prefix;
14pub mod compress;
15pub mod compress_api;
16pub mod cost;
17pub mod effort;
18pub mod forward;
19pub mod google;
20pub mod history_prune;
21pub mod holdout;
22pub mod introspect;
23pub mod metrics;
24pub mod openai;
25pub mod openai_responses;
26pub mod openai_responses_ws;
27pub mod output_savings;
28pub mod prose;
29pub mod prose_ranker;
30pub mod tool_kind;
31pub mod tool_output;
32pub mod usage;
33pub mod usage_meter;
34pub mod verbosity;
35
36use std::net::SocketAddr;
37use std::sync::Arc;
38use std::sync::atomic::{AtomicU64, Ordering};
39
40use crate::core::config::Upstreams;
41
42use axum::{
43    Router,
44    body::Body,
45    extract::State,
46    http::{Request, StatusCode},
47    response::{IntoResponse, Response},
48    routing::{any, get, post},
49};
50
51#[derive(Clone)]
52pub struct ProxyState {
53    pub client: reqwest::Client,
54    pub port: u16,
55    pub stats: Arc<ProxyStats>,
56    pub introspect: Arc<introspect::IntrospectState>,
57    /// Live provider upstreams, refreshed from config.toml without a proxy
58    /// restart (#449). Read per request via [`ProxyState::openai_upstream`] etc.
59    pub upstreams: tokio::sync::watch::Receiver<Arc<Upstreams>>,
60    /// Shared Cloudflare cookie jar (also wired into `client`), so the Codex
61    /// ChatGPT WebSocket passthrough replays the same clearance to chatgpt.com
62    /// that the reqwest rail accumulated (#597).
63    pub(crate) chatgpt_cookies: Arc<chatgpt_cookies::ChatGptCloudflareCookieStore>,
64}
65
66impl ProxyState {
67    /// Consistent snapshot of all upstreams for the current request/response.
68    pub fn upstream_snapshot(&self) -> Arc<Upstreams> {
69        self.upstreams.borrow().clone()
70    }
71
72    /// Current Anthropic upstream (live).
73    pub fn anthropic_upstream(&self) -> String {
74        self.upstreams.borrow().anthropic.clone()
75    }
76
77    /// Current OpenAI upstream (live).
78    pub fn openai_upstream(&self) -> String {
79        self.upstreams.borrow().openai.clone()
80    }
81
82    /// Current ChatGPT upstream (live).
83    pub fn chatgpt_upstream(&self) -> String {
84        self.upstreams.borrow().chatgpt.clone()
85    }
86
87    /// Current Gemini upstream (live).
88    pub fn gemini_upstream(&self) -> String {
89        self.upstreams.borrow().gemini.clone()
90    }
91
92    /// Cloudflare `Cookie` header for the current ChatGPT upstream, used by the
93    /// WebSocket passthrough handshake (#597). `None` until a request on the
94    /// reqwest rail has seen Cloudflare clearance.
95    pub fn chatgpt_cookie_header(&self) -> Option<String> {
96        let url = reqwest::Url::parse(&self.chatgpt_upstream()).ok()?;
97        self.chatgpt_cookies
98            .cookie_header(&url)
99            .and_then(|v| v.to_str().ok().map(str::to_owned))
100    }
101}
102
103pub struct ProxyStats {
104    pub requests_total: AtomicU64,
105    pub requests_compressed: AtomicU64,
106    pub tokens_saved: AtomicU64,
107    pub bytes_original: AtomicU64,
108    pub bytes_compressed: AtomicU64,
109    pub anthropic: ProviderStats,
110    pub openai: ProviderStats,
111    pub chatgpt: ProviderStats,
112    pub gemini: ProviderStats,
113}
114
115#[derive(Default)]
116pub struct ProviderStats {
117    pub requests_total: AtomicU64,
118    pub requests_compressed: AtomicU64,
119    pub tokens_saved: AtomicU64,
120    pub bytes_original: AtomicU64,
121    pub bytes_compressed: AtomicU64,
122}
123
124impl Default for ProxyStats {
125    fn default() -> Self {
126        Self {
127            requests_total: AtomicU64::new(0),
128            requests_compressed: AtomicU64::new(0),
129            tokens_saved: AtomicU64::new(0),
130            bytes_original: AtomicU64::new(0),
131            bytes_compressed: AtomicU64::new(0),
132            anthropic: ProviderStats::default(),
133            openai: ProviderStats::default(),
134            chatgpt: ProviderStats::default(),
135            gemini: ProviderStats::default(),
136        }
137    }
138}
139
140impl ProxyStats {
141    pub fn record_request(&self, original: usize, compressed: usize) {
142        self.record_totals(original, compressed);
143    }
144
145    pub fn record_provider_request(
146        &self,
147        provider_label: &str,
148        original: usize,
149        compressed: usize,
150    ) {
151        let (effective_compressed, saved_tokens, compressed_request) =
152            self.record_totals(original, compressed);
153
154        if let Some(provider) = self.provider(provider_label) {
155            provider.record(
156                original,
157                effective_compressed,
158                compressed_request,
159                saved_tokens,
160            );
161        }
162    }
163
164    fn record_totals(&self, original: usize, compressed: usize) -> (usize, u64, bool) {
165        self.requests_total.fetch_add(1, Ordering::Relaxed);
166        self.bytes_original
167            .fetch_add(original as u64, Ordering::Relaxed);
168        let effective_compressed = compressed.min(original);
169        self.bytes_compressed
170            .fetch_add(effective_compressed as u64, Ordering::Relaxed);
171        if compressed < original {
172            self.requests_compressed.fetch_add(1, Ordering::Relaxed);
173        }
174        let saved_tokens = (original.saturating_sub(effective_compressed) / 4) as u64;
175        self.tokens_saved.fetch_add(saved_tokens, Ordering::Relaxed);
176        (effective_compressed, saved_tokens, compressed < original)
177    }
178
179    pub fn compression_ratio(&self) -> f64 {
180        let original = self.bytes_original.load(Ordering::Relaxed);
181        if original == 0 {
182            return 0.0;
183        }
184        let compressed = self.bytes_compressed.load(Ordering::Relaxed);
185        (1.0 - compressed as f64 / original as f64) * 100.0
186    }
187
188    /// Maps a proxy `provider_label` to its per-upstream bucket. Unknown labels
189    /// return `None` (still counted in the totals, never misattributed to a bucket);
190    /// every real upstream — Gemini included — passes an explicit label.
191    fn provider(&self, provider_label: &str) -> Option<&ProviderStats> {
192        match provider_label {
193            "Anthropic" => Some(&self.anthropic),
194            "OpenAI" => Some(&self.openai),
195            "ChatGPT" => Some(&self.chatgpt),
196            "Gemini" => Some(&self.gemini),
197            _ => None,
198        }
199    }
200
201    pub fn provider_summary(&self) -> serde_json::Value {
202        serde_json::json!({
203            "anthropic": self.anthropic.summary(),
204            "openai": self.openai.summary(),
205            "chatgpt": self.chatgpt.summary(),
206            "gemini": self.gemini.summary(),
207        })
208    }
209}
210
211impl ProviderStats {
212    fn record(
213        &self,
214        original: usize,
215        effective_compressed: usize,
216        compressed_request: bool,
217        saved_tokens: u64,
218    ) {
219        self.requests_total.fetch_add(1, Ordering::Relaxed);
220        if compressed_request {
221            self.requests_compressed.fetch_add(1, Ordering::Relaxed);
222        }
223        self.tokens_saved.fetch_add(saved_tokens, Ordering::Relaxed);
224        self.bytes_original
225            .fetch_add(original as u64, Ordering::Relaxed);
226        self.bytes_compressed
227            .fetch_add(effective_compressed as u64, Ordering::Relaxed);
228    }
229
230    fn compression_ratio(&self) -> f64 {
231        let original = self.bytes_original.load(Ordering::Relaxed);
232        if original == 0 {
233            return 0.0;
234        }
235        let compressed = self.bytes_compressed.load(Ordering::Relaxed);
236        (1.0 - compressed as f64 / original as f64) * 100.0
237    }
238
239    fn summary(&self) -> serde_json::Value {
240        serde_json::json!({
241            "requests_total": self.requests_total.load(Ordering::Relaxed),
242            "requests_compressed": self.requests_compressed.load(Ordering::Relaxed),
243            "tokens_saved": self.tokens_saved.load(Ordering::Relaxed),
244            "bytes_original": self.bytes_original.load(Ordering::Relaxed),
245            "bytes_compressed": self.bytes_compressed.load(Ordering::Relaxed),
246            "compression_ratio_pct": format!("{:.1}", self.compression_ratio()),
247        })
248    }
249}
250
251#[cfg(test)]
252mod stats_tests {
253    use super::*;
254    use std::sync::atomic::Ordering;
255
256    #[test]
257    fn compression_ratio_includes_uncompressed_requests() {
258        let stats = ProxyStats::default();
259
260        stats.record_request(1_000, 500);
261        stats.record_request(1_000, 1_000);
262
263        assert_eq!(stats.requests_total.load(Ordering::Relaxed), 2);
264        assert_eq!(stats.requests_compressed.load(Ordering::Relaxed), 1);
265        assert_eq!(stats.tokens_saved.load(Ordering::Relaxed), 125);
266        assert_eq!(stats.compression_ratio(), 25.0);
267    }
268
269    #[test]
270    fn expanded_requests_count_as_zero_savings() {
271        let stats = ProxyStats::default();
272
273        stats.record_request(1_000, 1_500);
274
275        assert_eq!(stats.requests_total.load(Ordering::Relaxed), 1);
276        assert_eq!(stats.requests_compressed.load(Ordering::Relaxed), 0);
277        assert_eq!(stats.tokens_saved.load(Ordering::Relaxed), 0);
278        assert_eq!(stats.compression_ratio(), 0.0);
279    }
280
281    #[test]
282    fn provider_stats_are_separate() {
283        let stats = ProxyStats::default();
284
285        stats.record_provider_request("OpenAI", 1_000, 500);
286        stats.record_provider_request("ChatGPT", 2_000, 1_000);
287
288        assert_eq!(stats.requests_total.load(Ordering::Relaxed), 2);
289        assert_eq!(stats.openai.requests_total.load(Ordering::Relaxed), 1);
290        assert_eq!(stats.chatgpt.requests_total.load(Ordering::Relaxed), 1);
291        assert_eq!(stats.openai.tokens_saved.load(Ordering::Relaxed), 125);
292        assert_eq!(stats.chatgpt.tokens_saved.load(Ordering::Relaxed), 250);
293        assert_eq!(stats.openai.compression_ratio(), 50.0);
294        assert_eq!(stats.chatgpt.compression_ratio(), 50.0);
295    }
296
297    #[test]
298    fn unlabelled_requests_do_not_count_as_gemini() {
299        let stats = ProxyStats::default();
300
301        stats.record_request(1_000, 500);
302
303        assert_eq!(stats.requests_total.load(Ordering::Relaxed), 1);
304        assert_eq!(stats.gemini.requests_total.load(Ordering::Relaxed), 0);
305    }
306
307    #[test]
308    fn unknown_label_is_not_recorded_to_any_bucket() {
309        let stats = ProxyStats::default();
310
311        stats.record_provider_request("Mystery", 1_000, 500);
312
313        // Totals still count it; no per-upstream bucket is touched.
314        assert_eq!(stats.requests_total.load(Ordering::Relaxed), 1);
315        assert_eq!(stats.anthropic.requests_total.load(Ordering::Relaxed), 0);
316        assert_eq!(stats.openai.requests_total.load(Ordering::Relaxed), 0);
317        assert_eq!(stats.chatgpt.requests_total.load(Ordering::Relaxed), 0);
318        assert_eq!(stats.gemini.requests_total.load(Ordering::Relaxed), 0);
319    }
320}
321
322/// TCP connect timeout (seconds). Configurable via `LEAN_CTX_PROXY_CONNECT_TIMEOUT_SECS`.
323fn connect_timeout_secs() -> u64 {
324    std::env::var("LEAN_CTX_PROXY_CONNECT_TIMEOUT_SECS")
325        .ok()
326        .and_then(|v| v.trim().parse::<u64>().ok())
327        .filter(|s| *s > 0)
328        .unwrap_or(15)
329}
330
331/// Idle read timeout (seconds) between bytes from upstream. Generous by default
332/// so long extended-thinking phases (which still emit SSE keepalives) are never
333/// cut, while a truly dead connection eventually fails. Configurable via
334/// `LEAN_CTX_PROXY_READ_TIMEOUT_SECS`.
335fn read_idle_timeout_secs() -> u64 {
336    std::env::var("LEAN_CTX_PROXY_READ_TIMEOUT_SECS")
337        .ok()
338        .and_then(|v| v.trim().parse::<u64>().ok())
339        .filter(|s| *s > 0)
340        .unwrap_or(300)
341}
342
343/// How often (seconds) a running proxy re-reads config.toml for upstream
344/// changes. `LEAN_CTX_PROXY_RELOAD_SECS` overrides; default 5s.
345fn upstream_reload_secs() -> u64 {
346    std::env::var("LEAN_CTX_PROXY_RELOAD_SECS")
347        .ok()
348        .and_then(|v| v.trim().parse::<u64>().ok())
349        .filter(|s| *s > 0)
350        .unwrap_or(5)
351}
352
353/// Background task: re-resolves the provider upstreams from config.toml on an
354/// interval and publishes any change to the live request handlers (#449). Ends
355/// once every receiver (the proxy itself) has been dropped.
356///
357/// `Config::load()` already keeps an internal content-hash cache, so re-reading
358/// an unchanged `config.toml` skips the TOML parse + merge and costs only a small
359/// file read; combined with the relaxed default interval (#453) the idle steady
360/// state is negligible without needing a separate stat pre-check.
361fn spawn_upstream_refresh(tx: tokio::sync::watch::Sender<Arc<Upstreams>>, initial: Upstreams) {
362    let interval = std::time::Duration::from_secs(upstream_reload_secs());
363    tokio::spawn(async move {
364        let mut last = initial;
365        loop {
366            tokio::time::sleep(interval).await;
367            let next = crate::core::config::Config::load()
368                .proxy
369                .refresh_upstreams(&last);
370            if next != last {
371                log_upstream_change(&last, &next);
372                last = next.clone();
373                if tx.send(Arc::new(next)).is_err() {
374                    break;
375                }
376            }
377        }
378    });
379}
380
381/// One stdout line per changed provider, matching the startup banner style so a
382/// running proxy's log shows when (and to what) an upstream switched.
383fn log_upstream_change(old: &Upstreams, new: &Upstreams) {
384    if old.anthropic != new.anthropic {
385        println!("  ↻ Anthropic upstream → {}", new.anthropic);
386    }
387    if old.openai != new.openai {
388        println!("  ↻ OpenAI upstream → {}", new.openai);
389    }
390    if old.chatgpt != new.chatgpt {
391        println!("  ↻ ChatGPT upstream → {}", new.chatgpt);
392    }
393    if old.gemini != new.gemini {
394        println!("  ↻ Gemini upstream → {}", new.gemini);
395    }
396}
397
398pub async fn start_proxy(port: u16) -> anyhow::Result<()> {
399    let token = crate::core::session_token::resolve_proxy_token("LEAN_CTX_PROXY_TOKEN");
400    start_proxy_with_token(port, Some(token)).await
401}
402
403/// Security invariant: the proxy NEVER runs unauthenticated. `None` does not
404/// mean "no auth" — it means "resolve the session token for me". Provider
405/// routes additionally accept provider API keys (see `proxy_auth_guard`), so
406/// IDE clients keep working without any setup.
407fn effective_auth_token(auth_token: Option<String>) -> String {
408    auth_token
409        .filter(|t| !t.trim().is_empty())
410        .unwrap_or_else(|| crate::core::session_token::resolve_proxy_token("LEAN_CTX_PROXY_TOKEN"))
411}
412
413/// Install the process-default rustls `CryptoProvider` for the Codex ChatGPT
414/// WebSocket passthrough (#597).
415///
416/// `tokio-tungstenite`'s rustls connector builds its `ClientConfig` from the
417/// process-default provider. Our tree pulls *both* aws-lc-rs (reqwest) and ring
418/// (lettre/ureq), so rustls cannot auto-pick one and the `wss://chatgpt.com`
419/// handshake aborts with *"Could not automatically determine the process-level
420/// CryptoProvider"*. reqwest is unaffected (it configures aws-lc-rs explicitly),
421/// so we match it here. Idempotent: a prior install just returns the provider
422/// back as `Err`, which we ignore.
423fn install_default_crypto_provider() {
424    let _ = rustls::crypto::aws_lc_rs::default_provider().install_default();
425}
426
427pub async fn start_proxy_with_token(port: u16, auth_token: Option<String>) -> anyhow::Result<()> {
428    use crate::core::config::{Config, is_local_proxy_url};
429
430    // Must run before any WebSocket passthrough opens a wss:// upstream (#597).
431    install_default_crypto_provider();
432
433    let auth_token = effective_auth_token(auth_token);
434
435    // A single total timeout aborts long streaming generations (e.g. Opus doing
436    // a big refactor) mid-response. Use a connect timeout plus a read (idle)
437    // timeout instead: a genuinely hung upstream still fails, but a slow-but-
438    // alive stream is never cut off. Both are configurable for edge networks.
439    let chatgpt_cookies = chatgpt_cookies::shared_chatgpt_cloudflare_cookie_store();
440    let client = chatgpt_cookies::with_chatgpt_cloudflare_cookie_store(
441        reqwest::Client::builder(),
442        chatgpt_cookies.clone(),
443    )
444    .connect_timeout(std::time::Duration::from_secs(connect_timeout_secs()))
445    .read_timeout(std::time::Duration::from_secs(read_idle_timeout_secs()))
446    .build()?;
447
448    // Seed the measured-spend meter from disk so a proxy restart never zeroes
449    // the user's cumulative real provider bill.
450    usage_meter::resume_from_disk();
451    // Seed the cold-prefix baselines too so a long idle gap that straddles a
452    // proxy restart is still detected and the repack can fire (#499).
453    cold_prefix::resume_from_disk();
454
455    let cfg = Config::load();
456    // Read once at startup — avoids a Config::load() on every proxied request.
457    let require_token = cfg.proxy_require_token;
458    let initial = cfg.proxy.resolve_all();
459
460    // The proxy reads its upstreams live from a watch channel: a background task
461    // re-resolves them from config.toml on an interval and publishes any change,
462    // so `lean-ctx config set proxy.*_upstream` (or any config.toml edit) takes
463    // effect on the running proxy within seconds, without a restart (#449).
464    let (upstream_tx, upstream_rx) = tokio::sync::watch::channel(Arc::new(initial.clone()));
465    spawn_upstream_refresh(upstream_tx, initial.clone());
466
467    let Upstreams {
468        anthropic: anthropic_upstream,
469        openai: openai_upstream,
470        chatgpt: chatgpt_upstream,
471        gemini: gemini_upstream,
472    } = initial;
473
474    let state = ProxyState {
475        client,
476        port,
477        stats: Arc::new(ProxyStats::default()),
478        introspect: Arc::new(introspect::IntrospectState::default()),
479        upstreams: upstream_rx,
480        chatgpt_cookies,
481    };
482
483    let mut app = Router::new()
484        .route("/health", get(health))
485        .route("/status", get(status_handler))
486        .route("/v1/messages", any(anthropic::handler))
487        .route("/v1/messages/{*rest}", any(anthropic::handler))
488        .route("/v1/chat/completions", any(openai::handler))
489        // POST → HTTP/SSE forwarder; GET → Codex/OpenAI WebSocket bridge (#440).
490        .route(
491            "/v1/responses",
492            post(openai_responses::handler).get(openai_responses::ws_handler),
493        )
494        .route("/v1/responses/{*rest}", any(openai_responses::handler))
495        // Bare provider endpoints (no `/v1` prefix). Clients whose base URL points
496        // at the proxy root — notably OpenCode via `@ai-sdk/openai`, whose
497        // Responses-API requests hit `/responses` — dispatch here. The
498        // `normalize_provider_path` layer rewrites the URI to its canonical
499        // `/v1/...` form before the handler forwards upstream (#353).
500        .route("/messages", any(anthropic::handler))
501        .route("/messages/{*rest}", any(anthropic::handler))
502        .route("/chat/completions", any(openai::handler))
503        .route(
504            "/responses",
505            post(openai_responses::handler).get(openai_responses::ws_handler),
506        )
507        .route("/responses/{*rest}", any(openai_responses::handler))
508        .route(
509            "/backend-api/codex/responses",
510            post(chatgpt::codex_responses_handler).get(chatgpt::codex_responses_ws_handler),
511        )
512        .route(
513            "/backend-api/codex/responses/{*rest}",
514            any(chatgpt::codex_responses_handler),
515        )
516        // Non-model ChatGPT backend calls (including codex_apps MCP) are not
517        // prompt JSON. Keep them as credential-preserving passthrough traffic.
518        .route("/backend-api", any(chatgpt::backend_api_handler))
519        .route("/backend-api/{*rest}", any(chatgpt::backend_api_handler))
520        .route("/v1/references/{id}", get(v1_resolve_reference))
521        // Drop-in `compress(messages, model)` contract (#739): deterministic
522        // messages-in / messages-out compression for SDK clients.
523        .route("/v1/compress", post(compress_api::handler))
524        .fallback(fallback_router)
525        .layer(axum::middleware::from_fn(host_guard))
526        .with_state(state);
527
528    {
529        let expected = auth_token.clone();
530        app = app.layer(axum::middleware::from_fn(move |req, next| {
531            let expected = expected.clone();
532            proxy_auth_guard(req, next, expected, require_token)
533        }));
534    }
535
536    // Outermost layer (runs first): normalize bare provider endpoints to their
537    // canonical `/v1/...` form so auth, routing and upstream forwarding all agree,
538    // regardless of whether the client's base URL includes `/v1` (#353).
539    app = app.layer(axum::middleware::from_fn(normalize_provider_path));
540
541    let addr = SocketAddr::from(([127, 0, 0, 1], port));
542    println!("lean-ctx proxy listening on http://{addr} (token auth enabled)");
543    println!("  Anthropic: POST /v1/messages → {anthropic_upstream}");
544    println!("  OpenAI:    POST /v1/chat/completions → {openai_upstream}");
545    println!(
546        "  OpenAI:    POST /v1/responses → {openai_upstream}  (bare /responses also accepted)"
547    );
548    println!("  ChatGPT:   POST /backend-api/codex/responses → {chatgpt_upstream}");
549    println!("  ChatGPT:   any  /backend-api/* → {chatgpt_upstream}");
550    println!("  Gemini:    POST /v1beta/models/... → {gemini_upstream}");
551    println!("  Compress:  POST /v1/compress (deterministic messages-in/out, local)");
552    // Codex defaults to a WebSocket Responses transport (ws://…/responses). The
553    // proxy now bridges it to the HTTP/SSE upstream (#440), so Codex works as a
554    // drop-in without a `supports_websockets = false` workaround.
555    println!(
556        "  Codex:     WS  ws://{addr}/responses → bridged to {openai_upstream} (HTTP/SSE, #440)"
557    );
558    if openai_upstream.starts_with("http://") && !is_local_proxy_url(&openai_upstream) {
559        println!(
560            "  ⚠ OpenAI upstream is plaintext HTTP to a non-loopback host \
561             (allow_insecure_http_upstream) — use only on a trusted local network"
562        );
563    }
564
565    let listener = tokio::net::TcpListener::bind(addr).await?;
566    axum::serve(listener, app)
567        .with_graceful_shutdown(shutdown_signal())
568        .await?;
569
570    println!("lean-ctx proxy shut down cleanly.");
571    Ok(())
572}
573
574async fn shutdown_signal() {
575    let ctrl_c = tokio::signal::ctrl_c();
576
577    #[cfg(unix)]
578    {
579        // Fall back to Ctrl-C only if the SIGTERM handler cannot be installed,
580        // rather than panicking the proxy on startup.
581        match tokio::signal::unix::signal(tokio::signal::unix::SignalKind::terminate()) {
582            Ok(mut sigterm) => {
583                tokio::select! {
584                    _ = ctrl_c => {},
585                    _ = sigterm.recv() => {},
586                }
587            }
588            Err(e) => {
589                tracing::warn!("lean-ctx proxy: SIGTERM handler unavailable ({e}); Ctrl-C only");
590                ctrl_c.await.ok();
591            }
592        }
593    }
594
595    #[cfg(not(unix))]
596    {
597        ctrl_c.await.ok();
598    }
599
600    println!("lean-ctx proxy: received shutdown signal, draining…");
601}
602
603async fn health() -> impl IntoResponse {
604    let body = serde_json::json!({
605        "status": "ok",
606        "pid": std::process::id(),
607    });
608    (StatusCode::OK, axum::Json(body))
609}
610
611async fn v1_resolve_reference(
612    axum::extract::Path(id): axum::extract::Path<String>,
613) -> impl IntoResponse {
614    match crate::server::reference_store::resolve(&id) {
615        Some(content) => (StatusCode::OK, content),
616        None => (
617            StatusCode::NOT_FOUND,
618            "Reference expired or not found".to_string(),
619        ),
620    }
621}
622
623async fn status_handler(State(state): State<ProxyState>) -> impl IntoResponse {
624    use std::sync::atomic::Ordering::Relaxed;
625    let s = &state.stats;
626    let i = &state.introspect;
627
628    let last_breakdown = i
629        .last_breakdown
630        .lock()
631        .ok()
632        .and_then(|guard| guard.as_ref().map(|b| serde_json::to_value(b).ok()))
633        .flatten();
634
635    let spend = usage_meter::snapshot();
636    let spend_total: f64 = spend.iter().map(|m| m.cost_usd).sum();
637
638    // Live upstreams the proxy is forwarding to right now (#449). This is the
639    // single source of truth for "where is my traffic actually going" — it
640    // reflects config.toml hot-reloads and any start-time env override.
641    let up = state.upstream_snapshot();
642
643    // Resolve the effort level fresh so /status reflects config.toml hot-reloads
644    // and env overrides, matching the upstream snapshot above (#834).
645    let active_effort = crate::core::config::Config::load().proxy.resolved_effort();
646
647    let body = serde_json::json!({
648        "status": "running",
649        "port": state.port,
650        "upstreams": {
651            "anthropic": up.anthropic.clone(),
652            "openai": up.openai.clone(),
653            "chatgpt": up.chatgpt.clone(),
654            "gemini": up.gemini.clone(),
655        },
656        "requests_total": s.requests_total.load(Relaxed),
657        "requests_compressed": s.requests_compressed.load(Relaxed),
658        "tokens_saved": s.tokens_saved.load(Relaxed),
659        "tokens_saved_estimated": true,
660        "bytes_original": s.bytes_original.load(Relaxed),
661        "bytes_compressed": s.bytes_compressed.load(Relaxed),
662        "compression_ratio_pct": format!("{:.1}", s.compression_ratio()),
663        "per_upstream": s.provider_summary(),
664        "cache_safety": cache_safety::snapshot(),
665        "cache_attribution": cache_attribution::snapshot(),
666        "effort": effort::snapshot(active_effort),
667        "per_model": cost::snapshot(),
668        "spend": {
669            "source": "measured",
670            "total_usd": spend_total,
671            "per_model": spend,
672            "note": "Actual provider bill: real model + billed tokens (incl. cache reads/writes & reasoning) read from upstream responses for proxy-routed clients."
673        },
674        "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.",
675        "introspect": {
676            "total_requests_analyzed": i.total_requests.load(Relaxed),
677            "total_system_prompt_tokens": i.total_system_prompt_tokens.load(Relaxed),
678            "last_breakdown": last_breakdown,
679        }
680    });
681    (StatusCode::OK, axum::Json(body))
682}
683
684#[allow(clippy::result_large_err)]
685async fn proxy_auth_guard(
686    req: axum::extract::Request,
687    next: axum::middleware::Next,
688    expected_token: String,
689    require_token: bool,
690) -> Result<Response, Response> {
691    let path = req.uri().path();
692    if path == "/health" {
693        return Ok(next.run(req).await);
694    }
695
696    if let Some(auth) = req
697        .headers()
698        .get("authorization")
699        .and_then(|v| v.to_str().ok())
700        && let Some(token) = auth.strip_prefix("Bearer ")
701        && constant_time_eq(token.as_bytes(), expected_token.as_bytes())
702    {
703        return Ok(next.run(req).await);
704    }
705
706    // Accept provider API keys on provider routes (loopback-only, host_guard runs first).
707    // AI tools like Claude Code send x-api-key, not Bearer tokens. Since the proxy
708    // only binds to 127.0.0.1, the presence of a valid API key header is sufficient
709    // to authenticate the request as coming from a local AI tool. Disabled when
710    // `proxy_require_token` is set — strict hosts then require the Bearer token.
711    if provider_key_fallback_allowed(
712        require_token,
713        has_provider_api_key(&req),
714        is_provider_route(path),
715    ) {
716        return Ok(next.run(req).await);
717    }
718
719    let cfg = crate::core::config::Config::load();
720    let hint = match cfg.proxy_enabled {
721        Some(true) => {
722            "lean-ctx proxy requires authentication. Use a Bearer token (LEAN_CTX_PROXY_TOKEN) or configure your AI tool's API key."
723        }
724        Some(false) => "lean-ctx proxy is disabled but still running. Run: lean-ctx proxy cleanup",
725        None => {
726            "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"
727        }
728    };
729
730    let body = serde_json::json!({
731        "type": "error",
732        "error": {
733            "type": "authentication_error",
734            "message": format!("401 Unauthorized — {hint}")
735        }
736    });
737
738    Err((StatusCode::UNAUTHORIZED, axum::Json(body)).into_response())
739}
740
741fn has_provider_api_key(req: &axum::extract::Request) -> bool {
742    let headers = req.headers();
743    // Provider-specific key headers: Anthropic `x-api-key`, Google
744    // `x-goog-api-key`, Azure `api-key`. Any non-empty value authenticates.
745    for key in ["x-api-key", "x-goog-api-key", "api-key"] {
746        if headers
747            .get(key)
748            .and_then(|v| v.to_str().ok())
749            .is_some_and(|v| !v.trim().is_empty())
750        {
751            return true;
752        }
753    }
754    // OpenAI-style `Authorization` auth. Accept ANY non-empty credential, not
755    // just `Bearer sk-`/`gsk_`: OpenAI-*compatible* providers driven through
756    // OpenCode/Codex (Azure, OpenRouter, Groq, vLLM/Ollama gateways, project &
757    // service-account keys) issue keys that don't carry those prefixes. The proxy
758    // binds to loopback only and never injects upstream credentials — it forwards
759    // this header verbatim, so an invalid key is rejected by the real upstream,
760    // never silently honoured. Gating provider routes on key *shape* only ever
761    // produced false 401s for those clients (#362).
762    if let Some(auth) = headers.get("authorization").and_then(|v| v.to_str().ok()) {
763        let auth = auth.trim();
764        let credential = auth
765            .strip_prefix("Bearer ")
766            .or_else(|| auth.strip_prefix("bearer "))
767            .unwrap_or(auth)
768            .trim();
769        // Reject an empty value or a bare scheme keyword carrying no token.
770        return !credential.is_empty() && !credential.eq_ignore_ascii_case("bearer");
771    }
772    false
773}
774
775fn is_provider_route(path: &str) -> bool {
776    path.starts_with("/v1/")
777        || path.starts_with("/v1beta/")
778        || path.starts_with("/chat/completions")
779        || path.starts_with("/responses")
780        || path.starts_with("/messages")
781        || path.starts_with("/backend-api")
782}
783
784/// Decides whether a request authenticates via a provider API key alone, without
785/// the lean-ctx Bearer token. True only in the default, loopback-friendly mode
786/// where a local AI tool's own provider key is accepted on a provider route. When
787/// `require_token` is set (strict, shared-host mode) the fallback is disabled and
788/// the Bearer token becomes mandatory. Pure, so the policy is unit-testable
789/// without axum middleware plumbing.
790fn provider_key_fallback_allowed(
791    require_token: bool,
792    has_provider_key: bool,
793    is_provider_route: bool,
794) -> bool {
795    !require_token && has_provider_key && is_provider_route
796}
797
798/// Maps a bare provider endpoint to its canonical `/v1/...` form, preserving any
799/// sub-path. Returns `None` when the path is already canonical or not a known
800/// provider endpoint.
801///
802/// Some OpenAI-compatible clients treat the configured base URL as the API root
803/// and append the bare endpoint, so they send `POST /responses` or
804/// `/chat/completions` instead of `/v1/responses` — notably OpenCode via
805/// `@ai-sdk/openai`, whose Responses-API requests land on `/responses`. The proxy
806/// and every upstream only know the `/v1/...` paths, so an un-prefixed request
807/// would 401 (not a provider route) and then 404 (no handler). (#353)
808fn canonical_provider_path(path: &str) -> Option<String> {
809    // Inverse case of the bare-endpoint rewrite below: the advertised
810    // OPENAI_BASE_URL includes `/v1` (#366), so a client that treats the base URL
811    // as an origin and appends `/v1/...` itself produces `/v1/v1/...`.
812    if let Some(rest) = path.strip_prefix("/v1/v1/") {
813        return Some(format!("/v1/{rest}"));
814    }
815    const BARE_TO_CANONICAL: &[(&str, &str, &str)] = &[
816        ("/responses", "/v1/responses", "/responses/"),
817        (
818            "/chat/completions",
819            "/v1/chat/completions",
820            "/chat/completions/",
821        ),
822        ("/messages", "/v1/messages", "/messages/"),
823    ];
824    for (bare, canonical, bare_with_slash) in BARE_TO_CANONICAL {
825        if path == *bare {
826            return Some((*canonical).to_string());
827        }
828        if let Some(rest) = path.strip_prefix(bare_with_slash) {
829            return Some(format!("{canonical}/{rest}"));
830        }
831    }
832    None
833}
834
835/// Returns the canonicalized URI for a bare provider endpoint (query preserved),
836/// or `None` when no rewrite is needed. Pure, so the rewrite is unit-testable
837/// without constructing axum middleware plumbing.
838fn normalized_provider_uri(uri: &axum::http::Uri) -> Option<axum::http::Uri> {
839    let canonical = canonical_provider_path(uri.path())?;
840    let new_path_and_query = match uri.query() {
841        Some(q) => format!("{canonical}?{q}"),
842        None => canonical,
843    };
844    new_path_and_query.parse::<axum::http::Uri>().ok()
845}
846
847/// Rewrites the request URI in place when it targets a bare provider endpoint, so
848/// downstream auth (`is_provider_route`), routing and upstream forwarding all see
849/// the canonical `/v1/...` path. (#353)
850async fn normalize_provider_path(
851    mut req: axum::extract::Request,
852    next: axum::middleware::Next,
853) -> Response {
854    if let Some(uri) = normalized_provider_uri(req.uri()) {
855        *req.uri_mut() = uri;
856    }
857    next.run(req).await
858}
859
860fn constant_time_eq(a: &[u8], b: &[u8]) -> bool {
861    use subtle::ConstantTimeEq;
862    if a.len() != b.len() {
863        return false;
864    }
865    bool::from(a.ct_eq(b))
866}
867
868async fn host_guard(
869    req: axum::extract::Request,
870    next: axum::middleware::Next,
871) -> Result<Response, StatusCode> {
872    if let Some(host) = req.headers().get("host").and_then(|v| v.to_str().ok()) {
873        let h = host.split(':').next().unwrap_or(host);
874        if matches!(h, "127.0.0.1" | "localhost" | "[::1]") {
875            return Ok(next.run(req).await);
876        }
877    }
878    Err(StatusCode::FORBIDDEN)
879}
880
881async fn fallback_router(State(state): State<ProxyState>, req: Request<Body>) -> Response {
882    let path = req.uri().path().to_string();
883
884    if path.starts_with("/v1beta/models/") || path.starts_with("/v1/models/") {
885        match google::handler(State(state), req).await {
886            Ok(resp) => resp,
887            Err(status) => Response::builder()
888                .status(status)
889                .body(Body::from("proxy error"))
890                .expect("BUG: building error response with valid status should never fail"),
891        }
892    } else {
893        let method = req.method().to_string();
894        eprintln!("lean-ctx proxy: unmatched {method} {path}");
895        Response::builder()
896            .status(StatusCode::NOT_FOUND)
897            .body(Body::from(format!(
898                "lean-ctx proxy: no handler for {method} {path}"
899            )))
900            .expect("BUG: building 404 response should never fail")
901    }
902}
903
904#[cfg(test)]
905mod auth_tests {
906    use super::*;
907
908    // P0-4 (#416): the proxy must never run unauthenticated — `None` means
909    // "resolve the session token", not "no auth".
910    #[test]
911    fn effective_auth_token_never_yields_empty() {
912        let _env = crate::core::data_dir::test_env_lock();
913        let tmp = tempfile::tempdir().unwrap();
914        crate::test_env::set_var("LEAN_CTX_DATA_DIR", tmp.path());
915
916        assert_eq!(effective_auth_token(Some("tok".into())), "tok");
917        let auto = effective_auth_token(None);
918        assert!(!auto.trim().is_empty(), "None must auto-resolve a token");
919        let blank = effective_auth_token(Some("   ".into()));
920        assert!(!blank.trim().is_empty(), "blank tokens must be replaced");
921
922        crate::test_env::remove_var("LEAN_CTX_DATA_DIR");
923    }
924
925    // #597: the Codex ChatGPT WS passthrough opens a wss://chatgpt.com socket via
926    // tokio-tungstenite, which needs a process-default rustls CryptoProvider. The
927    // tree has both aws-lc-rs and ring, so one must be installed explicitly or the
928    // handshake aborts. Guards against that regression.
929    #[test]
930    fn installs_default_crypto_provider_for_ws_passthrough() {
931        install_default_crypto_provider();
932        assert!(
933            rustls::crypto::CryptoProvider::get_default().is_some(),
934            "WS passthrough needs a process-default CryptoProvider"
935        );
936    }
937
938    #[test]
939    fn is_provider_route_v1() {
940        assert!(is_provider_route("/v1/chat/completions"));
941        assert!(is_provider_route("/v1/messages"));
942        assert!(is_provider_route("/v1/completions"));
943    }
944
945    #[test]
946    fn is_provider_route_anthropic_subpaths() {
947        assert!(is_provider_route("/v1/messages/count_tokens"));
948        assert!(is_provider_route("/v1/messages/batches"));
949        assert!(is_provider_route("/v1/messages/batches/batch_123"));
950    }
951
952    #[test]
953    fn is_provider_route_v1beta() {
954        assert!(is_provider_route("/v1beta/models"));
955    }
956
957    #[test]
958    fn is_provider_route_chat() {
959        assert!(is_provider_route("/chat/completions"));
960    }
961
962    #[test]
963    fn is_provider_route_chatgpt_backend_api() {
964        assert!(is_provider_route("/backend-api/codex/responses"));
965        assert!(is_provider_route("/backend-api/codex/responses/resp_123"));
966        assert!(is_provider_route("/backend-api/wham/session"));
967        assert!(is_provider_route("/backend-api/ps/mcp"));
968        assert!(is_provider_route("/backend-api/codex_apps"));
969        assert!(is_provider_route("/backend-api/codex_apps/mcp"));
970        assert!(is_provider_route("/backend-api/mcp/codex_apps"));
971        assert!(is_provider_route("/backend-api/apps/codex_apps/mcp"));
972    }
973
974    #[test]
975    fn is_provider_route_rejects_non_provider() {
976        assert!(!is_provider_route("/health"));
977        assert!(!is_provider_route("/api/v2/test"));
978        assert!(!is_provider_route("/"));
979    }
980
981    fn build_request(headers: &[(&str, &str)], path: &str) -> axum::extract::Request {
982        let mut builder = axum::http::Request::builder().uri(path);
983        for (k, v) in headers {
984            builder = builder.header(*k, *v);
985        }
986        builder.body(axum::body::Body::empty()).unwrap()
987    }
988
989    #[test]
990    fn has_provider_api_key_x_api_key() {
991        let req = build_request(&[("x-api-key", "sk-ant-abc123")], "/v1/messages");
992        assert!(has_provider_api_key(&req));
993    }
994
995    #[test]
996    fn has_provider_api_key_x_goog() {
997        let req = build_request(&[("x-goog-api-key", "AIzaSyAbc")], "/v1beta/models");
998        assert!(has_provider_api_key(&req));
999    }
1000
1001    #[test]
1002    fn has_provider_api_key_azure() {
1003        let req = build_request(&[("api-key", "deadbeef")], "/v1/completions");
1004        assert!(has_provider_api_key(&req));
1005    }
1006
1007    #[test]
1008    fn has_provider_api_key_bearer_sk() {
1009        let req = build_request(
1010            &[("authorization", "Bearer sk-proj-abc123")],
1011            "/v1/chat/completions",
1012        );
1013        assert!(has_provider_api_key(&req));
1014    }
1015
1016    #[test]
1017    fn has_provider_api_key_empty_rejected() {
1018        let req = build_request(&[("x-api-key", "  ")], "/v1/messages");
1019        assert!(!has_provider_api_key(&req));
1020    }
1021
1022    #[test]
1023    fn has_provider_api_key_no_headers() {
1024        let req = build_request(&[], "/v1/messages");
1025        assert!(!has_provider_api_key(&req));
1026    }
1027
1028    #[test]
1029    fn has_provider_api_key_accepts_non_sk_bearer() {
1030        // #362: OpenAI-*compatible* providers (Azure, OpenRouter, Groq, vLLM/
1031        // Ollama gateways, project/service keys) issue keys without the sk-/gsk_
1032        // prefix. OpenCode (@ai-sdk/openai) forwards them as `Bearer <key>`; they
1033        // must authenticate on a loopback provider route. The upstream validates
1034        // the real key — the proxy never injects one.
1035        for key in [
1036            "Bearer or-v1-9f8e7d6c", // OpenRouter
1037            "Bearer gsk_live_1234",  // (still works)
1038            "Bearer abc.def.ghi",    // gateway/service token
1039            "Bearer 0123456789",     // opaque
1040        ] {
1041            let req = build_request(&[("authorization", key)], "/v1/responses");
1042            assert!(
1043                has_provider_api_key(&req),
1044                "non-sk Bearer must count as a provider credential: {key}"
1045            );
1046        }
1047    }
1048
1049    #[test]
1050    fn has_provider_api_key_empty_bearer_rejected() {
1051        // A blank credential — or a bare scheme word with no token (some HTTP
1052        // stacks trim trailing whitespace down to just "Bearer") — is not auth.
1053        for bad in ["Bearer    ", "", "Bearer", "bearer", "   "] {
1054            let req = build_request(&[("authorization", bad)], "/responses");
1055            assert!(
1056                !has_provider_api_key(&req),
1057                "blank/scheme-only Authorization must not authenticate: {bad:?}"
1058            );
1059        }
1060    }
1061
1062    // --- #334: opt-in strict proxy auth (proxy_require_token) ---
1063
1064    #[test]
1065    fn provider_key_fallback_allowed_in_default_mode() {
1066        // Default (require_token = false): a provider key on a provider route is
1067        // sufficient. This is what lets a local AI tool authenticate with its own
1068        // key and no lean-ctx Bearer token (the loopback-friendly behavior).
1069        assert!(provider_key_fallback_allowed(false, true, true));
1070    }
1071
1072    #[test]
1073    fn provider_key_fallback_denied_in_strict_mode() {
1074        // Strict (require_token = true, e.g. shared/multi-user host): the
1075        // provider-key fallback is disabled, so even a valid provider key on a
1076        // provider route is not enough — the Bearer token becomes mandatory.
1077        assert!(!provider_key_fallback_allowed(true, true, true));
1078    }
1079
1080    #[test]
1081    fn provider_key_fallback_requires_key_and_provider_route() {
1082        // The fallback never fires without a provider key, nor off a provider
1083        // route — regardless of mode.
1084        assert!(!provider_key_fallback_allowed(false, false, true));
1085        assert!(!provider_key_fallback_allowed(false, true, false));
1086        assert!(!provider_key_fallback_allowed(true, false, true));
1087    }
1088
1089    #[test]
1090    fn proxy_require_token_defaults_off() {
1091        // The strict mode must be opt-in: a fresh config keeps the loopback
1092        // behavior so existing local setups (Claude Code, OpenCode, Codex) keep
1093        // working without a token.
1094        assert!(!crate::core::config::Config::default().proxy_require_token);
1095    }
1096
1097    // --- #353: bare provider endpoints (OpenCode / @ai-sdk/openai) ---
1098
1099    #[test]
1100    fn is_provider_route_bare_responses_and_messages() {
1101        // Clients that point their base URL at the proxy root (no `/v1`) send the
1102        // bare endpoint; auth must still recognise it as a provider route.
1103        assert!(is_provider_route("/responses"));
1104        assert!(is_provider_route("/responses/resp_123/input_items"));
1105        assert!(is_provider_route("/messages"));
1106    }
1107
1108    #[test]
1109    fn canonical_provider_path_rewrites_bare_endpoints() {
1110        assert_eq!(
1111            canonical_provider_path("/responses").as_deref(),
1112            Some("/v1/responses")
1113        );
1114        assert_eq!(
1115            canonical_provider_path("/chat/completions").as_deref(),
1116            Some("/v1/chat/completions")
1117        );
1118        assert_eq!(
1119            canonical_provider_path("/messages").as_deref(),
1120            Some("/v1/messages")
1121        );
1122    }
1123
1124    #[test]
1125    fn canonical_provider_path_preserves_subpaths() {
1126        assert_eq!(
1127            canonical_provider_path("/responses/resp_abc/cancel").as_deref(),
1128            Some("/v1/responses/resp_abc/cancel")
1129        );
1130        assert_eq!(
1131            canonical_provider_path("/messages/batches/batch_1").as_deref(),
1132            Some("/v1/messages/batches/batch_1")
1133        );
1134    }
1135
1136    #[test]
1137    fn canonical_provider_path_ignores_already_canonical_and_unknown() {
1138        // Already canonical → no rewrite (avoids `/v1/v1/...`).
1139        assert_eq!(canonical_provider_path("/v1/responses"), None);
1140        assert_eq!(canonical_provider_path("/v1/chat/completions"), None);
1141        // Unrelated paths are untouched.
1142        assert_eq!(canonical_provider_path("/health"), None);
1143        assert_eq!(canonical_provider_path("/responsesx"), None);
1144        assert_eq!(canonical_provider_path("/"), None);
1145    }
1146
1147    #[test]
1148    fn canonical_provider_path_collapses_double_v1_prefix() {
1149        // OPENAI_BASE_URL now advertises `/v1` (#366); a client treating it as an
1150        // origin and appending `/v1/...` itself produces a double prefix.
1151        assert_eq!(
1152            canonical_provider_path("/v1/v1/responses").as_deref(),
1153            Some("/v1/responses")
1154        );
1155        assert_eq!(
1156            canonical_provider_path("/v1/v1/chat/completions").as_deref(),
1157            Some("/v1/chat/completions")
1158        );
1159    }
1160
1161    #[test]
1162    fn normalized_provider_uri_rewrites_path_and_preserves_query() {
1163        use axum::http::Uri;
1164        let uri: Uri = "/responses?stream=true".parse().unwrap();
1165        let rewritten = normalized_provider_uri(&uri).expect("bare /responses must rewrite");
1166        assert_eq!(rewritten.path(), "/v1/responses");
1167        assert_eq!(rewritten.query(), Some("stream=true"));
1168        assert_eq!(
1169            rewritten
1170                .path_and_query()
1171                .map(axum::http::uri::PathAndQuery::as_str),
1172            Some("/v1/responses?stream=true")
1173        );
1174    }
1175
1176    #[test]
1177    fn normalized_provider_uri_noop_for_canonical() {
1178        use axum::http::Uri;
1179        let uri: Uri = "/v1/responses".parse().unwrap();
1180        assert!(normalized_provider_uri(&uri).is_none());
1181    }
1182}
1183
1184#[cfg(test)]
1185mod upstream_tests {
1186    use super::*;
1187
1188    fn upstreams_with_openai(openai: &str) -> Upstreams {
1189        Upstreams {
1190            anthropic: "https://api.anthropic.com".into(),
1191            openai: openai.into(),
1192            chatgpt: "https://chatgpt.com".into(),
1193            gemini: "https://generativelanguage.googleapis.com".into(),
1194        }
1195    }
1196
1197    /// The #449 core wiring: provider handlers read the upstream per request from
1198    /// the watch channel, so a published change is served immediately, without
1199    /// rebuilding the `ProxyState`.
1200    #[tokio::test]
1201    async fn proxy_state_reads_upstream_live_from_watch() {
1202        let (tx, rx) =
1203            tokio::sync::watch::channel(Arc::new(upstreams_with_openai("https://old.example")));
1204        let state = ProxyState {
1205            client: reqwest::Client::new(),
1206            port: 0,
1207            stats: Arc::new(ProxyStats::default()),
1208            introspect: Arc::new(introspect::IntrospectState::default()),
1209            upstreams: rx,
1210            chatgpt_cookies: chatgpt_cookies::shared_chatgpt_cloudflare_cookie_store(),
1211        };
1212        assert_eq!(state.openai_upstream(), "https://old.example");
1213
1214        tx.send(Arc::new(upstreams_with_openai("https://new.example")))
1215            .unwrap();
1216        assert_eq!(
1217            state.openai_upstream(),
1218            "https://new.example",
1219            "a live handler read must reflect the published change"
1220        );
1221        assert_eq!(state.upstream_snapshot().openai, "https://new.example");
1222    }
1223
1224    /// End-to-end #449 repro (in-process, no network): a `config set`-style edit
1225    /// to config.toml is picked up by a *running* proxy's refresh task within the
1226    /// reload interval — without any restart. Before the fix this value stayed
1227    /// frozen at the start-time upstream forever.
1228    ///
1229    /// The process-global env lock is intentionally held across the polling
1230    /// `.await`s to keep `LEAN_CTX_*` isolated for the whole test; safe because
1231    /// each `#[tokio::test]` owns its current-thread runtime, so this std guard
1232    /// only makes *other* test threads wait — it can never deadlock this one.
1233    #[tokio::test]
1234    #[allow(clippy::await_holding_lock)]
1235    async fn config_change_is_picked_up_live_without_restart() {
1236        use crate::core::config::Config;
1237
1238        let _lock = crate::core::data_dir::test_env_lock();
1239        let tmp = tempfile::tempdir().unwrap();
1240        crate::test_env::set_var("LEAN_CTX_DATA_DIR", tmp.path());
1241        // Isolate from a developer shell that exports the env override (#449),
1242        // and make the reload fast + deterministic.
1243        crate::test_env::remove_var("LEAN_CTX_OPENAI_UPSTREAM");
1244        crate::test_env::set_var("LEAN_CTX_PROXY_RELOAD_SECS", "1");
1245
1246        // Start state: config.toml points OpenAI at a loopback upstream.
1247        Config::update_global(|c| {
1248            c.proxy.openai_upstream = Some("http://127.0.0.1:19101".into());
1249        })
1250        .unwrap();
1251        let initial = Config::load().proxy.resolve_all();
1252        assert_eq!(initial.openai, "http://127.0.0.1:19101");
1253
1254        let (tx, rx) = tokio::sync::watch::channel(Arc::new(initial.clone()));
1255        spawn_upstream_refresh(tx, initial);
1256
1257        // `lean-ctx config set proxy.openai_upstream …` (same safe write path).
1258        Config::update_global(|c| {
1259            c.proxy.openai_upstream = Some("http://127.0.0.1:19102".into());
1260        })
1261        .unwrap();
1262
1263        // Poll the live value the handlers would read — no restart in between.
1264        let mut live = rx.borrow().openai.clone();
1265        for _ in 0..80 {
1266            if live == "http://127.0.0.1:19102" {
1267                break;
1268            }
1269            tokio::time::sleep(std::time::Duration::from_millis(100)).await;
1270            live = rx.borrow().openai.clone();
1271        }
1272        assert_eq!(
1273            live, "http://127.0.0.1:19102",
1274            "running proxy must serve the new config.toml upstream without a restart"
1275        );
1276
1277        crate::test_env::remove_var("LEAN_CTX_PROXY_RELOAD_SECS");
1278        crate::test_env::remove_var("LEAN_CTX_DATA_DIR");
1279    }
1280}
lean_ctx/proxy/mod.rs

lean_ctx/proxy/
mod.rs
