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