feat: usage analytics, quota reset times, and zai balance support

src-tauri/Cargo.toml

@@ -22,6 +22,7 @@ tauri-plugin-updater = "2"
 tauri-plugin-process = "2"
 serde = { version = "1", features = ["derive"] }
 serde_json = "1"
+chrono = { version = "0.4", default-features = false, features = ["std"] }
 json5 = "1.3"
 tokio = { version = "1", features = ["full"] }
 axum = { version = "0.8", features = ["http2", "ws"] }

src-tauri/src/proxy.rs

@@ -225,6 +225,10 @@ struct Turn {
     key_id: Option<String>,
 }
 
+fn routed_stats_model(provider: &Provider, upstream_model: &str) -> String {
+    format!("{}/{}", provider.id, upstream_model)
+}
+
 impl Turn {
     fn new(
         provider: &str,
@@ -1063,6 +1067,11 @@ fn sanitize_stateless_responses_payload(payload: &mut Value) {
     let Some(input) = payload.get_mut("input").and_then(Value::as_array_mut) else {
         return;
     };
+    translate::repair_tool_exchange_items(input);
+    if input.is_empty() {
+        payload["input"] = Value::String(String::new());
+        return;
+    }
     for item in input.iter_mut() {
         if let Some(object) = item.as_object_mut() {
             object.remove("id");

src-tauri/src/proxy/auth.rs

@@ -1,22 +1,77 @@
 use axum::http::StatusCode;
 use axum::{body::Body, extract::Request, middleware::Next, response::Response};
+use std::path::Path;
 use std::sync::OnceLock;
 
 // The proxy accepts traffic from any local process, so every endpoint needs a
-// process-local secret before it can use a configured provider credential.
+// user-local secret before it can use a configured provider credential.
 static LOCAL_TOKEN: OnceLock<String> = OnceLock::new();
 
 /// Shared secret required on every request to the local proxy.
-/// Generated once per process from 32 random bytes (two UUIDv4s = 64 hex
-/// chars; `uuid` is the only RNG crate already in Cargo.toml).
+///
+/// The token is generated once and persisted under `~/.loomrouter`, so a
+/// running Codex keeps the same credentials after LoomRouter restarts.
+/// Regenerating it per process only worked while Codex reloaded
+/// `config.toml`; an already-open app kept sending the old provider token.
+#[cfg(not(test))]
 pub fn local_token() -> &'static str {
-    LOCAL_TOKEN.get_or_init(|| {
-        format!(
-            "{}{}",
-            uuid::Uuid::new_v4().simple(),
-            uuid::Uuid::new_v4().simple()
-        )
-    })
+    LOCAL_TOKEN.get_or_init(|| load_or_create_local_token(&local_token_path()))
+}
+
+#[cfg(test)]
+pub fn local_token() -> &'static str {
+    LOCAL_TOKEN.get_or_init(generate_local_token)
+}
+
+#[cfg(not(test))]
+fn local_token_path() -> std::path::PathBuf {
+    crate::config::config_dir().join("local-token")
+}
+
+fn managed_token_from(raw: &str) -> Option<String> {
+    let managed_start = raw.find(crate::codex::BEGIN_MARK)? + crate::codex::BEGIN_MARK.len();
+    let managed_end = raw[managed_start..].find(crate::codex::END_MARK)? + managed_start;
+    let managed = &raw[managed_start..managed_end];
+    let needle = "x-loomrouter-token\" = \"";
+    let start = managed.find(needle)? + needle.len();
+    let end = managed[start..].find('"')?;
+    Some(managed[start..start + end].to_string())
+}
+
+fn configured_token() -> Option<String> {
+    let raw = std::fs::read_to_string(crate::codex::codex_home().join("config.toml")).ok()?;
+    managed_token_from(&raw)
+}
+
+fn generate_local_token() -> String {
+    format!(
+        "{}{}",
+        uuid::Uuid::new_v4().simple(),
+        uuid::Uuid::new_v4().simple()
+    )
+}
+
+fn valid_local_token(token: &str) -> bool {
+    token.len() == 64 && token.chars().all(|c| c.is_ascii_hexdigit())
+}
+
+fn load_or_create_local_token(path: &Path) -> String {
+    if let Ok(token) = std::fs::read_to_string(path) {
+        if valid_local_token(&token) {
+            return token;
+        }
+    }
+    if let Some(token) = configured_token().filter(|token| valid_local_token(token)) {
+        if let Err(e) = crate::secure_fs::write_private(path, token.as_bytes()) {
+            tracing::warn!(path = %path.display(), error = %e, "failed to persist migrated local token");
+        }
+        return token;
+    }
+    let token = generate_local_token();
+    if let Err(e) = crate::secure_fs::write_private(path, token.as_bytes()) {
+        tracing::warn!(path = %path.display(), error = %e, "failed to persist local token");
+    }
+    token
 }
 
 /// Constant-time token comparison avoids leaking a valid token prefix through
@@ -80,3 +135,42 @@ pub(super) async fn auth_gate(req: Request, next: Next) -> Response {
         ))
         .expect("a static unauthorized response is valid")
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn persisted_local_token_is_reused() {
+        let dir = tempfile::tempdir().unwrap();
+        let path = dir.path().join("local-token");
+
+        let first = load_or_create_local_token(&path);
+        let second = load_or_create_local_token(&path);
+
+        assert_eq!(first, second);
+        assert_eq!(first.len(), 64);
+    }
+
+    #[test]
+    fn managed_block_token_is_read_without_leaking_adjacent_text() {
+        let raw = format!(
+            "# x-loomrouter-token\" = \"{}\"\n{}\nhttp_headers = {{ \"x-loomrouter-token\" = \"abc\", \"Authorization\" = \"Bearer abc\" }}\n{}",
+            "0".repeat(64),
+            crate::codex::BEGIN_MARK,
+            crate::codex::END_MARK,
+        );
+        assert_eq!(managed_token_from(&raw).as_deref(), Some("abc"));
+    }
+
+    #[test]
+    fn token_outside_managed_block_is_ignored() {
+        let raw = format!(
+            "x-loomrouter-token\" = \"{}\"\n{}\n# no token here\n{}",
+            "a".repeat(64),
+            crate::codex::BEGIN_MARK,
+            crate::codex::END_MARK,
+        );
+        assert_eq!(managed_token_from(&raw), None);
+    }
+}

src-tauri/src/proxy/dispatch.rs

@@ -88,12 +88,13 @@ pub(super) async fn dispatch_routed(
     wire: WireApi,
 ) -> anyhow::Result<Response> {
     if is_remote_compaction_v2(payload) {
-        return dispatch_routed_compaction(ctx, provider, upstream_model, model, payload).await;
+        return dispatch_routed_compaction(ctx, provider, upstream_model, payload).await;
     }
+    let stats_model = super::routed_stats_model(provider, upstream_model);
     if super::routing::codex_request_kind(payload).as_deref() == Some("compaction") {
         record_problem(
             &ctx.stats,
-            &Turn::new(&provider.id, upstream_model, "http", None),
+            &Turn::new(&provider.id, &stats_model, "http", None),
             "compaction",
             &format!(
                 "{BUILD_LABEL}: Codex sent a compaction call without a compaction_trigger item; treating it as a normal turn"
@@ -140,7 +141,7 @@ pub(super) async fn dispatch_routed(
                 return Err(anyhow::Error::new(visual_preparation_failure(
                     &ctx.stats,
                     &provider.id,
-                    model,
+                    &stats_model,
                     "http",
                     started,
                     &error,
@@ -168,7 +169,8 @@ pub(super) async fn dispatch_routed(
         build_upstream(provider, &prepared_payload, upstream_model, wire)?;
 
     let (upstream, key_id) = send(ctx, provider, path, &body).await?;
-    let turn = Turn::new(&provider.id, model, "http", Some(started)).with_key(key_id.as_deref());
+    let turn =
+        Turn::new(&provider.id, &stats_model, "http", Some(started)).with_key(key_id.as_deref());
     let status =
         StatusCode::from_u16(upstream.status().as_u16()).unwrap_or(StatusCode::BAD_GATEWAY);
 
@@ -345,7 +347,8 @@ async fn dispatch_claude_cli(
         .and_then(Value::as_bool)
         .unwrap_or(false);
     let started = std::time::Instant::now();
-    let turn = Turn::new(&provider.id, model, "http", Some(started));
+    let stats_model = super::routed_stats_model(provider, upstream_model);
+    let turn = Turn::new(&provider.id, &stats_model, "http", Some(started));
     let downstream_kind = wire.downstream();
     let (result, id) = super::run_claude_turn(payload, upstream_model, wire).await?;
     tracing::debug!(%model, input_tokens = result.input_tokens, output_tokens = result.output_tokens, "claude -p turn finished");
@@ -815,11 +818,11 @@ async fn dispatch_routed_compaction(
     ctx: &ProxyCtx,
     provider: &Provider,
     upstream_model: &str,
-    model: &str,
     payload: &Value,
 ) -> anyhow::Result<Response> {
     let started = std::time::Instant::now();
-    let turn = Turn::new(&provider.id, model, "http", Some(started));
+    let stats_model = super::routed_stats_model(provider, upstream_model);
+    let turn = Turn::new(&provider.id, &stats_model, "http", Some(started));
     let (summary, usage) = match summarize_compaction(ctx, provider, upstream_model, payload).await
     {
         Ok(ok) => ok,
@@ -875,17 +878,14 @@ pub(super) async fn routed_compaction_events(
     upstream_model: &str,
     payload: &Value,
 ) -> anyhow::Result<super::realtime::WsEvents> {
-    let model = payload
-        .get("model")
-        .and_then(Value::as_str)
-        .unwrap_or(upstream_model);
+    let stats_model = super::routed_stats_model(provider, upstream_model);
     let (summary, usage) = match summarize_compaction(ctx, provider, upstream_model, payload).await
     {
         Ok(ok) => ok,
         Err(error) => {
             record_problem(
                 &ctx.stats,
-                &Turn::new(&provider.id, model, "ws", None),
+                &Turn::new(&provider.id, &stats_model, "ws", None),
                 "compaction",
                 &format!("{BUILD_LABEL}: {error}"),
             );