alex_proxy/

lib.rs

use std::collections::{HashMap, HashSet};
use std::future::Future;
use std::path::PathBuf;
use std::pin::Pin;
use std::sync::Arc;
use std::time::Duration;

use alex_auth::{
    codex_reserve_blocked, codex_reserve_pct, codex_reset_selection, now_ms, Account,
    AccountPolicy, AccountPolicyMode, Vault,
};
use alex_core::{
    compute_cost, conversation_root, parse_grpc_web_response, parse_since, parse_sse_usage,
    parse_trace_tags, route_model, usage_from_json, validate_grpc_status_headers, window_label,
    ClientFormat, Provider, TraceRecord, GROK_CREDITS_ENDPOINT, GROK_CREDITS_REQUEST_BODY,
};
use alex_store::{Store, TraceFilter};
use axum::body::{Body, Bytes};
use axum::extract::{ConnectInfo, Path, Query, State};
use axum::http::{HeaderMap, HeaderValue, StatusCode};
use axum::response::{IntoResponse, Response};
use axum::routing::{get, post};
use axum::Router;
use futures_util::StreamExt;
use serde_json::{json, Value};
use tokio_stream::wrappers::ReceiverStream;

const ANTHROPIC_BASE: &str = "https://api.anthropic.com";
const OPENAI_BASE: &str = "https://api.openai.com";
const CODEX_BASE: &str = "https://chatgpt.com/backend-api/codex";
const XAI_BASE: &str = "https://cli-chat-proxy.grok.com/v1";
const GROK_CLIENT_VERSION: &str = "0.2.77";
const ANTHROPIC_OAUTH_BETA: &str = "oauth-2025-04-20";
const GEMINI_CODE_ASSIST_BASE: &str = "https://cloudcode-pa.googleapis.com";
const GEMINI_CODE_ASSIST_VERSION: &str = "v1internal";
const GEMINI_API_BASE: &str = "https://generativelanguage.googleapis.com";
const CODEX_AFFINITY_TTL_MS: i64 = 30 * 24 * 60 * 60 * 1000;
const CODEX_AFFINITY_MAX_ENTRIES: usize = 10_000;

#[derive(Debug, Clone)]
pub struct DarioActive {
    pub generation_id: String,
    pub base_url: String,
    pub api_key: String,
}

pub trait DarioRouter: Send + Sync {
    fn active(&self) -> Option<DarioActive>;
    fn begin(&self, generation_id: &str) -> Option<Box<dyn std::any::Any + Send>>;
    fn status(&self) -> Value;
    fn suspect(&self, generation_id: &str);
}

pub type UpdateApplyFuture =
    Pin<Box<dyn Future<Output = Result<Value, UpdateApplyError>> + Send + 'static>>;

#[derive(Debug)]
pub enum UpdateApplyError {
    Conflict(Value),
    Failed(String),
}

pub trait DaemonUpdater: Send + Sync {
    fn apply(&self) -> UpdateApplyFuture;
}

fn suspect_dario(state: &AppState, account: &Account) {
    if account.kind != "dario" {
        return;
    }
    if let (Some(dario), Some(gen)) = (&state.dario, account.id.strip_prefix("dario:")) {
        dario.suspect(gen);
    }
}

fn codex_limits_from_headers(headers: &reqwest::header::HeaderMap) -> Option<Value> {
    const LIMIT_HEADERS: &[&str] = &[
        "x-codex-primary-used-percent",
        "x-codex-primary-window-minutes",
        "x-codex-primary-reset-at",
        "x-codex-secondary-used-percent",
        "x-codex-secondary-window-minutes",
        "x-codex-secondary-reset-at",
        "x-codex-plan-type",
        "x-codex-active-limit",
        "x-codex-credits-balance",
        "x-codex-credits-has-credits",
        "x-codex-credits-unlimited",
    ];
    let mut safe = serde_json::Map::new();
    for name in LIMIT_HEADERS {
        if let Some(value) = headers.get(*name).and_then(|value| value.to_str().ok()) {
            safe.insert((*name).to_string(), json!(value));
        }
    }
    if safe.is_empty() {
        return None;
    }
    let parsed = alex_core::parse_limit_headers(Provider::Openai, &Value::Object(safe));
    let has_windows = parsed
        .get("windows")
        .and_then(Value::as_array)
        .map(|windows| !windows.is_empty())
        .unwrap_or(false);
    let has_plan = parsed.get("plan").is_some_and(|value| !value.is_null());
    (has_windows || has_plan).then_some(parsed)
}

pub struct AppState {
    pub local_key: String,
    pub vault: Arc<Vault>,
    pub store: Arc<Store>,
    pub http: reqwest::Client,
    pub dario: Option<Arc<dyn DarioRouter>>,
    pub in_flight: std::sync::atomic::AtomicI64,
    pub started_ms: i64,
    pub base_url: String,
    pub anthropic_usage: std::sync::Mutex<UsageCache>,
    pub xai_usage: std::sync::Mutex<UsageCache>,
    pub logins: alex_auth::sessions::LoginManager,
    pub run_keys: std::sync::RwLock<HashMap<String, CachedRunKey>>,
    pub update_status: Arc<tokio::sync::RwLock<Option<Value>>>,
    pub daemon_updater: std::sync::RwLock<Option<Arc<dyn DaemonUpdater>>>,
    codex_affinity: std::sync::Mutex<CodexAffinityCache>,
    codex_affinity_locks:
        std::sync::Mutex<HashMap<String, std::sync::Weak<tokio::sync::Mutex<()>>>>,
}

#[derive(Debug, Clone)]
pub struct CachedRunKey {
    pub run_id: Option<String>,
    pub tags_json: Option<String>,
    pub expires_ms: Option<i64>,
}

#[derive(Debug, Clone)]
struct CodexAffinityEntry {
    account_id: String,
    expires_at_ms: i64,
}

#[derive(Debug)]
struct CodexAffinityCache {
    entries: HashMap<String, CodexAffinityEntry>,
    ttl_ms: i64,
    max_entries: usize,
}

impl Default for CodexAffinityCache {
    fn default() -> Self {
        Self {
            entries: HashMap::new(),
            ttl_ms: CODEX_AFFINITY_TTL_MS,
            max_entries: CODEX_AFFINITY_MAX_ENTRIES,
        }
    }
}

impl CodexAffinityCache {
    fn preferred(&mut self, session_id: &str, now: i64) -> Option<String> {
        let entry = self.entries.get_mut(session_id)?;
        if entry.expires_at_ms <= now {
            self.entries.remove(session_id);
            return None;
        }
        // Sliding expiry keeps an active thread sticky without retaining
        // abandoned session IDs forever.
        entry.expires_at_ms = now.saturating_add(self.ttl_ms);
        Some(entry.account_id.clone())
    }

    fn bind(&mut self, session_id: &str, account_id: &str, now: i64) {
        self.entries.retain(|_, entry| entry.expires_at_ms > now);
        if !self.entries.contains_key(session_id) && self.entries.len() >= self.max_entries {
            if let Some(oldest) = self
                .entries
                .iter()
                .min_by_key(|(_, entry)| entry.expires_at_ms)
                .map(|(session, _)| session.clone())
            {
                self.entries.remove(&oldest);
            }
        }
        self.entries.insert(
            session_id.to_string(),
            CodexAffinityEntry {
                account_id: account_id.to_string(),
                expires_at_ms: now.saturating_add(self.ttl_ms),
            },
        );
    }

    fn unbind(&mut self, session_id: &str) {
        self.entries.remove(session_id);
    }
}

fn preferred_codex_account(state: &AppState, session_id: Option<&str>) -> Option<String> {
    let session_id = session_id.filter(|value| !value.is_empty())?;
    state
        .codex_affinity
        .lock()
        .unwrap_or_else(|poisoned| poisoned.into_inner())
        .preferred(session_id, now_ms())
}

fn codex_affinity_lock(state: &AppState, session_id: &str) -> Arc<tokio::sync::Mutex<()>> {
    let mut locks = state
        .codex_affinity_locks
        .lock()
        .unwrap_or_else(|poisoned| poisoned.into_inner());
    locks.retain(|_, lock| lock.strong_count() > 0);
    if let Some(lock) = locks.get(session_id).and_then(std::sync::Weak::upgrade) {
        return lock;
    }
    let lock = Arc::new(tokio::sync::Mutex::new(()));
    locks.insert(session_id.to_string(), Arc::downgrade(&lock));
    lock
}

fn bind_codex_account(state: &AppState, session_id: Option<&str>, account: &Account) {
    let Some(session_id) = session_id.filter(|value| !value.is_empty()) else {
        return;
    };
    if account.provider != Provider::Openai {
        return;
    }
    let mut affinity = state
        .codex_affinity
        .lock()
        .unwrap_or_else(|poisoned| poisoned.into_inner());
    if account.kind == "oauth" {
        affinity.bind(session_id, &account.id, now_ms());
    } else {
        // An API-key fallback is not a Codex subscription affinity. Forget a
        // stale subscription binding so the failed account is not restored as
        // soon as its cooldown ends.
        affinity.unbind(session_id);
    }
}

struct InFlight(Arc<AppState>);

impl InFlight {
    fn new(state: &Arc<AppState>) -> Self {
        state
            .in_flight
            .fetch_add(1, std::sync::atomic::Ordering::SeqCst);
        Self(state.clone())
    }
}

impl Drop for InFlight {
    fn drop(&mut self) {
        self.0
            .in_flight
            .fetch_sub(1, std::sync::atomic::Ordering::SeqCst);
    }
}

pub fn build_state(
    local_key: String,
    vault: Arc<Vault>,
    store: Arc<Store>,
    dario: Option<Arc<dyn DarioRouter>>,
    base_url: String,
) -> Arc<AppState> {
    let http = reqwest::Client::builder()
        .connect_timeout(Duration::from_secs(15))
        .build()
        .expect("reqwest client");
    Arc::new(AppState {
        local_key,
        vault,
        store,
        http,
        dario,
        in_flight: std::sync::atomic::AtomicI64::new(0),
        started_ms: now_ms(),
        base_url,
        anthropic_usage: std::sync::Mutex::new(UsageCache::default()),
        xai_usage: std::sync::Mutex::new(UsageCache::default()),
        logins: alex_auth::sessions::LoginManager::default(),
        run_keys: std::sync::RwLock::new(HashMap::new()),
        update_status: Arc::new(tokio::sync::RwLock::new(None)),
        daemon_updater: std::sync::RwLock::new(None),
        codex_affinity: std::sync::Mutex::new(CodexAffinityCache::default()),
        codex_affinity_locks: std::sync::Mutex::new(HashMap::new()),
    })
}

pub fn set_daemon_updater(state: &Arc<AppState>, updater: Arc<dyn DaemonUpdater>) {
    if let Ok(mut slot) = state.daemon_updater.write() {
        *slot = Some(updater);
    }
}

async fn require_local_key(
    State(state): State<Arc<AppState>>,
    req: axum::extract::Request,
    next: axum::middleware::Next,
) -> Response {
    let ok = client_key(req.headers())
        .map(|k| k == state.local_key)
        .unwrap_or(false);
    if !ok {
        return error_response(
            StatusCode::UNAUTHORIZED,
            "admin routes require x-api-key: <local_key>",
        );
    }
    next.run(req).await
}

pub fn router(state: Arc<AppState>) -> Router {
    // Control-plane + report routes: gated by the local key (sent as
    // `x-api-key: <key>` or `Authorization: Bearer <key>`) so a LAN/0.0.0.0
    // bind doesn't expose them. Run keys are NOT accepted here — a worker's
    // run key must not mint/revoke run keys or read the trace store.
    let gated = Router::new()
        .route(
            "/admin/run-keys",
            get(admin_run_keys_list).post(admin_run_keys_create),
        )
        .route(
            "/admin/run-keys/{id}",
            axum::routing::delete(admin_run_keys_revoke),
        )
        .route("/admin/storage", get(admin_storage))
        .route("/admin/storage/prune", post(admin_storage_prune))
        .route("/admin/traces", get(admin_traces))
        .route("/admin/accounts", get(admin_accounts))
        .route("/admin/accounts/analytics", get(admin_account_analytics))
        .route(
            "/admin/accounts/routing/openai",
            get(admin_codex_routing).put(admin_codex_routing_update),
        )
        .route(
            "/admin/accounts/{id}",
            axum::routing::delete(admin_account_remove).put(admin_account_update),
        )
        .route("/admin/auth/gemini-key", post(admin_auth_gemini_key))
        .route("/admin/health", get(admin_health))
        .route("/admin/analytics", get(admin_analytics))
        .route("/admin/limits", get(admin_limits))
        .route("/admin/update", get(admin_update).post(admin_update_apply))
        .route("/admin/dario", get(admin_dario))
        .route("/admin/dario/prompt-caches", get(admin_dario_prompt_caches))
        .route(
            "/admin/dario/prompt-caches/{key}",
            axum::routing::delete(admin_dario_prompt_cache_delete),
        )
        .route("/admin/auth/import", post(admin_auth_import))
        .route("/admin/auth/login/start", post(admin_auth_login_start))
        .route("/admin/auth/login/complete", post(admin_auth_login_complete))
        .route("/admin/auth/login/{id}", get(admin_auth_login_status))
        .route("/traces/search", get(traces_search))
        .route("/traces/export.ndjson", get(traces_export))
        .route("/traces/sessions", get(traces_sessions))
        .route(
            "/traces/sessions/{session_id}/transcript",
            get(traces_session_transcript),
        )
        .route("/traces/{id}", get(trace_get).delete(trace_delete))
        .route("/traces/{id}/reply.md", get(trace_reply_md))
        .route("/traces/{id}/body/{kind}", get(trace_body))
        .route("/traces/runs/{run_id}", get(traces_run_summary))
        .route("/traces/runs/{run_id}/events", get(traces_run_events))
        .route("/traces/runs/{run_id}/export.ndjson", get(traces_run_export))
        .route("/traces/runs/{run_id}/artifacts", get(traces_run_artifacts))
        .route_layer(axum::middleware::from_fn_with_state(
            state.clone(),
            require_local_key,
        ));

    Router::new()
        .route("/health", get(health))
        .route("/connect", get(connect_info))
        .route("/v1/models", get(models))
        .route("/v1/messages", post(anthropic_messages))
        .route("/v1/chat/completions", post(openai_chat))
        .route("/v1/responses", post(openai_responses))
        .route("/chat/completions", post(openai_chat))
        .route("/responses", post(openai_responses))
        .route("/v1beta/models/{model_action}", post(gemini_generate))
        .merge(gated)
        .layer(axum::extract::DefaultBodyLimit::max(64 * 1024 * 1024))
        .with_state(state)
}

async fn admin_auth_import(
    State(state): State<Arc<AppState>>,
    body: Option<axum::Json<Value>>,
) -> Response {
    let source = body
        .as_ref()
        .and_then(|b| b.0["source"].as_str().or(b.0["provider"].as_str()))
        .unwrap_or("all")
        .to_string();
    match alex_auth::import_all(&state.vault, &source).await {
        Ok(outcomes) => {
            let items: Vec<Value> = outcomes
                .iter()
                .map(|o| {
                    json!({
                        "source": o.source,
                        "imported": o.imported,
                        "note": o.note,
                    })
                })
                .collect();
            axum::Json(json!({"outcomes": items})).into_response()
        }
        Err(e) => error_response(StatusCode::BAD_REQUEST, &e.to_string()),
    }
}

async fn admin_account_remove(
    State(state): State<Arc<AppState>>,
    Path(id): Path<String>,
) -> Response {
    match state.vault.remove(&id).await {
        Ok(true) => axum::Json(json!({"removed": id})).into_response(),
        Ok(false) => error_response(StatusCode::NOT_FOUND, &format!("unknown account '{id}'")),
        Err(e) => error_response(StatusCode::INTERNAL_SERVER_ERROR, &e.to_string()),
    }
}

async fn admin_account_update(
    State(state): State<Arc<AppState>>,
    Path(id): Path<String>,
    body: axum::Json<Value>,
) -> Response {
    let Some(paused) = body.0["paused"].as_bool() else {
        return error_response(StatusCode::BAD_REQUEST, "missing boolean 'paused'");
    };
    match state.vault.set_paused(&id, paused).await {
        Ok(()) => axum::Json(json!({"id": id, "paused": paused})).into_response(),
        Err(e) if e.to_string().starts_with("unknown account") => {
            error_response(StatusCode::NOT_FOUND, &e.to_string())
        }
        Err(e) => error_response(StatusCode::INTERNAL_SERVER_ERROR, &e.to_string()),
    }
}

async fn admin_auth_gemini_key(
    State(state): State<Arc<AppState>>,
    body: axum::Json<Value>,
) -> Response {
    let Some(key) = body.0["key"].as_str().map(str::trim).filter(|k| !k.is_empty()) else {
        return error_response(StatusCode::BAD_REQUEST, "missing 'key'");
    };
    let account = Account {
        id: alex_auth::named_account_id(Provider::Gemini, "api_key", "default"),
        provider: Provider::Gemini,
        kind: "api_key".into(),
        name: "default".into(),
        description: None,
        paused: false,
        label: Some("gemini (AI Studio key)".into()),
        access_token: None,
        refresh_token: None,
        id_token: None,
        api_key: Some(key.to_string()),
        expires_at_ms: None,
        last_refresh_ms: None,
        account_meta: Value::Null,
        cooldown_until_ms: None,
        status: "active".into(),
        path: None,
    };
    match state.vault.upsert(account).await {
        Ok(()) => axum::Json(json!({"saved": "gemini-api-key"})).into_response(),
        Err(e) => error_response(StatusCode::INTERNAL_SERVER_ERROR, &e.to_string()),
    }
}

async fn admin_auth_login_start(
    State(state): State<Arc<AppState>>,
    body: axum::Json<Value>,
) -> Response {
    let Some(provider) = body.0["provider"].as_str() else {
        return error_response(StatusCode::BAD_REQUEST, "missing 'provider'");
    };
    if body.0["auto_identity"].as_bool() == Some(true) {
        return match state.logins.start_auto(state.vault.clone(), provider).await {
            Ok(snapshot) => axum::Json(snapshot).into_response(),
            Err(e) => error_response(StatusCode::BAD_REQUEST, &e.to_string()),
        };
    }
    let name = body.0["name"].as_str().unwrap_or("default");
    match state.logins.start(state.vault.clone(), provider, name).await {
        Ok(snapshot) => axum::Json(snapshot).into_response(),
        Err(e) => error_response(StatusCode::BAD_REQUEST, &e.to_string()),
    }
}

async fn admin_auth_login_complete(
    State(state): State<Arc<AppState>>,
    body: axum::Json<Value>,
) -> Response {
    let (Some(id), Some(input)) = (body.0["login_id"].as_str(), body.0["input"].as_str()) else {
        return error_response(StatusCode::BAD_REQUEST, "missing 'login_id' or 'input'");
    };
    match state.logins.complete(state.vault.clone(), id, input).await {
        Ok(snapshot) => axum::Json(snapshot).into_response(),
        Err(e) => error_response(StatusCode::BAD_REQUEST, &e.to_string()),
    }
}

async fn admin_auth_login_status(
    State(state): State<Arc<AppState>>,
    axum::extract::Path(id): axum::extract::Path<String>,
) -> Response {
    match state.logins.status(&id).await {
        Some(snapshot) => axum::Json(snapshot).into_response(),
        None => error_response(StatusCode::NOT_FOUND, "unknown or expired login session"),
    }
}

async fn health(State(state): State<Arc<AppState>>) -> impl IntoResponse {
    axum::Json(json!({
        "status": "ok",
        "service": "alexandria",
        "version": env!("CARGO_PKG_VERSION"),
        "in_flight": state.in_flight.load(std::sync::atomic::Ordering::SeqCst),
        "uptime_s": (now_ms() - state.started_ms) / 1000,
        "dario": state.dario.is_some(),
    }))
}

fn rewrite_host(base_url: &str, host: &str) -> String {
    let port = base_url.rsplit(':').next().unwrap_or("4100");
    format!("http://{host}:{port}")
}

pub fn connect_payload(base_url: &str, local_key: &str) -> (Value, String) {
    let base = base_url.trim_end_matches('/').to_string();
    let v1 = format!("{base}/v1");
    let exports = format!(
        "export ANTHROPIC_BASE_URL={base}\nexport ANTHROPIC_API_KEY={local_key}\n\
         export OPENAI_BASE_URL={v1}\nexport OPENAI_API_KEY={local_key}\n\
         export XAI_API_KEY={local_key}\nexport GROK_MODELS_BASE_URL={v1}\n\
         export GOOGLE_GEMINI_BASE_URL={base}\nexport GOOGLE_GENAI_API_VERSION=v1beta\n\
         export GEMINI_API_KEY={local_key}\nexport GEMINI_API_KEY_AUTH_MECHANISM=bearer\n\
         export GOOGLE_GENAI_USE_GCA=false\n"
    );
    let payload = json!({
        "service": "alexandria",
        "base_url": base,
        "api_key": local_key,
        "anthropic": {"base_url": base, "env": {"ANTHROPIC_BASE_URL": base, "ANTHROPIC_API_KEY": local_key}},
        "openai": {"base_url": v1, "env": {"OPENAI_BASE_URL": v1, "OPENAI_API_KEY": local_key}},
        "xai": {"base_url": v1, "env": {"XAI_API_KEY": local_key, "GROK_MODELS_BASE_URL": v1}},
        "gemini": {"base_url": base, "env": {
            "GOOGLE_GEMINI_BASE_URL": base,
            "GOOGLE_GENAI_API_VERSION": "v1beta",
            "GEMINI_API_KEY": local_key,
            "GEMINI_API_KEY_AUTH_MECHANISM": "bearer",
            "GOOGLE_GENAI_USE_GCA": "false",
        }},
        "exports": exports,
    });
    (payload, exports)
}

