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