firstpass_proxy/
run.rs

1//! Server bootstrap shared by the `firstpass` and `firstpass-proxy` binaries: build state from
2//! config, open the trace store, and serve until Ctrl-C. Keeping this in the lib means both the
3//! unified CLI (`firstpass up`) and the bare proxy binary start the server the exact same way.
4
5use std::sync::Arc;
6
7use crate::bandit::{ContextBucket, StartRungBandit};
8use crate::gate::GateHealthRegistry;
9use crate::provider::ProviderRegistry;
10use crate::{AppState, ProxyConfig, app, store};
11
12/// Initialize the global tracing subscriber from `RUST_LOG` (default `info`). Called by the
13/// binaries, not the library internals.
14pub fn init_tracing() {
15    tracing_subscriber::fmt()
16        .with_env_filter(
17            tracing_subscriber::EnvFilter::try_from_default_env()
18                .unwrap_or_else(|_| tracing_subscriber::EnvFilter::new("info")),
19        )
20        .init();
21}
22
23/// Register a default error budget for every gate named across enforce routes: auto-disable a gate
24/// whose abstain rate exceeds 25% over its last 50 runs (SPEC §7.2).
25#[must_use]
26pub fn build_gate_health(config: &ProxyConfig) -> GateHealthRegistry {
27    let mut gate_health = GateHealthRegistry::new();
28    if let Some(routing) = config.routing.as_ref() {
29        let mut seen = std::collections::HashSet::new();
30        for route in &routing.routes {
31            for gate in route.gates.iter().chain(&route.deferred_gates) {
32                if seen.insert(gate.clone()) {
33                    gate_health = gate_health.with_budget(gate.clone(), 50, 0.25);
34                }
35            }
36        }
37    }
38    gate_health
39}
40
41/// Open the trace store, build [`AppState`], and serve the HTTP proxy until Ctrl-C.
42///
43/// # Errors
44/// Returns any error from opening the store, building the HTTP client, binding the listener, or
45/// serving.
46pub async fn serve(config: ProxyConfig) -> Result<(), Box<dyn std::error::Error>> {
47    let (traces, spill, writer) = store::open_with_receipts(&config.db_path, config.receipts_mode)?;
48    let bind = config.bind.clone();
49    // Build the provider registry from any `[[provider]]` entries in the routing config (Groq,
50    // Together, Ollama, …), on top of the built-in anthropic/openai defaults. No config => defaults.
51    let provider_defs = config
52        .routing
53        .as_ref()
54        .map(|r| r.providers.as_slice())
55        .unwrap_or_default();
56    let providers = ProviderRegistry::from_config(
57        provider_defs,
58        &config.upstream_anthropic,
59        &config.upstream_openai,
60    );
61    let gate_health = build_gate_health(&config);
62    // Online adaptive conformal (opt-in): seed the live threshold from the fixed one (or 0.5) and
63    // let /v1/feedback track it. Absent config => None => fixed-threshold behavior, byte-identical.
64    let adaptive = config
65        .routing
66        .as_ref()
67        .and_then(|r| r.escalation.adaptive.as_ref())
68        .map(|a| {
69            let init = config
70                .routing
71                .as_ref()
72                .and_then(|r| r.escalation.serve_threshold)
73                .unwrap_or(0.5);
74            Arc::new(std::sync::Mutex::new(
75                firstpass_core::conformal::AdaptiveConformal::new(a.alpha, a.gamma, init),
76            ))
77        });
78    let tenant_rate_limiter = crate::proxy::build_tenant_rate_limiter(&config);
79
80    // UCB1 start-rung bandit (opt-in): warm-start from stored traces so learning survives
81    // restarts. Forgiving: an unreadable or absent store simply yields an empty bandit.
82    // ponytail: operator-wide load (load_all_traces) is correct here — the bandit is also
83    // operator-wide, keyed by context bucket, not by tenant.
84    let bandit = config
85        .routing
86        .as_ref()
87        .and_then(|r| r.escalation.bandit.as_ref())
88        .map(|bc| {
89            let algorithm = match bc.algorithm {
90                firstpass_core::BanditAlgorithm::Ucb1 => crate::bandit::Algorithm::Ucb1,
91                firstpass_core::BanditAlgorithm::Thompson => crate::bandit::Algorithm::Thompson,
92            };
93            // Seed from a fresh UUID: per-process randomness, no new deps (tests seed explicitly).
94            let seed = uuid::Uuid::now_v7().as_u128() as u64;
95            let mut b = StartRungBandit::with_algorithm(
96                bc.min_observations,
97                bc.exploration,
98                algorithm,
99                bc.discount,
100                seed,
101            );
102            if let Ok(traces_history) = store::load_all_traces(&config.db_path) {
103                for trace in &traces_history {
104                    let ctx = ContextBucket::from_features(&trace.request.features);
105                    b.feed_trace_attempts(&ctx, &trace.attempts);
106                }
107                tracing::info!(
108                    n = traces_history.len(),
109                    "bandit warm-started from trace store"
110                );
111            }
112            Arc::new(std::sync::Mutex::new(b))
113        });
114
115    let state = AppState {
116        config: Arc::new(config),
117        // Observe passthrough may stream SSE, so only bound the CONNECT phase here — a total or
118        // read timeout would sever a long-lived stream. (The enforce providers, which never stream
119        // through the adapter, carry a full request timeout — see `ProviderRegistry::new`.)
120        http: reqwest::Client::builder()
121            .connect_timeout(std::time::Duration::from_secs(10))
122            .build()?,
123        providers,
124        gate_health: Arc::new(gate_health),
125        traces,
126        adaptive,
127        bandit,
128        tenant_rate_limiter,
129        spill,
130    };
131
132    let listener = tokio::net::TcpListener::bind(&bind).await?;
133    tracing::info!(%bind, "firstpass listening");
134    tracing::info!("offboard: unset ANTHROPIC_BASE_URL");
135
136    axum::serve(listener, app(state)?)
137        .with_graceful_shutdown(shutdown_signal())
138        .await?;
139
140    // Dropping `traces` (via `state`, already out of scope) closes the channel; wait for the
141    // writer to flush and exit before the process ends.
142    drop(writer);
143    Ok(())
144}
145
146async fn shutdown_signal() {
147    if let Err(err) = tokio::signal::ctrl_c().await {
148        tracing::warn!(%err, "failed to install Ctrl-C handler; shutting down anyway");
149    }
150}
firstpass_proxy/run.rs

firstpass_proxy/
run.rs