async fn connect_info(
    State(state): State<Arc<AppState>>,
    axum::extract::ConnectInfo(peer): axum::extract::ConnectInfo<std::net::SocketAddr>,
    Query(q): Query<HashMap<String, String>>,
) -> Response {
    if !peer.ip().is_loopback() {
        return error_response(
            StatusCode::FORBIDDEN,
            "connection info is only served to loopback clients",
        );
    }
    let base = match q.get("host").map(String::as_str) {
        Some("docker") => rewrite_host(&state.base_url, "host.docker.internal"),
        Some(host) if !host.is_empty() => rewrite_host(&state.base_url, host),
        _ => state.base_url.clone(),
    };
    let (payload, exports) = connect_payload(&base, &state.local_key);
    if q.get("format").map(|f| f == "env").unwrap_or(false) {
        return Response::builder()
            .status(StatusCode::OK)
            .header("content-type", "text/plain")
            .body(Body::from(exports))
            .unwrap_or_else(|e| error_response(StatusCode::INTERNAL_SERVER_ERROR, &e.to_string()));
    }
    axum::Json(payload).into_response()
}

async fn models(State(state): State<Arc<AppState>>) -> impl IntoResponse {
    let mut ids = state.store.pricing_models();
    for (alias, _) in alex_core::model_aliases() {
        ids.push((*alias).to_string());
    }
    for id in ids.clone() {
        ids.push(format!("alexandria/{id}"));
    }
    let data: Vec<Value> = ids
        .into_iter()
        .map(|m| json!({"id": m, "object": "model", "owned_by": "alexandria"}))
        .collect();
    axum::Json(json!({"object": "list", "data": data}))
}

async fn admin_analytics(
    State(state): State<Arc<AppState>>,
    Query(q): Query<HashMap<String, String>>,
) -> Response {
    let minutes: i64 = q
        .get("since_minutes")
        .and_then(|s| s.parse().ok())
        .unwrap_or(60);
    match state.store.analytics(now_ms() - minutes * 60_000) {
        Ok(v) => axum::Json(v).into_response(),
        Err(e) => error_response(StatusCode::INTERNAL_SERVER_ERROR, &e.to_string()),
    }
}

async fn admin_account_analytics(
    State(state): State<Arc<AppState>>,
    Query(q): Query<HashMap<String, String>>,
) -> Response {
    let minutes: i64 = q
        .get("since_minutes")
        .and_then(|s| s.parse().ok())
        .unwrap_or(24 * 60)
        .clamp(1, 30 * 24 * 60);
    let bucket_minutes: i64 = q
        .get("bucket_minutes")
        .and_then(|s| s.parse().ok())
        .unwrap_or(60)
        .clamp(1, 24 * 60);
    match state
        .store
        .account_analytics(now_ms() - minutes * 60_000, bucket_minutes * 60_000)
    {
        Ok(v) => axum::Json(v).into_response(),
        Err(e) => error_response(StatusCode::INTERNAL_SERVER_ERROR, &e.to_string()),
    }
}

const USAGE_CACHE_TTL_MS: i64 = 300_000;
const USAGE_BACKOFF_BASE_MS: i64 = 60_000;
const USAGE_BACKOFF_MAX_MS: i64 = 3_600_000;

#[derive(Default)]
pub struct UsageCache {
    fetched_at_ms: i64,
    entry: Option<Value>,
    cooldown_until_ms: i64,
    failures: u32,
}

fn usage_backoff_ms(failures: u32, retry_after_ms: Option<i64>) -> i64 {
    let exp = USAGE_BACKOFF_BASE_MS
        .saturating_mul(1i64 << failures.saturating_sub(1).min(6))
        .min(USAGE_BACKOFF_MAX_MS);
    exp.max(retry_after_ms.unwrap_or(0))
}

async fn anthropic_usage_entry(state: &Arc<AppState>) -> Option<Value> {
    let account = state.vault.account_for(Provider::Anthropic, true).await.ok()?;
    if account.kind != "oauth" {
        return None;
    }
    let token = account.access_token.as_deref()?.to_string();
    {
        let cache = state.anthropic_usage.lock().unwrap();
        if cache.entry.is_some() && now_ms() < cache.fetched_at_ms + USAGE_CACHE_TTL_MS {
            return cache.entry.clone();
        }
        if now_ms() < cache.cooldown_until_ms {
            return cache.entry.clone();
        }
    }
    let result = state
        .http
        .get(format!("{ANTHROPIC_BASE}/api/oauth/usage"))
        .header("authorization", format!("Bearer {token}"))
        .header("anthropic-beta", ANTHROPIC_OAUTH_BETA)
        .header("accept", "application/json")
        .header("user-agent", "claude-cli/2.1.202 (external, cli)")
        .send()
        .await;
    match result {
        Ok(resp) if resp.status().is_success() => {
            let raw: Value = resp.json().await.unwrap_or(Value::Null);
            let mut windows = Vec::new();
            for (name, key) in [
                ("5h", "five_hour"),
                ("7d", "seven_day"),
                ("7d opus", "seven_day_opus"),
                ("7d sonnet", "seven_day_sonnet"),
            ] {
                let w = &raw[key];
                if w.is_object() {
                    windows.push(json!({
                        "window": name,
                        "used_pct": w["utilization"],
                        "resets_at": w["resets_at"],
                    }));
                }
            }
            let entry = json!({
                "provider": "anthropic",
                "source": "oauth usage endpoint",
                "plan": account.label,
                "windows": windows,
                "extra_usage": raw["extra_usage"],
            });
            let mut cache = state.anthropic_usage.lock().unwrap();
            cache.fetched_at_ms = now_ms();
            cache.cooldown_until_ms = 0;
            cache.failures = 0;
            cache.entry = Some(entry.clone());
            Some(entry)
        }
        Ok(resp) => {
            let retry_after = resp
                .headers()
                .get("retry-after")
                .and_then(|v| v.to_str().ok())
                .and_then(|v| v.parse::<i64>().ok())
                .map(|s| s.clamp(30, 3600) * 1000);
            let mut cache = state.anthropic_usage.lock().unwrap();
            cache.failures += 1;
            let cooldown = usage_backoff_ms(cache.failures, retry_after);
            cache.cooldown_until_ms = now_ms() + cooldown;
            tracing::debug!(
                status = resp.status().as_u16(),
                failures = cache.failures,
                cooldown_ms = cooldown,
                "anthropic usage endpoint unavailable; backing off"
            );
            cache.entry.clone()
        }
        Err(_) => {
            let mut cache = state.anthropic_usage.lock().unwrap();
            cache.failures += 1;
            let cooldown = usage_backoff_ms(cache.failures, None);
            cache.cooldown_until_ms = now_ms() + cooldown;
            cache.entry.clone()
        }
    }
}

/// Fetch SuperGrok weekly credits from grok.com gRPC-web billing RPC.
/// Uses the vault's xAI OAuth access token. Degrades gracefully on any failure.
async fn xai_usage_entry(state: &Arc<AppState>) -> Option<Value> {
    let account = state.vault.account_for(Provider::Xai, true).await.ok()?;
    if account.kind != "oauth" {
        return None;
    }
    let token = account.access_token.as_deref()?.to_string();
    {
        let cache = state.xai_usage.lock().unwrap();
        if cache.entry.is_some() && now_ms() < cache.fetched_at_ms + USAGE_CACHE_TTL_MS {
            return cache.entry.clone();
        }
        if now_ms() < cache.cooldown_until_ms {
            return cache.entry.clone();
        }
    }

    let result = state
        .http
        .post(GROK_CREDITS_ENDPOINT)
        .header("authorization", format!("Bearer {token}"))
        .header("origin", "https://grok.com")
        .header("referer", "https://grok.com/?_s=usage")
        .header("accept", "*/*")
        .header("content-type", "application/grpc-web+proto")
        .header("x-grpc-web", "1")
        .header("x-user-agent", "connect-es/2.1.1")
        .header("user-agent", "Alexandria")
        .body(GROK_CREDITS_REQUEST_BODY.to_vec())
        .timeout(Duration::from_secs(15))
        .send()
        .await;

    match result {
        Ok(resp) if resp.status().is_success() => {
            let headers_for_grpc: Vec<(String, String)> = resp
                .headers()
                .iter()
                .filter_map(|(k, v)| {
                    let key = k.as_str();
                    if key.starts_with("grpc-") {
                        v.to_str().ok().map(|val| (key.to_string(), val.to_string()))
                    } else {
                        None
                    }
                })
                .collect();
            if let Err(e) = validate_grpc_status_headers(headers_for_grpc) {
                tracing::debug!(error = %e, "xai grok credits grpc header status failed");
                let mut cache = state.xai_usage.lock().unwrap();
                cache.failures += 1;
                let cooldown = usage_backoff_ms(cache.failures, None);
                cache.cooldown_until_ms = now_ms() + cooldown;
                return cache.entry.clone();
            }
            let body = match resp.bytes().await {
                Ok(b) => b,
                Err(e) => {
                    tracing::debug!(error = %e, "xai grok credits body read failed");
                    let mut cache = state.xai_usage.lock().unwrap();
                    cache.failures += 1;
                    let cooldown = usage_backoff_ms(cache.failures, None);
                    cache.cooldown_until_ms = now_ms() + cooldown;
                    return cache.entry.clone();
                }
            };
            let now_s = now_ms() / 1000;
            match parse_grpc_web_response(&body, now_s) {
                Ok(snap) => {
                    let label = window_label(snap.resets_at_s, now_s);
                    let mut window = json!({
                        "window": label,
                        "used_pct": snap.used_percent,
                    });
                    if let Some(ts) = snap.resets_at_s {
                        window["resets_at_s"] = json!(ts);
                    }
                    let entry = json!({
                        "provider": "xai",
                        "source": "grok web billing",
                        "plan": account.label,
                        "windows": [window],
                    });
                    let mut cache = state.xai_usage.lock().unwrap();
                    cache.fetched_at_ms = now_ms();
                    cache.cooldown_until_ms = 0;
                    cache.failures = 0;
                    cache.entry = Some(entry.clone());
                    Some(entry)
                }
                Err(e) => {
                    tracing::debug!(error = %e, "xai grok credits parse failed");
                    let mut cache = state.xai_usage.lock().unwrap();
                    cache.failures += 1;
                    let cooldown = usage_backoff_ms(cache.failures, None);
                    cache.cooldown_until_ms = now_ms() + cooldown;
                    cache.entry.clone()
                }
            }
        }
        Ok(resp) => {
            let status = resp.status().as_u16();
            let retry_after = resp
                .headers()
                .get("retry-after")
                .and_then(|v| v.to_str().ok())
                .and_then(|v| v.parse::<i64>().ok())
                .map(|s| s.clamp(30, 3600) * 1000);
            let mut cache = state.xai_usage.lock().unwrap();
            cache.failures += 1;
            let cooldown = usage_backoff_ms(cache.failures, retry_after);
            cache.cooldown_until_ms = now_ms() + cooldown;
            tracing::debug!(
                status,
                failures = cache.failures,
                cooldown_ms = cooldown,
                "xai grok web billing unavailable; backing off"
            );
            cache.entry.clone()
        }
        Err(e) => {
            tracing::debug!(error = %e, "xai grok web billing request failed");
            let mut cache = state.xai_usage.lock().unwrap();
            cache.failures += 1;
            let cooldown = usage_backoff_ms(cache.failures, None);
            cache.cooldown_until_ms = now_ms() + cooldown;
            cache.entry.clone()
        }
    }
}

pub async fn limits_snapshot(state: &Arc<AppState>) -> Value {
    let mut providers: Vec<Value> = Vec::new();
    if let Some(entry) = anthropic_usage_entry(state).await {
        providers.push(entry);
    }
    if let Some(entry) = xai_usage_entry(state).await {
        providers.push(entry);
    }
    for (provider_str, ts_ms, headers_json) in
        state.store.latest_provider_headers().unwrap_or_default()
    {
        if providers
            .iter()
            .any(|p| p["provider"].as_str() == Some(&provider_str))
        {
            continue;
        }
        let Some(provider) = Provider::from_str_loose(&provider_str) else {
            continue;
        };
        let headers: Value = serde_json::from_str(&headers_json).unwrap_or(Value::Null);
        let mut parsed = alex_core::parse_limit_headers(provider, &headers);
        if let Some(o) = parsed.as_object_mut() {
            o.insert("provider".into(), json!(provider_str));
            o.insert("source".into(), json!("captured response headers"));
            o.insert("observed_at_ms".into(), json!(ts_ms));
            providers.push(parsed);
        }
    }
    providers.sort_by_key(|p| p["provider"].as_str().unwrap_or("").to_string());
    json!({"providers": providers})
}

async fn admin_limits(State(state): State<Arc<AppState>>) -> Response {
    axum::Json(limits_snapshot(&state).await).into_response()
}

async fn admin_update(State(state): State<Arc<AppState>>) -> Response {
    let stored = state.update_status.read().await.clone();
    let mut body = json!({
        "current": env!("CARGO_PKG_VERSION"),
        "latest": null,
        "update_available": false,
        "checked_at_ms": null,
    });
    if let Some(Value::Object(fields)) = stored {
        if let Some(obj) = body.as_object_mut() {
            obj.extend(fields);
            obj.insert("current".into(), json!(env!("CARGO_PKG_VERSION")));
        }
    }
    axum::Json(body).into_response()
}

async fn admin_update_apply(State(state): State<Arc<AppState>>) -> Response {
    let updater = state
        .daemon_updater
        .read()
        .ok()
        .and_then(|slot| slot.clone());
    let Some(updater) = updater else {
        return error_response(StatusCode::SERVICE_UNAVAILABLE, "daemon updater is not configured");
    };
    match updater.apply().await {
        Ok(body) => {
            let status = if body["applying"].as_bool() == Some(true) {
                StatusCode::ACCEPTED
            } else {
                StatusCode::OK
            };
            (status, axum::Json(body)).into_response()
        }
        Err(UpdateApplyError::Conflict(body)) => {
            (StatusCode::CONFLICT, axum::Json(body)).into_response()
        }
        Err(UpdateApplyError::Failed(message)) => {
            error_response(StatusCode::INTERNAL_SERVER_ERROR, &message)
        }
    }
}

async fn admin_dario(State(state): State<Arc<AppState>>) -> Response {
    match &state.dario {
        Some(d) => {
            let mut status = d.status();
            if let Some(obj) = status.as_object_mut() {
                obj.insert("prompt_caches".into(), json!(dario_prompt_caches(&state)));
            }
            axum::Json(status).into_response()
        }
        None => error_response(StatusCode::NOT_FOUND, "dario mode is not enabled"),
    }
}

fn dario_prompt_cache_dir(state: &AppState) -> PathBuf {
    state.store.data_dir.join("dario-prompt-cache")
}

fn dario_prompt_cache_key_valid(key: &str) -> bool {
    !key.is_empty()
        && key.len() <= 128
        && key
            .bytes()
            .all(|b| b.is_ascii_alphanumeric() || matches!(b, b'.' | b'_' | b'-'))
}

fn dario_prompt_cache_summary(path: PathBuf) -> Option<Value> {
    let raw = std::fs::read_to_string(&path).ok()?;
    let value: Value = serde_json::from_str(&raw).ok()?;
    let key = path.file_stem()?.to_string_lossy().to_string();
    let runs: Vec<Value> = value["runs"]
        .as_array()
        .map(|runs| runs.iter().rev().take(12).cloned().collect())
        .unwrap_or_default();
    Some(json!({
        "key": value["key"].as_str().unwrap_or(&key),
        "model": value["model"],
        "source": value["source"],
        "captured_at": value["captured_at"],
        "last_used_at": value["last_used_at"],
        "trace_id": value["trace_id"],
        "claude_bin": value["claude_bin"],
        "claude_version": value["claude_version"],
        "system_prompt_chars": value["system_prompt_chars"],
        "agent_identity_chars": value["agent_identity_chars"],
        "path": path.to_string_lossy(),
        "runs": runs,
    }))
}

fn dario_prompt_caches(state: &AppState) -> Vec<Value> {
    let dir = dario_prompt_cache_dir(state);
    let mut caches: Vec<Value> = std::fs::read_dir(dir)
        .ok()
        .into_iter()
        .flatten()
        .flatten()
        .filter_map(|entry| {
            let path = entry.path();
            (path.extension().and_then(|s| s.to_str()) == Some("json"))
                .then(|| dario_prompt_cache_summary(path))
                .flatten()
        })
        .collect();
    caches.sort_by(|a, b| {
        b["last_used_at"]
            .as_str()
            .unwrap_or("")
            .cmp(a["last_used_at"].as_str().unwrap_or(""))
    });
    caches
}

async fn admin_dario_prompt_caches(State(state): State<Arc<AppState>>) -> Response {
    axum::Json(json!({"prompt_caches": dario_prompt_caches(&state)})).into_response()
}

async fn admin_dario_prompt_cache_delete(
    State(state): State<Arc<AppState>>,
    Path(key): Path<String>,
) -> Response {
    if !dario_prompt_cache_key_valid(&key) {
        return error_response(StatusCode::BAD_REQUEST, "invalid prompt cache key");
    }
    let path = dario_prompt_cache_dir(&state).join(format!("{key}.json"));
    match std::fs::remove_file(&path) {
        Ok(()) => axum::Json(json!({"deleted": true, "key": key})).into_response(),
        Err(e) if e.kind() == std::io::ErrorKind::NotFound => {
            error_response(StatusCode::NOT_FOUND, "unknown prompt cache")
        }
        Err(e) => error_response(StatusCode::INTERNAL_SERVER_ERROR, &e.to_string()),
    }
}

async fn admin_storage(State(state): State<Arc<AppState>>) -> Response {
    match state.store.disk_usage() {
        Ok(v) => axum::Json(v).into_response(),
        Err(e) => error_response(StatusCode::INTERNAL_SERVER_ERROR, &e.to_string()),
    }
}

async fn admin_storage_prune(
    State(state): State<Arc<AppState>>,
    body: axum::Json<Value>,
) -> Response {
    let now = now_ms();
    let cutoff = match &body["older_than_ms"] {
        Value::Null => match body["older_than"].as_str() {
            Some(s) => match parse_since(s, now) {
                Some(ms) => ms,
                None => {
                    return error_response(
                        StatusCode::BAD_REQUEST,
                        &format!("invalid 'older_than' '{s}' (use 45s, 30m, 24h, 7d, or RFC3339)"),
                    )
                }
            },
            None => {
                return error_response(
                    StatusCode::BAD_REQUEST,
                    "missing 'older_than_ms' or 'older_than'",
                )
            }
        },
        v => match v.as_i64() {
            Some(ms) => ms,
            None => {
                return error_response(StatusCode::BAD_REQUEST, "'older_than_ms' must be an integer")
            }
        },
    };
    if cutoff > now {
        return error_response(StatusCode::BAD_REQUEST, "cutoff is in the future");
    }
    let bodies_only = body["bodies_only"].as_bool().unwrap_or(true);
    let dry_run = body["dry_run"].as_bool().unwrap_or(false);
    let store = state.store.clone();
    let report =
        tokio::task::spawn_blocking(move || store.prune(cutoff, bodies_only, dry_run)).await;
    match report {
        Ok(Ok(r)) => axum::Json(serde_json::to_value(r).unwrap_or_default()).into_response(),
        Ok(Err(e)) => error_response(StatusCode::INTERNAL_SERVER_ERROR, &e.to_string()),
        Err(e) => error_response(StatusCode::INTERNAL_SERVER_ERROR, &e.to_string()),
    }
}

async fn admin_traces(
    State(state): State<Arc<AppState>>,
    Query(q): Query<HashMap<String, String>>,
) -> Response {
    let limit = q.get("limit").and_then(|s| s.parse().ok()).unwrap_or(50);
    match state.store.list_traces(
        limit,
        q.get("session").map(|s| s.as_str()),
        q.get("model").map(|s| s.as_str()),
    ) {
        Ok(rows) => axum::Json(json!({"traces": rows})).into_response(),
        Err(e) => error_response(StatusCode::INTERNAL_SERVER_ERROR, &e.to_string()),
    }
}

fn filter_from_query(q: &HashMap<String, String>) -> TraceFilter {
    let now = now_ms();
    TraceFilter {
        since_ms: q.get("since").and_then(|s| parse_since(s, now)),
        until_ms: q.get("until").and_then(|s| parse_since(s, now)),
        run_id: q.get("run_id").cloned(),
        session: q.get("session").cloned(),
        model: q.get("model").cloned(),
        provider: q.get("provider").cloned(),
        path: q.get("path").cloned(),
        harness: q.get("harness").cloned(),
        status: q.get("status").and_then(|s| s.parse().ok()),
        errors_only: q
            .get("errors")
            .map(|v| v == "1" || v == "true")
            .unwrap_or(false),
        key_fingerprint: q.get("key_fingerprint").cloned(),
        reasoning_effort: q.get("effort").cloned(),
        limit: q.get("limit").and_then(|s| s.parse().ok()).unwrap_or(200),
    }
}

