Skip to main content

alex_proxy/
lib.rs

1use std::collections::HashMap;
2use std::future::Future;
3use std::path::PathBuf;
4use std::pin::Pin;
5use std::sync::Arc;
6use std::time::Duration;
7
8use alex_auth::{now_ms, Account, Vault};
9use alex_core::{
10    compute_cost, conversation_root, parse_grpc_web_response, parse_since, parse_sse_usage,
11    parse_trace_tags, route_model, usage_from_json, validate_grpc_status_headers, window_label,
12    ClientFormat, Provider, TraceRecord, GROK_CREDITS_ENDPOINT, GROK_CREDITS_REQUEST_BODY,
13};
14use alex_store::{Store, TraceFilter};
15use axum::body::{Body, Bytes};
16use axum::extract::{ConnectInfo, Path, Query, State};
17use axum::http::{HeaderMap, HeaderValue, StatusCode};
18use axum::response::{IntoResponse, Response};
19use axum::routing::{get, post};
20use axum::Router;
21use futures_util::StreamExt;
22use serde_json::{json, Value};
23use tokio_stream::wrappers::ReceiverStream;
24
25const ANTHROPIC_BASE: &str = "https://api.anthropic.com";
26const OPENAI_BASE: &str = "https://api.openai.com";
27const CODEX_BASE: &str = "https://chatgpt.com/backend-api/codex";
28const XAI_BASE: &str = "https://cli-chat-proxy.grok.com/v1";
29const GROK_CLIENT_VERSION: &str = "0.2.77";
30const ANTHROPIC_OAUTH_BETA: &str = "oauth-2025-04-20";
31const GEMINI_CODE_ASSIST_BASE: &str = "https://cloudcode-pa.googleapis.com";
32const GEMINI_CODE_ASSIST_VERSION: &str = "v1internal";
33const GEMINI_API_BASE: &str = "https://generativelanguage.googleapis.com";
34
35#[derive(Debug, Clone)]
36pub struct DarioActive {
37    pub generation_id: String,
38    pub base_url: String,
39    pub api_key: String,
40}
41
42pub trait DarioRouter: Send + Sync {
43    fn active(&self) -> Option<DarioActive>;
44    fn begin(&self, generation_id: &str) -> Option<Box<dyn std::any::Any + Send>>;
45    fn status(&self) -> Value;
46    fn suspect(&self, generation_id: &str);
47}
48
49pub type UpdateApplyFuture =
50    Pin<Box<dyn Future<Output = Result<Value, UpdateApplyError>> + Send + 'static>>;
51
52#[derive(Debug)]
53pub enum UpdateApplyError {
54    Conflict(Value),
55    Failed(String),
56}
57
58pub trait DaemonUpdater: Send + Sync {
59    fn apply(&self) -> UpdateApplyFuture;
60}
61
62fn suspect_dario(state: &AppState, account: &Account) {
63    if account.kind != "dario" {
64        return;
65    }
66    if let (Some(dario), Some(gen)) = (&state.dario, account.id.strip_prefix("dario:")) {
67        dario.suspect(gen);
68    }
69}
70
71pub struct AppState {
72    pub local_key: String,
73    pub vault: Arc<Vault>,
74    pub store: Arc<Store>,
75    pub http: reqwest::Client,
76    pub dario: Option<Arc<dyn DarioRouter>>,
77    pub in_flight: std::sync::atomic::AtomicI64,
78    pub started_ms: i64,
79    pub base_url: String,
80    pub anthropic_usage: std::sync::Mutex<UsageCache>,
81    pub xai_usage: std::sync::Mutex<UsageCache>,
82    pub logins: alex_auth::sessions::LoginManager,
83    pub run_keys: std::sync::RwLock<HashMap<String, CachedRunKey>>,
84    pub update_status: Arc<tokio::sync::RwLock<Option<Value>>>,
85    pub daemon_updater: std::sync::RwLock<Option<Arc<dyn DaemonUpdater>>>,
86}
87
88#[derive(Debug, Clone)]
89pub struct CachedRunKey {
90    pub run_id: Option<String>,
91    pub tags_json: Option<String>,
92    pub expires_ms: Option<i64>,
93}
94
95struct InFlight(Arc<AppState>);
96
97impl InFlight {
98    fn new(state: &Arc<AppState>) -> Self {
99        state
100            .in_flight
101            .fetch_add(1, std::sync::atomic::Ordering::SeqCst);
102        Self(state.clone())
103    }
104}
105
106impl Drop for InFlight {
107    fn drop(&mut self) {
108        self.0
109            .in_flight
110            .fetch_sub(1, std::sync::atomic::Ordering::SeqCst);
111    }
112}
113
114pub fn build_state(
115    local_key: String,
116    vault: Arc<Vault>,
117    store: Arc<Store>,
118    dario: Option<Arc<dyn DarioRouter>>,
119    base_url: String,
120) -> Arc<AppState> {
121    let http = reqwest::Client::builder()
122        .connect_timeout(Duration::from_secs(15))
123        .build()
124        .expect("reqwest client");
125    Arc::new(AppState {
126        local_key,
127        vault,
128        store,
129        http,
130        dario,
131        in_flight: std::sync::atomic::AtomicI64::new(0),
132        started_ms: now_ms(),
133        base_url,
134        anthropic_usage: std::sync::Mutex::new(UsageCache::default()),
135        xai_usage: std::sync::Mutex::new(UsageCache::default()),
136        logins: alex_auth::sessions::LoginManager::default(),
137        run_keys: std::sync::RwLock::new(HashMap::new()),
138        update_status: Arc::new(tokio::sync::RwLock::new(None)),
139        daemon_updater: std::sync::RwLock::new(None),
140    })
141}
142
143pub fn set_daemon_updater(state: &Arc<AppState>, updater: Arc<dyn DaemonUpdater>) {
144    if let Ok(mut slot) = state.daemon_updater.write() {
145        *slot = Some(updater);
146    }
147}
148
149async fn require_local_key(
150    State(state): State<Arc<AppState>>,
151    req: axum::extract::Request,
152    next: axum::middleware::Next,
153) -> Response {
154    let ok = client_key(req.headers())
155        .map(|k| k == state.local_key)
156        .unwrap_or(false);
157    if !ok {
158        return error_response(
159            StatusCode::UNAUTHORIZED,
160            "admin routes require x-api-key: <local_key>",
161        );
162    }
163    next.run(req).await
164}
165
166pub fn router(state: Arc<AppState>) -> Router {
167    // Control-plane + report routes: gated by the local key (sent as
168    // `x-api-key: <key>` or `Authorization: Bearer <key>`) so a LAN/0.0.0.0
169    // bind doesn't expose them. Run keys are NOT accepted here — a worker's
170    // run key must not mint/revoke run keys or read the trace store.
171    let gated = Router::new()
172        .route(
173            "/admin/run-keys",
174            get(admin_run_keys_list).post(admin_run_keys_create),
175        )
176        .route(
177            "/admin/run-keys/{id}",
178            axum::routing::delete(admin_run_keys_revoke),
179        )
180        .route("/admin/storage", get(admin_storage))
181        .route("/admin/storage/prune", post(admin_storage_prune))
182        .route("/admin/traces", get(admin_traces))
183        .route("/admin/accounts", get(admin_accounts))
184        .route(
185            "/admin/accounts/{id}",
186            axum::routing::delete(admin_account_remove),
187        )
188        .route("/admin/auth/gemini-key", post(admin_auth_gemini_key))
189        .route("/admin/health", get(admin_health))
190        .route("/admin/analytics", get(admin_analytics))
191        .route("/admin/limits", get(admin_limits))
192        .route("/admin/update", get(admin_update).post(admin_update_apply))
193        .route("/admin/dario", get(admin_dario))
194        .route("/admin/dario/prompt-caches", get(admin_dario_prompt_caches))
195        .route(
196            "/admin/dario/prompt-caches/{key}",
197            axum::routing::delete(admin_dario_prompt_cache_delete),
198        )
199        .route("/admin/auth/import", post(admin_auth_import))
200        .route("/admin/auth/login/start", post(admin_auth_login_start))
201        .route("/admin/auth/login/complete", post(admin_auth_login_complete))
202        .route("/admin/auth/login/{id}", get(admin_auth_login_status))
203        .route("/traces/search", get(traces_search))
204        .route("/traces/export.ndjson", get(traces_export))
205        .route("/traces/sessions", get(traces_sessions))
206        .route(
207            "/traces/sessions/{session_id}/transcript",
208            get(traces_session_transcript),
209        )
210        .route("/traces/{id}", get(trace_get).delete(trace_delete))
211        .route("/traces/{id}/reply.md", get(trace_reply_md))
212        .route("/traces/{id}/body/{kind}", get(trace_body))
213        .route("/traces/runs/{run_id}", get(traces_run_summary))
214        .route("/traces/runs/{run_id}/events", get(traces_run_events))
215        .route("/traces/runs/{run_id}/export.ndjson", get(traces_run_export))
216        .route("/traces/runs/{run_id}/artifacts", get(traces_run_artifacts))
217        .route_layer(axum::middleware::from_fn_with_state(
218            state.clone(),
219            require_local_key,
220        ));
221
222    Router::new()
223        .route("/health", get(health))
224        .route("/connect", get(connect_info))
225        .route("/v1/models", get(models))
226        .route("/v1/messages", post(anthropic_messages))
227        .route("/v1/chat/completions", post(openai_chat))
228        .route("/v1/responses", post(openai_responses))
229        .route("/chat/completions", post(openai_chat))
230        .route("/responses", post(openai_responses))
231        .route("/v1beta/models/{model_action}", post(gemini_generate))
232        .merge(gated)
233        .layer(axum::extract::DefaultBodyLimit::max(64 * 1024 * 1024))
234        .with_state(state)
235}
236
237async fn admin_auth_import(
238    State(state): State<Arc<AppState>>,
239    body: Option<axum::Json<Value>>,
240) -> Response {
241    let source = body
242        .as_ref()
243        .and_then(|b| b.0["source"].as_str().or(b.0["provider"].as_str()))
244        .unwrap_or("all")
245        .to_string();
246    match alex_auth::import_all(&state.vault, &source).await {
247        Ok(outcomes) => {
248            let items: Vec<Value> = outcomes
249                .iter()
250                .map(|o| {
251                    json!({
252                        "source": o.source,
253                        "imported": o.imported,
254                        "note": o.note,
255                    })
256                })
257                .collect();
258            axum::Json(json!({"outcomes": items})).into_response()
259        }
260        Err(e) => error_response(StatusCode::BAD_REQUEST, &e.to_string()),
261    }
262}
263
264async fn admin_account_remove(
265    State(state): State<Arc<AppState>>,
266    Path(id): Path<String>,
267) -> Response {
268    match state.vault.remove(&id).await {
269        Ok(true) => axum::Json(json!({"removed": id})).into_response(),
270        Ok(false) => error_response(StatusCode::NOT_FOUND, &format!("unknown account '{id}'")),
271        Err(e) => error_response(StatusCode::INTERNAL_SERVER_ERROR, &e.to_string()),
272    }
273}
274
275async fn admin_auth_gemini_key(
276    State(state): State<Arc<AppState>>,
277    body: axum::Json<Value>,
278) -> Response {
279    let Some(key) = body.0["key"].as_str().map(str::trim).filter(|k| !k.is_empty()) else {
280        return error_response(StatusCode::BAD_REQUEST, "missing 'key'");
281    };
282    let account = Account {
283        id: alex_auth::named_account_id(Provider::Gemini, "api_key", "default"),
284        provider: Provider::Gemini,
285        kind: "api_key".into(),
286        name: "default".into(),
287        description: None,
288        paused: false,
289        label: Some("gemini (AI Studio key)".into()),
290        access_token: None,
291        refresh_token: None,
292        id_token: None,
293        api_key: Some(key.to_string()),
294        expires_at_ms: None,
295        last_refresh_ms: None,
296        account_meta: Value::Null,
297        cooldown_until_ms: None,
298        status: "active".into(),
299        path: None,
300    };
301    match state.vault.upsert(account).await {
302        Ok(()) => axum::Json(json!({"saved": "gemini-api-key"})).into_response(),
303        Err(e) => error_response(StatusCode::INTERNAL_SERVER_ERROR, &e.to_string()),
304    }
305}
306
307async fn admin_auth_login_start(
308    State(state): State<Arc<AppState>>,
309    body: axum::Json<Value>,
310) -> Response {
311    let Some(provider) = body.0["provider"].as_str() else {
312        return error_response(StatusCode::BAD_REQUEST, "missing 'provider'");
313    };
314    match state.logins.start(state.vault.clone(), provider).await {
315        Ok(snapshot) => axum::Json(snapshot).into_response(),
316        Err(e) => error_response(StatusCode::BAD_REQUEST, &e.to_string()),
317    }
318}
319
320async fn admin_auth_login_complete(
321    State(state): State<Arc<AppState>>,
322    body: axum::Json<Value>,
323) -> Response {
324    let (Some(id), Some(input)) = (body.0["login_id"].as_str(), body.0["input"].as_str()) else {
325        return error_response(StatusCode::BAD_REQUEST, "missing 'login_id' or 'input'");
326    };
327    match state.logins.complete(state.vault.clone(), id, input).await {
328        Ok(snapshot) => axum::Json(snapshot).into_response(),
329        Err(e) => error_response(StatusCode::BAD_REQUEST, &e.to_string()),
330    }
331}
332
333async fn admin_auth_login_status(
334    State(state): State<Arc<AppState>>,
335    axum::extract::Path(id): axum::extract::Path<String>,
336) -> Response {
337    match state.logins.status(&id).await {
338        Some(snapshot) => axum::Json(snapshot).into_response(),
339        None => error_response(StatusCode::NOT_FOUND, "unknown or expired login session"),
340    }
341}
342
343async fn health(State(state): State<Arc<AppState>>) -> impl IntoResponse {
344    axum::Json(json!({
345        "status": "ok",
346        "service": "alexandria",
347        "version": env!("CARGO_PKG_VERSION"),
348        "in_flight": state.in_flight.load(std::sync::atomic::Ordering::SeqCst),
349        "uptime_s": (now_ms() - state.started_ms) / 1000,
350        "dario": state.dario.is_some(),
351    }))
352}
353
354fn rewrite_host(base_url: &str, host: &str) -> String {
355    let port = base_url.rsplit(':').next().unwrap_or("4100");
356    format!("http://{host}:{port}")
357}
358
359pub fn connect_payload(base_url: &str, local_key: &str) -> (Value, String) {
360    let base = base_url.trim_end_matches('/').to_string();
361    let v1 = format!("{base}/v1");
362    let exports = format!(
363        "export ANTHROPIC_BASE_URL={base}\nexport ANTHROPIC_API_KEY={local_key}\n\
364         export OPENAI_BASE_URL={v1}\nexport OPENAI_API_KEY={local_key}\n\
365         export XAI_API_KEY={local_key}\nexport GROK_MODELS_BASE_URL={v1}\n\
366         export GOOGLE_GEMINI_BASE_URL={base}\nexport GOOGLE_GENAI_API_VERSION=v1beta\n\
367         export GEMINI_API_KEY={local_key}\nexport GEMINI_API_KEY_AUTH_MECHANISM=bearer\n\
368         export GOOGLE_GENAI_USE_GCA=false\n"
369    );
370    let payload = json!({
371        "service": "alexandria",
372        "base_url": base,
373        "api_key": local_key,
374        "anthropic": {"base_url": base, "env": {"ANTHROPIC_BASE_URL": base, "ANTHROPIC_API_KEY": local_key}},
375        "openai": {"base_url": v1, "env": {"OPENAI_BASE_URL": v1, "OPENAI_API_KEY": local_key}},
376        "xai": {"base_url": v1, "env": {"XAI_API_KEY": local_key, "GROK_MODELS_BASE_URL": v1}},
377        "gemini": {"base_url": base, "env": {
378            "GOOGLE_GEMINI_BASE_URL": base,
379            "GOOGLE_GENAI_API_VERSION": "v1beta",
380            "GEMINI_API_KEY": local_key,
381            "GEMINI_API_KEY_AUTH_MECHANISM": "bearer",
382            "GOOGLE_GENAI_USE_GCA": "false",
383        }},
384        "exports": exports,
385    });
386    (payload, exports)
387}
388
389async fn connect_info(
390    State(state): State<Arc<AppState>>,
391    axum::extract::ConnectInfo(peer): axum::extract::ConnectInfo<std::net::SocketAddr>,
392    Query(q): Query<HashMap<String, String>>,
393) -> Response {
394    if !peer.ip().is_loopback() {
395        return error_response(
396            StatusCode::FORBIDDEN,
397            "connection info is only served to loopback clients",
398        );
399    }
400    let base = match q.get("host").map(String::as_str) {
401        Some("docker") => rewrite_host(&state.base_url, "host.docker.internal"),
402        Some(host) if !host.is_empty() => rewrite_host(&state.base_url, host),
403        _ => state.base_url.clone(),
404    };
405    let (payload, exports) = connect_payload(&base, &state.local_key);
406    if q.get("format").map(|f| f == "env").unwrap_or(false) {
407        return Response::builder()
408            .status(StatusCode::OK)
409            .header("content-type", "text/plain")
410            .body(Body::from(exports))
411            .unwrap_or_else(|e| error_response(StatusCode::INTERNAL_SERVER_ERROR, &e.to_string()));
412    }
413    axum::Json(payload).into_response()
414}
415
416async fn models(State(state): State<Arc<AppState>>) -> impl IntoResponse {
417    let mut ids = state.store.pricing_models();
418    for (alias, _) in alex_core::model_aliases() {
419        ids.push((*alias).to_string());
420    }
421    for id in ids.clone() {
422        ids.push(format!("alexandria/{id}"));
423    }
424    let data: Vec<Value> = ids
425        .into_iter()
426        .map(|m| json!({"id": m, "object": "model", "owned_by": "alexandria"}))
427        .collect();
428    axum::Json(json!({"object": "list", "data": data}))
429}
430
431async fn admin_analytics(
432    State(state): State<Arc<AppState>>,
433    Query(q): Query<HashMap<String, String>>,
434) -> Response {
435    let minutes: i64 = q
436        .get("since_minutes")
437        .and_then(|s| s.parse().ok())
438        .unwrap_or(60);
439    match state.store.analytics(now_ms() - minutes * 60_000) {
440        Ok(v) => axum::Json(v).into_response(),
441        Err(e) => error_response(StatusCode::INTERNAL_SERVER_ERROR, &e.to_string()),
442    }
443}
444
445const USAGE_CACHE_TTL_MS: i64 = 300_000;
446const USAGE_BACKOFF_BASE_MS: i64 = 60_000;
447const USAGE_BACKOFF_MAX_MS: i64 = 3_600_000;
448
449#[derive(Default)]
450pub struct UsageCache {
451    fetched_at_ms: i64,
452    entry: Option<Value>,
453    cooldown_until_ms: i64,
454    failures: u32,
455}
456
457fn usage_backoff_ms(failures: u32, retry_after_ms: Option<i64>) -> i64 {
458    let exp = USAGE_BACKOFF_BASE_MS
459        .saturating_mul(1i64 << failures.saturating_sub(1).min(6))
460        .min(USAGE_BACKOFF_MAX_MS);
461    exp.max(retry_after_ms.unwrap_or(0))
462}
463
464async fn anthropic_usage_entry(state: &Arc<AppState>) -> Option<Value> {
465    let account = state.vault.account_for(Provider::Anthropic, true).await.ok()?;
466    if account.kind != "oauth" {
467        return None;
468    }
469    let token = account.access_token.as_deref()?.to_string();
470    {
471        let cache = state.anthropic_usage.lock().unwrap();
472        if cache.entry.is_some() && now_ms() < cache.fetched_at_ms + USAGE_CACHE_TTL_MS {
473            return cache.entry.clone();
474        }
475        if now_ms() < cache.cooldown_until_ms {
476            return cache.entry.clone();
477        }
478    }
479    let result = state
480        .http
481        .get(format!("{ANTHROPIC_BASE}/api/oauth/usage"))
482        .header("authorization", format!("Bearer {token}"))
483        .header("anthropic-beta", ANTHROPIC_OAUTH_BETA)
484        .header("accept", "application/json")
485        .header("user-agent", "claude-cli/2.1.202 (external, cli)")
486        .send()
487        .await;
488    match result {
489        Ok(resp) if resp.status().is_success() => {
490            let raw: Value = resp.json().await.unwrap_or(Value::Null);
491            let mut windows = Vec::new();
492            for (name, key) in [
493                ("5h", "five_hour"),
494                ("7d", "seven_day"),
495                ("7d opus", "seven_day_opus"),
496                ("7d sonnet", "seven_day_sonnet"),
497            ] {
498                let w = &raw[key];
499                if w.is_object() {
500                    windows.push(json!({
501                        "window": name,
502                        "used_pct": w["utilization"],
503                        "resets_at": w["resets_at"],
504                    }));
505                }
506            }
507            let entry = json!({
508                "provider": "anthropic",
509                "source": "oauth usage endpoint",
510                "plan": account.label,
511                "windows": windows,
512                "extra_usage": raw["extra_usage"],
513            });
514            let mut cache = state.anthropic_usage.lock().unwrap();
515            cache.fetched_at_ms = now_ms();
516            cache.cooldown_until_ms = 0;
517            cache.failures = 0;
518            cache.entry = Some(entry.clone());
519            Some(entry)
520        }
521        Ok(resp) => {
522            let retry_after = resp
523                .headers()
524                .get("retry-after")
525                .and_then(|v| v.to_str().ok())
526                .and_then(|v| v.parse::<i64>().ok())
527                .map(|s| s.clamp(30, 3600) * 1000);
528            let mut cache = state.anthropic_usage.lock().unwrap();
529            cache.failures += 1;
530            let cooldown = usage_backoff_ms(cache.failures, retry_after);
531            cache.cooldown_until_ms = now_ms() + cooldown;
532            tracing::debug!(
533                status = resp.status().as_u16(),
534                failures = cache.failures,
535                cooldown_ms = cooldown,
536                "anthropic usage endpoint unavailable; backing off"
537            );
538            cache.entry.clone()
539        }
540        Err(_) => {
541            let mut cache = state.anthropic_usage.lock().unwrap();
542            cache.failures += 1;
543            let cooldown = usage_backoff_ms(cache.failures, None);
544            cache.cooldown_until_ms = now_ms() + cooldown;
545            cache.entry.clone()
546        }
547    }
548}
549
550/// Fetch SuperGrok weekly credits from grok.com gRPC-web billing RPC.
551/// Uses the vault's xAI OAuth access token. Degrades gracefully on any failure.
552async fn xai_usage_entry(state: &Arc<AppState>) -> Option<Value> {
553    let account = state.vault.account_for(Provider::Xai, true).await.ok()?;
554    if account.kind != "oauth" {
555        return None;
556    }
557    let token = account.access_token.as_deref()?.to_string();
558    {
559        let cache = state.xai_usage.lock().unwrap();
560        if cache.entry.is_some() && now_ms() < cache.fetched_at_ms + USAGE_CACHE_TTL_MS {
561            return cache.entry.clone();
562        }
563        if now_ms() < cache.cooldown_until_ms {
564            return cache.entry.clone();
565        }
566    }
567
568    let result = state
569        .http
570        .post(GROK_CREDITS_ENDPOINT)
571        .header("authorization", format!("Bearer {token}"))
572        .header("origin", "https://grok.com")
573        .header("referer", "https://grok.com/?_s=usage")
574        .header("accept", "*/*")
575        .header("content-type", "application/grpc-web+proto")
576        .header("x-grpc-web", "1")
577        .header("x-user-agent", "connect-es/2.1.1")
578        .header("user-agent", "Alexandria")
579        .body(GROK_CREDITS_REQUEST_BODY.to_vec())
580        .timeout(Duration::from_secs(15))
581        .send()
582        .await;
583
584    match result {
585        Ok(resp) if resp.status().is_success() => {
586            let headers_for_grpc: Vec<(String, String)> = resp
587                .headers()
588                .iter()
589                .filter_map(|(k, v)| {
590                    let key = k.as_str();
591                    if key.starts_with("grpc-") {
592                        v.to_str().ok().map(|val| (key.to_string(), val.to_string()))
593                    } else {
594                        None
595                    }
596                })
597                .collect();
598            if let Err(e) = validate_grpc_status_headers(headers_for_grpc) {
599                tracing::debug!(error = %e, "xai grok credits grpc header status failed");
600                let mut cache = state.xai_usage.lock().unwrap();
601                cache.failures += 1;
602                let cooldown = usage_backoff_ms(cache.failures, None);
603                cache.cooldown_until_ms = now_ms() + cooldown;
604                return cache.entry.clone();
605            }
606            let body = match resp.bytes().await {
607                Ok(b) => b,
608                Err(e) => {
609                    tracing::debug!(error = %e, "xai grok credits body read failed");
610                    let mut cache = state.xai_usage.lock().unwrap();
611                    cache.failures += 1;
612                    let cooldown = usage_backoff_ms(cache.failures, None);
613                    cache.cooldown_until_ms = now_ms() + cooldown;
614                    return cache.entry.clone();
615                }
616            };
617            let now_s = now_ms() / 1000;
618            match parse_grpc_web_response(&body, now_s) {
619                Ok(snap) => {
620                    let label = window_label(snap.resets_at_s, now_s);
621                    let mut window = json!({
622                        "window": label,
623                        "used_pct": snap.used_percent,
624                    });
625                    if let Some(ts) = snap.resets_at_s {
626                        window["resets_at_s"] = json!(ts);
627                    }
628                    let entry = json!({
629                        "provider": "xai",
630                        "source": "grok web billing",
631                        "plan": account.label,
632                        "windows": [window],
633                    });
634                    let mut cache = state.xai_usage.lock().unwrap();
635                    cache.fetched_at_ms = now_ms();
636                    cache.cooldown_until_ms = 0;
637                    cache.failures = 0;
638                    cache.entry = Some(entry.clone());
639                    Some(entry)
640                }
641                Err(e) => {
642                    tracing::debug!(error = %e, "xai grok credits parse failed");
643                    let mut cache = state.xai_usage.lock().unwrap();
644                    cache.failures += 1;
645                    let cooldown = usage_backoff_ms(cache.failures, None);
646                    cache.cooldown_until_ms = now_ms() + cooldown;
647                    cache.entry.clone()
648                }
649            }
650        }
651        Ok(resp) => {
652            let status = resp.status().as_u16();
653            let retry_after = resp
654                .headers()
655                .get("retry-after")
656                .and_then(|v| v.to_str().ok())
657                .and_then(|v| v.parse::<i64>().ok())
658                .map(|s| s.clamp(30, 3600) * 1000);
659            let mut cache = state.xai_usage.lock().unwrap();
660            cache.failures += 1;
661            let cooldown = usage_backoff_ms(cache.failures, retry_after);
662            cache.cooldown_until_ms = now_ms() + cooldown;
663            tracing::debug!(
664                status,
665                failures = cache.failures,
666                cooldown_ms = cooldown,
667                "xai grok web billing unavailable; backing off"
668            );
669            cache.entry.clone()
670        }
671        Err(e) => {
672            tracing::debug!(error = %e, "xai grok web billing request failed");
673            let mut cache = state.xai_usage.lock().unwrap();
674            cache.failures += 1;
675            let cooldown = usage_backoff_ms(cache.failures, None);
676            cache.cooldown_until_ms = now_ms() + cooldown;
677            cache.entry.clone()
678        }
679    }
680}
681
682pub async fn limits_snapshot(state: &Arc<AppState>) -> Value {
683    let mut providers: Vec<Value> = Vec::new();
684    if let Some(entry) = anthropic_usage_entry(state).await {
685        providers.push(entry);
686    }
687    if let Some(entry) = xai_usage_entry(state).await {
688        providers.push(entry);
689    }
690    for (provider_str, ts_ms, headers_json) in
691        state.store.latest_provider_headers().unwrap_or_default()
692    {
693        if providers
694            .iter()
695            .any(|p| p["provider"].as_str() == Some(&provider_str))
696        {
697            continue;
698        }
699        let Some(provider) = Provider::from_str_loose(&provider_str) else {
700            continue;
701        };
702        let headers: Value = serde_json::from_str(&headers_json).unwrap_or(Value::Null);
703        let mut parsed = alex_core::parse_limit_headers(provider, &headers);
704        if let Some(o) = parsed.as_object_mut() {
705            o.insert("provider".into(), json!(provider_str));
706            o.insert("source".into(), json!("captured response headers"));
707            o.insert("observed_at_ms".into(), json!(ts_ms));
708            providers.push(parsed);
709        }
710    }
711    providers.sort_by_key(|p| p["provider"].as_str().unwrap_or("").to_string());
712    json!({"providers": providers})
713}
714
715async fn admin_limits(State(state): State<Arc<AppState>>) -> Response {
716    axum::Json(limits_snapshot(&state).await).into_response()
717}
718
719async fn admin_update(State(state): State<Arc<AppState>>) -> Response {
720    let stored = state.update_status.read().await.clone();
721    let mut body = json!({
722        "current": env!("CARGO_PKG_VERSION"),
723        "latest": null,
724        "update_available": false,
725        "checked_at_ms": null,
726    });
727    if let Some(Value::Object(fields)) = stored {
728        if let Some(obj) = body.as_object_mut() {
729            obj.extend(fields);
730            obj.insert("current".into(), json!(env!("CARGO_PKG_VERSION")));
731        }
732    }
733    axum::Json(body).into_response()
734}
735
736async fn admin_update_apply(State(state): State<Arc<AppState>>) -> Response {
737    let updater = state
738        .daemon_updater
739        .read()
740        .ok()
741        .and_then(|slot| slot.clone());
742    let Some(updater) = updater else {
743        return error_response(StatusCode::SERVICE_UNAVAILABLE, "daemon updater is not configured");
744    };
745    match updater.apply().await {
746        Ok(body) => {
747            let status = if body["applying"].as_bool() == Some(true) {
748                StatusCode::ACCEPTED
749            } else {
750                StatusCode::OK
751            };
752            (status, axum::Json(body)).into_response()
753        }
754        Err(UpdateApplyError::Conflict(body)) => {
755            (StatusCode::CONFLICT, axum::Json(body)).into_response()
756        }
757        Err(UpdateApplyError::Failed(message)) => {
758            error_response(StatusCode::INTERNAL_SERVER_ERROR, &message)
759        }
760    }
761}
762
763async fn admin_dario(State(state): State<Arc<AppState>>) -> Response {
764    match &state.dario {
765        Some(d) => {
766            let mut status = d.status();
767            if let Some(obj) = status.as_object_mut() {
768                obj.insert("prompt_caches".into(), json!(dario_prompt_caches(&state)));
769            }
770            axum::Json(status).into_response()
771        }
772        None => error_response(StatusCode::NOT_FOUND, "dario mode is not enabled"),
773    }
774}
775
776fn dario_prompt_cache_dir(state: &AppState) -> PathBuf {
777    state.store.data_dir.join("dario-prompt-cache")
778}
779
780fn dario_prompt_cache_key_valid(key: &str) -> bool {
781    !key.is_empty()
782        && key.len() <= 128
783        && key
784            .bytes()
785            .all(|b| b.is_ascii_alphanumeric() || matches!(b, b'.' | b'_' | b'-'))
786}
787
788fn dario_prompt_cache_summary(path: PathBuf) -> Option<Value> {
789    let raw = std::fs::read_to_string(&path).ok()?;
790    let value: Value = serde_json::from_str(&raw).ok()?;
791    let key = path.file_stem()?.to_string_lossy().to_string();
792    let runs: Vec<Value> = value["runs"]
793        .as_array()
794        .map(|runs| runs.iter().rev().take(12).cloned().collect())
795        .unwrap_or_default();
796    Some(json!({
797        "key": value["key"].as_str().unwrap_or(&key),
798        "model": value["model"],
799        "source": value["source"],
800        "captured_at": value["captured_at"],
801        "last_used_at": value["last_used_at"],
802        "trace_id": value["trace_id"],
803        "claude_bin": value["claude_bin"],
804        "claude_version": value["claude_version"],
805        "system_prompt_chars": value["system_prompt_chars"],
806        "agent_identity_chars": value["agent_identity_chars"],
807        "path": path.to_string_lossy(),
808        "runs": runs,
809    }))
810}
811
812fn dario_prompt_caches(state: &AppState) -> Vec<Value> {
813    let dir = dario_prompt_cache_dir(state);
814    let mut caches: Vec<Value> = std::fs::read_dir(dir)
815        .ok()
816        .into_iter()
817        .flatten()
818        .flatten()
819        .filter_map(|entry| {
820            let path = entry.path();
821            (path.extension().and_then(|s| s.to_str()) == Some("json"))
822                .then(|| dario_prompt_cache_summary(path))
823                .flatten()
824        })
825        .collect();
826    caches.sort_by(|a, b| {
827        b["last_used_at"]
828            .as_str()
829            .unwrap_or("")
830            .cmp(a["last_used_at"].as_str().unwrap_or(""))
831    });
832    caches
833}
834
835async fn admin_dario_prompt_caches(State(state): State<Arc<AppState>>) -> Response {
836    axum::Json(json!({"prompt_caches": dario_prompt_caches(&state)})).into_response()
837}
838
839async fn admin_dario_prompt_cache_delete(
840    State(state): State<Arc<AppState>>,
841    Path(key): Path<String>,
842) -> Response {
843    if !dario_prompt_cache_key_valid(&key) {
844        return error_response(StatusCode::BAD_REQUEST, "invalid prompt cache key");
845    }
846    let path = dario_prompt_cache_dir(&state).join(format!("{key}.json"));
847    match std::fs::remove_file(&path) {
848        Ok(()) => axum::Json(json!({"deleted": true, "key": key})).into_response(),
849        Err(e) if e.kind() == std::io::ErrorKind::NotFound => {
850            error_response(StatusCode::NOT_FOUND, "unknown prompt cache")
851        }
852        Err(e) => error_response(StatusCode::INTERNAL_SERVER_ERROR, &e.to_string()),
853    }
854}
855
856async fn admin_storage(State(state): State<Arc<AppState>>) -> Response {
857    match state.store.disk_usage() {
858        Ok(v) => axum::Json(v).into_response(),
859        Err(e) => error_response(StatusCode::INTERNAL_SERVER_ERROR, &e.to_string()),
860    }
861}
862
863async fn admin_storage_prune(
864    State(state): State<Arc<AppState>>,
865    body: axum::Json<Value>,
866) -> Response {
867    let now = now_ms();
868    let cutoff = match &body["older_than_ms"] {
869        Value::Null => match body["older_than"].as_str() {
870            Some(s) => match parse_since(s, now) {
871                Some(ms) => ms,
872                None => {
873                    return error_response(
874                        StatusCode::BAD_REQUEST,
875                        &format!("invalid 'older_than' '{s}' (use 45s, 30m, 24h, 7d, or RFC3339)"),
876                    )
877                }
878            },
879            None => {
880                return error_response(
881                    StatusCode::BAD_REQUEST,
882                    "missing 'older_than_ms' or 'older_than'",
883                )
884            }
885        },
886        v => match v.as_i64() {
887            Some(ms) => ms,
888            None => {
889                return error_response(StatusCode::BAD_REQUEST, "'older_than_ms' must be an integer")
890            }
891        },
892    };
893    if cutoff > now {
894        return error_response(StatusCode::BAD_REQUEST, "cutoff is in the future");
895    }
896    let bodies_only = body["bodies_only"].as_bool().unwrap_or(true);
897    let dry_run = body["dry_run"].as_bool().unwrap_or(false);
898    let store = state.store.clone();
899    let report =
900        tokio::task::spawn_blocking(move || store.prune(cutoff, bodies_only, dry_run)).await;
901    match report {
902        Ok(Ok(r)) => axum::Json(serde_json::to_value(r).unwrap_or_default()).into_response(),
903        Ok(Err(e)) => error_response(StatusCode::INTERNAL_SERVER_ERROR, &e.to_string()),
904        Err(e) => error_response(StatusCode::INTERNAL_SERVER_ERROR, &e.to_string()),
905    }
906}
907
908async fn admin_traces(
909    State(state): State<Arc<AppState>>,
910    Query(q): Query<HashMap<String, String>>,
911) -> Response {
912    let limit = q.get("limit").and_then(|s| s.parse().ok()).unwrap_or(50);
913    match state.store.list_traces(
914        limit,
915        q.get("session").map(|s| s.as_str()),
916        q.get("model").map(|s| s.as_str()),
917    ) {
918        Ok(rows) => axum::Json(json!({"traces": rows})).into_response(),
919        Err(e) => error_response(StatusCode::INTERNAL_SERVER_ERROR, &e.to_string()),
920    }
921}
922
923fn filter_from_query(q: &HashMap<String, String>) -> TraceFilter {
924    let now = now_ms();
925    TraceFilter {
926        since_ms: q.get("since").and_then(|s| parse_since(s, now)),
927        until_ms: q.get("until").and_then(|s| parse_since(s, now)),
928        run_id: q.get("run_id").cloned(),
929        session: q.get("session").cloned(),
930        model: q.get("model").cloned(),
931        provider: q.get("provider").cloned(),
932        path: q.get("path").cloned(),
933        harness: q.get("harness").cloned(),
934        status: q.get("status").and_then(|s| s.parse().ok()),
935        errors_only: q
936            .get("errors")
937            .map(|v| v == "1" || v == "true")
938            .unwrap_or(false),
939        key_fingerprint: q.get("key_fingerprint").cloned(),
940        reasoning_effort: q.get("effort").cloned(),
941        limit: q.get("limit").and_then(|s| s.parse().ok()).unwrap_or(200),
942    }
943}
944
945fn wants_bodies(q: &HashMap<String, String>) -> bool {
946    q.get("bodies").map(|v| v == "1" || v == "true").unwrap_or(false)
947}
948
949fn read_gz_file(path: &str) -> Option<Vec<u8>> {
950    let file = std::fs::File::open(path).ok()?;
951    let mut decoder = flate2::read::GzDecoder::new(file);
952    let mut buf = Vec::new();
953    std::io::Read::read_to_end(&mut decoder, &mut buf).ok()?;
954    Some(buf)
955}
956
957fn inline_row_bodies(row: &mut Value) {
958    use base64::Engine;
959    for (path_key, out_key) in [
960        ("req_body_path", "req_body_b64"),
961        ("upstream_req_body_path", "upstream_req_body_b64"),
962        ("resp_body_path", "resp_body_b64"),
963    ] {
964        let Some(buf) = row[path_key].as_str().and_then(read_gz_file) else {
965            continue;
966        };
967        row[out_key] = json!(base64::engine::general_purpose::STANDARD.encode(&buf));
968    }
969}
970
971fn ndjson_response(mut rows: Vec<Value>, inline_bodies: bool) -> Response {
972    rows.sort_by_key(|r| r["ts_request_ms"].as_i64().unwrap_or(0));
973    let mut out = String::new();
974    for mut row in rows {
975        if inline_bodies {
976            inline_row_bodies(&mut row);
977        }
978        out.push_str(&serde_json::to_string(&row).unwrap_or_default());
979        out.push('\n');
980    }
981    Response::builder()
982        .status(StatusCode::OK)
983        .header("content-type", "application/x-ndjson")
984        .body(Body::from(out))
985        .unwrap_or_else(|e| error_response(StatusCode::INTERNAL_SERVER_ERROR, &e.to_string()))
986}
987
988const TEXT_SCAN_CAP: usize = 300;
989
990fn trace_matches_text(row: &Value, needle: &str) -> bool {
991    for key in ["req_body_path", "resp_body_path"] {
992        if let Some(path) = row.get(key).and_then(|v| v.as_str()) {
993            if let Some(bytes) = read_gz_file(path) {
994                if String::from_utf8_lossy(&bytes)
995                    .to_lowercase()
996                    .contains(needle)
997                {
998                    return true;
999                }
1000            }
1001        }
1002    }
1003    false
1004}
1005
1006async fn traces_search(
1007    State(state): State<Arc<AppState>>,
1008    Query(q): Query<HashMap<String, String>>,
1009) -> Response {
1010    let mut filter = filter_from_query(&q);
1011    let text = q
1012        .get("text")
1013        .or_else(|| q.get("q"))
1014        .map(|t| t.trim().to_lowercase())
1015        .filter(|t| !t.is_empty());
1016    if text.is_some() {
1017        filter.limit = TEXT_SCAN_CAP;
1018    }
1019    match state.store.search_traces(&filter) {
1020        Ok(rows) => match text {
1021            Some(needle) => {
1022                let scanned = rows.len();
1023                let rows = tokio::task::spawn_blocking(move || {
1024                    rows.into_iter()
1025                        .filter(|r| trace_matches_text(r, &needle))
1026                        .collect::<Vec<_>>()
1027                })
1028                .await
1029                .unwrap_or_default();
1030                axum::Json(json!({"traces": rows, "scanned": scanned, "scan_cap": TEXT_SCAN_CAP}))
1031                    .into_response()
1032            }
1033            None => axum::Json(json!({"traces": rows})).into_response(),
1034        },
1035        Err(e) => error_response(StatusCode::INTERNAL_SERVER_ERROR, &e.to_string()),
1036    }
1037}
1038
1039async fn traces_export(
1040    State(state): State<Arc<AppState>>,
1041    Query(q): Query<HashMap<String, String>>,
1042) -> Response {
1043    match state.store.search_traces(&filter_from_query(&q)) {
1044        Ok(rows) => ndjson_response(rows, wants_bodies(&q)),
1045        Err(e) => error_response(StatusCode::INTERNAL_SERVER_ERROR, &e.to_string()),
1046    }
1047}
1048
1049fn truncate_chars(s: String, max: usize) -> String {
1050    if s.chars().count() <= max {
1051        s
1052    } else {
1053        s.chars().take(max).collect()
1054    }
1055}
1056
1057fn read_gz_json(path: Option<&str>) -> Option<Value> {
1058    let buf = path.and_then(read_gz_file)?;
1059    serde_json::from_slice(&buf).ok()
1060}
1061
1062fn read_gz_text(path: Option<&str>) -> Option<String> {
1063    let buf = path.and_then(read_gz_file)?;
1064    Some(String::from_utf8_lossy(&buf).to_string())
1065}
1066
1067fn body_date_dir_name(name: &str) -> bool {
1068    let bytes = name.as_bytes();
1069    bytes.len() == 10
1070        && bytes[4] == b'-'
1071        && bytes[7] == b'-'
1072        && bytes
1073            .iter()
1074            .enumerate()
1075            .all(|(i, b)| i == 4 || i == 7 || b.is_ascii_digit())
1076}
1077
1078fn dario_capture_suffix(kind: &str) -> Option<&'static str> {
1079    match kind {
1080        "dario-upstream-request" => Some("dario-upstream-request.json.gz"),
1081        "dario-upstream-response" => Some("dario-upstream-response.json.gz"),
1082        _ => None,
1083    }
1084}
1085
1086fn is_dario_trace(row: &Value) -> bool {
1087    row["account_id"]
1088        .as_str()
1089        .map(|id| id.starts_with("dario:"))
1090        .unwrap_or(false)
1091}
1092
1093fn find_dario_capture_path(state: &AppState, row: &Value, kind: &str) -> Option<String> {
1094    if !is_dario_trace(row) {
1095        return None;
1096    }
1097    let trace_id = row["id"].as_str()?;
1098    let suffix = dario_capture_suffix(kind)?;
1099    let filename = format!("{trace_id}.{suffix}");
1100    let bodies = state.store.data_dir.join("bodies");
1101    let mut days: Vec<PathBuf> = std::fs::read_dir(&bodies)
1102        .ok()?
1103        .flatten()
1104        .filter_map(|entry| {
1105            let name = entry.file_name().to_string_lossy().to_string();
1106            (entry.path().is_dir() && body_date_dir_name(&name)).then(|| entry.path())
1107        })
1108        .collect();
1109    days.sort_by(|a, b| b.file_name().cmp(&a.file_name()));
1110    for day in days {
1111        let path = day.join(&filename);
1112        if path.is_file() {
1113            return Some(path.to_string_lossy().to_string());
1114        }
1115    }
1116    None
1117}
1118
1119fn dario_capture_summary(state: &AppState, row: &Value) -> Option<Value> {
1120    let request_path = find_dario_capture_path(state, row, "dario-upstream-request");
1121    let response_path = find_dario_capture_path(state, row, "dario-upstream-response");
1122    if request_path.is_none() && response_path.is_none() {
1123        return None;
1124    }
1125    let prompt_cache = request_path
1126        .as_deref()
1127        .and_then(|path| read_gz_json(Some(path)))
1128        .and_then(|body| {
1129            body["prompt_cache"]
1130                .as_object()
1131                .map(|_| body["prompt_cache"].clone())
1132        });
1133    Some(json!({
1134        "request_available": request_path.is_some(),
1135        "response_available": response_path.is_some(),
1136        "request_path": request_path,
1137        "response_path": response_path,
1138        "prompt_cache": prompt_cache,
1139    }))
1140}
1141
1142async fn traces_sessions(
1143    State(state): State<Arc<AppState>>,
1144    Query(q): Query<HashMap<String, String>>,
1145) -> Response {
1146    let since = q.get("since").and_then(|s| parse_since(s, now_ms()));
1147    let limit = q.get("limit").and_then(|s| s.parse().ok()).unwrap_or(0);
1148    match state.store.sessions(since, limit) {
1149        Ok(rows) => axum::Json(json!({"sessions": rows})).into_response(),
1150        Err(e) => error_response(StatusCode::INTERNAL_SERVER_ERROR, &e.to_string()),
1151    }
1152}
1153
1154fn transcript_turn(row: &Value) -> Value {
1155    use alex_core::translate;
1156    let user = read_gz_json(row["req_body_path"].as_str())
1157        .and_then(|req| {
1158            translate::last_user_text(row["client_format"].as_str().unwrap_or(""), &req)
1159        })
1160        .map(|s| truncate_chars(s, 8000));
1161    let resp_text = read_gz_text(row["resp_body_path"].as_str());
1162    let fmt = row["upstream_format"]
1163        .as_str()
1164        .or(row["client_format"].as_str())
1165        .unwrap_or("")
1166        .to_string();
1167    let assistant = resp_text
1168        .as_deref()
1169        .and_then(|text| translate::assistant_reply_text(&fmt, text))
1170        .map(|s| truncate_chars(s, 8000));
1171    let tool_calls: Vec<Value> = resp_text
1172        .as_deref()
1173        .map(|text| translate::assistant_tool_calls(&fmt, text))
1174        .unwrap_or_default()
1175        .into_iter()
1176        .take(24)
1177        .map(|mut c| {
1178            if let Some(a) = c["arguments"].as_str() {
1179                let t = truncate_chars(a.to_string(), 600);
1180                c["arguments"] = json!(t);
1181            }
1182            c
1183        })
1184        .collect();
1185    json!({
1186        "trace_id": row["id"],
1187        "ts_request_ms": row["ts_request_ms"],
1188        "ts_response_ms": row["ts_response_ms"],
1189        "model": row["routed_model"],
1190        "status": row["status"],
1191        "input_tokens": row["input_tokens"],
1192        "output_tokens": row["output_tokens"],
1193        "reasoning_effort": row["reasoning_effort"],
1194        "thinking_budget": row["thinking_budget"],
1195        "cost_usd": row["cost_usd"],
1196        "error": row["error"],
1197        "user": user,
1198        "assistant": assistant,
1199        "tool_calls": tool_calls,
1200    })
1201}
1202
1203async fn traces_session_transcript(
1204    State(state): State<Arc<AppState>>,
1205    Path(session_id): Path<String>,
1206    Query(q): Query<HashMap<String, String>>,
1207) -> Response {
1208    let since = q
1209        .get("since_ms")
1210        .and_then(|s| s.parse::<i64>().ok())
1211        .or_else(|| q.get("since").and_then(|s| parse_since(s, now_ms())));
1212    let limit: usize = q.get("limit").and_then(|s| s.parse().ok()).unwrap_or(500);
1213    let rows = match state.store.session_traces(&session_id, since) {
1214        Ok(rows) => rows,
1215        Err(e) => return error_response(StatusCode::INTERNAL_SERVER_ERROR, &e.to_string()),
1216    };
1217    let turns: Vec<Value> = rows.iter().take(limit).map(transcript_turn).collect();
1218    axum::Json(json!({"session_id": session_id, "turns": turns})).into_response()
1219}
1220
1221fn trace_reasoning_fields(req: &Value) -> (Option<String>, Option<i64>) {
1222    (
1223        req["reasoning"]["effort"].as_str().map(String::from),
1224        req["thinking"]["budget_tokens"].as_i64(),
1225    )
1226}
1227
1228fn trace_extras(req: &Value) -> Value {
1229    let system_text: Option<String> = match &req["system"] {
1230        Value::String(s) => Some(s.clone()),
1231        Value::Array(parts) => {
1232            let joined: Vec<&str> = parts.iter().filter_map(|p| p["text"].as_str()).collect();
1233            (!joined.is_empty()).then(|| joined.join("\n\n"))
1234        }
1235        _ => req["instructions"].as_str().map(String::from).or_else(|| {
1236            let joined: Vec<String> = req["messages"]
1237                .as_array()
1238                .into_iter()
1239                .flatten()
1240                .filter(|m| matches!(m["role"].as_str(), Some("system") | Some("developer")))
1241                .map(|m| match &m["content"] {
1242                    Value::String(s) => s.clone(),
1243                    Value::Array(parts) => parts
1244                        .iter()
1245                        .filter_map(|p| p["text"].as_str())
1246                        .collect::<Vec<_>>()
1247                        .join("\n"),
1248                    _ => String::new(),
1249                })
1250                .filter(|s| !s.is_empty())
1251                .collect();
1252            (!joined.is_empty()).then(|| joined.join("\n\n"))
1253        }),
1254    };
1255    let system_chars = system_text.as_ref().map(|s| s.chars().count());
1256    let system_prompt = system_text.map(|s| truncate_chars(s, 64_000));
1257    let (reasoning_effort, thinking_budget) = trace_reasoning_fields(req);
1258    json!({
1259        "reasoning_effort": reasoning_effort,
1260        "thinking_budget": thinking_budget,
1261        "max_tokens": req["max_tokens"].as_i64().or(req["max_output_tokens"].as_i64()),
1262        "temperature": req["temperature"],
1263        "message_count": req["messages"].as_array().or(req["input"].as_array()).map(|a| a.len()),
1264        "system_chars": system_chars,
1265        "system_prompt": system_prompt,
1266    })
1267}
1268
1269async fn trace_get(State(state): State<Arc<AppState>>, Path(id): Path<String>) -> Response {
1270    match state.store.get_trace(&id) {
1271        Ok(Some(row)) => {
1272            let mut extras = read_gz_json(row["req_body_path"].as_str())
1273                .map(|req| trace_extras(&req))
1274                .unwrap_or_else(|| json!({}));
1275            if let Some(summary) = dario_capture_summary(&state, &row) {
1276                if !extras.is_object() {
1277                    extras = json!({});
1278                }
1279                if let Some(obj) = extras.as_object_mut() {
1280                    obj.insert("dario_capture".into(), summary);
1281                }
1282            }
1283            axum::Json(json!({"trace": row, "extras": extras})).into_response()
1284        }
1285        Ok(None) => error_response(StatusCode::NOT_FOUND, &format!("unknown trace '{id}'")),
1286        Err(e) => error_response(StatusCode::INTERNAL_SERVER_ERROR, &e.to_string()),
1287    }
1288}
1289
1290async fn trace_body(
1291    State(state): State<Arc<AppState>>,
1292    Path((id, kind)): Path<(String, String)>,
1293) -> Response {
1294    let row = match state.store.get_trace(&id) {
1295        Ok(Some(row)) => row,
1296        Ok(None) => {
1297            return error_response(StatusCode::NOT_FOUND, &format!("unknown trace '{id}'"))
1298        }
1299        Err(e) => return error_response(StatusCode::INTERNAL_SERVER_ERROR, &e.to_string()),
1300    };
1301    let path = match kind.as_str() {
1302        "request" => row["req_body_path"].as_str().map(String::from),
1303        "upstream-request" => row["upstream_req_body_path"].as_str().map(String::from),
1304        "response" => row["resp_body_path"].as_str().map(String::from),
1305        "dario-upstream-request" | "dario-upstream-response" => {
1306            find_dario_capture_path(&state, &row, &kind)
1307        }
1308        _ => {
1309            return error_response(
1310                StatusCode::BAD_REQUEST,
1311                "kind must be request|upstream-request|response|dario-upstream-request|dario-upstream-response",
1312            )
1313        }
1314    };
1315    match read_gz_text(path.as_deref()) {
1316        Some(text) => {
1317            let ct = if text.trim_start().starts_with('{') || text.trim_start().starts_with('[') {
1318                "application/json; charset=utf-8"
1319            } else {
1320                "text/plain; charset=utf-8"
1321            };
1322            Response::builder()
1323                .status(StatusCode::OK)
1324                .header("content-type", ct)
1325                .header("x-alexandria-body-path", path.as_deref().unwrap_or(""))
1326                .body(Body::from(text))
1327                .unwrap_or_else(|e| error_response(StatusCode::INTERNAL_SERVER_ERROR, &e.to_string()))
1328        }
1329        None => error_response(
1330            StatusCode::NOT_FOUND,
1331            &format!("no {kind} body stored for trace '{id}'"),
1332        ),
1333    }
1334}
1335
1336async fn trace_reply_md(State(state): State<Arc<AppState>>, Path(id): Path<String>) -> Response {
1337    use alex_core::translate;
1338    let row = match state.store.get_trace(&id) {
1339        Ok(Some(row)) => row,
1340        Ok(None) => {
1341            return error_response(StatusCode::NOT_FOUND, &format!("unknown trace '{id}'"))
1342        }
1343        Err(e) => return error_response(StatusCode::INTERNAL_SERVER_ERROR, &e.to_string()),
1344    };
1345    let fmt = row["upstream_format"]
1346        .as_str()
1347        .or(row["client_format"].as_str())
1348        .unwrap_or("");
1349    let reply = read_gz_text(row["resp_body_path"].as_str())
1350        .and_then(|text| translate::assistant_reply_text(fmt, &text));
1351    match reply {
1352        Some(md) => Response::builder()
1353            .status(StatusCode::OK)
1354            .header("content-type", "text/markdown; charset=utf-8")
1355            .body(Body::from(md))
1356            .unwrap_or_else(|e| error_response(StatusCode::INTERNAL_SERVER_ERROR, &e.to_string())),
1357        None => error_response(
1358            StatusCode::NOT_FOUND,
1359            &format!("no assistant reply available for trace '{id}'"),
1360        ),
1361    }
1362}
1363
1364async fn trace_delete(State(state): State<Arc<AppState>>, Path(id): Path<String>) -> Response {
1365    match state.store.get_trace(&id) {
1366        Ok(Some(_)) => {}
1367        Ok(None) => {
1368            return error_response(StatusCode::NOT_FOUND, &format!("unknown trace '{id}'"))
1369        }
1370        Err(e) => return error_response(StatusCode::INTERNAL_SERVER_ERROR, &e.to_string()),
1371    }
1372    match state.store.delete_trace(&id) {
1373        Ok(paths) => {
1374            let removed = paths
1375                .iter()
1376                .filter(|p| std::fs::remove_file(p).is_ok())
1377                .count();
1378            axum::Json(json!({"deleted": true, "files_removed": removed})).into_response()
1379        }
1380        Err(e) => error_response(StatusCode::INTERNAL_SERVER_ERROR, &e.to_string()),
1381    }
1382}
1383
1384async fn traces_run_summary(
1385    State(state): State<Arc<AppState>>,
1386    Path(run_id): Path<String>,
1387) -> Response {
1388    match state.store.run_summary(&run_id) {
1389        Ok(v) if v["trace_count"].as_i64().unwrap_or(0) == 0 => error_response(
1390            StatusCode::NOT_FOUND,
1391            &format!("no traces for run '{run_id}'"),
1392        ),
1393        Ok(v) => axum::Json(v).into_response(),
1394        Err(e) => error_response(StatusCode::INTERNAL_SERVER_ERROR, &e.to_string()),
1395    }
1396}
1397
1398async fn traces_run_events(
1399    State(state): State<Arc<AppState>>,
1400    Path(run_id): Path<String>,
1401    Query(q): Query<HashMap<String, String>>,
1402) -> Response {
1403    let filter = TraceFilter {
1404        run_id: Some(run_id),
1405        limit: q.get("limit").and_then(|s| s.parse().ok()).unwrap_or(1000),
1406        ..Default::default()
1407    };
1408    match state.store.search_traces(&filter) {
1409        Ok(rows) => axum::Json(json!({"traces": rows})).into_response(),
1410        Err(e) => error_response(StatusCode::INTERNAL_SERVER_ERROR, &e.to_string()),
1411    }
1412}
1413
1414async fn traces_run_export(
1415    State(state): State<Arc<AppState>>,
1416    Path(run_id): Path<String>,
1417    Query(q): Query<HashMap<String, String>>,
1418) -> Response {
1419    let filter = TraceFilter {
1420        run_id: Some(run_id),
1421        limit: q.get("limit").and_then(|s| s.parse().ok()).unwrap_or(5000),
1422        ..Default::default()
1423    };
1424    match state.store.search_traces(&filter) {
1425        Ok(rows) => ndjson_response(rows, wants_bodies(&q)),
1426        Err(e) => error_response(StatusCode::INTERNAL_SERVER_ERROR, &e.to_string()),
1427    }
1428}
1429
1430async fn traces_run_artifacts(
1431    State(state): State<Arc<AppState>>,
1432    Path(run_id): Path<String>,
1433) -> Response {
1434    match state.store.run_artifacts(&run_id) {
1435        Ok(artifacts) => axum::Json(json!({
1436            "run_id": run_id,
1437            "data_dir": state.store.data_dir.to_string_lossy(),
1438            "artifacts": artifacts,
1439        }))
1440        .into_response(),
1441        Err(e) => error_response(StatusCode::INTERNAL_SERVER_ERROR, &e.to_string()),
1442    }
1443}
1444
1445async fn admin_accounts(State(state): State<Arc<AppState>>) -> impl IntoResponse {
1446    let accounts: Vec<Value> = state
1447        .vault
1448        .list()
1449        .await
1450        .into_iter()
1451        .map(|a| {
1452            json!({
1453                "id": a.id,
1454                "provider": a.provider.as_str(),
1455                "name": a.name,
1456                "kind": a.kind,
1457                "label": a.label,
1458                "description": a.description,
1459                "paused": a.paused,
1460                "path": a.path.map(|p| p.display().to_string()),
1461                "status": a.status,
1462                "expires_at_ms": a.expires_at_ms,
1463                "expires_in_s": a.expires_at_ms.map(|e| (e - now_ms()) / 1000),
1464            })
1465        })
1466        .collect();
1467    axum::Json(json!({"accounts": accounts}))
1468}
1469
1470async fn admin_health(State(state): State<Arc<AppState>>) -> Response {
1471    let heartbeats = state.store.last_heartbeats().unwrap_or_default();
1472    let accounts: Vec<Value> = state
1473        .vault
1474        .list()
1475        .await
1476        .into_iter()
1477        .map(|a| {
1478            let last = heartbeats
1479                .iter()
1480                .find(|h| h["provider"].as_str() == Some(a.provider.as_str()))
1481                .cloned();
1482            json!({
1483                "id": a.id,
1484                "provider": a.provider.as_str(),
1485                "name": a.name,
1486                "kind": a.kind,
1487                "status": a.status,
1488                "paused": a.paused,
1489                "path": a.path.map(|p| p.display().to_string()),
1490                "token_expires_in_s": a.expires_at_ms.map(|e| (e - now_ms()) / 1000),
1491                "last_heartbeat": last,
1492            })
1493        })
1494        .collect();
1495    axum::Json(json!({"accounts": accounts})).into_response()
1496}
1497
1498#[derive(Debug, Clone, serde::Serialize)]
1499pub struct PingResult {
1500    pub provider: &'static str,
1501    pub account_id: Option<String>,
1502    pub ok: bool,
1503    pub status: Option<u16>,
1504    pub latency_ms: i64,
1505    pub message: String,
1506}
1507
1508#[derive(Debug, Clone)]
1509pub struct PingModels {
1510    pub anthropic: String,
1511    pub openai: String,
1512    pub xai: String,
1513    pub gemini: String,
1514}
1515
1516pub async fn ping_provider(
1517    state: &Arc<AppState>,
1518    provider: Provider,
1519    models: &PingModels,
1520) -> PingResult {
1521    let start = now_ms();
1522    let prompt = "Health check: what time is it? If you cannot know, just reply: creds ok";
1523    let (format, path, body) = match provider {
1524        Provider::Anthropic => (
1525            ClientFormat::AnthropicMessages,
1526            "/v1/messages",
1527            json!({
1528                "model": models.anthropic,
1529                "max_tokens": 64,
1530                "system": "You are Claude Code, Anthropic's official CLI for Claude.",
1531                "messages": [{"role": "user", "content": prompt}],
1532            }),
1533        ),
1534        Provider::Openai => (
1535            ClientFormat::OpenaiResponses,
1536            "/v1/responses",
1537            json!({
1538                "model": models.openai,
1539                "stream": true,
1540                "store": false,
1541                "instructions": "You are a helpful assistant.",
1542                "input": [{"type": "message", "role": "user",
1543                           "content": [{"type": "input_text", "text": prompt}]}],
1544            }),
1545        ),
1546        Provider::Xai => (
1547            ClientFormat::OpenaiChat,
1548            "/v1/chat/completions",
1549            json!({
1550                "model": models.xai,
1551                "stream": false,
1552                "messages": [{"role": "user", "content": prompt}],
1553            }),
1554        ),
1555        Provider::Gemini => (
1556            ClientFormat::GeminiGenerate,
1557            "/v1beta/models/:generateContent",
1558            json!({
1559                "model": models.gemini,
1560                "contents": [{"role": "user", "parts": [{"text": prompt}]}],
1561                "generationConfig": {"maxOutputTokens": 64},
1562            }),
1563        ),
1564    };
1565    let account_id = state
1566        .vault
1567        .account_for(provider, true)
1568        .await
1569        .ok()
1570        .map(|a| a.id);
1571    let mut headers = HeaderMap::new();
1572    headers.insert(
1573        "x-api-key",
1574        HeaderValue::from_str(&state.local_key).expect("key header"),
1575    );
1576    headers.insert("user-agent", HeaderValue::from_static("alexandria-ping"));
1577    let resp = proxy(
1578        state.clone(),
1579        format,
1580        path,
1581        headers,
1582        Bytes::from(serde_json::to_vec(&body).expect("ping body")),
1583        None,
1584    )
1585    .await;
1586    let status = resp.status().as_u16();
1587    let bytes = axum::body::to_bytes(resp.into_body(), 8 * 1024 * 1024)
1588        .await
1589        .unwrap_or_default();
1590    let text = String::from_utf8_lossy(&bytes);
1591    let message = extract_reply(&text).unwrap_or_else(|| snippet(&text));
1592    PingResult {
1593        provider: provider.as_str(),
1594        account_id,
1595        ok: (200..300).contains(&status),
1596        status: Some(status),
1597        latency_ms: now_ms() - start,
1598        message,
1599    }
1600}
1601
1602fn extract_reply(text: &str) -> Option<String> {
1603    if let Ok(v) = serde_json::from_str::<Value>(text) {
1604        if let Some(t) = v["content"][0]["text"].as_str() {
1605            return Some(t.to_string());
1606        }
1607        if let Some(t) = v["choices"][0]["message"]["content"].as_str() {
1608            return Some(t.to_string());
1609        }
1610        if let Some(t) = v["error"]["message"].as_str() {
1611            return Some(t.to_string());
1612        }
1613        if let Some(t) = v["detail"].as_str() {
1614            return Some(t.to_string());
1615        }
1616    }
1617    for line in text.lines() {
1618        let Some(data) = line.strip_prefix("data:") else {
1619            continue;
1620        };
1621        let Ok(v) = serde_json::from_str::<Value>(data.trim()) else {
1622            continue;
1623        };
1624        if v["type"].as_str() == Some("response.output_text.done") {
1625            if let Some(t) = v["text"].as_str() {
1626                return Some(t.to_string());
1627            }
1628        }
1629    }
1630    None
1631}
1632
1633fn snippet(text: &str) -> String {
1634    let t: String = text.chars().take(200).collect();
1635    t.replace('\n', " ")
1636}
1637
1638pub async fn heartbeat_once(state: &Arc<AppState>, models: &PingModels) -> Vec<PingResult> {
1639    let providers: Vec<Provider> = {
1640        let mut seen = Vec::new();
1641        for a in state.vault.list().await {
1642            if a.status == "active"
1643                && matches!(
1644                    a.provider,
1645                    Provider::Anthropic | Provider::Openai | Provider::Xai | Provider::Gemini
1646                )
1647                && !seen.contains(&a.provider)
1648            {
1649                seen.push(a.provider);
1650            }
1651        }
1652        seen
1653    };
1654    let mut results = Vec::new();
1655    for provider in providers {
1656        let r = ping_provider(state, provider, models).await;
1657        if let Err(e) = state.store.insert_heartbeat(
1658            now_ms(),
1659            r.provider,
1660            r.account_id.as_deref(),
1661            r.ok,
1662            r.status.map(|s| s as i64),
1663            r.latency_ms,
1664            &r.message,
1665        ) {
1666            tracing::warn!("failed to record heartbeat: {e}");
1667        }
1668        tracing::info!(
1669            provider = r.provider,
1670            ok = r.ok,
1671            status = r.status,
1672            latency_ms = r.latency_ms,
1673            reply = %r.message,
1674            "heartbeat"
1675        );
1676        results.push(r);
1677    }
1678    results
1679}
1680
1681async fn anthropic_messages(
1682    State(state): State<Arc<AppState>>,
1683    ConnectInfo(peer): ConnectInfo<std::net::SocketAddr>,
1684    headers: HeaderMap,
1685    body: Bytes,
1686) -> Response {
1687    proxy(
1688        state,
1689        ClientFormat::AnthropicMessages,
1690        "/v1/messages",
1691        headers,
1692        body,
1693        Some(peer),
1694    )
1695    .await
1696}
1697
1698async fn openai_chat(
1699    State(state): State<Arc<AppState>>,
1700    ConnectInfo(peer): ConnectInfo<std::net::SocketAddr>,
1701    headers: HeaderMap,
1702    body: Bytes,
1703) -> Response {
1704    proxy(
1705        state,
1706        ClientFormat::OpenaiChat,
1707        "/v1/chat/completions",
1708        headers,
1709        body,
1710        Some(peer),
1711    )
1712    .await
1713}
1714
1715async fn openai_responses(
1716    State(state): State<Arc<AppState>>,
1717    ConnectInfo(peer): ConnectInfo<std::net::SocketAddr>,
1718    headers: HeaderMap,
1719    body: Bytes,
1720) -> Response {
1721    proxy(
1722        state,
1723        ClientFormat::OpenaiResponses,
1724        "/v1/responses",
1725        headers,
1726        body,
1727        Some(peer),
1728    )
1729    .await
1730}
1731
1732async fn gemini_generate(
1733    State(state): State<Arc<AppState>>,
1734    ConnectInfo(peer): ConnectInfo<std::net::SocketAddr>,
1735    axum::extract::Path(model_action): axum::extract::Path<String>,
1736    headers: HeaderMap,
1737    body: Bytes,
1738) -> Response {
1739    let (model, method) = model_action
1740        .split_once(':')
1741        .unwrap_or((model_action.as_str(), "generateContent"));
1742    let stream = method == "streamGenerateContent";
1743    if method != "generateContent" && !stream {
1744        return error_response(
1745            StatusCode::NOT_FOUND,
1746            &format!("unsupported gemini method '{method}' (expected generateContent or streamGenerateContent)"),
1747        );
1748    }
1749    if model.is_empty() {
1750        return error_response(StatusCode::BAD_REQUEST, "missing model in path");
1751    }
1752    let mut v: Value = match serde_json::from_slice(&body) {
1753        Ok(v) => v,
1754        Err(e) => {
1755            return error_response(StatusCode::BAD_REQUEST, &format!("invalid JSON body: {e}"))
1756        }
1757    };
1758    v["model"] = json!(model);
1759    if stream {
1760        v["stream"] = json!(true);
1761    }
1762    let body = Bytes::from(serde_json::to_vec(&v).unwrap_or_default());
1763    let path = if stream {
1764        "/v1beta/models/:streamGenerateContent"
1765    } else {
1766        "/v1beta/models/:generateContent"
1767    };
1768    proxy(
1769        state,
1770        ClientFormat::GeminiGenerate,
1771        path,
1772        headers,
1773        body,
1774        Some(peer),
1775    )
1776    .await
1777}
1778
1779fn error_response(status: StatusCode, message: &str) -> Response {
1780    (
1781        status,
1782        axum::Json(json!({"error": {"type": "alexandria", "message": message}})),
1783    )
1784        .into_response()
1785}
1786
1787fn client_key(headers: &HeaderMap) -> Option<String> {
1788    if let Some(v) = headers.get("x-api-key").and_then(|v| v.to_str().ok()) {
1789        return Some(v.to_string());
1790    }
1791    if let Some(v) = headers.get("x-goog-api-key").and_then(|v| v.to_str().ok()) {
1792        return Some(v.to_string());
1793    }
1794    headers
1795        .get("authorization")
1796        .and_then(|v| v.to_str().ok())
1797        .and_then(|v| v.strip_prefix("Bearer "))
1798        .map(String::from)
1799}
1800
1801fn redacted_headers(headers: &HeaderMap) -> String {
1802    let map: HashMap<String, String> = headers
1803        .iter()
1804        .map(|(k, v)| {
1805            let key = k.as_str().to_lowercase();
1806            let val = if ["authorization", "x-api-key", "cookie", "chatgpt-account-id"]
1807                .contains(&key.as_str())
1808            {
1809                "<redacted>".to_string()
1810            } else {
1811                v.to_str().unwrap_or("<binary>").to_string()
1812            };
1813            (key, val)
1814        })
1815        .collect();
1816    serde_json::to_string(&map).unwrap_or_default()
1817}
1818
1819#[derive(Debug, Clone, Copy, PartialEq)]
1820enum RespondAs {
1821    Anthropic,
1822    OpenaiChat,
1823    OpenaiResponses,
1824    Gemini,
1825}
1826
1827struct UpstreamPlan {
1828    url: String,
1829    account: Account,
1830    body: Vec<u8>,
1831    upstream_format: &'static str,
1832    destream: bool,
1833    respond_as: Option<RespondAs>,
1834    client_stream: bool,
1835    extra_headers: Vec<(String, String)>,
1836    dario_guard: Option<Box<dyn std::any::Any + Send>>,
1837}
1838
1839fn dario_account(active: &DarioActive) -> Account {
1840    Account {
1841        id: format!("dario:{}", active.generation_id),
1842        provider: Provider::Anthropic,
1843        kind: "dario".into(),
1844        name: active.generation_id.clone(),
1845        description: None,
1846        paused: false,
1847        label: Some("dario generational proxy".into()),
1848        access_token: None,
1849        refresh_token: None,
1850        id_token: None,
1851        api_key: Some(active.api_key.clone()),
1852        expires_at_ms: None,
1853        last_refresh_ms: None,
1854        account_meta: Value::Null,
1855        cooldown_until_ms: None,
1856        status: "active".into(),
1857        path: None,
1858    }
1859}
1860
1861async fn ensure_gemini_project(
1862    state: &AppState,
1863    account: &Account,
1864) -> Result<String, (StatusCode, String)> {
1865    if let Some(p) = account.account_meta.get("project_id").and_then(|v| v.as_str()) {
1866        if !p.is_empty() {
1867            return Ok(p.to_string());
1868        }
1869    }
1870    if let Ok(env_p) = std::env::var("GOOGLE_CLOUD_PROJECT") {
1871        if !env_p.is_empty() {
1872            let _ = state
1873                .vault
1874                .set_account_meta(&account.id, "project_id", json!(env_p))
1875                .await;
1876            return Ok(env_p);
1877        }
1878    }
1879    let token = account
1880        .access_token
1881        .as_deref()
1882        .ok_or_else(|| (StatusCode::BAD_GATEWAY, "gemini account has no access token".into()))?;
1883    let load_url = format!("{GEMINI_CODE_ASSIST_BASE}/{GEMINI_CODE_ASSIST_VERSION}:loadCodeAssist");
1884    let load_body = json!({
1885        "cloudaicompanionProject": null,
1886        "metadata": {"ideType": "IDE_UNSPECIFIED", "platform": "PLATFORM_UNSPECIFIED", "pluginType": "GEMINI"},
1887    });
1888    let resp = state
1889        .http
1890        .post(&load_url)
1891        .header("authorization", format!("Bearer {token}"))
1892        .json(&load_body)
1893        .send()
1894        .await
1895        .map_err(|e| (StatusCode::BAD_GATEWAY, format!("loadCodeAssist failed: {e}")))?;
1896    let load: Value = resp.json().await.unwrap_or(Value::Null);
1897    let extract = |v: &Value| -> Option<String> {
1898        for key in ["cloudaicompanionProject", "projectId", "project"] {
1899            match &v[key] {
1900                Value::String(s) if !s.is_empty() => return Some(s.clone()),
1901                obj if obj["id"].is_string() => {
1902                    return obj["id"].as_str().map(String::from)
1903                }
1904                _ => {}
1905            }
1906        }
1907        None
1908    };
1909    if let Some(p) = extract(&load) {
1910        let _ = state
1911            .vault
1912            .set_account_meta(&account.id, "project_id", json!(p))
1913            .await;
1914        return Ok(p);
1915    }
1916    if let Some(free) = load["ineligibleTiers"]
1917        .as_array()
1918        .and_then(|t| t.iter().find(|t| t["tierId"] == json!("free-tier")))
1919    {
1920        let reason = free["reasonCode"].as_str().unwrap_or("");
1921        let msg = free["reasonMessage"].as_str().unwrap_or("");
1922        let hint = if let Some(url) = free["validationUrl"].as_str() {
1923            format!(
1924                "gemini account needs verification before the free Code Assist tier works — \
1925                 open {url} to verify, then retry (or sign in with another personal Google \
1926                 account, or set gemini_project to a GCP project)"
1927            )
1928        } else if reason == "DASHER_USER" {
1929            "this is a Google Workspace account, which cannot use the free Code Assist tier — \
1930             set gemini_project to a GCP project (Code Assist API enabled), or sign in with a \
1931             personal Google account"
1932                .to_string()
1933        } else {
1934            format!("gemini account is not eligible for the free Code Assist tier: {msg}")
1935        };
1936        return Err((StatusCode::BAD_GATEWAY, hint));
1937    }
1938    let tier = load["allowedTiers"]
1939        .as_array()
1940        .and_then(|tiers| tiers.iter().find(|t| t["isDefault"] == json!(true)))
1941        .and_then(|t| t["id"].as_str())
1942        .unwrap_or("free-tier")
1943        .to_string();
1944    let onboard_url =
1945        format!("{GEMINI_CODE_ASSIST_BASE}/{GEMINI_CODE_ASSIST_VERSION}:onboardUser");
1946    let onboard_body = json!({
1947        "tierId": tier,
1948        "cloudaicompanionProject": load["cloudaicompanionProject"],
1949        "metadata": {"ideType": "IDE_UNSPECIFIED", "platform": "PLATFORM_UNSPECIFIED", "pluginType": "GEMINI"},
1950    });
1951    for _ in 0..5 {
1952        let resp = state
1953            .http
1954            .post(&onboard_url)
1955            .header("authorization", format!("Bearer {token}"))
1956            .json(&onboard_body)
1957            .send()
1958            .await
1959            .map_err(|e| (StatusCode::BAD_GATEWAY, format!("onboardUser failed: {e}")))?;
1960        let lro: Value = resp.json().await.unwrap_or(Value::Null);
1961        if lro["done"] == json!(true) {
1962            if let Some(p) = extract(&lro["response"]) {
1963                let _ = state
1964                    .vault
1965                    .set_account_meta(&account.id, "project_id", json!(p))
1966                    .await;
1967                return Ok(p);
1968            }
1969        }
1970        tokio::time::sleep(Duration::from_secs(2)).await;
1971    }
1972    Err((
1973        StatusCode::BAD_GATEWAY,
1974        "gemini Code Assist onboarding did not return a project id (set GOOGLE_CLOUD_PROJECT to override)".into(),
1975    ))
1976}
1977
1978async fn plan_upstream(
1979    state: &AppState,
1980    format: ClientFormat,
1981    provider: Provider,
1982    routed_model: &str,
1983    body_json: &mut Value,
1984    original_body: &[u8],
1985    trace_id: &str,
1986) -> Result<UpstreamPlan, (StatusCode, String)> {
1987    use alex_core::translate;
1988    let client_stream = body_json["stream"].as_bool().unwrap_or(false);
1989    match provider {
1990        Provider::Anthropic => {
1991            let converted = match format {
1992                ClientFormat::AnthropicMessages => None,
1993                ClientFormat::OpenaiChat => Some((
1994                    translate::openai_chat_to_anthropic(body_json),
1995                    RespondAs::OpenaiChat,
1996                )),
1997                ClientFormat::OpenaiResponses => Some((
1998                    translate::openai_responses_to_anthropic(body_json),
1999                    RespondAs::OpenaiResponses,
2000                )),
2001                ClientFormat::GeminiGenerate => Some((
2002                    translate::gemini_to_anthropic(body_json),
2003                    RespondAs::Gemini,
2004                )),
2005            };
2006            let dario_active = state.dario.as_ref().and_then(|d| d.active());
2007            let (base, account, dario_guard, dario_capture) = match (&state.dario, dario_active) {
2008                (Some(dario), Some(active)) => (
2009                    active.base_url.trim_end_matches('/').to_string(),
2010                    dario_account(&active),
2011                    dario.begin(&active.generation_id),
2012                    true,
2013                ),
2014                _ => {
2015                    let account = state
2016                        .vault
2017                        .account_for(provider, true)
2018                        .await
2019                        .map_err(|e| (StatusCode::BAD_GATEWAY, e.to_string()))?;
2020                    (ANTHROPIC_BASE.to_string(), account, None, false)
2021                }
2022            };
2023            let (body, respond_as) = match converted {
2024                None => {
2025                    body_json["model"] = json!(routed_model);
2026                    (
2027                        serde_json::to_vec(body_json).unwrap_or_else(|_| original_body.to_vec()),
2028                        None,
2029                    )
2030                }
2031                Some((mut converted, respond_as)) => {
2032                    converted["model"] = json!(routed_model);
2033                    converted["stream"] = json!(false);
2034                    (
2035                        serde_json::to_vec(&converted).unwrap_or_else(|_| original_body.to_vec()),
2036                        Some(respond_as),
2037                    )
2038                }
2039            };
2040            Ok(UpstreamPlan {
2041                url: format!("{base}/v1/messages"),
2042                account,
2043                body,
2044                upstream_format: "anthropic",
2045                destream: false,
2046                respond_as,
2047                client_stream,
2048                extra_headers: if dario_capture {
2049                    vec![
2050                        ("x-dario-capture-id".into(), trace_id.into()),
2051                        ("x-dario-capture-model".into(), routed_model.into()),
2052                    ]
2053                } else {
2054                    vec![]
2055                },
2056                dario_guard,
2057            })
2058        }
2059        Provider::Openai => {
2060            let prefer_oauth = format != ClientFormat::OpenaiChat;
2061            let account = state
2062                .vault
2063                .account_for(provider, prefer_oauth)
2064                .await
2065                .map_err(|e| (StatusCode::BAD_GATEWAY, e.to_string()))?;
2066            let oauth = account.kind == "oauth";
2067            match format {
2068                ClientFormat::OpenaiChat if !oauth => {
2069                    body_json["model"] = json!(routed_model);
2070                    let body =
2071                        serde_json::to_vec(body_json).unwrap_or_else(|_| original_body.to_vec());
2072                    Ok(UpstreamPlan {
2073                        url: format!("{OPENAI_BASE}/v1/chat/completions"),
2074                        account,
2075                        body,
2076                        upstream_format: "openai-chat",
2077                        destream: false,
2078                        respond_as: None,
2079                        client_stream,
2080                        extra_headers: vec![],
2081                        dario_guard: None,
2082                    })
2083                }
2084                ClientFormat::OpenaiChat => {
2085                    let pivot = translate::openai_chat_to_anthropic(body_json);
2086                    let mut converted = translate::anthropic_to_openai_responses(&pivot);
2087                    converted["model"] = json!(routed_model);
2088                    translate::normalize_codex_request(&mut converted);
2089                    let body = serde_json::to_vec(&converted)
2090                        .unwrap_or_else(|_| original_body.to_vec());
2091                    Ok(UpstreamPlan {
2092                        url: format!("{CODEX_BASE}/responses"),
2093                        account,
2094                        body,
2095                        upstream_format: "openai-responses",
2096                        destream: false,
2097                        respond_as: Some(RespondAs::OpenaiChat),
2098                        client_stream,
2099                        extra_headers: vec![],
2100                        dario_guard: None,
2101                    })
2102                }
2103                ClientFormat::OpenaiResponses => {
2104                    body_json["model"] = json!(routed_model);
2105                    let mut destream = false;
2106                    let url = if oauth {
2107                        if body_json["stream"].as_bool() != Some(true) {
2108                            destream = true;
2109                        }
2110                        translate::normalize_codex_request(body_json);
2111                        format!("{CODEX_BASE}/responses")
2112                    } else {
2113                        format!("{OPENAI_BASE}/v1/responses")
2114                    };
2115                    let body =
2116                        serde_json::to_vec(body_json).unwrap_or_else(|_| original_body.to_vec());
2117                    Ok(UpstreamPlan {
2118                        url,
2119                        account,
2120                        body,
2121                        upstream_format: "openai-responses",
2122                        destream,
2123                        respond_as: None,
2124                        client_stream,
2125                        extra_headers: vec![],
2126                        dario_guard: None,
2127                    })
2128                }
2129                ClientFormat::AnthropicMessages => {
2130                    let mut converted = translate::anthropic_to_openai_responses(body_json);
2131                    converted["model"] = json!(routed_model);
2132                    let url = if oauth {
2133                        translate::normalize_codex_request(&mut converted);
2134                        format!("{CODEX_BASE}/responses")
2135                    } else {
2136                        converted["stream"] = json!(false);
2137                        format!("{OPENAI_BASE}/v1/responses")
2138                    };
2139                    let body = serde_json::to_vec(&converted)
2140                        .unwrap_or_else(|_| original_body.to_vec());
2141                    Ok(UpstreamPlan {
2142                        url,
2143                        account,
2144                        body,
2145                        upstream_format: "openai-responses",
2146                        destream: false,
2147                        respond_as: Some(RespondAs::Anthropic),
2148                        client_stream,
2149                        extra_headers: vec![],
2150                        dario_guard: None,
2151                    })
2152                }
2153                ClientFormat::GeminiGenerate => {
2154                    let pivot = translate::gemini_to_anthropic(body_json);
2155                    let mut converted = translate::anthropic_to_openai_responses(&pivot);
2156                    converted["model"] = json!(routed_model);
2157                    let url = if oauth {
2158                        translate::normalize_codex_request(&mut converted);
2159                        format!("{CODEX_BASE}/responses")
2160                    } else {
2161                        converted["stream"] = json!(false);
2162                        format!("{OPENAI_BASE}/v1/responses")
2163                    };
2164                    let body = serde_json::to_vec(&converted)
2165                        .unwrap_or_else(|_| original_body.to_vec());
2166                    Ok(UpstreamPlan {
2167                        url,
2168                        account,
2169                        body,
2170                        upstream_format: "openai-responses",
2171                        destream: false,
2172                        respond_as: Some(RespondAs::Gemini),
2173                        client_stream,
2174                        extra_headers: vec![],
2175                        dario_guard: None,
2176                    })
2177                }
2178            }
2179        }
2180        Provider::Xai => {
2181            let account = state
2182                .vault
2183                .account_for(provider, true)
2184                .await
2185                .map_err(|e| (StatusCode::BAD_GATEWAY, e.to_string()))?;
2186            let extra_headers = vec![
2187                ("x-grok-model-override".into(), routed_model.to_string()),
2188                ("x-grok-conv-id".into(), trace_id.to_string()),
2189            ];
2190            match format {
2191                ClientFormat::OpenaiChat => {
2192                    body_json["model"] = json!(routed_model);
2193                    let body =
2194                        serde_json::to_vec(body_json).unwrap_or_else(|_| original_body.to_vec());
2195                    Ok(UpstreamPlan {
2196                        url: format!("{XAI_BASE}/chat/completions"),
2197                        account,
2198                        body,
2199                        upstream_format: "openai-chat",
2200                        destream: false,
2201                        respond_as: None,
2202                        client_stream,
2203                        extra_headers,
2204                        dario_guard: None,
2205                    })
2206                }
2207                ClientFormat::AnthropicMessages => {
2208                    // pivot (anthropic) → chat completions for xAI/Grok upstream
2209                    let mut converted = translate::anthropic_to_openai_chat(body_json);
2210                    converted["model"] = json!(routed_model);
2211                    // buffer full upstream body then re-synth client dialect
2212                    converted["stream"] = json!(false);
2213                    let body = serde_json::to_vec(&converted)
2214                        .unwrap_or_else(|_| original_body.to_vec());
2215                    Ok(UpstreamPlan {
2216                        url: format!("{XAI_BASE}/chat/completions"),
2217                        account,
2218                        body,
2219                        upstream_format: "openai-chat",
2220                        destream: false,
2221                        respond_as: Some(RespondAs::Anthropic),
2222                        client_stream,
2223                        extra_headers,
2224                        dario_guard: None,
2225                    })
2226                }
2227                ClientFormat::OpenaiResponses | ClientFormat::GeminiGenerate => Err((
2228                    StatusCode::NOT_IMPLEMENTED,
2229                    "the xai/grok upstream speaks OpenAI chat completions; POST to /v1/chat/completions or /v1/messages"
2230                        .to_string(),
2231                )),
2232            }
2233        }
2234        Provider::Gemini => {
2235            // Prefer an AI Studio API key over the OAuth/Code-Assist path.
2236            let account = state
2237                .vault
2238                .account_for(provider, false)
2239                .await
2240                .map_err(|e| (StatusCode::BAD_GATEWAY, e.to_string()))?;
2241            let (gemini_req, respond_as) = match format {
2242                ClientFormat::GeminiGenerate => {
2243                    let mut g = body_json.clone();
2244                    if let Some(o) = g.as_object_mut() {
2245                        o.remove("model");
2246                        o.remove("stream");
2247                    }
2248                    (g, None)
2249                }
2250                ClientFormat::AnthropicMessages => (
2251                    translate::anthropic_to_gemini_request(body_json),
2252                    Some(RespondAs::Anthropic),
2253                ),
2254                ClientFormat::OpenaiChat => (
2255                    translate::anthropic_to_gemini_request(&translate::openai_chat_to_anthropic(
2256                        body_json,
2257                    )),
2258                    Some(RespondAs::OpenaiChat),
2259                ),
2260                ClientFormat::OpenaiResponses => (
2261                    translate::anthropic_to_gemini_request(
2262                        &translate::openai_responses_to_anthropic(body_json),
2263                    ),
2264                    Some(RespondAs::OpenaiResponses),
2265                ),
2266            };
2267            let method = if client_stream {
2268                "streamGenerateContent"
2269            } else {
2270                "generateContent"
2271            };
2272            let (url, body) = if account.kind == "api_key" {
2273                // AI Studio: plain request, model in the path, x-goog-api-key header.
2274                let mut url = format!(
2275                    "{GEMINI_API_BASE}/v1beta/models/{routed_model}:{method}"
2276                );
2277                if client_stream {
2278                    url.push_str("?alt=sse");
2279                }
2280                let body =
2281                    serde_json::to_vec(&gemini_req).unwrap_or_else(|_| original_body.to_vec());
2282                (url, body)
2283            } else {
2284                // Code Assist (OAuth): wrapped envelope + project.
2285                let project = ensure_gemini_project(state, &account).await?;
2286                let mut url = format!(
2287                    "{GEMINI_CODE_ASSIST_BASE}/{GEMINI_CODE_ASSIST_VERSION}:{method}"
2288                );
2289                if client_stream {
2290                    url.push_str("?alt=sse");
2291                }
2292                let envelope = json!({
2293                    "model": routed_model,
2294                    "project": project,
2295                    "request": gemini_req,
2296                });
2297                let body =
2298                    serde_json::to_vec(&envelope).unwrap_or_else(|_| original_body.to_vec());
2299                (url, body)
2300            };
2301            Ok(UpstreamPlan {
2302                url,
2303                account,
2304                body,
2305                upstream_format: "gemini",
2306                destream: false,
2307                respond_as: respond_as.or(Some(RespondAs::Gemini)),
2308                client_stream,
2309                extra_headers: vec![],
2310                dario_guard: None,
2311            })
2312        }
2313    }
2314}
2315
2316#[cfg(test)]
2317mod trace_api_tests {
2318    use super::{
2319        session_from_metadata, trace_extras, trace_reasoning_fields, transcript_turn,
2320        truncate_chars,
2321    };
2322    use serde_json::json;
2323
2324    #[test]
2325    fn metadata_session_variants() {
2326        let claude = json!({"metadata": {"user_id":
2327            "{\"device_id\":\"d1\",\"session_id\":\"ses_inner\"}"}});
2328        assert_eq!(session_from_metadata(&claude), Some("ses_inner".into()));
2329        let plain = json!({"metadata": {"user_id": "user-123"}});
2330        assert_eq!(session_from_metadata(&plain), Some("user-123".into()));
2331        let json_no_session = json!({"metadata": {"user_id": "{\"device_id\":\"d1\"}"}});
2332        assert_eq!(
2333            session_from_metadata(&json_no_session),
2334            Some("{\"device_id\":\"d1\"}".into())
2335        );
2336        assert_eq!(session_from_metadata(&json!({})), None);
2337    }
2338
2339    #[test]
2340    fn extras_per_format() {
2341        let anthropic = json!({
2342            "system": [{"type": "text", "text": "abcd"}],
2343            "messages": [{"role": "user", "content": "hi"}],
2344            "max_tokens": 100,
2345            "temperature": 0.5,
2346            "thinking": {"type": "enabled", "budget_tokens": 4096},
2347        });
2348        let e = trace_extras(&anthropic);
2349        assert_eq!(e["thinking_budget"], 4096);
2350        assert_eq!(trace_reasoning_fields(&anthropic), (None, Some(4096)));
2351        assert_eq!(e["max_tokens"], 100);
2352        assert_eq!(e["temperature"], 0.5);
2353        assert_eq!(e["message_count"], 1);
2354        assert_eq!(e["system_chars"], 4);
2355        assert_eq!(e["reasoning_effort"], serde_json::Value::Null);
2356        let responses = json!({
2357            "instructions": "abc",
2358            "input": [{"type": "message"}, {"type": "message"}],
2359            "max_output_tokens": 200,
2360            "reasoning": {"effort": "high"},
2361        });
2362        let e = trace_extras(&responses);
2363        assert_eq!(e["reasoning_effort"], "high");
2364        assert_eq!(
2365            trace_reasoning_fields(&responses),
2366            (Some("high".into()), None)
2367        );
2368        assert_eq!(e["max_tokens"], 200);
2369        assert_eq!(e["message_count"], 2);
2370        assert_eq!(e["system_chars"], 3);
2371        assert_eq!(e["thinking_budget"], serde_json::Value::Null);
2372    }
2373
2374    #[test]
2375    fn truncates_on_char_boundaries() {
2376        assert_eq!(truncate_chars("abc".into(), 8000), "abc");
2377        assert_eq!(truncate_chars("héllo".repeat(2000), 8000).chars().count(), 8000);
2378    }
2379
2380    #[test]
2381    fn transcript_turn_missing_bodies_are_null() {
2382        let row = json!({
2383            "id": "t1", "ts_request_ms": 1, "ts_response_ms": 2,
2384            "routed_model": "m", "status": 200,
2385            "input_tokens": 10, "output_tokens": 5, "cost_usd": 0.01, "error": null,
2386            "reasoning_effort": "minimal", "thinking_budget": null,
2387            "req_body_path": "/nonexistent/x.gz", "resp_body_path": null,
2388            "client_format": "anthropic", "upstream_format": "anthropic",
2389        });
2390        let turn = transcript_turn(&row);
2391        assert_eq!(turn["trace_id"], "t1");
2392        assert_eq!(turn["user"], serde_json::Value::Null);
2393        assert_eq!(turn["assistant"], serde_json::Value::Null);
2394        assert_eq!(turn["model"], "m");
2395        assert_eq!(turn["reasoning_effort"], "minimal");
2396        assert_eq!(turn["thinking_budget"], serde_json::Value::Null);
2397    }
2398}
2399
2400#[cfg(test)]
2401mod run_key_tests {
2402    use super::{generate_run_key, key_fingerprint, key_hash_hex, merge_run_key_tags};
2403
2404    #[test]
2405    fn merge_header_tags_win_per_key() {
2406        let merged = merge_run_key_tags(
2407            Some(r#"{"task":"demo","suite":"swebench"}"#),
2408            Some(r#"{"task":"override","case":"x1"}"#),
2409        )
2410        .unwrap();
2411        let v: serde_json::Value = serde_json::from_str(&merged).unwrap();
2412        assert_eq!(v["task"], "override");
2413        assert_eq!(v["suite"], "swebench");
2414        assert_eq!(v["case"], "x1");
2415    }
2416
2417    #[test]
2418    fn merge_handles_missing_sides() {
2419        assert_eq!(merge_run_key_tags(None, None), None);
2420        assert_eq!(
2421            merge_run_key_tags(Some(r#"{"a":"1"}"#), None).unwrap(),
2422            r#"{"a":"1"}"#
2423        );
2424        assert_eq!(
2425            merge_run_key_tags(None, Some(r#"{"b":"2"}"#)).unwrap(),
2426            r#"{"b":"2"}"#
2427        );
2428        assert_eq!(merge_run_key_tags(Some("not json"), None), None);
2429    }
2430
2431    #[test]
2432    fn run_key_shape_and_fingerprint() {
2433        let key = generate_run_key();
2434        assert!(key.starts_with("alxk-"));
2435        assert_eq!(key.len(), 5 + 64);
2436        assert_ne!(key, generate_run_key());
2437        let hash = key_hash_hex(&key);
2438        assert_eq!(hash.len(), 64);
2439        assert_eq!(hash[..16], key_fingerprint(&key));
2440    }
2441}
2442
2443#[cfg(test)]
2444mod usage_tests {
2445    use super::{key_fingerprint, usage_backoff_ms};
2446
2447    #[test]
2448    fn fingerprint_is_first_8_sha256_bytes_hex() {
2449        assert_eq!(key_fingerprint("abc"), "ba7816bf8f01cfea");
2450        assert_eq!(key_fingerprint("abc").len(), 16);
2451        assert_ne!(key_fingerprint("abc"), key_fingerprint("abd"));
2452    }
2453
2454    #[test]
2455    fn backoff_grows_and_caps() {
2456        assert_eq!(usage_backoff_ms(1, None), 60_000);
2457        assert_eq!(usage_backoff_ms(2, None), 120_000);
2458        assert_eq!(usage_backoff_ms(3, None), 240_000);
2459        assert_eq!(usage_backoff_ms(10, None), 3_600_000);
2460    }
2461
2462    #[test]
2463    fn retry_after_wins_when_larger() {
2464        assert_eq!(usage_backoff_ms(1, Some(600_000)), 600_000);
2465        assert_eq!(usage_backoff_ms(5, Some(1_000)), 960_000);
2466    }
2467}
2468
2469fn upstream_headers(
2470    account: &Account,
2471    client_headers: &HeaderMap,
2472) -> Result<reqwest::header::HeaderMap, (StatusCode, String)> {
2473    let mut h = reqwest::header::HeaderMap::new();
2474    h.insert("content-type", HeaderValue::from_static("application/json"));
2475    h.insert("accept", HeaderValue::from_static("*/*"));
2476    h.insert("accept-encoding", HeaderValue::from_static("identity"));
2477    if let Some(ua) = client_headers.get("user-agent") {
2478        h.insert("user-agent", ua.clone());
2479    }
2480    match (account.provider, account.kind.as_str()) {
2481        (Provider::Anthropic, "oauth") => {
2482            let token = account.access_token.as_deref().ok_or((
2483                StatusCode::BAD_GATEWAY,
2484                "anthropic oauth account has no access token".to_string(),
2485            ))?;
2486            h.insert(
2487                "authorization",
2488                HeaderValue::from_str(&format!("Bearer {token}"))
2489                    .map_err(|e| (StatusCode::BAD_GATEWAY, e.to_string()))?,
2490            );
2491            let mut beta = ANTHROPIC_OAUTH_BETA.to_string();
2492            if let Some(client_beta) = client_headers
2493                .get("anthropic-beta")
2494                .and_then(|v| v.to_str().ok())
2495            {
2496                if !client_beta.contains(ANTHROPIC_OAUTH_BETA) {
2497                    beta = format!("{ANTHROPIC_OAUTH_BETA},{client_beta}");
2498                } else {
2499                    beta = client_beta.to_string();
2500                }
2501            }
2502            h.insert(
2503                "anthropic-beta",
2504                HeaderValue::from_str(&beta)
2505                    .map_err(|e| (StatusCode::BAD_GATEWAY, e.to_string()))?,
2506            );
2507            insert_anthropic_version(&mut h, client_headers);
2508        }
2509        (Provider::Anthropic, _) => {
2510            let key = account
2511                .api_key
2512                .as_deref()
2513                .or(account.access_token.as_deref())
2514                .ok_or((
2515                    StatusCode::BAD_GATEWAY,
2516                    "anthropic account has no api key".to_string(),
2517                ))?;
2518            h.insert(
2519                "x-api-key",
2520                HeaderValue::from_str(key).map_err(|e| (StatusCode::BAD_GATEWAY, e.to_string()))?,
2521            );
2522            if let Some(client_beta) = client_headers.get("anthropic-beta") {
2523                h.insert("anthropic-beta", client_beta.clone());
2524            }
2525            insert_anthropic_version(&mut h, client_headers);
2526        }
2527        (Provider::Openai, "oauth") => {
2528            let token = account.access_token.as_deref().ok_or((
2529                StatusCode::BAD_GATEWAY,
2530                "openai oauth account has no access token".to_string(),
2531            ))?;
2532            h.insert(
2533                "authorization",
2534                HeaderValue::from_str(&format!("Bearer {token}"))
2535                    .map_err(|e| (StatusCode::BAD_GATEWAY, e.to_string()))?,
2536            );
2537            if let Some(acct) = account.chatgpt_account_id() {
2538                h.insert(
2539                    "chatgpt-account-id",
2540                    HeaderValue::from_str(&acct)
2541                        .map_err(|e| (StatusCode::BAD_GATEWAY, e.to_string()))?,
2542                );
2543            }
2544            h.insert(
2545                "openai-beta",
2546                HeaderValue::from_static("responses=experimental"),
2547            );
2548            h.insert("originator", HeaderValue::from_static("codex_cli_rs"));
2549            let session = client_headers
2550                .get("session_id")
2551                .and_then(|v| v.to_str().ok())
2552                .map(String::from)
2553                .unwrap_or_else(|| uuid::Uuid::new_v4().to_string());
2554            h.insert(
2555                "session_id",
2556                HeaderValue::from_str(&session)
2557                    .map_err(|e| (StatusCode::BAD_GATEWAY, e.to_string()))?,
2558            );
2559        }
2560        (Provider::Openai, _) => {
2561            let key = account.api_key.as_deref().ok_or((
2562                StatusCode::BAD_GATEWAY,
2563                "openai account has no api key".to_string(),
2564            ))?;
2565            h.insert(
2566                "authorization",
2567                HeaderValue::from_str(&format!("Bearer {key}"))
2568                    .map_err(|e| (StatusCode::BAD_GATEWAY, e.to_string()))?,
2569            );
2570        }
2571        (Provider::Xai, _) => {
2572            let token = account
2573                .access_token
2574                .as_deref()
2575                .or(account.api_key.as_deref())
2576                .ok_or((
2577                    StatusCode::BAD_GATEWAY,
2578                    "xai account has no access token".to_string(),
2579                ))?;
2580            h.insert(
2581                "authorization",
2582                HeaderValue::from_str(&format!("Bearer {token}"))
2583                    .map_err(|e| (StatusCode::BAD_GATEWAY, e.to_string()))?,
2584            );
2585            h.insert("x-xai-token-auth", HeaderValue::from_static("xai-grok-cli"));
2586            h.insert(
2587                "x-grok-client-version",
2588                HeaderValue::from_static(GROK_CLIENT_VERSION),
2589            );
2590            h.insert("user-agent", HeaderValue::from_static("xai-grok-workspace"));
2591            h.insert("accept", HeaderValue::from_static("text/event-stream"));
2592        }
2593        (Provider::Gemini, "api_key") => {
2594            let key = account.api_key.as_deref().ok_or((
2595                StatusCode::BAD_GATEWAY,
2596                "gemini api_key account has no key".to_string(),
2597            ))?;
2598            h.insert(
2599                "x-goog-api-key",
2600                HeaderValue::from_str(key)
2601                    .map_err(|e| (StatusCode::BAD_GATEWAY, e.to_string()))?,
2602            );
2603        }
2604        (Provider::Gemini, _) => {
2605            let token = account.access_token.as_deref().ok_or((
2606                StatusCode::BAD_GATEWAY,
2607                "gemini account has no access token".to_string(),
2608            ))?;
2609            h.insert(
2610                "authorization",
2611                HeaderValue::from_str(&format!("Bearer {token}"))
2612                    .map_err(|e| (StatusCode::BAD_GATEWAY, e.to_string()))?,
2613            );
2614        }
2615    }
2616    Ok(h)
2617}
2618
2619fn insert_anthropic_version(h: &mut reqwest::header::HeaderMap, client_headers: &HeaderMap) {
2620    match client_headers.get("anthropic-version") {
2621        Some(v) => {
2622            h.insert("anthropic-version", v.clone());
2623        }
2624        None => {
2625            h.insert("anthropic-version", HeaderValue::from_static("2023-06-01"));
2626        }
2627    }
2628}
2629
2630fn key_fingerprint(key: &str) -> String {
2631    use sha2::{Digest, Sha256};
2632    let digest = Sha256::digest(key.as_bytes());
2633    digest[..8].iter().map(|b| format!("{b:02x}")).collect()
2634}
2635
2636const RUN_KEY_PREFIX: &str = "alxk-";
2637const RUN_KEY_DEFAULT_TTL_S: i64 = 86_400;
2638const RUN_KEY_MAX_TTL_S: i64 = 604_800;
2639
2640pub fn generate_run_key() -> String {
2641    use rand::RngCore;
2642    let mut bytes = [0u8; 32];
2643    rand::thread_rng().fill_bytes(&mut bytes);
2644    format!(
2645        "{RUN_KEY_PREFIX}{}",
2646        bytes.iter().map(|b| format!("{b:02x}")).collect::<String>()
2647    )
2648}
2649
2650fn key_hash_hex(key: &str) -> String {
2651    use sha2::{Digest, Sha256};
2652    let digest = Sha256::digest(key.as_bytes());
2653    digest.iter().map(|b| format!("{b:02x}")).collect()
2654}
2655
2656fn merge_run_key_tags(key_tags: Option<&str>, header_tags: Option<&str>) -> Option<String> {
2657    let parse = |s: Option<&str>| {
2658        s.and_then(|s| serde_json::from_str::<Value>(s).ok())
2659            .and_then(|v| v.as_object().cloned())
2660            .unwrap_or_default()
2661    };
2662    let mut merged = parse(key_tags);
2663    merged.extend(parse(header_tags));
2664    if merged.is_empty() {
2665        None
2666    } else {
2667        serde_json::to_string(&Value::Object(merged)).ok()
2668    }
2669}
2670
2671fn run_key_entry(state: &AppState, key_hash: &str) -> Option<CachedRunKey> {
2672    let now = now_ms();
2673    let cached = state.run_keys.read().unwrap().get(key_hash).cloned();
2674    if let Some(entry) = cached {
2675        if entry.expires_ms.map(|e| e > now).unwrap_or(true) {
2676            return Some(entry);
2677        }
2678        state.run_keys.write().unwrap().remove(key_hash);
2679        return None;
2680    }
2681    let row = state.store.lookup_run_key(key_hash, now).ok().flatten()?;
2682    let entry = CachedRunKey {
2683        run_id: row["run_id"].as_str().map(String::from),
2684        tags_json: row["tags"]
2685            .as_object()
2686            .filter(|o| !o.is_empty())
2687            .and_then(|o| serde_json::to_string(o).ok()),
2688        expires_ms: row["expires_ms"].as_i64(),
2689    };
2690    state
2691        .run_keys
2692        .write()
2693        .unwrap()
2694        .insert(key_hash.to_string(), entry.clone());
2695    Some(entry)
2696}
2697
2698async fn admin_run_keys_create(
2699    State(state): State<Arc<AppState>>,
2700    body: Option<axum::Json<Value>>,
2701) -> Response {
2702    let body = body.map(|b| b.0).unwrap_or_else(|| json!({}));
2703    let tags = match &body["tags"] {
2704        Value::Null => None,
2705        Value::Object(o) => Some(o.clone()),
2706        _ => return error_response(StatusCode::BAD_REQUEST, "'tags' must be an object"),
2707    };
2708    let kind = body["kind"].as_str().unwrap_or("run");
2709    if kind != "run" && kind != "harness" {
2710        return error_response(StatusCode::BAD_REQUEST, "'kind' must be 'run' or 'harness'");
2711    }
2712    let run_id = body["run_id"].as_str().map(String::from);
2713    let label = body["label"].as_str().map(String::from);
2714    if kind == "harness" && label.as_deref().unwrap_or("").trim().is_empty() {
2715        return error_response(StatusCode::BAD_REQUEST, "'label' is required for harness keys");
2716    }
2717    let created_ms = now_ms();
2718    let expires_ms = if kind == "harness" {
2719        None
2720    } else {
2721        let ttl_seconds = body["ttl_seconds"].as_i64().unwrap_or(RUN_KEY_DEFAULT_TTL_S);
2722        if ttl_seconds <= 0 {
2723            return error_response(StatusCode::BAD_REQUEST, "'ttl_seconds' must be positive");
2724        }
2725        Some(created_ms + ttl_seconds.min(RUN_KEY_MAX_TTL_S) * 1000)
2726    };
2727    let key = generate_run_key();
2728    let key_hash = key_hash_hex(&key);
2729    let id = format!("rk-{}", &key_hash[..8]);
2730    let tags_json = tags
2731        .as_ref()
2732        .filter(|o| !o.is_empty())
2733        .and_then(|o| serde_json::to_string(o).ok());
2734    match state.store.insert_run_key(
2735        &id,
2736        &key_hash,
2737        kind,
2738        run_id.as_deref(),
2739        tags_json.as_deref(),
2740        label.as_deref(),
2741        created_ms,
2742        expires_ms,
2743    ) {
2744        Ok(()) => {
2745            let (_, exports) = connect_payload(&state.base_url, &key);
2746            (
2747                StatusCode::CREATED,
2748                axum::Json(json!({
2749                    "id": id,
2750                    "key": key,
2751                    "kind": kind,
2752                    "run_id": run_id,
2753                    "label": label,
2754                    "tags": tags.map(Value::Object).unwrap_or_else(|| json!({})),
2755                    "expires_ms": expires_ms,
2756                    "exports": exports,
2757                })),
2758            )
2759                .into_response()
2760        }
2761        Err(e) => error_response(StatusCode::INTERNAL_SERVER_ERROR, &e.to_string()),
2762    }
2763}
2764
2765async fn admin_run_keys_list(
2766    State(state): State<Arc<AppState>>,
2767    Query(q): Query<HashMap<String, String>>,
2768) -> Response {
2769    let all = q.get("all").map(|v| v == "1" || v == "true").unwrap_or(false);
2770    match state.store.list_run_keys(all) {
2771        Ok(rows) => axum::Json(json!({"run_keys": rows})).into_response(),
2772        Err(e) => error_response(StatusCode::INTERNAL_SERVER_ERROR, &e.to_string()),
2773    }
2774}
2775
2776async fn admin_run_keys_revoke(
2777    State(state): State<Arc<AppState>>,
2778    Path(id): Path<String>,
2779) -> Response {
2780    match state.store.revoke_run_key(&id) {
2781        Ok(true) => {
2782            state.run_keys.write().unwrap().clear();
2783            axum::Json(json!({"revoked": true})).into_response()
2784        }
2785        Ok(false) => error_response(StatusCode::NOT_FOUND, &format!("unknown run key '{id}'")),
2786        Err(e) => error_response(StatusCode::INTERNAL_SERVER_ERROR, &e.to_string()),
2787    }
2788}
2789
2790fn session_from_metadata(body_json: &Value) -> Option<String> {
2791    let raw = body_json["metadata"]["user_id"].as_str()?;
2792    if let Ok(v) = serde_json::from_str::<Value>(raw) {
2793        if let Some(inner) = v["session_id"].as_str() {
2794            return Some(inner.to_string());
2795        }
2796    }
2797    Some(raw.to_string())
2798}
2799
2800const METADATA_HEADERS: &[(&str, &str)] = &[
2801    ("x-alexandria-harness", "harness"),
2802    ("x-alexandria-harness-version", "harness_version"),
2803    ("x-alexandria-task", "task"),
2804    ("x-alexandria-model", "model"),
2805    ("x-alexandria-job", "job"),
2806    ("x-alexandria-phase", "phase"),
2807    ("x-alexandria-kind", "kind"),
2808];
2809
2810fn trace_tags_json(headers: &HeaderMap) -> Option<String> {
2811    let values: Vec<&str> = headers
2812        .get_all("x-alexandria-trace-tag")
2813        .iter()
2814        .filter_map(|v| v.to_str().ok())
2815        .collect();
2816    let mut tags = parse_trace_tags(&values);
2817    if let Some(o) = tags.as_object_mut() {
2818        for (header, key) in METADATA_HEADERS {
2819            if o.contains_key(*key) {
2820                continue;
2821            }
2822            if let Some(v) = headers.get(*header).and_then(|v| v.to_str().ok()) {
2823                let v = v.trim();
2824                if !v.is_empty() {
2825                    o.insert((*key).to_string(), json!(v));
2826                }
2827            }
2828        }
2829    }
2830    if tags.as_object().map(|o| o.is_empty()).unwrap_or(true) {
2831        return None;
2832    }
2833    serde_json::to_string(&tags).ok()
2834}
2835
2836async fn proxy(
2837    state: Arc<AppState>,
2838    format: ClientFormat,
2839    path: &'static str,
2840    headers: HeaderMap,
2841    body: Bytes,
2842    peer: Option<std::net::SocketAddr>,
2843) -> Response {
2844    let mut run_key: Option<CachedRunKey> = None;
2845    let client_fingerprint = match client_key(&headers) {
2846        Some(k) if k == state.local_key => key_fingerprint(&k),
2847        Some(k) => {
2848            let key_hash = key_hash_hex(&k);
2849            match run_key_entry(&state, &key_hash) {
2850                Some(entry) => {
2851                    if let Err(e) = state.store.touch_run_key(&key_hash, now_ms()) {
2852                        tracing::warn!("failed to touch run key: {e}");
2853                    }
2854                    run_key = Some(entry);
2855                    key_hash.chars().take(16).collect()
2856                }
2857                None if k.starts_with(RUN_KEY_PREFIX) => {
2858                    return error_response(
2859                        StatusCode::UNAUTHORIZED,
2860                        "run key expired or revoked",
2861                    )
2862                }
2863                None => {
2864                    return error_response(
2865                        StatusCode::UNAUTHORIZED,
2866                        "bad or missing local key (x-api-key / Authorization: Bearer)",
2867                    )
2868                }
2869            }
2870        }
2871        None => {
2872            return error_response(
2873                StatusCode::UNAUTHORIZED,
2874                "bad or missing local key (x-api-key / Authorization: Bearer)",
2875            )
2876        }
2877    };
2878
2879    let in_flight = InFlight::new(&state);
2880    let trace_id = uuid::Uuid::new_v4().to_string();
2881    let mut trace = TraceRecord {
2882        id: trace_id.clone(),
2883        ts_request_ms: now_ms(),
2884        method: Some("POST".into()),
2885        path: Some(path.into()),
2886        client_format: Some(format.as_str().into()),
2887        harness: headers
2888            .get("user-agent")
2889            .and_then(|v| v.to_str().ok())
2890            .map(String::from),
2891        req_headers_json: Some(redacted_headers(&headers)),
2892        run_id: headers
2893            .get("x-alexandria-run-id")
2894            .and_then(|v| v.to_str().ok())
2895            .map(String::from)
2896            .or_else(|| run_key.as_ref().and_then(|k| k.run_id.clone())),
2897        tags: merge_run_key_tags(
2898            run_key.as_ref().and_then(|k| k.tags_json.as_deref()),
2899            trace_tags_json(&headers).as_deref(),
2900        ),
2901        client_ip: peer.map(|p| p.ip().to_string()),
2902        key_fingerprint: Some(client_fingerprint),
2903        ..Default::default()
2904    };
2905
2906    let mut body_json: Value = match serde_json::from_slice(&body) {
2907        Ok(v) => v,
2908        Err(e) => {
2909            return error_response(StatusCode::BAD_REQUEST, &format!("body is not JSON: {e}"))
2910        }
2911    };
2912    let (reasoning_effort, thinking_budget) = trace_reasoning_fields(&body_json);
2913    trace.reasoning_effort = reasoning_effort;
2914    trace.thinking_budget = thinking_budget;
2915
2916    trace.session_id = headers
2917        .get("x-session-id")
2918        .and_then(|v| v.to_str().ok())
2919        .map(String::from)
2920        .or_else(|| {
2921            headers
2922                .get("session_id")
2923                .and_then(|v| v.to_str().ok())
2924                .map(String::from)
2925        })
2926        .or_else(|| session_from_metadata(&body_json))
2927        .or_else(|| {
2928            conversation_root(format, &body_json).map(|root| {
2929                let ua = headers
2930                    .get("user-agent")
2931                    .and_then(|v| v.to_str().ok())
2932                    .unwrap_or("");
2933                let ip = peer.map(|p| p.ip().to_string()).unwrap_or_default();
2934                format!("auto-{}", key_fingerprint(&format!("{root}{ua}{ip}")))
2935            })
2936        });
2937
2938    let requested_model = body_json["model"].as_str().unwrap_or("").to_string();
2939    if requested_model.is_empty() {
2940        return error_response(StatusCode::BAD_REQUEST, "missing 'model' in request body");
2941    }
2942    let (routed_provider, routed_model) = route_model(&requested_model);
2943    let provider = routed_provider.unwrap_or_else(|| format.default_provider());
2944    trace.requested_model = Some(requested_model.clone());
2945    trace.routed_model = Some(routed_model.clone());
2946    trace.upstream_provider = Some(provider.as_str().into());
2947    trace.streamed = Some(body_json["stream"].as_bool().unwrap_or(false));
2948
2949    let mut plan = match plan_upstream(
2950        &state,
2951        format,
2952        provider,
2953        &routed_model,
2954        &mut body_json,
2955        &body,
2956        &trace_id,
2957    )
2958    .await
2959    {
2960        Ok(p) => p,
2961        Err((status, msg)) => {
2962            trace.status = Some(status.as_u16() as i64);
2963            trace.error = Some(msg.clone());
2964            finalize_trace(&state, trace, &body, None, None);
2965            return error_response(status, &msg);
2966        }
2967    };
2968    trace.upstream_format = Some(plan.upstream_format.into());
2969    trace.account_id = Some(plan.account.id.clone());
2970    trace.billing_bucket = Some(
2971        if plan.account.kind == "oauth" || plan.account.kind == "dario" {
2972            "subscription"
2973        } else {
2974            "api"
2975        }
2976        .into(),
2977    );
2978
2979    tracing::info!(
2980        trace_id,
2981        model = %routed_model,
2982        provider = provider.as_str(),
2983        account = %plan.account.id,
2984        url = %plan.url,
2985        "proxying request"
2986    );
2987
2988    let mut account = plan.account.clone();
2989    let mut upstream_resp = None;
2990    for attempt in 0..2 {
2991        let mut up_headers = match upstream_headers(&account, &headers) {
2992            Ok(h) => h,
2993            Err((status, msg)) => {
2994                trace.status = Some(status.as_u16() as i64);
2995                trace.error = Some(msg.clone());
2996                finalize_trace(&state, trace, &body, Some(&plan.body), None);
2997                return error_response(status, &msg);
2998            }
2999        };
3000        for (k, v) in &plan.extra_headers {
3001            if let (Ok(name), Ok(value)) = (
3002                reqwest::header::HeaderName::from_bytes(k.as_bytes()),
3003                HeaderValue::from_str(v),
3004            ) {
3005                up_headers.insert(name, value);
3006            }
3007        }
3008        let resp = match state
3009            .http
3010            .post(&plan.url)
3011            .headers(up_headers)
3012            .body(plan.body.clone())
3013            .send()
3014            .await
3015        {
3016            Ok(r) => r,
3017            Err(e) => {
3018                let msg = format!("upstream request failed: {e}");
3019                suspect_dario(&state, &account);
3020                trace.status = Some(502);
3021                trace.error = Some(msg.clone());
3022                finalize_trace(&state, trace, &body, Some(&plan.body), None);
3023                return error_response(StatusCode::BAD_GATEWAY, &msg);
3024            }
3025        };
3026        let retryable_status = resp.status() == reqwest::StatusCode::TOO_MANY_REQUESTS
3027            || resp.status() == reqwest::StatusCode::UNAUTHORIZED
3028            || resp.status() == reqwest::StatusCode::FORBIDDEN
3029            || resp.status().is_server_error();
3030        if retryable_status && account.kind != "dario" {
3031            let retry_after_s = resp
3032                .headers()
3033                .get("retry-after")
3034                .and_then(|v| v.to_str().ok())
3035                .and_then(|v| v.parse::<i64>().ok())
3036                .unwrap_or(60)
3037                .clamp(1, 3600);
3038            tracing::warn!(account = %account.id, retry_after_s, status = %resp.status(), "upstream failed; cooling account down");
3039            if let Err(e) = state.vault.mark_cooldown(&account.id, now_ms() + retry_after_s * 1000).await {
3040                tracing::warn!("failed to mark cooldown: {e}");
3041            }
3042            if attempt == 0 {
3043                match plan_upstream(&state, format, provider, &routed_model, &mut body_json, &body, &trace_id).await {
3044                    Ok(p) if p.account.id != account.id => {
3045                        plan = p;
3046                        account = plan.account.clone();
3047                        trace.account_id = Some(account.id.clone());
3048                        trace.upstream_format = Some(plan.upstream_format.into());
3049                        tracing::warn!(account = %account.id, "retrying once with failover account");
3050                        continue;
3051                    }
3052                    Ok(_) => {}
3053                    Err((status, msg)) => tracing::warn!(status = %status, error = %msg, "no failover account available"),
3054                }
3055            }
3056        }
3057        if resp.status() == reqwest::StatusCode::UNAUTHORIZED
3058            && account.kind == "oauth"
3059            && attempt == 0
3060        {
3061            tracing::warn!(
3062                account = %account.id,
3063                "upstream returned 401 for oauth account; forcing token refresh and retrying"
3064            );
3065            match state.vault.refresh(&account.id, true).await {
3066                Ok(fresh) => {
3067                    account = fresh;
3068                    continue;
3069                }
3070                Err(e) => {
3071                    tracing::warn!("forced refresh failed: {e}");
3072                }
3073            }
3074        }
3075        upstream_resp = Some(resp);
3076        break;
3077    }
3078    let upstream_resp = upstream_resp.expect("upstream response after retry loop");
3079    trace.account_id = Some(account.id.clone());
3080
3081    let status = upstream_resp.status();
3082    trace.status = Some(status.as_u16() as i64);
3083    let resp_headers = upstream_resp.headers().clone();
3084    let content_type = resp_headers
3085        .get("content-type")
3086        .and_then(|v| v.to_str().ok())
3087        .unwrap_or("application/octet-stream")
3088        .to_string();
3089    trace.resp_headers_json = Some(redacted_headers(&resp_headers));
3090    let is_sse = content_type.starts_with("text/event-stream");
3091
3092    if let Some(target) = plan.respond_as {
3093        use alex_core::translate;
3094        let buf = match upstream_resp.bytes().await {
3095            Ok(b) => b.to_vec(),
3096            Err(e) => {
3097                let msg = format!("upstream body read failed: {e}");
3098                suspect_dario(&state, &account);
3099                trace.status = Some(502);
3100                trace.error = Some(msg.clone());
3101                finalize_trace(&state, trace, &body, Some(&plan.body), None);
3102                return error_response(StatusCode::BAD_GATEWAY, &msg);
3103            }
3104        };
3105        drop(plan.dario_guard.take());
3106        trace.ts_response_ms = Some(now_ms());
3107        fill_usage_and_cost(&state, &mut trace, &buf, is_sse);
3108        let text = String::from_utf8_lossy(&buf).to_string();
3109        if !status.is_success() {
3110            finalize_trace(&state, trace, &body, Some(&plan.body), Some(&buf));
3111            return Response::builder()
3112                .status(status)
3113                .header("content-type", "application/json")
3114                .header("x-alexandria-trace-id", &trace_id)
3115                .body(Body::from(buf))
3116                .unwrap_or_else(|e| {
3117                    error_response(StatusCode::INTERNAL_SERVER_ERROR, &e.to_string())
3118                });
3119        }
3120        let trimmed = text.trim_start();
3121        let looks_sse = trimmed.starts_with("event:") || trimmed.starts_with("data:");
3122        let upstream_final = if plan.upstream_format == "anthropic" {
3123            if is_sse || looks_sse {
3124                translate::parse_anthropic_sse_to_message(&text)
3125            } else {
3126                serde_json::from_str(&text).ok()
3127            }
3128        } else if plan.upstream_format == "gemini" {
3129            translate::parse_gemini_upstream_final(&text)
3130        } else if plan.upstream_format == "openai-chat" {
3131            if is_sse || looks_sse {
3132                translate::parse_openai_chat_sse_final(&text)
3133            } else {
3134                serde_json::from_str(&text).ok()
3135            }
3136        } else if is_sse || looks_sse {
3137            translate::parse_responses_sse_final(&text)
3138        } else {
3139            serde_json::from_str(&text).ok()
3140        };
3141        let Some(upstream_final) = upstream_final else {
3142            let msg = "could not reassemble upstream response for translation";
3143            trace.error = Some(msg.to_string());
3144            finalize_trace(&state, trace, &body, Some(&plan.body), Some(&buf));
3145            return error_response(StatusCode::BAD_GATEWAY, msg);
3146        };
3147        let out = match (target, plan.upstream_format) {
3148            (RespondAs::Gemini, "gemini") => upstream_final.clone(),
3149            (RespondAs::Anthropic, "gemini") => {
3150                translate::gemini_response_to_anthropic(&upstream_final, &requested_model)
3151            }
3152            (RespondAs::OpenaiChat, "gemini") => translate::anthropic_response_to_openai_chat(
3153                &translate::gemini_response_to_anthropic(&upstream_final, &requested_model),
3154                &requested_model,
3155            ),
3156            (RespondAs::OpenaiResponses, "gemini") => {
3157                translate::anthropic_response_to_openai_responses(
3158                    &translate::gemini_response_to_anthropic(&upstream_final, &requested_model),
3159                    &requested_model,
3160                )
3161            }
3162            (RespondAs::OpenaiChat, "anthropic") => {
3163                translate::anthropic_response_to_openai_chat(&upstream_final, &requested_model)
3164            }
3165            (RespondAs::OpenaiResponses, "anthropic") => translate::anthropic_response_to_openai_responses(
3166                &upstream_final,
3167                &requested_model,
3168            ),
3169            (RespondAs::Gemini, "anthropic") => {
3170                translate::anthropic_response_to_gemini(&upstream_final, &requested_model)
3171            }
3172            (RespondAs::Anthropic, "openai-chat") => {
3173                translate::openai_chat_response_to_anthropic(&upstream_final, &requested_model)
3174            }
3175            (RespondAs::OpenaiChat, "openai-chat") => upstream_final.clone(),
3176            (RespondAs::OpenaiResponses, "openai-chat") => {
3177                translate::anthropic_response_to_openai_responses(
3178                    &translate::openai_chat_response_to_anthropic(
3179                        &upstream_final,
3180                        &requested_model,
3181                    ),
3182                    &requested_model,
3183                )
3184            }
3185            (RespondAs::Gemini, "openai-chat") => translate::anthropic_response_to_gemini(
3186                &translate::openai_chat_response_to_anthropic(&upstream_final, &requested_model),
3187                &requested_model,
3188            ),
3189            (RespondAs::Gemini, _) => translate::anthropic_response_to_gemini(
3190                &translate::responses_final_to_anthropic(&upstream_final, &requested_model),
3191                &requested_model,
3192            ),
3193            (RespondAs::Anthropic, "anthropic") | (RespondAs::OpenaiResponses, _) => {
3194                upstream_final.clone()
3195            }
3196            (RespondAs::Anthropic, _) => {
3197                translate::responses_final_to_anthropic(&upstream_final, &requested_model)
3198            }
3199            (RespondAs::OpenaiChat, _) => {
3200                translate::responses_final_to_openai_chat(&upstream_final, &requested_model)
3201            }
3202        };
3203        let (out_ct, out_body) = if plan.client_stream {
3204            let sse = match target {
3205                RespondAs::Anthropic => translate::synth_anthropic_sse(&out),
3206                RespondAs::OpenaiChat => translate::synth_openai_chat_sse(&out),
3207                RespondAs::OpenaiResponses => translate::synth_openai_responses_sse(&out),
3208                RespondAs::Gemini => translate::synth_gemini_sse(&out),
3209            };
3210            ("text/event-stream", sse.into_bytes())
3211        } else {
3212            (
3213                "application/json",
3214                serde_json::to_vec(&out).unwrap_or_default(),
3215            )
3216        };
3217        finalize_trace(&state, trace, &body, Some(&plan.body), Some(&buf));
3218        return Response::builder()
3219            .status(StatusCode::OK)
3220            .header("content-type", out_ct)
3221            .header("x-alexandria-trace-id", &trace_id)
3222            .body(Body::from(out_body))
3223            .unwrap_or_else(|e| error_response(StatusCode::INTERNAL_SERVER_ERROR, &e.to_string()));
3224    }
3225
3226    if plan.destream {
3227        let buf = match upstream_resp.bytes().await {
3228            Ok(b) => b.to_vec(),
3229            Err(e) => {
3230                let msg = format!("upstream body read failed: {e}");
3231                trace.status = Some(502);
3232                trace.error = Some(msg.clone());
3233                finalize_trace(&state, trace, &body, Some(&plan.body), None);
3234                return error_response(StatusCode::BAD_GATEWAY, &msg);
3235            }
3236        };
3237        trace.ts_response_ms = Some(now_ms());
3238        fill_usage_and_cost(&state, &mut trace, &buf, is_sse);
3239        let text = String::from_utf8_lossy(&buf).to_string();
3240        let (out_status, out_body) = if status.is_success() {
3241            match extract_final_response(&text) {
3242                Some(v) => (StatusCode::OK, serde_json::to_vec(&v).unwrap_or_default()),
3243                None => {
3244                    let msg = "upstream stream ended without a response.completed event";
3245                    trace.error = Some(msg.to_string());
3246                    (
3247                        StatusCode::BAD_GATEWAY,
3248                        serde_json::to_vec(
3249                            &json!({"error": {"type": "alexandria", "message": msg}}),
3250                        )
3251                        .unwrap_or_default(),
3252                    )
3253                }
3254            }
3255        } else {
3256            (status, buf.clone())
3257        };
3258        finalize_trace(&state, trace, &body, Some(&plan.body), Some(&buf));
3259        return Response::builder()
3260            .status(out_status)
3261            .header("content-type", "application/json")
3262            .header("x-alexandria-trace-id", &trace_id)
3263            .body(Body::from(out_body))
3264            .unwrap_or_else(|e| error_response(StatusCode::INTERNAL_SERVER_ERROR, &e.to_string()));
3265    }
3266
3267    let (tx, rx) = tokio::sync::mpsc::channel::<Result<Bytes, std::io::Error>>(64);
3268    let mut upstream_stream = upstream_resp.bytes_stream();
3269    let state2 = state.clone();
3270    let client_body = body.clone();
3271    let upstream_body = plan.body.clone();
3272    let dario_guard = plan.dario_guard.take();
3273    tokio::spawn(async move {
3274        let _dario_guard = dario_guard;
3275        let _in_flight = in_flight;
3276        let mut buf: Vec<u8> = Vec::new();
3277        let mut stream_error: Option<String> = None;
3278        while let Some(chunk) = upstream_stream.next().await {
3279            match chunk {
3280                Ok(b) => {
3281                    buf.extend_from_slice(&b);
3282                    let _ = tx.send(Ok(b)).await;
3283                }
3284                Err(e) => {
3285                    stream_error = Some(format!("upstream stream error: {e}"));
3286                    let _ = tx
3287                        .send(Err(std::io::Error::new(
3288                            std::io::ErrorKind::Other,
3289                            e.to_string(),
3290                        )))
3291                        .await;
3292                    break;
3293                }
3294            }
3295        }
3296        drop(tx);
3297        trace.ts_response_ms = Some(now_ms());
3298        if trace.error.is_none() {
3299            trace.error = stream_error;
3300        }
3301        if trace.error.is_some() {
3302            if let (Some(dario), Some(gen)) = (
3303                &state2.dario,
3304                trace
3305                    .account_id
3306                    .as_deref()
3307                    .and_then(|id| id.strip_prefix("dario:")),
3308            ) {
3309                dario.suspect(gen);
3310            }
3311        }
3312
3313        fill_usage_and_cost(&state2, &mut trace, &buf, is_sse);
3314        finalize_trace(
3315            &state2,
3316            trace,
3317            &client_body,
3318            Some(&upstream_body),
3319            Some(&buf),
3320        );
3321    });
3322
3323    let mut response = Response::builder().status(status);
3324    for (k, v) in resp_headers.iter() {
3325        let key = k.as_str().to_lowercase();
3326        if [
3327            "transfer-encoding",
3328            "connection",
3329            "content-encoding",
3330            "content-length",
3331        ]
3332        .contains(&key.as_str())
3333        {
3334            continue;
3335        }
3336        response = response.header(k, v);
3337    }
3338    response = response.header("x-alexandria-trace-id", &trace_id);
3339    response
3340        .body(Body::from_stream(ReceiverStream::new(rx)))
3341        .unwrap_or_else(|e| error_response(StatusCode::INTERNAL_SERVER_ERROR, &e.to_string()))
3342}
3343
3344fn fill_usage_and_cost(state: &AppState, trace: &mut TraceRecord, buf: &[u8], is_sse: bool) {
3345    let text = String::from_utf8_lossy(buf);
3346    let trimmed = text.trim_start();
3347    let looks_sse = trimmed.starts_with("event:") || trimmed.starts_with("data:");
3348    trace.usage = if is_sse || looks_sse {
3349        parse_sse_usage(&text)
3350    } else {
3351        serde_json::from_str::<Value>(&text)
3352            .map(|v| usage_from_json(&v))
3353            .unwrap_or_default()
3354    };
3355    if !trace.usage.is_empty() {
3356        if let Some(pricing) = trace
3357            .routed_model
3358            .as_deref()
3359            .and_then(|m| state.store.pricing_for(m))
3360        {
3361            let input_includes_cached = trace
3362                .upstream_format
3363                .as_deref()
3364                .map(|f| f.starts_with("openai"))
3365                .unwrap_or(false);
3366            trace.cost_usd = Some(compute_cost(&trace.usage, &pricing, input_includes_cached));
3367        }
3368    }
3369}
3370
3371fn extract_final_response(text: &str) -> Option<Value> {
3372    let mut last: Option<Value> = None;
3373    for line in text.lines() {
3374        let Some(data) = line.strip_prefix("data:") else {
3375            continue;
3376        };
3377        let Ok(v) = serde_json::from_str::<Value>(data.trim()) else {
3378            continue;
3379        };
3380        if matches!(
3381            v["type"].as_str(),
3382            Some("response.completed") | Some("response.incomplete") | Some("response.failed")
3383        ) && v["response"].is_object()
3384        {
3385            last = Some(v["response"].clone());
3386        }
3387    }
3388    last
3389}
3390
3391fn finalize_trace(
3392    state: &AppState,
3393    mut trace: TraceRecord,
3394    client_body: &[u8],
3395    upstream_body: Option<&[u8]>,
3396    resp_body: Option<&[u8]>,
3397) {
3398    let store = &state.store;
3399    match store.write_body(&trace.id, "request.json", client_body) {
3400        Ok(p) => trace.req_body_path = Some(p),
3401        Err(e) => tracing::warn!("failed to write request body: {e}"),
3402    }
3403    if let Some(up) = upstream_body {
3404        if up != client_body {
3405            match store.write_body(&trace.id, "upstream-request.json", up) {
3406                Ok(p) => trace.upstream_req_body_path = Some(p),
3407                Err(e) => tracing::warn!("failed to write upstream request body: {e}"),
3408            }
3409        }
3410    }
3411    if let Some(resp) = resp_body {
3412        match store.write_body(&trace.id, "response.body", resp) {
3413            Ok(p) => trace.resp_body_path = Some(p),
3414            Err(e) => tracing::warn!("failed to write response body: {e}"),
3415        }
3416    }
3417    if let Err(e) = store.insert_trace(&trace) {
3418        tracing::error!("failed to insert trace {}: {e}", trace.id);
3419    } else {
3420        tracing::info!(
3421            trace_id = %trace.id,
3422            status = trace.status,
3423            input = trace.usage.input_tokens,
3424            output = trace.usage.output_tokens,
3425            cost = trace.cost_usd,
3426            "trace recorded"
3427        );
3428    }
3429}
3430
3431#[cfg(test)]
3432mod tests {
3433    use super::*;
3434    use std::path::PathBuf;
3435
3436    fn tmpdir(name: &str) -> PathBuf {
3437        let dir = std::env::temp_dir().join(format!(
3438            "alex-proxy-{name}-{}-{}",
3439            std::process::id(),
3440            std::time::SystemTime::now()
3441                .duration_since(std::time::UNIX_EPOCH)
3442                .unwrap()
3443                .as_nanos()
3444        ));
3445        std::fs::create_dir_all(&dir).unwrap();
3446        dir
3447    }
3448
3449    fn test_state(name: &str) -> Arc<AppState> {
3450        let dir = tmpdir(name);
3451        let store = Arc::new(Store::open(dir.join("store")).unwrap());
3452        let vault = Arc::new(Vault::open(dir.join("vault")).unwrap());
3453        build_state(
3454            "alx-local".into(),
3455            vault,
3456            store,
3457            None,
3458            "http://127.0.0.1:4100".into(),
3459        )
3460    }
3461
3462    async fn response_json(resp: Response) -> (StatusCode, Value) {
3463        let status = resp.status();
3464        let body = axum::body::to_bytes(resp.into_body(), usize::MAX)
3465            .await
3466            .unwrap();
3467        let value = serde_json::from_slice(&body).unwrap_or(Value::Null);
3468        (status, value)
3469    }
3470
3471    struct FakeUpdater {
3472        result: Result<Value, UpdateApplyError>,
3473    }
3474
3475    impl DaemonUpdater for FakeUpdater {
3476        fn apply(&self) -> UpdateApplyFuture {
3477            let result = match &self.result {
3478                Ok(body) => Ok(body.clone()),
3479                Err(UpdateApplyError::Conflict(body)) => {
3480                    Err(UpdateApplyError::Conflict(body.clone()))
3481                }
3482                Err(UpdateApplyError::Failed(message)) => {
3483                    Err(UpdateApplyError::Failed(message.clone()))
3484                }
3485            };
3486            Box::pin(async move { result })
3487        }
3488    }
3489
3490    #[tokio::test]
3491    async fn admin_update_apply_response_shapes() {
3492        let state = test_state("admin-update-apply");
3493        set_daemon_updater(
3494            &state,
3495            Arc::new(FakeUpdater {
3496                result: Ok(json!({
3497                    "applying": true,
3498                    "current": "0.1.0",
3499                    "latest": "0.2.0",
3500                    "update_available": true,
3501                })),
3502            }),
3503        );
3504        let (status, body) = response_json(admin_update_apply(State(state.clone())).await).await;
3505        assert_eq!(status, StatusCode::ACCEPTED);
3506        assert_eq!(body["applying"], true);
3507        assert_eq!(body["latest"], "0.2.0");
3508
3509        set_daemon_updater(
3510            &state,
3511            Arc::new(FakeUpdater {
3512                result: Ok(json!({
3513                    "applying": false,
3514                    "current": "0.2.0",
3515                    "latest": "0.2.0",
3516                    "update_available": false,
3517                })),
3518            }),
3519        );
3520        let (status, body) = response_json(admin_update_apply(State(state.clone())).await).await;
3521        assert_eq!(status, StatusCode::OK);
3522        assert_eq!(body["update_available"], false);
3523
3524        set_daemon_updater(
3525            &state,
3526            Arc::new(FakeUpdater {
3527                result: Err(UpdateApplyError::Conflict(json!({
3528                    "applying": false,
3529                    "current": "0.1.0",
3530                    "latest": "0.2.0",
3531                    "update_available": true,
3532                    "reason": "alex is managed by Homebrew - run `brew upgrade alex`",
3533                }))),
3534            }),
3535        );
3536        let (status, body) = response_json(admin_update_apply(State(state)).await).await;
3537        assert_eq!(status, StatusCode::CONFLICT);
3538        assert_eq!(
3539            body["reason"],
3540            "alex is managed by Homebrew - run `brew upgrade alex`"
3541        );
3542    }
3543
3544    #[tokio::test]
3545    async fn harness_run_key_mints_lists_validates_and_revokes() {
3546        let state = test_state("harness-run-key");
3547        let (_status, created) = response_json(
3548            admin_run_keys_create(
3549                State(state.clone()),
3550                Some(axum::Json(json!({
3551                    "kind": "harness",
3552                    "label": "pi",
3553                    "ttl_seconds": 1,
3554                    "tags": {"harness": "pi"},
3555                }))),
3556            )
3557            .await,
3558        )
3559        .await;
3560        assert_eq!(_status, StatusCode::CREATED);
3561        assert_eq!(created["kind"], "harness");
3562        assert_eq!(created["label"], "pi");
3563        assert_eq!(created["expires_ms"], Value::Null);
3564        let id = created["id"].as_str().unwrap().to_string();
3565        let key = created["key"].as_str().unwrap();
3566        let hash = key_hash_hex(key);
3567
3568        let entry = run_key_entry(&state, &hash).unwrap();
3569        assert_eq!(entry.expires_ms, None);
3570        assert_eq!(entry.tags_json.as_deref(), Some(r#"{"harness":"pi"}"#));
3571
3572        let (_status, listed) = response_json(
3573            admin_run_keys_list(State(state.clone()), Query(HashMap::new())).await,
3574        )
3575        .await;
3576        assert_eq!(_status, StatusCode::OK);
3577        let rows = listed["run_keys"].as_array().unwrap();
3578        assert_eq!(rows.len(), 1);
3579        assert_eq!(rows[0]["id"], id);
3580        assert_eq!(rows[0]["kind"], "harness");
3581        assert_eq!(rows[0]["label"], "pi");
3582        assert_eq!(rows[0]["expires_ms"], Value::Null);
3583
3584        let (_status, revoked) =
3585            response_json(admin_run_keys_revoke(State(state.clone()), Path(id)).await).await;
3586        assert_eq!(_status, StatusCode::OK);
3587        assert_eq!(revoked["revoked"], true);
3588        assert!(run_key_entry(&state, &hash).is_none());
3589    }
3590
3591    #[tokio::test]
3592    async fn harness_run_key_requires_label() {
3593        let state = test_state("harness-run-key-label");
3594        let (status, body) = response_json(
3595            admin_run_keys_create(
3596                State(state),
3597                Some(axum::Json(json!({
3598                    "kind": "harness",
3599                    "tags": {"harness": "pi"},
3600                }))),
3601            )
3602            .await,
3603        )
3604        .await;
3605        assert_eq!(status, StatusCode::BAD_REQUEST);
3606        assert!(body["error"]["message"]
3607            .as_str()
3608            .unwrap()
3609            .contains("label"));
3610    }
3611}