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