fn wants_bodies(q: &HashMap<String, String>) -> bool {
    q.get("bodies").map(|v| v == "1" || v == "true").unwrap_or(false)
}

fn read_gz_file(path: &str) -> Option<Vec<u8>> {
    let file = std::fs::File::open(path).ok()?;
    let mut decoder = flate2::read::GzDecoder::new(file);
    let mut buf = Vec::new();
    std::io::Read::read_to_end(&mut decoder, &mut buf).ok()?;
    Some(buf)
}

fn inline_row_bodies(row: &mut Value) {
    use base64::Engine;
    for (path_key, out_key) in [
        ("req_body_path", "req_body_b64"),
        ("upstream_req_body_path", "upstream_req_body_b64"),
        ("resp_body_path", "resp_body_b64"),
    ] {
        let Some(buf) = row[path_key].as_str().and_then(read_gz_file) else {
            continue;
        };
        row[out_key] = json!(base64::engine::general_purpose::STANDARD.encode(&buf));
    }
}

fn ndjson_response(mut rows: Vec<Value>, inline_bodies: bool) -> Response {
    rows.sort_by_key(|r| r["ts_request_ms"].as_i64().unwrap_or(0));
    let mut out = String::new();
    for mut row in rows {
        if inline_bodies {
            inline_row_bodies(&mut row);
        }
        out.push_str(&serde_json::to_string(&row).unwrap_or_default());
        out.push('\n');
    }
    Response::builder()
        .status(StatusCode::OK)
        .header("content-type", "application/x-ndjson")
        .body(Body::from(out))
        .unwrap_or_else(|e| error_response(StatusCode::INTERNAL_SERVER_ERROR, &e.to_string()))
}

const TEXT_SCAN_CAP: usize = 300;

fn trace_matches_text(row: &Value, needle: &str) -> bool {
    for key in ["req_body_path", "resp_body_path"] {
        if let Some(path) = row.get(key).and_then(|v| v.as_str()) {
            if let Some(bytes) = read_gz_file(path) {
                if String::from_utf8_lossy(&bytes)
                    .to_lowercase()
                    .contains(needle)
                {
                    return true;
                }
            }
        }
    }
    false
}

async fn traces_search(
    State(state): State<Arc<AppState>>,
    Query(q): Query<HashMap<String, String>>,
) -> Response {
    let mut filter = filter_from_query(&q);
    let text = q
        .get("text")
        .or_else(|| q.get("q"))
        .map(|t| t.trim().to_lowercase())
        .filter(|t| !t.is_empty());
    if text.is_some() {
        filter.limit = TEXT_SCAN_CAP;
    }
    match state.store.search_traces(&filter) {
        Ok(rows) => match text {
            Some(needle) => {
                let scanned = rows.len();
                let rows = tokio::task::spawn_blocking(move || {
                    rows.into_iter()
                        .filter(|r| trace_matches_text(r, &needle))
                        .collect::<Vec<_>>()
                })
                .await
                .unwrap_or_default();
                axum::Json(json!({"traces": rows, "scanned": scanned, "scan_cap": TEXT_SCAN_CAP}))
                    .into_response()
            }
            None => axum::Json(json!({"traces": rows})).into_response(),
        },
        Err(e) => error_response(StatusCode::INTERNAL_SERVER_ERROR, &e.to_string()),
    }
}

async fn traces_export(
    State(state): State<Arc<AppState>>,
    Query(q): Query<HashMap<String, String>>,
) -> Response {
    match state.store.search_traces(&filter_from_query(&q)) {
        Ok(rows) => ndjson_response(rows, wants_bodies(&q)),
        Err(e) => error_response(StatusCode::INTERNAL_SERVER_ERROR, &e.to_string()),
    }
}

fn truncate_chars(s: String, max: usize) -> String {
    if s.chars().count() <= max {
        s
    } else {
        s.chars().take(max).collect()
    }
}

fn read_gz_json(path: Option<&str>) -> Option<Value> {
    let buf = path.and_then(read_gz_file)?;
    serde_json::from_slice(&buf).ok()
}

fn read_gz_text(path: Option<&str>) -> Option<String> {
    let buf = path.and_then(read_gz_file)?;
    Some(String::from_utf8_lossy(&buf).to_string())
}

fn body_date_dir_name(name: &str) -> bool {
    let bytes = name.as_bytes();
    bytes.len() == 10
        && bytes[4] == b'-'
        && bytes[7] == b'-'
        && bytes
            .iter()
            .enumerate()
            .all(|(i, b)| i == 4 || i == 7 || b.is_ascii_digit())
}

fn dario_capture_suffix(kind: &str) -> Option<&'static str> {
    match kind {
        "dario-upstream-request" => Some("dario-upstream-request.json.gz"),
        "dario-upstream-response" => Some("dario-upstream-response.json.gz"),
        _ => None,
    }
}

fn is_dario_trace(row: &Value) -> bool {
    row["account_id"]
        .as_str()
        .map(|id| id.starts_with("dario:"))
        .unwrap_or(false)
}

fn find_dario_capture_path(state: &AppState, row: &Value, kind: &str) -> Option<String> {
    if !is_dario_trace(row) {
        return None;
    }
    let trace_id = row["id"].as_str()?;
    let suffix = dario_capture_suffix(kind)?;
    let filename = format!("{trace_id}.{suffix}");
    let bodies = state.store.data_dir.join("bodies");
    let mut days: Vec<PathBuf> = std::fs::read_dir(&bodies)
        .ok()?
        .flatten()
        .filter_map(|entry| {
            let name = entry.file_name().to_string_lossy().to_string();
            (entry.path().is_dir() && body_date_dir_name(&name)).then(|| entry.path())
        })
        .collect();
    days.sort_by(|a, b| b.file_name().cmp(&a.file_name()));
    for day in days {
        let path = day.join(&filename);
        if path.is_file() {
            return Some(path.to_string_lossy().to_string());
        }
    }
    None
}

fn dario_capture_summary(state: &AppState, row: &Value) -> Option<Value> {
    let request_path = find_dario_capture_path(state, row, "dario-upstream-request");
    let response_path = find_dario_capture_path(state, row, "dario-upstream-response");
    if request_path.is_none() && response_path.is_none() {
        return None;
    }
    let prompt_cache = request_path
        .as_deref()
        .and_then(|path| read_gz_json(Some(path)))
        .and_then(|body| {
            body["prompt_cache"]
                .as_object()
                .map(|_| body["prompt_cache"].clone())
        });
    Some(json!({
        "request_available": request_path.is_some(),
        "response_available": response_path.is_some(),
        "request_path": request_path,
        "response_path": response_path,
        "prompt_cache": prompt_cache,
    }))
}

async fn traces_sessions(
    State(state): State<Arc<AppState>>,
    Query(q): Query<HashMap<String, String>>,
) -> Response {
    let since = q.get("since").and_then(|s| parse_since(s, now_ms()));
    let limit = q.get("limit").and_then(|s| s.parse().ok()).unwrap_or(0);
    match state.store.sessions(since, limit) {
        Ok(rows) => axum::Json(json!({"sessions": rows})).into_response(),
        Err(e) => error_response(StatusCode::INTERNAL_SERVER_ERROR, &e.to_string()),
    }
}

fn transcript_turn(row: &Value) -> Value {
    use alex_core::translate;
    let user = read_gz_json(row["req_body_path"].as_str())
        .and_then(|req| {
            translate::last_user_text(row["client_format"].as_str().unwrap_or(""), &req)
        })
        .map(|s| truncate_chars(s, 8000));
    let resp_text = read_gz_text(row["resp_body_path"].as_str());
    let fmt = row["upstream_format"]
        .as_str()
        .or(row["client_format"].as_str())
        .unwrap_or("")
        .to_string();
    let assistant = resp_text
        .as_deref()
        .and_then(|text| translate::assistant_reply_text(&fmt, text))
        .map(|s| truncate_chars(s, 8000));
    let tool_calls: Vec<Value> = resp_text
        .as_deref()
        .map(|text| translate::assistant_tool_calls(&fmt, text))
        .unwrap_or_default()
        .into_iter()
        .take(24)
        .map(|mut c| {
            if let Some(a) = c["arguments"].as_str() {
                let t = truncate_chars(a.to_string(), 600);
                c["arguments"] = json!(t);
            }
            c
        })
        .collect();
    json!({
        "trace_id": row["id"],
        "ts_request_ms": row["ts_request_ms"],
        "ts_response_ms": row["ts_response_ms"],
        "model": row["routed_model"],
        "status": row["status"],
        "input_tokens": row["input_tokens"],
        "output_tokens": row["output_tokens"],
        "reasoning_effort": row["reasoning_effort"],
        "thinking_budget": row["thinking_budget"],
        "cost_usd": row["cost_usd"],
        "account_id": row["account_id"],
        "error": row["error"],
        "user": user,
        "assistant": assistant,
        "tool_calls": tool_calls,
    })
}

async fn traces_session_transcript(
    State(state): State<Arc<AppState>>,
    Path(session_id): Path<String>,
    Query(q): Query<HashMap<String, String>>,
) -> Response {
    let since = q
        .get("since_ms")
        .and_then(|s| s.parse::<i64>().ok())
        .or_else(|| q.get("since").and_then(|s| parse_since(s, now_ms())));
    let limit: usize = q.get("limit").and_then(|s| s.parse().ok()).unwrap_or(500);
    let rows = match state.store.session_traces(&session_id, since) {
        Ok(rows) => rows,
        Err(e) => return error_response(StatusCode::INTERNAL_SERVER_ERROR, &e.to_string()),
    };
    let turns: Vec<Value> = rows.iter().take(limit).map(transcript_turn).collect();
    axum::Json(json!({"session_id": session_id, "turns": turns})).into_response()
}

fn trace_reasoning_fields(req: &Value) -> (Option<String>, Option<i64>) {
    (
        req["reasoning"]["effort"].as_str().map(String::from),
        req["thinking"]["budget_tokens"].as_i64(),
    )
}

fn trace_extras(req: &Value) -> Value {
    let system_text: Option<String> = match &req["system"] {
        Value::String(s) => Some(s.clone()),
        Value::Array(parts) => {
            let joined: Vec<&str> = parts.iter().filter_map(|p| p["text"].as_str()).collect();
            (!joined.is_empty()).then(|| joined.join("\n\n"))
        }
        _ => req["instructions"].as_str().map(String::from).or_else(|| {
            let joined: Vec<String> = req["messages"]
                .as_array()
                .into_iter()
                .flatten()
                .filter(|m| matches!(m["role"].as_str(), Some("system") | Some("developer")))
                .map(|m| match &m["content"] {
                    Value::String(s) => s.clone(),
                    Value::Array(parts) => parts
                        .iter()
                        .filter_map(|p| p["text"].as_str())
                        .collect::<Vec<_>>()
                        .join("\n"),
                    _ => String::new(),
                })
                .filter(|s| !s.is_empty())
                .collect();
            (!joined.is_empty()).then(|| joined.join("\n\n"))
        }),
    };
    let system_chars = system_text.as_ref().map(|s| s.chars().count());
    let system_prompt = system_text.map(|s| truncate_chars(s, 64_000));
    let (reasoning_effort, thinking_budget) = trace_reasoning_fields(req);
    json!({
        "reasoning_effort": reasoning_effort,
        "thinking_budget": thinking_budget,
        "max_tokens": req["max_tokens"].as_i64().or(req["max_output_tokens"].as_i64()),
        "temperature": req["temperature"],
        "message_count": req["messages"].as_array().or(req["input"].as_array()).map(|a| a.len()),
        "system_chars": system_chars,
        "system_prompt": system_prompt,
    })
}

async fn trace_get(State(state): State<Arc<AppState>>, Path(id): Path<String>) -> Response {
    match state.store.get_trace(&id) {
        Ok(Some(row)) => {
            let mut extras = read_gz_json(row["req_body_path"].as_str())
                .map(|req| trace_extras(&req))
                .unwrap_or_else(|| json!({}));
            if let Some(summary) = dario_capture_summary(&state, &row) {
                if !extras.is_object() {
                    extras = json!({});
                }
                if let Some(obj) = extras.as_object_mut() {
                    obj.insert("dario_capture".into(), summary);
                }
            }
            axum::Json(json!({"trace": row, "extras": extras})).into_response()
        }
        Ok(None) => error_response(StatusCode::NOT_FOUND, &format!("unknown trace '{id}'")),
        Err(e) => error_response(StatusCode::INTERNAL_SERVER_ERROR, &e.to_string()),
    }
}

async fn trace_body(
    State(state): State<Arc<AppState>>,
    Path((id, kind)): Path<(String, String)>,
) -> Response {
    let row = match state.store.get_trace(&id) {
        Ok(Some(row)) => row,
        Ok(None) => {
            return error_response(StatusCode::NOT_FOUND, &format!("unknown trace '{id}'"))
        }
        Err(e) => return error_response(StatusCode::INTERNAL_SERVER_ERROR, &e.to_string()),
    };
    let path = match kind.as_str() {
        "request" => row["req_body_path"].as_str().map(String::from),
        "upstream-request" => row["upstream_req_body_path"].as_str().map(String::from),
        "response" => row["resp_body_path"].as_str().map(String::from),
        "dario-upstream-request" | "dario-upstream-response" => {
            find_dario_capture_path(&state, &row, &kind)
        }
        _ => {
            return error_response(
                StatusCode::BAD_REQUEST,
                "kind must be request|upstream-request|response|dario-upstream-request|dario-upstream-response",
            )
        }
    };
    match read_gz_text(path.as_deref()) {
        Some(text) => {
            let ct = if text.trim_start().starts_with('{') || text.trim_start().starts_with('[') {
                "application/json; charset=utf-8"
            } else {
                "text/plain; charset=utf-8"
            };
            Response::builder()
                .status(StatusCode::OK)
                .header("content-type", ct)
                .header("x-alexandria-body-path", path.as_deref().unwrap_or(""))
                .body(Body::from(text))
                .unwrap_or_else(|e| error_response(StatusCode::INTERNAL_SERVER_ERROR, &e.to_string()))
        }
        None => error_response(
            StatusCode::NOT_FOUND,
            &format!("no {kind} body stored for trace '{id}'"),
        ),
    }
}

async fn trace_reply_md(State(state): State<Arc<AppState>>, Path(id): Path<String>) -> Response {
    use alex_core::translate;
    let row = match state.store.get_trace(&id) {
        Ok(Some(row)) => row,
        Ok(None) => {
            return error_response(StatusCode::NOT_FOUND, &format!("unknown trace '{id}'"))
        }
        Err(e) => return error_response(StatusCode::INTERNAL_SERVER_ERROR, &e.to_string()),
    };
    let fmt = row["upstream_format"]
        .as_str()
        .or(row["client_format"].as_str())
        .unwrap_or("");
    let reply = read_gz_text(row["resp_body_path"].as_str())
        .and_then(|text| translate::assistant_reply_text(fmt, &text));
    match reply {
        Some(md) => Response::builder()
            .status(StatusCode::OK)
            .header("content-type", "text/markdown; charset=utf-8")
            .body(Body::from(md))
            .unwrap_or_else(|e| error_response(StatusCode::INTERNAL_SERVER_ERROR, &e.to_string())),
        None => error_response(
            StatusCode::NOT_FOUND,
            &format!("no assistant reply available for trace '{id}'"),
        ),
    }
}

async fn trace_delete(State(state): State<Arc<AppState>>, Path(id): Path<String>) -> Response {
    match state.store.get_trace(&id) {
        Ok(Some(_)) => {}
        Ok(None) => {
            return error_response(StatusCode::NOT_FOUND, &format!("unknown trace '{id}'"))
        }
        Err(e) => return error_response(StatusCode::INTERNAL_SERVER_ERROR, &e.to_string()),
    }
    match state.store.delete_trace(&id) {
        Ok(paths) => {
            let removed = paths
                .iter()
                .filter(|p| std::fs::remove_file(p).is_ok())
                .count();
            axum::Json(json!({"deleted": true, "files_removed": removed})).into_response()
        }
        Err(e) => error_response(StatusCode::INTERNAL_SERVER_ERROR, &e.to_string()),
    }
}

async fn traces_run_summary(
    State(state): State<Arc<AppState>>,
    Path(run_id): Path<String>,
) -> Response {
    match state.store.run_summary(&run_id) {
        Ok(v) if v["trace_count"].as_i64().unwrap_or(0) == 0 => error_response(
            StatusCode::NOT_FOUND,
            &format!("no traces for run '{run_id}'"),
        ),
        Ok(v) => axum::Json(v).into_response(),
        Err(e) => error_response(StatusCode::INTERNAL_SERVER_ERROR, &e.to_string()),
    }
}

async fn traces_run_events(
    State(state): State<Arc<AppState>>,
    Path(run_id): Path<String>,
    Query(q): Query<HashMap<String, String>>,
) -> Response {
    let filter = TraceFilter {
        run_id: Some(run_id),
        limit: q.get("limit").and_then(|s| s.parse().ok()).unwrap_or(1000),
        ..Default::default()
    };
    match state.store.search_traces(&filter) {
        Ok(rows) => axum::Json(json!({"traces": rows})).into_response(),
        Err(e) => error_response(StatusCode::INTERNAL_SERVER_ERROR, &e.to_string()),
    }
}

async fn traces_run_export(
    State(state): State<Arc<AppState>>,
    Path(run_id): Path<String>,
    Query(q): Query<HashMap<String, String>>,
) -> Response {
    let filter = TraceFilter {
        run_id: Some(run_id),
        limit: q.get("limit").and_then(|s| s.parse().ok()).unwrap_or(5000),
        ..Default::default()
    };
    match state.store.search_traces(&filter) {
        Ok(rows) => ndjson_response(rows, wants_bodies(&q)),
        Err(e) => error_response(StatusCode::INTERNAL_SERVER_ERROR, &e.to_string()),
    }
}

async fn traces_run_artifacts(
    State(state): State<Arc<AppState>>,
    Path(run_id): Path<String>,
) -> Response {
    match state.store.run_artifacts(&run_id) {
        Ok(artifacts) => axum::Json(json!({
            "run_id": run_id,
            "data_dir": state.store.data_dir.to_string_lossy(),
            "artifacts": artifacts,
        }))
        .into_response(),
        Err(e) => error_response(StatusCode::INTERNAL_SERVER_ERROR, &e.to_string()),
    }
}

async fn admin_accounts(State(state): State<Arc<AppState>>) -> impl IntoResponse {
    let openai_policy = state.vault.policy(Provider::Openai);
    let accounts: Vec<Value> = state
        .vault
        .list()
        .await
        .into_iter()
        .map(|a| {
            let email = a.email();
            let routing = if a.provider == Provider::Openai {
                let reserve_pct = codex_reserve_pct(&a, &openai_policy);
                json!({
                    "eligible": !openai_policy.disabled.iter().any(|name| name == &a.name || name == &a.id),
                    "priority": openai_policy.order.iter().position(|name| name == &a.name),
                    "reserve_pct": reserve_pct,
                    "reserve_blocked": codex_reserve_blocked(&a, reserve_pct, now_ms() / 1000),
                    "reset_selection": codex_reset_selection(&a, now_ms() / 1000),
                })
            } else {
                Value::Null
            };
            let limits = if a.provider == Provider::Openai {
                a.account_meta
                    .get("codex_limits")
                    .cloned()
                    .unwrap_or(Value::Null)
            } else {
                Value::Null
            };
            json!({
                "id": a.id,
                "provider": a.provider.as_str(),
                "name": a.name,
                "kind": a.kind,
                "label": a.label,
                "description": a.description,
                "email": email,
                "paused": a.paused,
                "path": a.path.map(|p| p.display().to_string()),
                "status": a.status,
                "expires_at_ms": a.expires_at_ms,
                "expires_in_s": a.expires_at_ms.map(|e| (e - now_ms()) / 1000),
                "routing": routing,
                "limits": limits,
            })
        })
        .collect();
    axum::Json(json!({"accounts": accounts}))
}

fn codex_strategy_name(mode: &AccountPolicyMode) -> &'static str {
    match mode {
        AccountPolicyMode::ResetFirst => "reset_first",
        AccountPolicyMode::RoundRobin => "round_robin",
        AccountPolicyMode::Priority | AccountPolicyMode::Threshold => "priority",
    }
}

