timeline bug patch #154

Author: pacifio

crates/atlas-checkpoint/src/import.rs

@@ -54,7 +54,7 @@ use serde::{Deserialize, Serialize};
 use crate::blobs;
 use crate::capture::{Capture, SessionKey, ToolCallContent, TurnContent};
 use crate::error::{Error, Result};
-use crate::model::{Mode, Role, Source, ToolStatus, WorkspaceMode};
+use crate::model::{Mode, Role, Source, TokenTotals, ToolStatus, WorkspaceMode};
 use crate::store::Store;
 use crate::tools::{canonical_name, ToolName};
 
@@ -319,6 +319,14 @@ fn import_file(
     // matching `tool_result` carries just the id. Remembered per file so the
     // result's upsert does not downgrade the stored name to `Other`.
     let mut call_meta: HashMap<String, (ToolName, i64)> = HashMap::new();
+    // Usage, deduplicated by request — see `read_usage`. Accumulated for the
+    // whole file and written once at the end, because the total is only true
+    // once the file has been read.
+    let mut usage_by_request: HashMap<String, TokenTotals> = HashMap::new();
+    // `read_turn` only sees `message.model` on lines that carry text, so a
+    // transcript whose assistant lines are all tool calls would never name its
+    // model. Taken from any line that has one.
+    let mut model_seen: Option<String> = None;
     let no_locations = serde_json::json!([]);
 
     {
@@ -355,6 +363,34 @@ fn import_file(
                 continue;
             }
 
+            // Read before the "nothing usable here" filter below: a tool-only
+            // assistant line produces no turn and no tool_result, and dropping
+            // it would silently lose the usage of every tool-calling request.
+            if let Some(line_usage) = read_usage(&value) {
+                let slot = usage_by_request.entry(line_usage.key).or_default();
+                // Per-field max rather than overwrite: two copies of one
+                // request disagree when the first was written mid-stream, and
+                // the finished one is the larger.
+                slot.input_tokens = slot.input_tokens.max(line_usage.totals.input_tokens);
+                slot.output_tokens = slot.output_tokens.max(line_usage.totals.output_tokens);
+                slot.cache_creation_tokens =
+                    slot.cache_creation_tokens.max(line_usage.totals.cache_creation_tokens);
+                slot.cache_read_tokens =
+                    slot.cache_read_tokens.max(line_usage.totals.cache_read_tokens);
+            }
+            // Assistant lines only: the model is a property of what answered,
+            // and a user line that happens to carry one is echoing the client's
+            // request rather than reporting what ran.
+            if model_seen.is_none()
+                && value.get("type").and_then(serde_json::Value::as_str) == Some("assistant")
+            {
+                model_seen = value
+                    .get("message")
+                    .and_then(|m| m.get("model"))
+                    .and_then(serde_json::Value::as_str)
+                    .map(str::to_string);
+            }
+
             let created_at = line_timestamp(&value);
             let turn = read_turn(&value).filter(|t| !is_envelope(t));
             let tool_uses = read_tool_uses(&value);
@@ -504,6 +540,30 @@ fn import_file(
 
     if let Some(id) = &session_id {
         outcome.sessions_imported += imported_here;
+
+        // The whole-file total, deduplicated by request and written as a
+        // replace. After an I/O error this covers only the prefix that was
+        // read, which is correct: the progress marker below is not advanced, so
+        // the next pass recomputes the file from the top.
+        let mut usage = TokenTotals::default();
+        for request in usage_by_request.values() {
+            usage.input_tokens = usage.input_tokens.saturating_add(request.input_tokens);
+            usage.output_tokens = usage.output_tokens.saturating_add(request.output_tokens);
+            usage.cache_creation_tokens =
+                usage.cache_creation_tokens.saturating_add(request.cache_creation_tokens);
+            usage.cache_read_tokens =
+                usage.cache_read_tokens.saturating_add(request.cache_read_tokens);
+        }
+        if usage != TokenTotals::default() {
+            store.replace_usage_totals(id, &usage)?;
+        }
+        // The model, when the Session did not already have one — `COALESCE` in
+        // the upsert takes the new value and leaves an existing one alone.
+        if let Some(model) = &model_seen {
+            let mut capture = Capture::new(store, mode);
+            capture.ensure_session(&session_key, None, Some(model.as_str()), None, None)?;
+        }
+
         // Live capture stamped "now" at first sighting; the transcript knows
         // when the conversation really began. Without this, a year of imported
         // history all dates from the day the import ran.
@@ -522,6 +582,47 @@ fn import_file(
     Ok(())
 }
 
+/// The usage one assistant line reports, with the request it belongs to.
+struct UsageLine {
+    key: String,
+    totals: TokenTotals,
+}
+
+/// Read `message.usage` off an assistant line.
+///
+/// The transcript writes the same logical assistant message more than once — a
+/// streaming record, then its rewrite — and every copy repeats the same usage
+/// block. On a real transcript here, 18 usage lines collapsed to 8 requests:
+/// summing lines instead of requests over-counts by 2.2x. `requestId` is the
+/// key that collapses them, falling back to the message id and then the line's
+/// own uuid so a line without one is at least counted once rather than merged
+/// into a neighbour.
+fn read_usage(value: &serde_json::Value) -> Option<UsageLine> {
+    let message = value.get("message")?;
+    let usage = message.get("usage")?;
+    let key = value
+        .get("requestId")
+        .and_then(serde_json::Value::as_str)
+        .or_else(|| message.get("id").and_then(serde_json::Value::as_str))
+        .or_else(|| value.get("uuid").and_then(serde_json::Value::as_str))?
+        .to_string();
+
+    let field = |name: &str| usage.get(name).and_then(serde_json::Value::as_u64).unwrap_or(0);
+    let totals = TokenTotals {
+        input_tokens: field("input_tokens"),
+        output_tokens: field("output_tokens"),
+        cache_creation_tokens: field("cache_creation_input_tokens"),
+        cache_read_tokens: field("cache_read_input_tokens"),
+        ..Default::default()
+    };
+    // A usage block of all zeros carries no information and would otherwise
+    // occupy a request slot, so an empty total stays empty.
+    if totals == TokenTotals::default() {
+        return None;
+    }
+    Some(UsageLine { key, totals })
+}
+
 /// One usable line of a transcript.
 struct ImportedTurn {
     role: Role,

crates/atlas-checkpoint/src/model.rs

@@ -277,6 +277,13 @@ pub struct Session {
     pub summary: Option<String>,
     pub started_at: DateTime<Utc>,
     pub updated_at: DateTime<Utc>,
+    /// When this Session last did work — its newest message stamp or turn end.
+    ///
+    /// Deliberately not `updated_at`, which is a row-mutation clock: deriving a
+    /// title, recording usage or re-importing a transcript all bump that and
+    /// none of them is work happening now. Null only for a Session with no
+    /// message and no closed turn.
+    pub last_activity_at: Option<DateTime<Utc>>,
     /// Something went wrong recording this Session and a human should know.
     pub needs_attention: bool,
     pub attention_reason: Option<String>,

crates/atlas-checkpoint/src/schema.rs

@@ -16,7 +16,7 @@ use rusqlite::Connection;
 use crate::error::{Error, Result};
 
 /// Bump when adding a migration, and add the matching arm in [`migrate`].
-pub const SCHEMA_VERSION: i64 = 7;
+pub const SCHEMA_VERSION: i64 = 8;
 
 pub fn migrate(conn: &Connection) -> Result<()> {
     // Fast path, outside any transaction: the overwhelmingly common case is a
@@ -70,6 +70,11 @@ pub fn migrate(conn: &Connection) -> Result<()> {
             // a store that half-applied it must still be openable.
             apply_tolerant(conn, V7)?;
         }
+        if found < 8 {
+            // Same tolerance again: one ALTER TABLE, two indexes, and three
+            // repair statements that are all safe to re-run.
+            apply_tolerant(conn, V8)?;
+        }
         conn.pragma_update(None, "user_version", SCHEMA_VERSION)?;
         Ok(())
     })();
@@ -132,6 +137,48 @@ const V7: &str = r#"
 ALTER TABLE agent_session ADD COLUMN branch TEXT;
 "#;
 
+const V8: &str = r#"
+-- When the Session last did work, as opposed to when its row was last written.
+--
+-- `updated_at` is stamped with the wall clock by every write there is: an
+-- upsert, a title derivation, a token update, an import. The Timeline read it
+-- as the end of the Session and got `updated_at - started_at` as a duration,
+-- so a transcript that ran in June and was imported in July reported 1395
+-- hours, and the day grouping filed a year of history under Today. Activity
+-- and mutation are two different facts and now live in two different columns.
+ALTER TABLE agent_session ADD COLUMN last_activity_at TEXT;
+
+-- Backfill from the only timestamps in the store that were never the wall clock
+-- at write time: a message carries the transcript's own stamp, and a turn ends
+-- when the turn ended.
+UPDATE agent_session SET last_activity_at = (
+    SELECT MAX(stamp) FROM (
+        SELECT MAX(created_at) AS stamp FROM agent_message WHERE session_id = agent_session.id
+        UNION ALL
+        SELECT MAX(ended_at) AS stamp FROM turn WHERE session_id = agent_session.id
+    )
+);
+
+-- A Session with no message and no closed turn has nothing better to offer.
+UPDATE agent_session SET last_activity_at = updated_at WHERE last_activity_at IS NULL;
+
+-- Board ordering and day bucketing.
+CREATE INDEX IF NOT EXISTS idx_session_activity
+    ON agent_session (workspace_id, last_activity_at);
+
+-- Covers the gap-capped active-time scan, which reads (session_id, created_at)
+-- and nothing else.
+CREATE INDEX IF NOT EXISTS idx_message_activity
+    ON agent_message (session_id, created_at);
+
+-- Re-read every transcript once. Token usage was never parsed out of the JSONL,
+-- so every imported Session carries an empty total — and progress is keyed on
+-- file size, which means a naive re-run skips the entire corpus. Clearing this
+-- is the whole token backfill. Re-reading is idempotent: every line carries the
+-- agent's own message id, and the usage write is a replace rather than a sum.
+DELETE FROM import_progress;
+"#;
+
 const V1: &str = r#"
 -- One row per Session. Named `agent_session`, never `session`: the server's
 -- `session` table belongs to Better Auth, and the local schema mirrors the
@@ -514,4 +561,6 @@ pub const REQUIRED_INDEXES: &[&str] = &[
     "idx_checkpoint_patch",
     "idx_checkpoint_outbox",
     "idx_file_touch_unconsumed",
+    "idx_session_activity",
+    "idx_message_activity",
 ];