async fn codex_routing_snapshot(state: &Arc<AppState>) -> Value {
    let policy = state.vault.policy(Provider::Openai);
    let mut accounts: Vec<Account> = state
        .vault
        .list()
        .await
        .into_iter()
        .filter(|account| account.provider == Provider::Openai)
        .collect();
    accounts.sort_by_key(|account| {
        (
            policy
                .order
                .iter()
                .position(|name| name == &account.name)
                .unwrap_or(usize::MAX / 2),
            account.name.clone(),
            account.id.clone(),
        )
    });
    let accounts: Vec<Value> = accounts
        .into_iter()
        .enumerate()
        .map(|(fallback_priority, account)| {
            let eligible = !policy
                .disabled
                .iter()
                .any(|name| name == &account.name || name == &account.id);
            let priority = policy
                .order
                .iter()
                .position(|name| name == &account.name)
                .unwrap_or(fallback_priority);
            let limits = account
                .account_meta
                .get("codex_limits")
                .cloned()
                .unwrap_or_else(|| json!({}));
            let reserve_pct = codex_reserve_pct(&account, &policy);
            let reset_selection = codex_reset_selection(&account, now_ms() / 1000);
            json!({
                "account_id": account.id,
                "eligible": eligible,
                "priority": priority,
                "reserve_pct": reserve_pct,
                "reserve_blocked": codex_reserve_blocked(&account, reserve_pct, now_ms() / 1000),
                "reset_selection": reset_selection,
                "observed_at_ms": limits.get("observed_at_ms"),
                "plan": limits.get("plan"),
                "active_limit": limits.get("active_limit"),
                "windows": limits.get("windows").cloned().unwrap_or_else(|| json!([])),
                "credits": limits.get("credits").cloned().unwrap_or(Value::Null),
            })
        })
        .collect();
    json!({
        "provider": "openai",
        "strategy": codex_strategy_name(&policy.mode),
        "reserve_pct": policy.reserve_pct.unwrap_or(10).min(100),
        "allow_mid_thread_failover": policy.allow_mid_thread_failover,
        "reset_selection_rule": "highest_used_pct_then_earliest_reset",
        "accounts": accounts,
    })
}

async fn admin_codex_routing(State(state): State<Arc<AppState>>) -> Response {
    axum::Json(codex_routing_snapshot(&state).await).into_response()
}

async fn admin_codex_routing_update(
    State(state): State<Arc<AppState>>,
    body: axum::Json<Value>,
) -> Response {
    let current_policy = state.vault.policy(Provider::Openai);
    let mode = match body.0.get("strategy").and_then(Value::as_str) {
        Some("reset_first") => AccountPolicyMode::ResetFirst,
        Some("priority") => AccountPolicyMode::Priority,
        Some("round_robin") => AccountPolicyMode::RoundRobin,
        _ => {
            return error_response(
                StatusCode::BAD_REQUEST,
                "strategy must be reset_first, priority, or round_robin",
            )
        }
    };
    let reserve_pct = match body.0.get("reserve_pct") {
        Some(value) => match value.as_u64() {
            Some(value) => value,
            None => return error_response(StatusCode::BAD_REQUEST, "reserve_pct must be an integer"),
        },
        None => current_policy.reserve_pct.unwrap_or(10) as u64,
    };
    if reserve_pct > 100 {
        return error_response(StatusCode::BAD_REQUEST, "reserve_pct must be between 0 and 100");
    }
    let allow_mid_thread_failover = match body.0.get("allow_mid_thread_failover") {
        Some(value) => match value.as_bool() {
            Some(value) => value,
            None => return error_response(
                StatusCode::BAD_REQUEST,
                "allow_mid_thread_failover must be boolean",
            ),
        },
        None => current_policy.allow_mid_thread_failover,
    };
    let Some(requested) = body.0.get("accounts").and_then(Value::as_array) else {
        return error_response(StatusCode::BAD_REQUEST, "accounts must be an array");
    };
    let openai_accounts: Vec<Account> = state
        .vault
        .list()
        .await
        .into_iter()
        .filter(|account| account.provider == Provider::Openai)
        .collect();
    let by_id: HashMap<String, String> = openai_accounts
        .iter()
        .map(|account| (account.id.clone(), account.name.clone()))
        .collect();
    let mut seen = std::collections::HashSet::new();
    let mut ordered = Vec::new();
    for item in requested {
        let Some(account_id) = item.get("account_id").and_then(Value::as_str) else {
            return error_response(StatusCode::BAD_REQUEST, "each account needs account_id");
        };
        let Some(eligible) = item.get("eligible").and_then(Value::as_bool) else {
            return error_response(StatusCode::BAD_REQUEST, "each account needs boolean eligible");
        };
        let Some(priority) = item.get("priority").and_then(Value::as_u64) else {
            return error_response(StatusCode::BAD_REQUEST, "each account needs integer priority");
        };
        let account_reserve_pct = match item.get("reserve_pct") {
            Some(value) => match value.as_u64() {
                Some(value) if value <= 100 => value as u8,
                Some(_) => return error_response(
                    StatusCode::BAD_REQUEST,
                    "account reserve_pct must be between 0 and 100",
                ),
                None => return error_response(
                    StatusCode::BAD_REQUEST,
                    "account reserve_pct must be an integer",
                ),
            },
            None => reserve_pct as u8,
        };
        let Some(name) = by_id.get(account_id) else {
            return error_response(
                StatusCode::BAD_REQUEST,
                &format!("unknown OpenAI account '{account_id}'"),
            );
        };
        if !seen.insert(account_id.to_string()) {
            return error_response(
                StatusCode::BAD_REQUEST,
                &format!("duplicate OpenAI account '{account_id}'"),
            );
        }
        ordered.push((priority, name.clone(), eligible, account_reserve_pct));
    }
    if seen.len() != by_id.len() {
        return error_response(
            StatusCode::CONFLICT,
            "account list changed; refresh Settings and save again",
        );
    }
    ordered.sort_by_key(|(priority, name, _, _)| (*priority, name.clone()));
    let policy = AccountPolicy {
        order: ordered.iter().map(|(_, name, _, _)| name.clone()).collect(),
        mode,
        threshold_pct: None,
        reserve_pct: Some(reserve_pct as u8),
        account_reserve_pct: ordered
            .iter()
            .map(|(_, name, _, reserve_pct)| (name.clone(), *reserve_pct))
            .collect(),
        allow_mid_thread_failover,
        disabled: ordered
            .iter()
            .filter(|(_, _, eligible, _)| !*eligible)
            .map(|(_, name, _, _)| name.clone())
            .collect(),
    };
    match state
        .vault
        .set_policy_persisted(Provider::Openai, policy)
        .await
    {
        Ok(()) => axum::Json(codex_routing_snapshot(&state).await).into_response(),
        Err(error) => error_response(StatusCode::INTERNAL_SERVER_ERROR, &error.to_string()),
    }
}

async fn admin_health(State(state): State<Arc<AppState>>) -> Response {
    let heartbeats = state.store.last_heartbeats().unwrap_or_default();
    let accounts: Vec<Value> = state
        .vault
        .list()
        .await
        .into_iter()
        .map(|a| {
            let last = heartbeats
                .iter()
                .find(|h| h["provider"].as_str() == Some(a.provider.as_str()))
                .cloned();
            json!({
                "id": a.id,
                "provider": a.provider.as_str(),
                "name": a.name,
                "kind": a.kind,
                "status": a.status,
                "paused": a.paused,
                "path": a.path.map(|p| p.display().to_string()),
                "token_expires_in_s": a.expires_at_ms.map(|e| (e - now_ms()) / 1000),
                "last_heartbeat": last,
            })
        })
        .collect();
    axum::Json(json!({"accounts": accounts})).into_response()
}

#[derive(Debug, Clone, serde::Serialize)]
pub struct PingResult {
    pub provider: &'static str,
    pub account_id: Option<String>,
    pub ok: bool,
    pub status: Option<u16>,
    pub latency_ms: i64,
    pub message: String,
}

#[derive(Debug, Clone)]
pub struct PingModels {
    pub anthropic: String,
    pub openai: String,
    pub xai: String,
    pub gemini: String,
}

pub async fn ping_provider(
    state: &Arc<AppState>,
    provider: Provider,
    models: &PingModels,
) -> PingResult {
    let start = now_ms();
    let prompt = "Health check: what time is it? If you cannot know, just reply: creds ok";
    let (format, path, body) = match provider {
        Provider::Anthropic => (
            ClientFormat::AnthropicMessages,
            "/v1/messages",
            json!({
                "model": models.anthropic,
                "max_tokens": 64,
                "system": "You are Claude Code, Anthropic's official CLI for Claude.",
                "messages": [{"role": "user", "content": prompt}],
            }),
        ),
        Provider::Openai => (
            ClientFormat::OpenaiResponses,
            "/v1/responses",
            json!({
                "model": models.openai,
                "stream": true,
                "store": false,
                "instructions": "You are a helpful assistant.",
                "input": [{"type": "message", "role": "user",
                           "content": [{"type": "input_text", "text": prompt}]}],
            }),
        ),
        Provider::Xai => (
            ClientFormat::OpenaiChat,
            "/v1/chat/completions",
            json!({
                "model": models.xai,
                "stream": false,
                "messages": [{"role": "user", "content": prompt}],
            }),
        ),
        Provider::Gemini => (
            ClientFormat::GeminiGenerate,
            "/v1beta/models/:generateContent",
            json!({
                "model": models.gemini,
                "contents": [{"role": "user", "parts": [{"text": prompt}]}],
                "generationConfig": {"maxOutputTokens": 64},
            }),
        ),
    };
    let account_id = state
        .vault
        .account_for(provider, true)
        .await
        .ok()
        .map(|a| a.id);
    let mut headers = HeaderMap::new();
    headers.insert(
        "x-api-key",
        HeaderValue::from_str(&state.local_key).expect("key header"),
    );
    headers.insert("user-agent", HeaderValue::from_static("alexandria-ping"));
    let resp = proxy(
        state.clone(),
        format,
        path,
        headers,
        Bytes::from(serde_json::to_vec(&body).expect("ping body")),
        None,
    )
    .await;
    let status = resp.status().as_u16();
    let bytes = axum::body::to_bytes(resp.into_body(), 8 * 1024 * 1024)
        .await
        .unwrap_or_default();
    let text = String::from_utf8_lossy(&bytes);
    let message = extract_reply(&text).unwrap_or_else(|| snippet(&text));
    PingResult {
        provider: provider.as_str(),
        account_id,
        ok: (200..300).contains(&status),
        status: Some(status),
        latency_ms: now_ms() - start,
        message,
    }
}

fn extract_reply(text: &str) -> Option<String> {
    if let Ok(v) = serde_json::from_str::<Value>(text) {
        if let Some(t) = v["content"][0]["text"].as_str() {
            return Some(t.to_string());
        }
        if let Some(t) = v["choices"][0]["message"]["content"].as_str() {
            return Some(t.to_string());
        }
        if let Some(t) = v["error"]["message"].as_str() {
            return Some(t.to_string());
        }
        if let Some(t) = v["detail"].as_str() {
            return Some(t.to_string());
        }
    }
    for line in text.lines() {
        let Some(data) = line.strip_prefix("data:") else {
            continue;
        };
        let Ok(v) = serde_json::from_str::<Value>(data.trim()) else {
            continue;
        };
        if v["type"].as_str() == Some("response.output_text.done") {
            if let Some(t) = v["text"].as_str() {
                return Some(t.to_string());
            }
        }
    }
    None
}

fn snippet(text: &str) -> String {
    let t: String = text.chars().take(200).collect();
    t.replace('\n', " ")
}

pub async fn heartbeat_once(state: &Arc<AppState>, models: &PingModels) -> Vec<PingResult> {
    let providers: Vec<Provider> = {
        let mut seen = Vec::new();
        for a in state.vault.list().await {
            if a.status == "active"
                && matches!(
                    a.provider,
                    Provider::Anthropic | Provider::Openai | Provider::Xai | Provider::Gemini
                )
                && !seen.contains(&a.provider)
            {
                seen.push(a.provider);
            }
        }
        seen
    };
    let mut results = Vec::new();
    for provider in providers {
        let r = ping_provider(state, provider, models).await;
        if let Err(e) = state.store.insert_heartbeat(
            now_ms(),
            r.provider,
            r.account_id.as_deref(),
            r.ok,
            r.status.map(|s| s as i64),
            r.latency_ms,
            &r.message,
        ) {
            tracing::warn!("failed to record heartbeat: {e}");
        }
        tracing::info!(
            provider = r.provider,
            ok = r.ok,
            status = r.status,
            latency_ms = r.latency_ms,
            reply = %r.message,
            "heartbeat"
        );
        results.push(r);
    }
    results
}

async fn anthropic_messages(
    State(state): State<Arc<AppState>>,
    ConnectInfo(peer): ConnectInfo<std::net::SocketAddr>,
    headers: HeaderMap,
    body: Bytes,
) -> Response {
    proxy(
        state,
        ClientFormat::AnthropicMessages,
        "/v1/messages",
        headers,
        body,
        Some(peer),
    )
    .await
}

async fn openai_chat(
    State(state): State<Arc<AppState>>,
    ConnectInfo(peer): ConnectInfo<std::net::SocketAddr>,
    headers: HeaderMap,
    body: Bytes,
) -> Response {
    proxy(
        state,
        ClientFormat::OpenaiChat,
        "/v1/chat/completions",
        headers,
        body,
        Some(peer),
    )
    .await
}

async fn openai_responses(
    State(state): State<Arc<AppState>>,
    ConnectInfo(peer): ConnectInfo<std::net::SocketAddr>,
    headers: HeaderMap,
    body: Bytes,
) -> Response {
    proxy(
        state,
        ClientFormat::OpenaiResponses,
        "/v1/responses",
        headers,
        body,
        Some(peer),
    )
    .await
}

async fn gemini_generate(
    State(state): State<Arc<AppState>>,
    ConnectInfo(peer): ConnectInfo<std::net::SocketAddr>,
    axum::extract::Path(model_action): axum::extract::Path<String>,
    headers: HeaderMap,
    body: Bytes,
) -> Response {
    let (model, method) = model_action
        .split_once(':')
        .unwrap_or((model_action.as_str(), "generateContent"));
    let stream = method == "streamGenerateContent";
    if method != "generateContent" && !stream {
        return error_response(
            StatusCode::NOT_FOUND,
            &format!("unsupported gemini method '{method}' (expected generateContent or streamGenerateContent)"),
        );
    }
    if model.is_empty() {
        return error_response(StatusCode::BAD_REQUEST, "missing model in path");
    }
    let mut v: Value = match serde_json::from_slice(&body) {
        Ok(v) => v,
        Err(e) => {
            return error_response(StatusCode::BAD_REQUEST, &format!("invalid JSON body: {e}"))
        }
    };
    v["model"] = json!(model);
    if stream {
        v["stream"] = json!(true);
    }
    let body = Bytes::from(serde_json::to_vec(&v).unwrap_or_default());
    let path = if stream {
        "/v1beta/models/:streamGenerateContent"
    } else {
        "/v1beta/models/:generateContent"
    };
    proxy(
        state,
        ClientFormat::GeminiGenerate,
        path,
        headers,
        body,
        Some(peer),
    )
    .await
}

fn error_response(status: StatusCode, message: &str) -> Response {
    (
        status,
        axum::Json(json!({"error": {"type": "alexandria", "message": message}})),
    )
        .into_response()
}

fn client_key(headers: &HeaderMap) -> Option<String> {
    if let Some(v) = headers.get("x-api-key").and_then(|v| v.to_str().ok()) {
        return Some(v.to_string());
    }
    if let Some(v) = headers.get("x-goog-api-key").and_then(|v| v.to_str().ok()) {
        return Some(v.to_string());
    }
    headers
        .get("authorization")
        .and_then(|v| v.to_str().ok())
        .and_then(|v| v.strip_prefix("Bearer "))
        .map(String::from)
}

const CLAUDE_CODE_PASSTHROUGH_HEADERS: &[&str] = &[
    "accept",
    "x-app",
    "x-claude-code-session-id",
    "x-stainless-arch",
    "x-stainless-lang",
    "x-stainless-os",
    "x-stainless-package-version",
    "x-stainless-retry-count",
    "x-stainless-runtime",
    "x-stainless-runtime-version",
    "x-stainless-timeout",
    "anthropic-dangerous-direct-browser-access",
];

fn is_genuine_claude_code_request(
    format: ClientFormat,
    headers: &HeaderMap,
    body: &Value,
) -> bool {
    if format != ClientFormat::AnthropicMessages {
        return false;
    }

    for harness in headers.get_all("x-alexandria-harness").iter() {
        let Ok(harness) = harness.to_str() else {
            return false;
        };
        let harness = harness.trim().to_ascii_lowercase();
        if !matches!(harness.as_str(), "claude" | "claude-code") {
            return false;
        }
    }

    let claude_user_agent = headers
        .get("user-agent")
        .and_then(|value| value.to_str().ok())
        .is_some_and(|value| value.to_ascii_lowercase().starts_with("claude-cli/"));
    let cli_app = headers
        .get("x-app")
        .and_then(|value| value.to_str().ok())
        .is_some_and(|value| value.eq_ignore_ascii_case("cli"));
    let has_session = headers
        .get("x-claude-code-session-id")
        .and_then(|value| value.to_str().ok())
        .is_some_and(|value| !value.trim().is_empty());
    let has_billing_header = body["system"]
        .as_array()
        .and_then(|blocks| blocks.first())
        .and_then(|block| block["text"].as_str())
        .is_some_and(|text| text.starts_with("x-anthropic-billing-header:"));

    claude_user_agent && cli_app && has_session && has_billing_header
}

fn redacted_headers(headers: &HeaderMap) -> String {
    let map: HashMap<String, String> = headers
        .iter()
        .map(|(k, v)| {
            let key = k.as_str().to_lowercase();
            let val = if ["authorization", "x-api-key", "cookie", "chatgpt-account-id"]
                .contains(&key.as_str())
            {
                "<redacted>".to_string()
            } else {
                v.to_str().unwrap_or("<binary>").to_string()
            };
            (key, val)
        })
        .collect();
    serde_json::to_string(&map).unwrap_or_default()
}

#[derive(Debug, Clone, Copy, PartialEq)]
enum RespondAs {
    Anthropic,
    OpenaiChat,
    OpenaiResponses,
    Gemini,
}

struct UpstreamPlan {
    url: String,
    account: Account,
    body: Vec<u8>,
    upstream_format: &'static str,
    destream: bool,
    respond_as: Option<RespondAs>,
    client_stream: bool,
    extra_headers: Vec<(String, String)>,
    dario_guard: Option<Box<dyn std::any::Any + Send>>,
}

fn dario_account(active: &DarioActive) -> Account {
    Account {
        id: format!("dario:{}", active.generation_id),
        provider: Provider::Anthropic,
        kind: "dario".into(),
        name: active.generation_id.clone(),
        description: None,
        paused: false,
        label: Some("dario generational proxy".into()),
        access_token: None,
        refresh_token: None,
        id_token: None,
        api_key: Some(active.api_key.clone()),
        expires_at_ms: None,
        last_refresh_ms: None,
        account_meta: Value::Null,
        cooldown_until_ms: None,
        status: "active".into(),
        path: None,
    }
}

async fn ensure_gemini_project(
    state: &AppState,
    account: &Account,
) -> Result<String, (StatusCode, String)> {
    if let Some(p) = account.account_meta.get("project_id").and_then(|v| v.as_str()) {
        if !p.is_empty() {
            return Ok(p.to_string());
        }
    }
    if let Ok(env_p) = std::env::var("GOOGLE_CLOUD_PROJECT") {
        if !env_p.is_empty() {
            let _ = state
                .vault
                .set_account_meta(&account.id, "project_id", json!(env_p))
                .await;
            return Ok(env_p);
        }
    }
    let token = account
        .access_token
        .as_deref()
        .ok_or_else(|| (StatusCode::BAD_GATEWAY, "gemini account has no access token".into()))?;
    let load_url = format!("{GEMINI_CODE_ASSIST_BASE}/{GEMINI_CODE_ASSIST_VERSION}:loadCodeAssist");
    let load_body = json!({
        "cloudaicompanionProject": null,
        "metadata": {"ideType": "IDE_UNSPECIFIED", "platform": "PLATFORM_UNSPECIFIED", "pluginType": "GEMINI"},
    });
    let resp = state
        .http
        .post(&load_url)
        .header("authorization", format!("Bearer {token}"))
        .json(&load_body)
        .send()
        .await
        .map_err(|e| (StatusCode::BAD_GATEWAY, format!("loadCodeAssist failed: {e}")))?;
    let load: Value = resp.json().await.unwrap_or(Value::Null);
    let extract = |v: &Value| -> Option<String> {
        for key in ["cloudaicompanionProject", "projectId", "project"] {
            match &v[key] {
                Value::String(s) if !s.is_empty() => return Some(s.clone()),
                obj if obj["id"].is_string() => {
                    return obj["id"].as_str().map(String::from)
                }
                _ => {}
            }
        }
        None
    };
    if let Some(p) = extract(&load) {
        let _ = state
            .vault
            .set_account_meta(&account.id, "project_id", json!(p))
            .await;
        return Ok(p);
    }
    if let Some(free) = load["ineligibleTiers"]
        .as_array()
        .and_then(|t| t.iter().find(|t| t["tierId"] == json!("free-tier")))
    {
        let reason = free["reasonCode"].as_str().unwrap_or("");
        let msg = free["reasonMessage"].as_str().unwrap_or("");
        let hint = if let Some(url) = free["validationUrl"].as_str() {
            format!(
                "gemini account needs verification before the free Code Assist tier works — \
                 open {url} to verify, then retry (or sign in with another personal Google \
                 account, or set gemini_project to a GCP project)"
            )
        } else if reason == "DASHER_USER" {
            "this is a Google Workspace account, which cannot use the free Code Assist tier — \
             set gemini_project to a GCP project (Code Assist API enabled), or sign in with a \
             personal Google account"
                .to_string()
        } else {
            format!("gemini account is not eligible for the free Code Assist tier: {msg}")
        };
        return Err((StatusCode::BAD_GATEWAY, hint));
    }
    let tier = load["allowedTiers"]
        .as_array()
        .and_then(|tiers| tiers.iter().find(|t| t["isDefault"] == json!(true)))
        .and_then(|t| t["id"].as_str())
        .unwrap_or("free-tier")
        .to_string();
    let onboard_url =
        format!("{GEMINI_CODE_ASSIST_BASE}/{GEMINI_CODE_ASSIST_VERSION}:onboardUser");
    let onboard_body = json!({
        "tierId": tier,
        "cloudaicompanionProject": load["cloudaicompanionProject"],
        "metadata": {"ideType": "IDE_UNSPECIFIED", "platform": "PLATFORM_UNSPECIFIED", "pluginType": "GEMINI"},
    });
    for _ in 0..5 {
        let resp = state
            .http
            .post(&onboard_url)
            .header("authorization", format!("Bearer {token}"))
            .json(&onboard_body)
            .send()
            .await
            .map_err(|e| (StatusCode::BAD_GATEWAY, format!("onboardUser failed: {e}")))?;
        let lro: Value = resp.json().await.unwrap_or(Value::Null);
        if lro["done"] == json!(true) {
            if let Some(p) = extract(&lro["response"]) {
                let _ = state
                    .vault
                    .set_account_meta(&account.id, "project_id", json!(p))
                    .await;
                return Ok(p);
            }
        }
        tokio::time::sleep(Duration::from_secs(2)).await;
    }
    Err((
        StatusCode::BAD_GATEWAY,
        "gemini Code Assist onboarding did not return a project id (set GOOGLE_CLOUD_PROJECT to override)".into(),
    ))
}

async fn plan_upstream(
    state: &AppState,
    format: ClientFormat,
    provider: Provider,
    routed_model: &str,
    body_json: &mut Value,
    original_body: &[u8],
    trace_id: &str,
    excluded_accounts: &HashSet<String>,
    affinity_session: Option<&str>,
    client_headers: &HeaderMap,
) -> Result<UpstreamPlan, (StatusCode, String)> {
    use alex_core::translate;
    let client_stream = body_json["stream"].as_bool().unwrap_or(false);
    let genuine_claude_code =
        is_genuine_claude_code_request(format, client_headers, body_json);
    match provider {
        Provider::Anthropic => {
            let converted = match format {
                ClientFormat::AnthropicMessages => None,
                ClientFormat::OpenaiChat => Some((
                    translate::openai_chat_to_anthropic(body_json),
                    RespondAs::OpenaiChat,
                )),
                ClientFormat::OpenaiResponses => Some((
                    translate::openai_responses_to_anthropic(body_json),
                    RespondAs::OpenaiResponses,
                )),
                ClientFormat::GeminiGenerate => Some((
                    translate::gemini_to_anthropic(body_json),
                    RespondAs::Gemini,
                )),
            };
            let (base, account, dario_guard, dario_capture) = if genuine_claude_code {
                let account = state
                    .vault
                    .account_for_excluding(provider, true, excluded_accounts)
                    .await
                    .map_err(|e| (StatusCode::BAD_GATEWAY, e.to_string()))?;
                (ANTHROPIC_BASE.to_string(), account, None, false)
            } else {
                let dario_active = state.dario.as_ref().and_then(|dario| dario.active());
                match (&state.dario, dario_active) {
                    (Some(dario), Some(active)) => {
                        let guard = dario.begin(&active.generation_id).ok_or_else(|| {
                            (
                                StatusCode::SERVICE_UNAVAILABLE,
                                "Dario became unavailable while routing the Anthropic request"
                                    .to_string(),
                            )
                        })?;
                        (
                            active.base_url.trim_end_matches('/').to_string(),
                            dario_account(&active),
                            Some(guard),
                            true,
                        )
                    }
                    (Some(_), None) => {
                        return Err((
                            StatusCode::SERVICE_UNAVAILABLE,
                            "Dario is configured but no healthy generation is available"
                                .to_string(),
                        ));
                    }
                    (None, None) => {
                        let account = state
                            .vault
                            .account_for_excluding(provider, true, excluded_accounts)
                            .await
                            .map_err(|e| (StatusCode::BAD_GATEWAY, e.to_string()))?;
                        (ANTHROPIC_BASE.to_string(), account, None, false)
                    }
                    (None, Some(_)) => unreachable!("Dario cannot be active when it is disabled"),
                }
            };
            let (body, respond_as) = match converted {
                None => {
                    body_json["model"] = json!(routed_model);
                    (
                        serde_json::to_vec(body_json).unwrap_or_else(|_| original_body.to_vec()),
                        None,
                    )
                }
                Some((mut converted, respond_as)) => {
                    converted["model"] = json!(routed_model);
                    converted["stream"] = json!(false);
                    (
                        serde_json::to_vec(&converted).unwrap_or_else(|_| original_body.to_vec()),
                        Some(respond_as),
                    )
                }
            };
            Ok(UpstreamPlan {
                url: format!("{base}/v1/messages"),
                account,
                body,
                upstream_format: "anthropic",
                destream: false,
                respond_as,
                client_stream,
                extra_headers: if dario_capture {
                    vec![
                        ("x-dario-capture-id".into(), trace_id.into()),
                        ("x-dario-capture-model".into(), routed_model.into()),
                    ]
                } else {
                    vec![]
                },
                dario_guard,
            })
        }
        Provider::Openai => {
            // Serialize only the account-planning step for the same thread.
            // This closes the first-request race without blocking unrelated
            // sessions or holding a lock during the upstream model request.
            let _affinity_guard = match affinity_session.filter(|value| !value.is_empty()) {
                Some(session_id) => Some(
                    codex_affinity_lock(state, session_id)
                        .lock_owned()
                        .await,
                ),
                None => None,
            };
            let prefer_oauth = format != ClientFormat::OpenaiChat;
            let policy = state.vault.policy(Provider::Openai);
            let preferred = preferred_codex_account(state, affinity_session);
            let account = state
                .vault
                .account_for_excluding_preferred_mode(
                    provider,
                    prefer_oauth,
                    excluded_accounts,
                    preferred.as_deref(),
                    preferred.is_some() && !policy.allow_mid_thread_failover,
                )
                .await
                .map_err(|e| (StatusCode::BAD_GATEWAY, e.to_string()))?;
            let oauth = account.kind == "oauth";
            bind_codex_account(state, affinity_session, &account);
            match format {
                ClientFormat::OpenaiChat if !oauth => {
                    body_json["model"] = json!(routed_model);
                    let body =
                        serde_json::to_vec(body_json).unwrap_or_else(|_| original_body.to_vec());
                    Ok(UpstreamPlan {
                        url: format!("{OPENAI_BASE}/v1/chat/completions"),
                        account,
                        body,
                        upstream_format: "openai-chat",
                        destream: false,
                        respond_as: None,
                        client_stream,
                        extra_headers: vec![],
                        dario_guard: None,
                    })
                }
                ClientFormat::OpenaiChat => {
                    let pivot = translate::openai_chat_to_anthropic(body_json);
                    let mut converted = translate::anthropic_to_openai_responses(&pivot);
                    converted["model"] = json!(routed_model);
                    translate::normalize_codex_request(&mut converted);
                    let body = serde_json::to_vec(&converted)
                        .unwrap_or_else(|_| original_body.to_vec());
                    Ok(UpstreamPlan {
                        url: format!("{CODEX_BASE}/responses"),
                        account,
                        body,
                        upstream_format: "openai-responses",
                        destream: false,
                        respond_as: Some(RespondAs::OpenaiChat),
                        client_stream,
                        extra_headers: vec![],
                        dario_guard: None,
                    })
                }
                ClientFormat::OpenaiResponses => {
                    body_json["model"] = json!(routed_model);
                    let mut destream = false;
                    let url = if oauth {
                        if body_json["stream"].as_bool() != Some(true) {
                            destream = true;
                        }
                        translate::normalize_codex_request(body_json);
                        format!("{CODEX_BASE}/responses")
                    } else {
                        format!("{OPENAI_BASE}/v1/responses")
                    };
                    let body =
                        serde_json::to_vec(body_json).unwrap_or_else(|_| original_body.to_vec());
                    Ok(UpstreamPlan {
                        url,
                        account,
                        body,
                        upstream_format: "openai-responses",
                        destream,
                        respond_as: None,
                        client_stream,
                        extra_headers: vec![],
                        dario_guard: None,
                    })
                }
                ClientFormat::AnthropicMessages => {
                    let mut converted = translate::anthropic_to_openai_responses(body_json);
                    converted["model"] = json!(routed_model);
                    let url = if oauth {
                        translate::normalize_codex_request(&mut converted);
                        format!("{CODEX_BASE}/responses")
                    } else {
                        converted["stream"] = json!(false);
                        format!("{OPENAI_BASE}/v1/responses")
                    };
                    let body = serde_json::to_vec(&converted)
                        .unwrap_or_else(|_| original_body.to_vec());
                    Ok(UpstreamPlan {
                        url,
                        account,
                        body,
                        upstream_format: "openai-responses",
                        destream: false,
                        respond_as: Some(RespondAs::Anthropic),
                        client_stream,
                        extra_headers: vec![],
                        dario_guard: None,
                    })
                }
                ClientFormat::GeminiGenerate => {
                    let pivot = translate::gemini_to_anthropic(body_json);
                    let mut converted = translate::anthropic_to_openai_responses(&pivot);
                    converted["model"] = json!(routed_model);
                    let url = if oauth {
                        translate::normalize_codex_request(&mut converted);
                        format!("{CODEX_BASE}/responses")
                    } else {
                        converted["stream"] = json!(false);
                        format!("{OPENAI_BASE}/v1/responses")
                    };
                    let body = serde_json::to_vec(&converted)
                        .unwrap_or_else(|_| original_body.to_vec());
                    Ok(UpstreamPlan {
                        url,
                        account,
                        body,
                        upstream_format: "openai-responses",
                        destream: false,
                        respond_as: Some(RespondAs::Gemini),
                        client_stream,
                        extra_headers: vec![],
                        dario_guard: None,
                    })
                }
            }
        }
        Provider::Xai => {
            let account = state
                .vault
                .account_for_excluding(provider, true, excluded_accounts)
                .await
                .map_err(|e| (StatusCode::BAD_GATEWAY, e.to_string()))?;
            let extra_headers = vec![
                ("x-grok-model-override".into(), routed_model.to_string()),
                ("x-grok-conv-id".into(), trace_id.to_string()),
            ];
            match format {
                ClientFormat::OpenaiChat => {
                    body_json["model"] = json!(routed_model);
                    let body =
                        serde_json::to_vec(body_json).unwrap_or_else(|_| original_body.to_vec());
                    Ok(UpstreamPlan {
                        url: format!("{XAI_BASE}/chat/completions"),
                        account,
                        body,
                        upstream_format: "openai-chat",
                        destream: false,
                        respond_as: None,
                        client_stream,
                        extra_headers,
                        dario_guard: None,
                    })
                }
                ClientFormat::AnthropicMessages => {
                    // pivot (anthropic) → chat completions for xAI/Grok upstream
                    let mut converted = translate::anthropic_to_openai_chat(body_json);
                    converted["model"] = json!(routed_model);
                    // buffer full upstream body then re-synth client dialect
                    converted["stream"] = json!(false);
                    let body = serde_json::to_vec(&converted)
                        .unwrap_or_else(|_| original_body.to_vec());
                    Ok(UpstreamPlan {
                        url: format!("{XAI_BASE}/chat/completions"),
                        account,
                        body,
                        upstream_format: "openai-chat",
                        destream: false,
                        respond_as: Some(RespondAs::Anthropic),
                        client_stream,
                        extra_headers,
                        dario_guard: None,
                    })
                }
                ClientFormat::OpenaiResponses | ClientFormat::GeminiGenerate => Err((
                    StatusCode::NOT_IMPLEMENTED,
                    "the xai/grok upstream speaks OpenAI chat completions; POST to /v1/chat/completions or /v1/messages"
                        .to_string(),
                )),
            }
        }
        Provider::Gemini => {
            // Prefer an AI Studio API key over the OAuth/Code-Assist path.
            let account = state
                .vault
                .account_for_excluding(provider, false, excluded_accounts)
                .await
                .map_err(|e| (StatusCode::BAD_GATEWAY, e.to_string()))?;
            let (gemini_req, respond_as) = match format {
                ClientFormat::GeminiGenerate => {
                    let mut g = body_json.clone();
                    if let Some(o) = g.as_object_mut() {
                        o.remove("model");
                        o.remove("stream");
                    }
                    (g, None)
                }
                ClientFormat::AnthropicMessages => (
                    translate::anthropic_to_gemini_request(body_json),
                    Some(RespondAs::Anthropic),
                ),
                ClientFormat::OpenaiChat => (
                    translate::anthropic_to_gemini_request(&translate::openai_chat_to_anthropic(
                        body_json,
                    )),
                    Some(RespondAs::OpenaiChat),
                ),
                ClientFormat::OpenaiResponses => (
                    translate::anthropic_to_gemini_request(
                        &translate::openai_responses_to_anthropic(body_json),
                    ),
                    Some(RespondAs::OpenaiResponses),
                ),
            };
            let method = if client_stream {
                "streamGenerateContent"
            } else {
                "generateContent"
            };
            let (url, body) = if account.kind == "api_key" {
                // AI Studio: plain request, model in the path, x-goog-api-key header.
                let mut url = format!(
                    "{GEMINI_API_BASE}/v1beta/models/{routed_model}:{method}"
                );
                if client_stream {
                    url.push_str("?alt=sse");
                }
                let body =
                    serde_json::to_vec(&gemini_req).unwrap_or_else(|_| original_body.to_vec());
                (url, body)
            } else {
                // Code Assist (OAuth): wrapped envelope + project.
                let project = ensure_gemini_project(state, &account).await?;
                let mut url = format!(
                    "{GEMINI_CODE_ASSIST_BASE}/{GEMINI_CODE_ASSIST_VERSION}:{method}"
                );
                if client_stream {
                    url.push_str("?alt=sse");
                }
                let envelope = json!({
                    "model": routed_model,
                    "project": project,
                    "request": gemini_req,
                });
                let body =
                    serde_json::to_vec(&envelope).unwrap_or_else(|_| original_body.to_vec());
                (url, body)
            };
            Ok(UpstreamPlan {
                url,
                account,
                body,
                upstream_format: "gemini",
                destream: false,
                respond_as: respond_as.or(Some(RespondAs::Gemini)),
                client_stream,
                extra_headers: vec![],
                dario_guard: None,
            })
        }
    }
}

#[cfg(test)]
mod trace_api_tests {
    use super::{
        session_from_metadata, trace_extras, trace_harness, trace_reasoning_fields, transcript_turn,
        truncate_chars,
    };
    use axum::http::{HeaderMap, HeaderValue};
    use serde_json::json;

    #[test]
    fn metadata_session_variants() {
        let claude = json!({"metadata": {"user_id":
            "{\"device_id\":\"d1\",\"session_id\":\"ses_inner\"}"}});
        assert_eq!(session_from_metadata(&claude), Some("ses_inner".into()));
        let plain = json!({"metadata": {"user_id": "user-123"}});
        assert_eq!(session_from_metadata(&plain), Some("user-123".into()));
        let json_no_session = json!({"metadata": {"user_id": "{\"device_id\":\"d1\"}"}});
        assert_eq!(
            session_from_metadata(&json_no_session),
            Some("{\"device_id\":\"d1\"}".into())
        );
        assert_eq!(session_from_metadata(&json!({})), None);
    }

    #[test]
    fn explicit_harness_header_wins_over_sdk_user_agent() {
        let mut headers = HeaderMap::new();
        headers.insert(
            "user-agent",
            HeaderValue::from_static("Anthropic/JS 0.91.1"),
        );
        assert_eq!(trace_harness(&headers).as_deref(), Some("Anthropic/JS 0.91.1"));
        headers.insert(
            "x-alexandria-harness",
            HeaderValue::from_static("pi"),
        );
        assert_eq!(trace_harness(&headers).as_deref(), Some("pi"));
    }

    #[test]
    fn extras_per_format() {
        let anthropic = json!({
            "system": [{"type": "text", "text": "abcd"}],
            "messages": [{"role": "user", "content": "hi"}],
            "max_tokens": 100,
            "temperature": 0.5,
            "thinking": {"type": "enabled", "budget_tokens": 4096},
        });
        let e = trace_extras(&anthropic);
        assert_eq!(e["thinking_budget"], 4096);
        assert_eq!(trace_reasoning_fields(&anthropic), (None, Some(4096)));
        assert_eq!(e["max_tokens"], 100);
        assert_eq!(e["temperature"], 0.5);
        assert_eq!(e["message_count"], 1);
        assert_eq!(e["system_chars"], 4);
        assert_eq!(e["reasoning_effort"], serde_json::Value::Null);
        let responses = json!({
            "instructions": "abc",
            "input": [{"type": "message"}, {"type": "message"}],
            "max_output_tokens": 200,
            "reasoning": {"effort": "high"},
        });
        let e = trace_extras(&responses);
        assert_eq!(e["reasoning_effort"], "high");
        assert_eq!(
            trace_reasoning_fields(&responses),
            (Some("high".into()), None)
        );
        assert_eq!(e["max_tokens"], 200);
        assert_eq!(e["message_count"], 2);
        assert_eq!(e["system_chars"], 3);
        assert_eq!(e["thinking_budget"], serde_json::Value::Null);
    }

    #[test]
    fn truncates_on_char_boundaries() {
        assert_eq!(truncate_chars("abc".into(), 8000), "abc");
        assert_eq!(truncate_chars("héllo".repeat(2000), 8000).chars().count(), 8000);
    }

    #[test]
    fn transcript_turn_missing_bodies_are_null() {
        let row = json!({
            "id": "t1", "ts_request_ms": 1, "ts_response_ms": 2,
            "routed_model": "m", "status": 200,
            "input_tokens": 10, "output_tokens": 5, "cost_usd": 0.01, "error": null,
            "reasoning_effort": "minimal", "thinking_budget": null,
            "req_body_path": "/nonexistent/x.gz", "resp_body_path": null,
            "client_format": "anthropic", "upstream_format": "anthropic",
        });
        let turn = transcript_turn(&row);
        assert_eq!(turn["trace_id"], "t1");
        assert_eq!(turn["user"], serde_json::Value::Null);
        assert_eq!(turn["assistant"], serde_json::Value::Null);
        assert_eq!(turn["model"], "m");
        assert_eq!(turn["reasoning_effort"], "minimal");
        assert_eq!(turn["thinking_budget"], serde_json::Value::Null);
    }
}

#[cfg(test)]
mod run_key_tests {
    use super::{generate_run_key, key_fingerprint, key_hash_hex, merge_run_key_tags};

    #[test]
    fn merge_header_tags_win_per_key() {
        let merged = merge_run_key_tags(
            Some(r#"{"task":"demo","suite":"swebench"}"#),
            Some(r#"{"task":"override","case":"x1"}"#),
        )
        .unwrap();
        let v: serde_json::Value = serde_json::from_str(&merged).unwrap();
        assert_eq!(v["task"], "override");
        assert_eq!(v["suite"], "swebench");
        assert_eq!(v["case"], "x1");
    }

    #[test]
    fn merge_handles_missing_sides() {
        assert_eq!(merge_run_key_tags(None, None), None);
        assert_eq!(
            merge_run_key_tags(Some(r#"{"a":"1"}"#), None).unwrap(),
            r#"{"a":"1"}"#
        );
        assert_eq!(
            merge_run_key_tags(None, Some(r#"{"b":"2"}"#)).unwrap(),
            r#"{"b":"2"}"#
        );
        assert_eq!(merge_run_key_tags(Some("not json"), None), None);
    }

    #[test]
    fn run_key_shape_and_fingerprint() {
        let key = generate_run_key();
        assert!(key.starts_with("alxk-"));
        assert_eq!(key.len(), 5 + 64);
        assert_ne!(key, generate_run_key());
        let hash = key_hash_hex(&key);
        assert_eq!(hash.len(), 64);
        assert_eq!(hash[..16], key_fingerprint(&key));
    }
}

#[cfg(test)]
mod usage_tests {
    use super::{key_fingerprint, usage_backoff_ms};

    #[test]
    fn fingerprint_is_first_8_sha256_bytes_hex() {
        assert_eq!(key_fingerprint("abc"), "ba7816bf8f01cfea");
        assert_eq!(key_fingerprint("abc").len(), 16);
        assert_ne!(key_fingerprint("abc"), key_fingerprint("abd"));
    }

    #[test]
    fn backoff_grows_and_caps() {
        assert_eq!(usage_backoff_ms(1, None), 60_000);
        assert_eq!(usage_backoff_ms(2, None), 120_000);
        assert_eq!(usage_backoff_ms(3, None), 240_000);
        assert_eq!(usage_backoff_ms(10, None), 3_600_000);
    }

    #[test]
    fn retry_after_wins_when_larger() {
        assert_eq!(usage_backoff_ms(1, Some(600_000)), 600_000);
        assert_eq!(usage_backoff_ms(5, Some(1_000)), 960_000);
    }
}

fn upstream_headers(
    account: &Account,
    client_headers: &HeaderMap,
    genuine_claude_code: bool,
) -> Result<reqwest::header::HeaderMap, (StatusCode, String)> {
    let mut h = reqwest::header::HeaderMap::new();
    h.insert("content-type", HeaderValue::from_static("application/json"));
    h.insert("accept", HeaderValue::from_static("*/*"));
    h.insert("accept-encoding", HeaderValue::from_static("identity"));
    if let Some(ua) = client_headers.get("user-agent") {
        h.insert("user-agent", ua.clone());
    }
    match (account.provider, account.kind.as_str()) {
        (Provider::Anthropic, "oauth") => {
            let token = account.access_token.as_deref().ok_or((
                StatusCode::BAD_GATEWAY,
                "anthropic oauth account has no access token".to_string(),
            ))?;
            h.insert(
                "authorization",
                HeaderValue::from_str(&format!("Bearer {token}"))
                    .map_err(|e| (StatusCode::BAD_GATEWAY, e.to_string()))?,
            );
            let mut beta = ANTHROPIC_OAUTH_BETA.to_string();
            if let Some(client_beta) = client_headers
                .get("anthropic-beta")
                .and_then(|v| v.to_str().ok())
            {
                if !client_beta.contains(ANTHROPIC_OAUTH_BETA) {
                    beta = format!("{ANTHROPIC_OAUTH_BETA},{client_beta}");
                } else {
                    beta = client_beta.to_string();
                }
            }
            h.insert(
                "anthropic-beta",
                HeaderValue::from_str(&beta)
                    .map_err(|e| (StatusCode::BAD_GATEWAY, e.to_string()))?,
            );
            insert_anthropic_version(&mut h, client_headers);
        }
        (Provider::Anthropic, _) => {
            let key = account
                .api_key
                .as_deref()
                .or(account.access_token.as_deref())
                .ok_or((
                    StatusCode::BAD_GATEWAY,
                    "anthropic account has no api key".to_string(),
                ))?;
            h.insert(
                "x-api-key",
                HeaderValue::from_str(key).map_err(|e| (StatusCode::BAD_GATEWAY, e.to_string()))?,
            );
            if let Some(client_beta) = client_headers.get("anthropic-beta") {
                h.insert("anthropic-beta", client_beta.clone());
            }
            insert_anthropic_version(&mut h, client_headers);
        }
        (Provider::Openai, "oauth") => {
            let token = account.access_token.as_deref().ok_or((
                StatusCode::BAD_GATEWAY,
                "openai oauth account has no access token".to_string(),
            ))?;
            h.insert(
                "authorization",
                HeaderValue::from_str(&format!("Bearer {token}"))
                    .map_err(|e| (StatusCode::BAD_GATEWAY, e.to_string()))?,
            );
            if let Some(acct) = account.chatgpt_account_id() {
                h.insert(
                    "chatgpt-account-id",
                    HeaderValue::from_str(&acct)
                        .map_err(|e| (StatusCode::BAD_GATEWAY, e.to_string()))?,
                );
            }
            h.insert(
                "openai-beta",
                HeaderValue::from_static("responses=experimental"),
            );
            h.insert("originator", HeaderValue::from_static("codex_cli_rs"));
            let session = client_headers
                .get("session_id")
                .and_then(|v| v.to_str().ok())
                .map(String::from)
                .unwrap_or_else(|| uuid::Uuid::new_v4().to_string());
            h.insert(
                "session_id",
                HeaderValue::from_str(&session)
                    .map_err(|e| (StatusCode::BAD_GATEWAY, e.to_string()))?,
            );
        }
        (Provider::Openai, _) => {
            let key = account.api_key.as_deref().ok_or((
                StatusCode::BAD_GATEWAY,
                "openai account has no api key".to_string(),
            ))?;
            h.insert(
                "authorization",
                HeaderValue::from_str(&format!("Bearer {key}"))
                    .map_err(|e| (StatusCode::BAD_GATEWAY, e.to_string()))?,
            );
        }
        (Provider::Xai, _) => {
            let token = account
                .access_token
                .as_deref()
                .or(account.api_key.as_deref())
                .ok_or((
                    StatusCode::BAD_GATEWAY,
                    "xai account has no access token".to_string(),
                ))?;
            h.insert(
                "authorization",
                HeaderValue::from_str(&format!("Bearer {token}"))
                    .map_err(|e| (StatusCode::BAD_GATEWAY, e.to_string()))?,
            );
            h.insert("x-xai-token-auth", HeaderValue::from_static("xai-grok-cli"));
            h.insert(
                "x-grok-client-version",
                HeaderValue::from_static(GROK_CLIENT_VERSION),
            );
            h.insert("user-agent", HeaderValue::from_static("xai-grok-workspace"));
            h.insert("accept", HeaderValue::from_static("text/event-stream"));
        }
        (Provider::Gemini, "api_key") => {
            let key = account.api_key.as_deref().ok_or((
                StatusCode::BAD_GATEWAY,
                "gemini api_key account has no key".to_string(),
            ))?;
            h.insert(
                "x-goog-api-key",
                HeaderValue::from_str(key)
                    .map_err(|e| (StatusCode::BAD_GATEWAY, e.to_string()))?,
            );
        }
        (Provider::Gemini, _) => {
            let token = account.access_token.as_deref().ok_or((
                StatusCode::BAD_GATEWAY,
                "gemini account has no access token".to_string(),
            ))?;
            h.insert(
                "authorization",
                HeaderValue::from_str(&format!("Bearer {token}"))
                    .map_err(|e| (StatusCode::BAD_GATEWAY, e.to_string()))?,
            );
        }
    }
    if genuine_claude_code && account.provider == Provider::Anthropic {
        for name in CLAUDE_CODE_PASSTHROUGH_HEADERS {
            if let Some(value) = client_headers.get(*name) {
                h.insert(*name, value.clone());
            }
        }
    }
    Ok(h)
}

fn insert_anthropic_version(h: &mut reqwest::header::HeaderMap, client_headers: &HeaderMap) {
    match client_headers.get("anthropic-version") {
        Some(v) => {
            h.insert("anthropic-version", v.clone());
        }
        None => {
            h.insert("anthropic-version", HeaderValue::from_static("2023-06-01"));
        }
    }
}

fn key_fingerprint(key: &str) -> String {
    use sha2::{Digest, Sha256};
    let digest = Sha256::digest(key.as_bytes());
    digest[..8].iter().map(|b| format!("{b:02x}")).collect()
}

const RUN_KEY_PREFIX: &str = "alxk-";
const RUN_KEY_DEFAULT_TTL_S: i64 = 86_400;
const RUN_KEY_MAX_TTL_S: i64 = 604_800;

pub fn generate_run_key() -> String {
    use rand::RngCore;
    let mut bytes = [0u8; 32];
    rand::thread_rng().fill_bytes(&mut bytes);
    format!(
        "{RUN_KEY_PREFIX}{}",
        bytes.iter().map(|b| format!("{b:02x}")).collect::<String>()
    )
}

fn key_hash_hex(key: &str) -> String {
    use sha2::{Digest, Sha256};
    let digest = Sha256::digest(key.as_bytes());
    digest.iter().map(|b| format!("{b:02x}")).collect()
}

fn merge_run_key_tags(key_tags: Option<&str>, header_tags: Option<&str>) -> Option<String> {
    let parse = |s: Option<&str>| {
        s.and_then(|s| serde_json::from_str::<Value>(s).ok())
            .and_then(|v| v.as_object().cloned())
            .unwrap_or_default()
    };
    let mut merged = parse(key_tags);
    merged.extend(parse(header_tags));
    if merged.is_empty() {
        None
    } else {
        serde_json::to_string(&Value::Object(merged)).ok()
    }
}

fn run_key_entry(state: &AppState, key_hash: &str) -> Option<CachedRunKey> {
    let now = now_ms();
    let cached = state.run_keys.read().unwrap().get(key_hash).cloned();
    if let Some(entry) = cached {
        if entry.expires_ms.map(|e| e > now).unwrap_or(true) {
            return Some(entry);
        }
        state.run_keys.write().unwrap().remove(key_hash);
        return None;
    }
    let row = state.store.lookup_run_key(key_hash, now).ok().flatten()?;
    let entry = CachedRunKey {
        run_id: row["run_id"].as_str().map(String::from),
        tags_json: row["tags"]
            .as_object()
            .filter(|o| !o.is_empty())
            .and_then(|o| serde_json::to_string(o).ok()),
        expires_ms: row["expires_ms"].as_i64(),
    };
    state
        .run_keys
        .write()
        .unwrap()
        .insert(key_hash.to_string(), entry.clone());
    Some(entry)
}

async fn admin_run_keys_create(
    State(state): State<Arc<AppState>>,
    body: Option<axum::Json<Value>>,
) -> Response {
    let body = body.map(|b| b.0).unwrap_or_else(|| json!({}));
    let tags = match &body["tags"] {
        Value::Null => None,
        Value::Object(o) => Some(o.clone()),
        _ => return error_response(StatusCode::BAD_REQUEST, "'tags' must be an object"),
    };
    let kind = body["kind"].as_str().unwrap_or("run");
    if kind != "run" && kind != "harness" {
        return error_response(StatusCode::BAD_REQUEST, "'kind' must be 'run' or 'harness'");
    }
    let run_id = body["run_id"].as_str().map(String::from);
    let label = body["label"].as_str().map(String::from);
    if kind == "harness" && label.as_deref().unwrap_or("").trim().is_empty() {
        return error_response(StatusCode::BAD_REQUEST, "'label' is required for harness keys");
    }
    let created_ms = now_ms();
    let expires_ms = if kind == "harness" {
        None
    } else {
        let ttl_seconds = body["ttl_seconds"].as_i64().unwrap_or(RUN_KEY_DEFAULT_TTL_S);
        if ttl_seconds <= 0 {
            return error_response(StatusCode::BAD_REQUEST, "'ttl_seconds' must be positive");
        }
        Some(created_ms + ttl_seconds.min(RUN_KEY_MAX_TTL_S) * 1000)
    };
    let key = generate_run_key();
    let key_hash = key_hash_hex(&key);
    let id = format!("rk-{}", &key_hash[..8]);
    let tags_json = tags
        .as_ref()
        .filter(|o| !o.is_empty())
        .and_then(|o| serde_json::to_string(o).ok());
    match state.store.insert_run_key(
        &id,
        &key_hash,
        kind,
        run_id.as_deref(),
        tags_json.as_deref(),
        label.as_deref(),
        created_ms,
        expires_ms,
    ) {
        Ok(()) => {
            let (_, exports) = connect_payload(&state.base_url, &key);
            (
                StatusCode::CREATED,
                axum::Json(json!({
                    "id": id,
                    "key": key,
                    "kind": kind,
                    "run_id": run_id,
                    "label": label,
                    "tags": tags.map(Value::Object).unwrap_or_else(|| json!({})),
                    "expires_ms": expires_ms,
                    "exports": exports,
                })),
            )
                .into_response()
        }
        Err(e) => error_response(StatusCode::INTERNAL_SERVER_ERROR, &e.to_string()),
    }
}

async fn admin_run_keys_list(
    State(state): State<Arc<AppState>>,
    Query(q): Query<HashMap<String, String>>,
) -> Response {
    let all = q.get("all").map(|v| v == "1" || v == "true").unwrap_or(false);
    match state.store.list_run_keys(all) {
        Ok(rows) => axum::Json(json!({"run_keys": rows})).into_response(),
        Err(e) => error_response(StatusCode::INTERNAL_SERVER_ERROR, &e.to_string()),
    }
}

async fn admin_run_keys_revoke(
    State(state): State<Arc<AppState>>,
    Path(id): Path<String>,
) -> Response {
    match state.store.revoke_run_key(&id) {
        Ok(true) => {
            state.run_keys.write().unwrap().clear();
            axum::Json(json!({"revoked": true})).into_response()
        }
        Ok(false) => error_response(StatusCode::NOT_FOUND, &format!("unknown run key '{id}'")),
        Err(e) => error_response(StatusCode::INTERNAL_SERVER_ERROR, &e.to_string()),
    }
}

fn session_from_metadata(body_json: &Value) -> Option<String> {
    let raw = body_json["metadata"]["user_id"].as_str()?;
    if let Ok(v) = serde_json::from_str::<Value>(raw) {
        if let Some(inner) = v["session_id"].as_str() {
            return Some(inner.to_string());
        }
    }
    Some(raw.to_string())
}

fn trace_harness(headers: &HeaderMap) -> Option<String> {
    headers
        .get("x-alexandria-harness")
        .and_then(|value| value.to_str().ok())
        .map(str::trim)
        .filter(|value| !value.is_empty())
        .map(String::from)
        .or_else(|| {
            headers
                .get("user-agent")
                .and_then(|value| value.to_str().ok())
                .map(String::from)
        })
}

const METADATA_HEADERS: &[(&str, &str)] = &[
    ("x-alexandria-harness", "harness"),
    ("x-alexandria-harness-version", "harness_version"),
    ("x-alexandria-task", "task"),
    ("x-alexandria-model", "model"),
    ("x-alexandria-job", "job"),
    ("x-alexandria-phase", "phase"),
    ("x-alexandria-kind", "kind"),
];

fn trace_tags_json(headers: &HeaderMap) -> Option<String> {
    let values: Vec<&str> = headers
        .get_all("x-alexandria-trace-tag")
        .iter()
        .filter_map(|v| v.to_str().ok())
        .collect();
    let mut tags = parse_trace_tags(&values);
    if let Some(o) = tags.as_object_mut() {
        for (header, key) in METADATA_HEADERS {
            if o.contains_key(*key) {
                continue;
            }
            if let Some(v) = headers.get(*header).and_then(|v| v.to_str().ok()) {
                let v = v.trim();
                if !v.is_empty() {
                    o.insert((*key).to_string(), json!(v));
                }
            }
        }
    }
    if tags.as_object().map(|o| o.is_empty()).unwrap_or(true) {
        return None;
    }
    serde_json::to_string(&tags).ok()
}

fn retryable_failover_status(status: reqwest::StatusCode) -> bool {
    status == reqwest::StatusCode::TOO_MANY_REQUESTS
        || status == reqwest::StatusCode::UNAUTHORIZED
        || status == reqwest::StatusCode::FORBIDDEN
        || status.is_server_error()
}

fn retry_failover_allowed(
    provider: Provider,
    thread_was_affined: bool,
    allow_mid_thread_failover: bool,
) -> bool {
    provider != Provider::Openai || !thread_was_affined || allow_mid_thread_failover
}

async fn proxy(
    state: Arc<AppState>,
    format: ClientFormat,
    path: &'static str,
    headers: HeaderMap,
    body: Bytes,
    peer: Option<std::net::SocketAddr>,
) -> Response {
    let mut run_key: Option<CachedRunKey> = None;
    let client_fingerprint = match client_key(&headers) {
        Some(k) if k == state.local_key => key_fingerprint(&k),
        Some(k) => {
            let key_hash = key_hash_hex(&k);
            match run_key_entry(&state, &key_hash) {
                Some(entry) => {
                    if let Err(e) = state.store.touch_run_key(&key_hash, now_ms()) {
                        tracing::warn!("failed to touch run key: {e}");
                    }
                    run_key = Some(entry);
                    key_hash.chars().take(16).collect()
                }
                None if k.starts_with(RUN_KEY_PREFIX) => {
                    return error_response(
                        StatusCode::UNAUTHORIZED,
                        "run key expired or revoked",
                    )
                }
                None => {
                    return error_response(
                        StatusCode::UNAUTHORIZED,
                        "bad or missing local key (x-api-key / Authorization: Bearer)",
                    )
                }
            }
        }
        None => {
            return error_response(
                StatusCode::UNAUTHORIZED,
                "bad or missing local key (x-api-key / Authorization: Bearer)",
            )
        }
    };

    let in_flight = InFlight::new(&state);
    let trace_id = uuid::Uuid::new_v4().to_string();
    let mut trace = TraceRecord {
        id: trace_id.clone(),
        ts_request_ms: now_ms(),
        method: Some("POST".into()),
        path: Some(path.into()),
        client_format: Some(format.as_str().into()),
        harness: trace_harness(&headers),
        req_headers_json: Some(redacted_headers(&headers)),
        run_id: headers
            .get("x-alexandria-run-id")
            .and_then(|v| v.to_str().ok())
            .map(String::from)
            .or_else(|| run_key.as_ref().and_then(|k| k.run_id.clone())),
        tags: merge_run_key_tags(
            run_key.as_ref().and_then(|k| k.tags_json.as_deref()),
            trace_tags_json(&headers).as_deref(),
        ),
        client_ip: peer.map(|p| p.ip().to_string()),
        key_fingerprint: Some(client_fingerprint),
        ..Default::default()
    };

    let mut body_json: Value = match serde_json::from_slice(&body) {
        Ok(v) => v,
        Err(e) => {
            return error_response(StatusCode::BAD_REQUEST, &format!("body is not JSON: {e}"))
        }
    };
    let (reasoning_effort, thinking_budget) = trace_reasoning_fields(&body_json);
    trace.reasoning_effort = reasoning_effort;
    trace.thinking_budget = thinking_budget;
    let genuine_claude_code = is_genuine_claude_code_request(format, &headers, &body_json);

    trace.session_id = headers
        .get("x-session-id")
        .and_then(|v| v.to_str().ok())
        .map(String::from)
        .or_else(|| {
            headers
                .get("session_id")
                .and_then(|v| v.to_str().ok())
                .map(String::from)
        })
        .or_else(|| {
            headers
                .get("x-claude-code-session-id")
                .and_then(|v| v.to_str().ok())
                .map(String::from)
        })
        .or_else(|| session_from_metadata(&body_json))
        .or_else(|| {
            body_json
                .get("prompt_cache_key")
                .and_then(Value::as_str)
                .filter(|value| !value.is_empty())
                .map(String::from)
        })
        .or_else(|| {
            conversation_root(format, &body_json).map(|root| {
                let ua = headers
                    .get("user-agent")
                    .and_then(|v| v.to_str().ok())
                    .unwrap_or("");
                let ip = peer.map(|p| p.ip().to_string()).unwrap_or_default();
                format!("auto-{}", key_fingerprint(&format!("{root}{ua}{ip}")))
            })
        });

    let requested_model = body_json["model"].as_str().unwrap_or("").to_string();
    if requested_model.is_empty() {
        return error_response(StatusCode::BAD_REQUEST, "missing 'model' in request body");
    }
    let (routed_provider, routed_model) = route_model(&requested_model);
    let provider = routed_provider.unwrap_or_else(|| format.default_provider());
    trace.requested_model = Some(requested_model.clone());
    trace.routed_model = Some(routed_model.clone());
    trace.upstream_provider = Some(provider.as_str().into());
    trace.streamed = Some(body_json["stream"].as_bool().unwrap_or(false));

    // Capture this before planning binds a brand-new session. The policy only
    // suppresses failover for a thread that arrived with an existing affinity;
    // an unaffined first request may still fail over and then pin its winner.
    let codex_thread_was_affined = provider == Provider::Openai
        && preferred_codex_account(&state, trace.session_id.as_deref()).is_some();
    let allow_mid_thread_failover = state
        .vault
        .policy(Provider::Openai)
        .allow_mid_thread_failover;

    let original_body_json = body_json.clone();
    let mut attempted_accounts = HashSet::new();
    let mut plan = match plan_upstream(
        &state,
        format,
        provider,
        &routed_model,
        &mut body_json,
        &body,
        &trace_id,
        &attempted_accounts,
        trace.session_id.as_deref(),
        &headers,
    )
    .await
    {
        Ok(p) => p,
        Err((status, msg)) => {
            trace.status = Some(status.as_u16() as i64);
            trace.error = Some(msg.clone());
            finalize_trace(&state, trace, &body, None, None);
            return error_response(status, &msg);
        }
    };
    trace.upstream_format = Some(plan.upstream_format.into());
    trace.account_id = Some(plan.account.id.clone());
    trace.billing_bucket = Some(
        if plan.account.kind == "oauth" || plan.account.kind == "dario" {
            "subscription"
        } else {
            "api"
        }
        .into(),
    );

    tracing::info!(
        trace_id,
        model = %routed_model,
        provider = provider.as_str(),
        account = %plan.account.id,
        url = %plan.url,
        genuine_claude_code,
        "proxying request"
    );

    let mut account = plan.account.clone();
    let mut upstream_resp = None;
    'accounts: loop {
        if !attempted_accounts.insert(account.id.clone()) {
            tracing::error!(account = %account.id, "refusing to retry an already-attempted account");
            break;
        }

        let mut forced_oauth_refresh = false;
        loop {
            let mut up_headers =
                match upstream_headers(&account, &headers, genuine_claude_code) {
                Ok(h) => h,
                Err((status, msg)) => {
                    trace.status = Some(status.as_u16() as i64);
                    trace.error = Some(msg.clone());
                    finalize_trace(&state, trace, &body, Some(&plan.body), None);
                    return error_response(status, &msg);
                }
            };
            for (k, v) in &plan.extra_headers {
                if let (Ok(name), Ok(value)) = (
                    reqwest::header::HeaderName::from_bytes(k.as_bytes()),
                    HeaderValue::from_str(v),
                ) {
                    up_headers.insert(name, value);
                }
            }
            let resp = match state
                .http
                .post(&plan.url)
                .headers(up_headers)
                .body(plan.body.clone())
                .send()
                .await
            {
                Ok(r) => r,
                Err(e) => {
                    let msg = format!("upstream request failed: {e}");
                    suspect_dario(&state, &account);
                    trace.status = Some(502);
                    trace.error = Some(msg.clone());
                    finalize_trace(&state, trace, &body, Some(&plan.body), None);
                    return error_response(StatusCode::BAD_GATEWAY, &msg);
                }
            };

            if account.provider == Provider::Openai && account.kind == "oauth" {
                if let Some(snapshot) = codex_limits_from_headers(resp.headers()) {
                    if let Err(error) = state
                        .vault
                        .record_codex_limits(&account.id, snapshot)
                        .await
                    {
                        tracing::warn!(account = %account.id, %error, "could not persist Codex limit snapshot");
                    }
                }
            }
            if resp.status() == reqwest::StatusCode::UNAUTHORIZED
                && account.kind == "oauth"
                && !forced_oauth_refresh
            {
                tracing::warn!(
                    account = %account.id,
                    "upstream returned 401 for oauth account; forcing token refresh and retrying"
                );
                forced_oauth_refresh = true;
                match state.vault.refresh(&account.id, true).await {
                    Ok(fresh) => {
                        account = fresh;
                        continue;
                    }
                    Err(e) => {
                        tracing::warn!("forced refresh failed: {e}");
                    }
                }
            }

            if retryable_failover_status(resp.status()) && account.kind != "dario" {
                let retry_after_s = resp
                    .headers()
                    .get("retry-after")
                    .and_then(|v| v.to_str().ok())
                    .and_then(|v| v.parse::<i64>().ok())
                    .unwrap_or(60)
                    .clamp(1, 3600);
                tracing::warn!(account = %account.id, retry_after_s, status = %resp.status(), "upstream failed; cooling account down");
                if let Err(e) = state
                    .vault
                    .mark_cooldown(&account.id, now_ms() + retry_after_s * 1000)
                    .await
                {
                    tracing::warn!("failed to mark cooldown: {e}");
                }

                if !retry_failover_allowed(
                    provider,
                    codex_thread_was_affined,
                    allow_mid_thread_failover,
                ) {
                    tracing::warn!(
                        account = %account.id,
                        status = %resp.status(),
                        "returning retryable Codex error without switching an affined thread"
                    );
                    upstream_resp = Some(resp);
                    break 'accounts;
                }

                let mut retry_body_json = original_body_json.clone();
                match plan_upstream(
                    &state,
                    format,
                    provider,
                    &routed_model,
                    &mut retry_body_json,
                    &body,
                    &trace_id,
                    &attempted_accounts,
                    trace.session_id.as_deref(),
                    &headers,
                )
                .await
                {
                    Ok(next_plan) if !attempted_accounts.contains(&next_plan.account.id) => {
                        plan = next_plan;
                        account = plan.account.clone();
                        trace.account_id = Some(account.id.clone());
                        trace.upstream_format = Some(plan.upstream_format.into());
                        trace.billing_bucket = Some(
                            if account.kind == "oauth" || account.kind == "dario" {
                                "subscription"
                            } else {
                                "api"
                            }
                            .into(),
                        );
                        tracing::warn!(
                            account = %account.id,
                            attempted = attempted_accounts.len(),
                            "retrying with failover account"
                        );
                        continue 'accounts;
                    }
                    Ok(next_plan) => {
                        tracing::error!(
                            account = %next_plan.account.id,
                            "account selector returned an already-attempted failover account"
                        );
                    }
                    Err((status, msg)) => tracing::warn!(
                        status = %status,
                        error = %msg,
                        attempted = attempted_accounts.len(),
                        "no untried failover account available"
                    ),
                }
            }

            upstream_resp = Some(resp);
            break 'accounts;
        }
    }
    let upstream_resp = upstream_resp.expect("upstream response after retry loop");
    trace.account_id = Some(account.id.clone());

    let status = upstream_resp.status();
    trace.status = Some(status.as_u16() as i64);
    let resp_headers = upstream_resp.headers().clone();
    let content_type = resp_headers
        .get("content-type")
        .and_then(|v| v.to_str().ok())
        .unwrap_or("application/octet-stream")
        .to_string();
    trace.resp_headers_json = Some(redacted_headers(&resp_headers));
    let is_sse = content_type.starts_with("text/event-stream");

    if let Some(target) = plan.respond_as {
        use alex_core::translate;
        let buf = match upstream_resp.bytes().await {
            Ok(b) => b.to_vec(),
            Err(e) => {
                let msg = format!("upstream body read failed: {e}");
                suspect_dario(&state, &account);
                trace.status = Some(502);
                trace.error = Some(msg.clone());
                finalize_trace(&state, trace, &body, Some(&plan.body), None);
                return error_response(StatusCode::BAD_GATEWAY, &msg);
            }
        };
        drop(plan.dario_guard.take());
        trace.ts_response_ms = Some(now_ms());
        fill_usage_and_cost(&state, &mut trace, &buf, is_sse);
        let text = String::from_utf8_lossy(&buf).to_string();
        if !status.is_success() {
            finalize_trace(&state, trace, &body, Some(&plan.body), Some(&buf));
            return Response::builder()
                .status(status)
                .header("content-type", "application/json")
                .header("x-alexandria-trace-id", &trace_id)
                .body(Body::from(buf))
                .unwrap_or_else(|e| {
                    error_response(StatusCode::INTERNAL_SERVER_ERROR, &e.to_string())
                });
        }
        let trimmed = text.trim_start();
        let looks_sse = trimmed.starts_with("event:") || trimmed.starts_with("data:");
        let upstream_final = if plan.upstream_format == "anthropic" {
            if is_sse || looks_sse {
                translate::parse_anthropic_sse_to_message(&text)
            } else {
                serde_json::from_str(&text).ok()
            }
        } else if plan.upstream_format == "gemini" {
            translate::parse_gemini_upstream_final(&text)
        } else if plan.upstream_format == "openai-chat" {
            if is_sse || looks_sse {
                translate::parse_openai_chat_sse_final(&text)
            } else {
                serde_json::from_str(&text).ok()
            }
        } else if is_sse || looks_sse {
            translate::parse_responses_sse_final(&text)
        } else {
            serde_json::from_str(&text).ok()
        };
        let Some(upstream_final) = upstream_final else {
            let msg = "could not reassemble upstream response for translation";
            trace.error = Some(msg.to_string());
            finalize_trace(&state, trace, &body, Some(&plan.body), Some(&buf));
            return error_response(StatusCode::BAD_GATEWAY, msg);
        };
        let out = match (target, plan.upstream_format) {
            (RespondAs::Gemini, "gemini") => upstream_final.clone(),
            (RespondAs::Anthropic, "gemini") => {
                translate::gemini_response_to_anthropic(&upstream_final, &requested_model)
            }
            (RespondAs::OpenaiChat, "gemini") => translate::anthropic_response_to_openai_chat(
                &translate::gemini_response_to_anthropic(&upstream_final, &requested_model),
                &requested_model,
            ),
            (RespondAs::OpenaiResponses, "gemini") => {
                translate::anthropic_response_to_openai_responses(
                    &translate::gemini_response_to_anthropic(&upstream_final, &requested_model),
                    &requested_model,
                )
            }
            (RespondAs::OpenaiChat, "anthropic") => {
                translate::anthropic_response_to_openai_chat(&upstream_final, &requested_model)
            }
            (RespondAs::OpenaiResponses, "anthropic") => translate::anthropic_response_to_openai_responses(
                &upstream_final,
                &requested_model,
            ),
            (RespondAs::Gemini, "anthropic") => {
                translate::anthropic_response_to_gemini(&upstream_final, &requested_model)
            }
            (RespondAs::Anthropic, "openai-chat") => {
                translate::openai_chat_response_to_anthropic(&upstream_final, &requested_model)
            }
            (RespondAs::OpenaiChat, "openai-chat") => upstream_final.clone(),
            (RespondAs::OpenaiResponses, "openai-chat") => {
                translate::anthropic_response_to_openai_responses(
                    &translate::openai_chat_response_to_anthropic(
                        &upstream_final,
                        &requested_model,
                    ),
                    &requested_model,
                )
            }
            (RespondAs::Gemini, "openai-chat") => translate::anthropic_response_to_gemini(
                &translate::openai_chat_response_to_anthropic(&upstream_final, &requested_model),
                &requested_model,
            ),
            (RespondAs::Gemini, _) => translate::anthropic_response_to_gemini(
                &translate::responses_final_to_anthropic(&upstream_final, &requested_model),
                &requested_model,
            ),
            (RespondAs::Anthropic, "anthropic") | (RespondAs::OpenaiResponses, _) => {
                upstream_final.clone()
            }
            (RespondAs::Anthropic, _) => {
                translate::responses_final_to_anthropic(&upstream_final, &requested_model)
            }
            (RespondAs::OpenaiChat, _) => {
                translate::responses_final_to_openai_chat(&upstream_final, &requested_model)
            }
        };
        let (out_ct, out_body) = if plan.client_stream {
            let sse = match target {
                RespondAs::Anthropic => translate::synth_anthropic_sse(&out),
                RespondAs::OpenaiChat => translate::synth_openai_chat_sse(&out),
                RespondAs::OpenaiResponses => translate::synth_openai_responses_sse(&out),
                RespondAs::Gemini => translate::synth_gemini_sse(&out),
            };
            ("text/event-stream", sse.into_bytes())
        } else {
            (
                "application/json",
                serde_json::to_vec(&out).unwrap_or_default(),
            )
        };
        finalize_trace(&state, trace, &body, Some(&plan.body), Some(&buf));
        return Response::builder()
            .status(StatusCode::OK)
            .header("content-type", out_ct)
            .header("x-alexandria-trace-id", &trace_id)
            .body(Body::from(out_body))
            .unwrap_or_else(|e| error_response(StatusCode::INTERNAL_SERVER_ERROR, &e.to_string()));
    }

    if plan.destream {
        let buf = match upstream_resp.bytes().await {
            Ok(b) => b.to_vec(),
            Err(e) => {
                let msg = format!("upstream body read failed: {e}");
                trace.status = Some(502);
                trace.error = Some(msg.clone());
                finalize_trace(&state, trace, &body, Some(&plan.body), None);
                return error_response(StatusCode::BAD_GATEWAY, &msg);
            }
        };
        trace.ts_response_ms = Some(now_ms());
        fill_usage_and_cost(&state, &mut trace, &buf, is_sse);
        let text = String::from_utf8_lossy(&buf).to_string();
        let (out_status, out_body) = if status.is_success() {
            match extract_final_response(&text) {
                Some(v) => (StatusCode::OK, serde_json::to_vec(&v).unwrap_or_default()),
                None => {
                    let msg = "upstream stream ended without a response.completed event";
                    trace.error = Some(msg.to_string());
                    (
                        StatusCode::BAD_GATEWAY,
                        serde_json::to_vec(
                            &json!({"error": {"type": "alexandria", "message": msg}}),
                        )
                        .unwrap_or_default(),
                    )
                }
            }
        } else {
            (status, buf.clone())
        };
        finalize_trace(&state, trace, &body, Some(&plan.body), Some(&buf));
        return Response::builder()
            .status(out_status)
            .header("content-type", "application/json")
            .header("x-alexandria-trace-id", &trace_id)
            .body(Body::from(out_body))
            .unwrap_or_else(|e| error_response(StatusCode::INTERNAL_SERVER_ERROR, &e.to_string()));
    }

    let (tx, rx) = tokio::sync::mpsc::channel::<Result<Bytes, std::io::Error>>(64);
    let mut upstream_stream = upstream_resp.bytes_stream();
    let state2 = state.clone();
    let client_body = body.clone();
    let upstream_body = plan.body.clone();
    let dario_guard = plan.dario_guard.take();
    tokio::spawn(async move {
        let _dario_guard = dario_guard;
        let _in_flight = in_flight;
        let mut buf: Vec<u8> = Vec::new();
        let mut stream_error: Option<String> = None;
        while let Some(chunk) = upstream_stream.next().await {
            match chunk {
                Ok(b) => {
                    buf.extend_from_slice(&b);
                    let _ = tx.send(Ok(b)).await;
                }
                Err(e) => {
                    stream_error = Some(format!("upstream stream error: {e}"));
                    let _ = tx
                        .send(Err(std::io::Error::new(
                            std::io::ErrorKind::Other,
                            e.to_string(),
                        )))
                        .await;
                    break;
                }
            }
        }
        drop(tx);
        trace.ts_response_ms = Some(now_ms());
        if trace.error.is_none() {
            trace.error = stream_error;
        }
        if trace.error.is_some() {
            if let (Some(dario), Some(gen)) = (
                &state2.dario,
                trace
                    .account_id
                    .as_deref()
                    .and_then(|id| id.strip_prefix("dario:")),
            ) {
                dario.suspect(gen);
            }
        }

        fill_usage_and_cost(&state2, &mut trace, &buf, is_sse);
        finalize_trace(
            &state2,
            trace,
            &client_body,
            Some(&upstream_body),
            Some(&buf),
        );
    });

    let mut response = Response::builder().status(status);
    for (k, v) in resp_headers.iter() {
        let key = k.as_str().to_lowercase();
        if [
            "transfer-encoding",
            "connection",
            "content-encoding",
            "content-length",
        ]
        .contains(&key.as_str())
        {
            continue;
        }
        response = response.header(k, v);
    }
    response = response.header("x-alexandria-trace-id", &trace_id);
    response
        .body(Body::from_stream(ReceiverStream::new(rx)))
        .unwrap_or_else(|e| error_response(StatusCode::INTERNAL_SERVER_ERROR, &e.to_string()))
}

fn fill_usage_and_cost(state: &AppState, trace: &mut TraceRecord, buf: &[u8], is_sse: bool) {
    let text = String::from_utf8_lossy(buf);
    let trimmed = text.trim_start();
    let looks_sse = trimmed.starts_with("event:") || trimmed.starts_with("data:");
    trace.usage = if is_sse || looks_sse {
        parse_sse_usage(&text)
    } else {
        serde_json::from_str::<Value>(&text)
            .map(|v| usage_from_json(&v))
            .unwrap_or_default()
    };
    if !trace.usage.is_empty() {
        if let Some(pricing) = trace
            .routed_model
            .as_deref()
            .and_then(|m| state.store.pricing_for(m))
        {
            let input_includes_cached = trace
                .upstream_format
                .as_deref()
                .map(|f| f.starts_with("openai"))
                .unwrap_or(false);
            trace.cost_usd = Some(compute_cost(&trace.usage, &pricing, input_includes_cached));
        }
    }
}

fn extract_final_response(text: &str) -> Option<Value> {
    let mut last: Option<Value> = None;
    for line in text.lines() {
        let Some(data) = line.strip_prefix("data:") else {
            continue;
        };
        let Ok(v) = serde_json::from_str::<Value>(data.trim()) else {
            continue;
        };
        if matches!(
            v["type"].as_str(),
            Some("response.completed") | Some("response.incomplete") | Some("response.failed")
        ) && v["response"].is_object()
        {
            last = Some(v["response"].clone());
        }
    }
    last
}

fn finalize_trace(
    state: &AppState,
    mut trace: TraceRecord,
    client_body: &[u8],
    upstream_body: Option<&[u8]>,
    resp_body: Option<&[u8]>,
) {
    let store = &state.store;
    match store.write_body(&trace.id, "request.json", client_body) {
        Ok(p) => trace.req_body_path = Some(p),
        Err(e) => tracing::warn!("failed to write request body: {e}"),
    }
    if let Some(up) = upstream_body {
        if up != client_body {
            match store.write_body(&trace.id, "upstream-request.json", up) {
                Ok(p) => trace.upstream_req_body_path = Some(p),
                Err(e) => tracing::warn!("failed to write upstream request body: {e}"),
            }
        }
    }
    if let Some(resp) = resp_body {
        match store.write_body(&trace.id, "response.body", resp) {
            Ok(p) => trace.resp_body_path = Some(p),
            Err(e) => tracing::warn!("failed to write response body: {e}"),
        }
    }
    if let Err(e) = store.insert_trace(&trace) {
        tracing::error!("failed to insert trace {}: {e}", trace.id);
    } else {
        tracing::info!(
            trace_id = %trace.id,
            status = trace.status,
            input = trace.usage.input_tokens,
            output = trace.usage.output_tokens,
            cost = trace.cost_usd,
            "trace recorded"
        );
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::path::PathBuf;

    fn tmpdir(name: &str) -> PathBuf {
        let dir = std::env::temp_dir().join(format!(
            "alex-proxy-{name}-{}-{}",
            std::process::id(),
            std::time::SystemTime::now()
                .duration_since(std::time::UNIX_EPOCH)
                .unwrap()
                .as_nanos()
        ));
        std::fs::create_dir_all(&dir).unwrap();
        dir
    }

    fn test_state(name: &str) -> Arc<AppState> {
        test_state_with_dario(name, None)
    }

    fn test_state_with_dario(
        name: &str,
        dario: Option<Arc<dyn DarioRouter>>,
    ) -> Arc<AppState> {
        let dir = tmpdir(name);
        let store = Arc::new(Store::open(dir.join("store")).unwrap());
        let vault = Arc::new(Vault::open(dir.join("vault")).unwrap());
        build_state(
            "alx-local".into(),
            vault,
            store,
            dario,
            "http://127.0.0.1:4100".into(),
        )
    }

    fn anthropic_account() -> Account {
        Account {
            id: "anthropic:test".into(),
            provider: Provider::Anthropic,
            kind: "oauth".into(),
            name: "test".into(),
            description: None,
            paused: false,
            label: None,
            access_token: Some("direct-token".into()),
            refresh_token: None,
            id_token: None,
            api_key: None,
            expires_at_ms: Some(now_ms() + 3_600_000),
            last_refresh_ms: None,
            account_meta: Value::Null,
            cooldown_until_ms: None,
            status: "active".into(),
            path: None,
        }
    }

    struct FakeDario {
        active: Option<DarioActive>,
        begin_succeeds: bool,
    }

    impl DarioRouter for FakeDario {
        fn active(&self) -> Option<DarioActive> {
            self.active.clone()
        }

        fn begin(&self, _generation_id: &str) -> Option<Box<dyn std::any::Any + Send>> {
            self.begin_succeeds
                .then(|| Box::new(()) as Box<dyn std::any::Any + Send>)
        }

        fn status(&self) -> Value {
            json!({"active": self.active.is_some()})
        }

        fn suspect(&self, _generation_id: &str) {}
    }

    fn active_dario() -> Arc<dyn DarioRouter> {
        Arc::new(FakeDario {
            active: Some(DarioActive {
                generation_id: "test-generation".into(),
                base_url: "http://127.0.0.1:9191".into(),
                api_key: "dario-key".into(),
            }),
            begin_succeeds: true,
        })
    }

    fn claude_code_request() -> (HeaderMap, Value) {
        let mut headers = HeaderMap::new();
        headers.insert("user-agent", HeaderValue::from_static("claude-cli/2.1.0"));
        headers.insert("x-app", HeaderValue::from_static("cli"));
        headers.insert(
            "x-claude-code-session-id",
            HeaderValue::from_static("session-123"),
        );
        let body = json!({
            "model": "claude-sonnet-4-5",
            "system": [{
                "type": "text",
                "text": "x-anthropic-billing-header: cc_version=2.1.0"
            }],
            "messages": []
        });
        (headers, body)
    }

    #[test]
    fn claude_code_detection_requires_the_complete_signature() {
        let (headers, body) = claude_code_request();
        assert!(is_genuine_claude_code_request(
            ClientFormat::AnthropicMessages,
            &headers,
            &body
        ));
        assert!(!is_genuine_claude_code_request(
            ClientFormat::OpenaiChat,
            &headers,
            &body
        ));

        for required in ["user-agent", "x-app", "x-claude-code-session-id"] {
            let mut missing = headers.clone();
            missing.remove(required);
            assert!(
                !is_genuine_claude_code_request(
                    ClientFormat::AnthropicMessages,
                    &missing,
                    &body
                ),
                "request without {required} must not bypass Dario"
            );
        }

        let mut missing_billing = body.clone();
        missing_billing["system"] = json!([]);
        assert!(!is_genuine_claude_code_request(
            ClientFormat::AnthropicMessages,
            &headers,
            &missing_billing
        ));

        let mut explicit_other_harness = headers;
        explicit_other_harness.insert(
            "x-alexandria-harness",
            HeaderValue::from_static("pi"),
        );
        assert!(!is_genuine_claude_code_request(
            ClientFormat::AnthropicMessages,
            &explicit_other_harness,
            &body
        ));

        let (mut conflicting_harness, body) = claude_code_request();
        conflicting_harness.append(
            "x-alexandria-harness",
            HeaderValue::from_static("claude-code"),
        );
        conflicting_harness.append(
            "x-alexandria-harness",
            HeaderValue::from_static("codex"),
        );
        assert!(!is_genuine_claude_code_request(
            ClientFormat::AnthropicMessages,
            &conflicting_harness,
            &body
        ));

        let (mut malformed_harness, body) = claude_code_request();
        malformed_harness.insert(
            "x-alexandria-harness",
            HeaderValue::from_bytes(b"\x80").unwrap(),
        );
        assert!(!is_genuine_claude_code_request(
            ClientFormat::AnthropicMessages,
            &malformed_harness,
            &body
        ));
    }

    #[tokio::test]
    async fn anthropic_routing_bypasses_dario_only_for_claude_code() {
        let state = test_state_with_dario("dario-routing", Some(active_dario()));
        state.vault.upsert(anthropic_account()).await.unwrap();
        let (claude_headers, request) = claude_code_request();
        let original = serde_json::to_vec(&request).unwrap();
        let excluded = HashSet::new();

        let mut direct_body = request.clone();
        let direct = plan_upstream(
            &state,
            ClientFormat::AnthropicMessages,
            Provider::Anthropic,
            "claude-sonnet-4-5",
            &mut direct_body,
            &original,
            "trace-direct",
            &excluded,
            None,
            &claude_headers,
        )
        .await
        .unwrap();
        assert_eq!(direct.url, "https://api.anthropic.com/v1/messages");
        assert_eq!(direct.account.id, "anthropic:test");
        assert!(direct.extra_headers.is_empty());

        let mut harness_headers = HeaderMap::new();
        harness_headers.insert("x-alexandria-harness", HeaderValue::from_static("pi"));
        let mut dario_body = request;
        let dario = plan_upstream(
            &state,
            ClientFormat::AnthropicMessages,
            Provider::Anthropic,
            "claude-sonnet-4-5",
            &mut dario_body,
            &original,
            "trace-dario",
            &excluded,
            None,
            &harness_headers,
        )
        .await
        .unwrap();
        assert_eq!(dario.url, "http://127.0.0.1:9191/v1/messages");
        assert_eq!(dario.account.kind, "dario");
        assert!(dario
            .extra_headers
            .contains(&("x-dario-capture-id".into(), "trace-dario".into())));
    }

    #[tokio::test]
    async fn configured_dario_fails_closed_when_unhealthy() {
        let state = test_state_with_dario(
            "dario-unhealthy",
            Some(Arc::new(FakeDario {
                active: None,
                begin_succeeds: false,
            })),
        );
        state.vault.upsert(anthropic_account()).await.unwrap();
        let (_, mut request) = claude_code_request();
        let client_headers = HeaderMap::new();
        let excluded = HashSet::new();
        let original = serde_json::to_vec(&request).unwrap();

        let result = plan_upstream(
            &state,
            ClientFormat::AnthropicMessages,
            Provider::Anthropic,
            "claude-sonnet-4-5",
            &mut request,
            &original,
            "trace-unhealthy",
            &excluded,
            None,
            &client_headers,
        )
        .await;
        match result {
            Err((status, message)) => {
                assert_eq!(status, StatusCode::SERVICE_UNAVAILABLE);
                assert!(message.contains("no healthy generation"));
            }
            Ok(_) => panic!("an unhealthy configured Dario must not fall back to direct traffic"),
        }
    }

    #[tokio::test]
    async fn dario_generation_race_fails_closed() {
        let state = test_state_with_dario(
            "dario-generation-race",
            Some(Arc::new(FakeDario {
                active: Some(DarioActive {
                    generation_id: "vanished-generation".into(),
                    base_url: "http://127.0.0.1:9191".into(),
                    api_key: "dario-key".into(),
                }),
                begin_succeeds: false,
            })),
        );
        let (_, mut request) = claude_code_request();
        let original = serde_json::to_vec(&request).unwrap();
        let client_headers = HeaderMap::new();
        let excluded = HashSet::new();

        let result = plan_upstream(
            &state,
            ClientFormat::AnthropicMessages,
            Provider::Anthropic,
            "claude-sonnet-4-5",
            &mut request,
            &original,
            "trace-race",
            &excluded,
            None,
            &client_headers,
        )
        .await;
        match result {
            Err((status, message)) => {
                assert_eq!(status, StatusCode::SERVICE_UNAVAILABLE);
                assert!(message.contains("became unavailable"));
            }
            Ok(_) => panic!("a vanished Dario generation must not fall back to direct traffic"),
        }
    }

    #[test]
    fn direct_claude_code_headers_are_allowlisted() {
        let (mut client_headers, _) = claude_code_request();
        for name in CLAUDE_CODE_PASSTHROUGH_HEADERS {
            client_headers.insert(*name, HeaderValue::from_static("safe-client-value"));
        }
        client_headers.insert(
            "authorization",
            HeaderValue::from_static("Bearer client-secret"),
        );
        client_headers.insert("x-api-key", HeaderValue::from_static("client-secret"));
        client_headers.insert("host", HeaderValue::from_static("attacker.invalid"));
        client_headers.insert("content-length", HeaderValue::from_static("999"));
        client_headers.insert("connection", HeaderValue::from_static("close"));
        client_headers.insert("accept-encoding", HeaderValue::from_static("br"));
        client_headers.insert(
            "x-alexandria-harness",
            HeaderValue::from_static("claude"),
        );
        client_headers.insert(
            "x-dario-capture-id",
            HeaderValue::from_static("spoofed"),
        );

        let direct = upstream_headers(&anthropic_account(), &client_headers, true).unwrap();
        for name in CLAUDE_CODE_PASSTHROUGH_HEADERS {
            assert_eq!(direct[*name], "safe-client-value", "missing {name}");
        }
        assert_eq!(direct["authorization"], "Bearer direct-token");
        assert!(direct.get("x-api-key").is_none());
        assert!(direct.get("host").is_none());
        assert!(direct.get("content-length").is_none());
        assert!(direct.get("connection").is_none());
        assert_eq!(direct["accept-encoding"], "identity");
        assert!(direct.get("x-alexandria-harness").is_none());
        assert!(direct.get("x-dario-capture-id").is_none());

        let non_claude = upstream_headers(&anthropic_account(), &client_headers, false).unwrap();
        assert!(non_claude.get("x-app").is_none());
        assert!(non_claude.get("x-claude-code-session-id").is_none());
        assert_eq!(non_claude["authorization"], "Bearer direct-token");
    }

    async fn response_json(resp: Response) -> (StatusCode, Value) {
        let status = resp.status();
        let body = axum::body::to_bytes(resp.into_body(), usize::MAX)
            .await
            .unwrap();
        let value = serde_json::from_slice(&body).unwrap_or(Value::Null);
        (status, value)
    }

    struct FakeUpdater {
        result: Result<Value, UpdateApplyError>,
    }

    impl DaemonUpdater for FakeUpdater {
        fn apply(&self) -> UpdateApplyFuture {
            let result = match &self.result {
                Ok(body) => Ok(body.clone()),
                Err(UpdateApplyError::Conflict(body)) => {
                    Err(UpdateApplyError::Conflict(body.clone()))
                }
                Err(UpdateApplyError::Failed(message)) => {
                    Err(UpdateApplyError::Failed(message.clone()))
                }
            };
            Box::pin(async move { result })
        }
    }

    fn test_openai_account(name: &str) -> Account {
        Account {
            id: if name == "default" {
                "openai-oauth".into()
            } else {
                format!("openai-oauth-{name}")
            },
            provider: Provider::Openai,
            kind: "oauth".into(),
            name: name.into(),
            description: None,
            paused: false,
            label: Some("codex (test)".into()),
            access_token: Some(format!("token-{name}")),
            refresh_token: Some(format!("refresh-{name}")),
            id_token: None,
            api_key: None,
            expires_at_ms: Some(now_ms() + 3_600_000),
            last_refresh_ms: Some(now_ms()),
            account_meta: json!({"account_id": format!("chatgpt-{name}")}),
            cooldown_until_ms: None,
            status: "active".into(),
            path: None,
        }
    }

    #[test]
    fn codex_affinity_cache_expires_and_evicts_oldest_session() {
        let mut cache = CodexAffinityCache {
            entries: HashMap::new(),
            ttl_ms: 100,
            max_entries: 2,
        };
        cache.bind("session-a", "account-a", 0);
        cache.bind("session-b", "account-b", 10);
        assert_eq!(cache.preferred("session-a", 20).as_deref(), Some("account-a"));

        // session-a's lookup extends its expiry, so session-b is oldest.
        cache.bind("session-c", "account-c", 30);
        assert!(cache.preferred("session-b", 30).is_none());
        assert_eq!(cache.preferred("session-a", 30).as_deref(), Some("account-a"));
        assert_eq!(cache.preferred("session-c", 30).as_deref(), Some("account-c"));
        assert!(cache.preferred("session-a", 131).is_none());
    }

    #[test]
    fn captures_only_safe_codex_limit_headers() {
        let mut headers = reqwest::header::HeaderMap::new();
        headers.insert("authorization", HeaderValue::from_static("Bearer secret"));
        headers.insert("x-codex-plan-type", HeaderValue::from_static("plus"));
        headers.insert(
            "x-codex-primary-used-percent",
            HeaderValue::from_static("42"),
        );
        headers.insert(
            "x-codex-primary-window-minutes",
            HeaderValue::from_static("300"),
        );
        headers.insert(
            "x-codex-primary-reset-at",
            HeaderValue::from_static("1800000000"),
        );
        let snapshot = codex_limits_from_headers(&headers).unwrap();
        assert_eq!(snapshot["plan"], "plus");
        assert_eq!(snapshot["windows"][0]["window"], "5h");
        assert_eq!(snapshot["windows"][0]["used_pct"], 42.0);
        assert!(!snapshot.to_string().contains("secret"));
    }

    #[tokio::test]
    async fn codex_routing_api_updates_eligibility_order_and_reserve() {
        let state = test_state("codex-routing");
        state
            .vault
            .upsert(test_openai_account("default"))
            .await
            .unwrap();
        state
            .vault
            .upsert(test_openai_account("work"))
            .await
            .unwrap();
        let (status, body) = response_json(
            admin_codex_routing_update(
                State(state.clone()),
                axum::Json(json!({
                    "strategy": "priority",
                    "reserve_pct": 15,
                    "allow_mid_thread_failover": false,
                    "accounts": [
                        {"account_id": "openai-oauth-work", "eligible": true, "priority": 0, "reserve_pct": 22},
                        {"account_id": "openai-oauth", "eligible": false, "priority": 1}
                    ]
                })),
            )
            .await,
        )
        .await;
        assert_eq!(status, StatusCode::OK);
        assert_eq!(body["strategy"], "priority");
        assert_eq!(body["reserve_pct"], 15);
        assert_eq!(body["allow_mid_thread_failover"], false);
        assert_eq!(body["accounts"][0]["account_id"], "openai-oauth-work");
        assert_eq!(body["accounts"][0]["eligible"], true);
        assert_eq!(body["accounts"][0]["reserve_pct"], 22);
        assert_eq!(body["accounts"][1]["eligible"], false);
        assert_eq!(body["accounts"][1]["reserve_pct"], 15);
        assert_eq!(
            state
                .vault
                .policy(Provider::Openai)
                .account_reserve_pct["work"],
            22
        );
        assert_eq!(
            state
                .vault
                .account_for(Provider::Openai, true)
                .await
                .unwrap()
                .name,
            "work"
        );
    }

    #[tokio::test]
    async fn admin_update_apply_response_shapes() {
        let state = test_state("admin-update-apply");
        set_daemon_updater(
            &state,
            Arc::new(FakeUpdater {
                result: Ok(json!({
                    "applying": true,
                    "current": "0.1.0",
                    "latest": "0.2.0",
                    "update_available": true,
                })),
            }),
        );
        let (status, body) = response_json(admin_update_apply(State(state.clone())).await).await;
        assert_eq!(status, StatusCode::ACCEPTED);
        assert_eq!(body["applying"], true);
        assert_eq!(body["latest"], "0.2.0");

        set_daemon_updater(
            &state,
            Arc::new(FakeUpdater {
                result: Ok(json!({
                    "applying": false,
                    "current": "0.2.0",
                    "latest": "0.2.0",
                    "update_available": false,
                })),
            }),
        );
        let (status, body) = response_json(admin_update_apply(State(state.clone())).await).await;
        assert_eq!(status, StatusCode::OK);
        assert_eq!(body["update_available"], false);

        set_daemon_updater(
            &state,
            Arc::new(FakeUpdater {
                result: Err(UpdateApplyError::Conflict(json!({
                    "applying": false,
                    "current": "0.1.0",
                    "latest": "0.2.0",
                    "update_available": true,
                    "reason": "alex is managed by Homebrew - run `brew upgrade alex`",
                }))),
            }),
        );
        let (status, body) = response_json(admin_update_apply(State(state)).await).await;
        assert_eq!(status, StatusCode::CONFLICT);
        assert_eq!(
            body["reason"],
            "alex is managed by Homebrew - run `brew upgrade alex`"
        );
    }

    #[tokio::test]
    async fn harness_run_key_mints_lists_validates_and_revokes() {
        let state = test_state("harness-run-key");
        let (_status, created) = response_json(
            admin_run_keys_create(
                State(state.clone()),
                Some(axum::Json(json!({
                    "kind": "harness",
                    "label": "pi",
                    "ttl_seconds": 1,
                    "tags": {"harness": "pi"},
                }))),
            )
            .await,
        )
        .await;
        assert_eq!(_status, StatusCode::CREATED);
        assert_eq!(created["kind"], "harness");
        assert_eq!(created["label"], "pi");
        assert_eq!(created["expires_ms"], Value::Null);
        let id = created["id"].as_str().unwrap().to_string();
        let key = created["key"].as_str().unwrap();
        let hash = key_hash_hex(key);

        let entry = run_key_entry(&state, &hash).unwrap();
        assert_eq!(entry.expires_ms, None);
        assert_eq!(entry.tags_json.as_deref(), Some(r#"{"harness":"pi"}"#));

        let (_status, listed) = response_json(
            admin_run_keys_list(State(state.clone()), Query(HashMap::new())).await,
        )
        .await;
        assert_eq!(_status, StatusCode::OK);
        let rows = listed["run_keys"].as_array().unwrap();
        assert_eq!(rows.len(), 1);
        assert_eq!(rows[0]["id"], id);
        assert_eq!(rows[0]["kind"], "harness");
        assert_eq!(rows[0]["label"], "pi");
        assert_eq!(rows[0]["expires_ms"], Value::Null);

        let (_status, revoked) =
            response_json(admin_run_keys_revoke(State(state.clone()), Path(id)).await).await;
        assert_eq!(_status, StatusCode::OK);
        assert_eq!(revoked["revoked"], true);
        assert!(run_key_entry(&state, &hash).is_none());
    }

    #[tokio::test]
    async fn harness_run_key_requires_label() {
        let state = test_state("harness-run-key-label");
        let (status, body) = response_json(
            admin_run_keys_create(
                State(state),
                Some(axum::Json(json!({
                    "kind": "harness",
                    "tags": {"harness": "pi"},
                }))),
            )
            .await,
        )
        .await;
        assert_eq!(status, StatusCode::BAD_REQUEST);
        assert!(body["error"]["message"]
            .as_str()
            .unwrap()
            .contains("label"));
    }

    #[test]
    fn failover_statuses_are_limited_to_auth_rate_limit_and_server_errors() {
        for status in [
            reqwest::StatusCode::UNAUTHORIZED,
            reqwest::StatusCode::FORBIDDEN,
            reqwest::StatusCode::TOO_MANY_REQUESTS,
            reqwest::StatusCode::INTERNAL_SERVER_ERROR,
            reqwest::StatusCode::BAD_GATEWAY,
            reqwest::StatusCode::SERVICE_UNAVAILABLE,
        ] {
            assert!(
                retryable_failover_status(status),
                "{status} should fail over"
            );
        }
        for status in [
            reqwest::StatusCode::OK,
            reqwest::StatusCode::BAD_REQUEST,
            reqwest::StatusCode::NOT_FOUND,
            reqwest::StatusCode::UNPROCESSABLE_ENTITY,
        ] {
            assert!(
                !retryable_failover_status(status),
                "{status} should be returned without failover"
            );
        }
    }

    #[test]
    fn disabled_mid_thread_failover_only_pins_existing_codex_threads() {
        assert!(!retry_failover_allowed(Provider::Openai, true, false));
        assert!(retry_failover_allowed(Provider::Openai, false, false));
        assert!(retry_failover_allowed(Provider::Openai, true, true));
        assert!(retry_failover_allowed(Provider::Anthropic, true, false));
    }

    #[tokio::test]
    async fn codex_planning_excludes_every_account_already_attempted() {
        let state = test_state("codex-plan-exclusions");
        for (id, name) in [
            ("openai-oauth-a", "a"),
            ("openai-oauth-b", "b"),
            ("openai-oauth-c", "c"),
        ] {
            state
                .vault
                .upsert(Account {
                    id: id.into(),
                    provider: Provider::Openai,
                    kind: "oauth".into(),
                    name: name.into(),
                    description: None,
                    paused: false,
                    label: None,
                    access_token: Some(format!("access-{name}")),
                    refresh_token: Some(format!("refresh-{name}")),
                    id_token: None,
                    api_key: None,
                    expires_at_ms: Some(now_ms() + 3_600_000),
                    last_refresh_ms: Some(now_ms()),
                    account_meta: json!({"account_id": format!("chatgpt-{name}")}),
                    cooldown_until_ms: None,
                    status: "active".into(),
                    path: None,
                })
                .await
                .unwrap();
        }

        let original = json!({"model": "gpt-5.3-codex", "stream": true});
        let raw = serde_json::to_vec(&original).unwrap();
        let client_headers = HeaderMap::new();
        let mut excluded = HashSet::new();
        let mut selected = Vec::new();
        for _ in 0..3 {
            let mut body = original.clone();
            let plan = plan_upstream(
                &state,
                ClientFormat::OpenaiResponses,
                Provider::Openai,
                "gpt-5.3-codex",
                &mut body,
                &raw,
                "trace-test",
                &excluded,
                None,
                &client_headers,
            )
            .await
            .unwrap();
            assert!(excluded.insert(plan.account.id.clone()));
            selected.push(plan.account.id);
        }

        assert_eq!(
            selected,
            vec!["openai-oauth-a", "openai-oauth-b", "openai-oauth-c"]
        );
        let mut body = original;
        assert!(plan_upstream(
            &state,
            ClientFormat::OpenaiResponses,
            Provider::Openai,
            "gpt-5.3-codex",
            &mut body,
            &raw,
            "trace-test",
            &excluded,
            None,
            &client_headers,
        )
        .await
        .is_err());
    }

    #[tokio::test]
    async fn codex_sessions_stay_affined_and_failover_rebinds() {
        let state = test_state("codex-session-affinity");
        state
            .vault
            .upsert(test_openai_account("a"))
            .await
            .unwrap();
        state
            .vault
            .upsert(test_openai_account("b"))
            .await
            .unwrap();
        state
            .vault
            .set_policies(vec![(
                Provider::Openai,
                AccountPolicy {
                    order: vec!["a".into(), "b".into()],
                    mode: AccountPolicyMode::RoundRobin,
                    reserve_pct: Some(0),
                    ..AccountPolicy::default()
                },
            )])
            .await;

        let request = json!({"model": "gpt-5.3-codex", "stream": true});
        let raw = serde_json::to_vec(&request).unwrap();
        let headers = HeaderMap::new();
        let none_excluded = HashSet::new();

        let mut body = request.clone();
        let session_a_first = plan_upstream(
            &state,
            ClientFormat::OpenaiResponses,
            Provider::Openai,
            "gpt-5.3-codex",
            &mut body,
            &raw,
            "trace-a-1",
            &none_excluded,
            Some("session-a"),
            &headers,
        )
        .await
        .unwrap()
        .account
        .id;

        let mut body = request.clone();
        let session_b_first = plan_upstream(
            &state,
            ClientFormat::OpenaiResponses,
            Provider::Openai,
            "gpt-5.3-codex",
            &mut body,
            &raw,
            "trace-b-1",
            &none_excluded,
            Some("session-b"),
            &headers,
        )
        .await
        .unwrap()
        .account
        .id;
        assert_ne!(session_a_first, session_b_first);

        let mut body = request.clone();
        let session_a_again = plan_upstream(
            &state,
            ClientFormat::OpenaiResponses,
            Provider::Openai,
            "gpt-5.3-codex",
            &mut body,
            &raw,
            "trace-a-2",
            &none_excluded,
            Some("session-a"),
            &headers,
        )
        .await
        .unwrap()
        .account
        .id;
        assert_eq!(session_a_again, session_a_first);

        let excluded = HashSet::from([session_a_first.clone()]);
        let mut body = request.clone();
        let failover = plan_upstream(
            &state,
            ClientFormat::OpenaiResponses,
            Provider::Openai,
            "gpt-5.3-codex",
            &mut body,
            &raw,
            "trace-a-failover",
            &excluded,
            Some("session-a"),
            &headers,
        )
        .await
        .unwrap()
        .account
        .id;
        assert_eq!(failover, session_b_first);

        let mut body = request;
        let after_failover = plan_upstream(
            &state,
            ClientFormat::OpenaiResponses,
            Provider::Openai,
            "gpt-5.3-codex",
            &mut body,
            &raw,
            "trace-a-3",
            &none_excluded,
            Some("session-a"),
            &headers,
        )
        .await
        .unwrap()
        .account
        .id;
        assert_eq!(after_failover, failover);
    }
}