muonroi-cli 1.8.0 → 1.8.3

package/dist/src/storage/transcript.js

@@ -1,5 +1,6 @@
 import { getCompactionSummaryText } from "../orchestrator/compaction.js";
 import { getResponseTaskType, isResponseTool } from "../pil/response-tools.js";
+import { logger } from "../utils/logger.js";
 import { getDatabase, withTransaction } from "./db.js";
 import { extractToolResultFromOutput, getOutputKind, isOutputSuccess } from "./tool-results.js";
 import { buildEffectiveTranscript } from "./transcript-view.js";
@@ -42,11 +43,11 @@ function isInternalCouncilMarker(content) {
 }
 function loadMessageRows(sessionId) {
     const rows = getDatabase()
-        .prepare(`
-    SELECT session_id, seq, role, message_json, created_at
-    FROM messages
-    WHERE session_id = ?
-    ORDER BY seq ASC
+        .prepare(`
+    SELECT session_id, seq, role, message_json, created_at
+    FROM messages
+    WHERE session_id = ?
+    ORDER BY seq ASC
   `)
         .all(sessionId);
     for (const row of rows) {
@@ -66,12 +67,12 @@ function toPersistedCompaction(row) {
 }
 export function loadLatestCompaction(sessionId) {
     const row = getDatabase()
-        .prepare(`
-    SELECT session_id, first_kept_seq, summary, tokens_before, created_at
-    FROM compactions
-    WHERE session_id = ?
-    ORDER BY id DESC
-    LIMIT 1
+        .prepare(`
+    SELECT session_id, first_kept_seq, summary, tokens_before, created_at
+    FROM compactions
+    WHERE session_id = ?
+    ORDER BY id DESC
+    LIMIT 1
   `)
         .get(sessionId);
     return toPersistedCompaction(row);
@@ -112,37 +113,37 @@ export function appendMessages(sessionId, messages) {
         // with `status='completed'` and the full message_json. Pre-A5 callers
         // hit this path identically: there is no pending row so the IGNORE
         // branch is unused and the INSERT wins.
-        const insertMessage = db.prepare(`
-      INSERT INTO messages (session_id, seq, role, message_json, created_at, status)
-      VALUES (?, ?, ?, ?, ?, 'completed')
-      ON CONFLICT(session_id, seq) DO UPDATE SET
-        role = excluded.role,
-        message_json = excluded.message_json,
-        status = 'completed'
+        const insertMessage = db.prepare(`
+      INSERT INTO messages (session_id, seq, role, message_json, created_at, status)
+      VALUES (?, ?, ?, ?, ?, 'completed')
+      ON CONFLICT(session_id, seq) DO UPDATE SET
+        role = excluded.role,
+        message_json = excluded.message_json,
+        status = 'completed'
     `);
-        const insertToolCall = db.prepare(`
-      INSERT OR IGNORE INTO tool_calls (
-        session_id, message_seq, tool_call_id, tool_name, args_json, status, started_at, completed_at
-      ) VALUES (?, ?, ?, ?, ?, ?, ?, ?)
+        const insertToolCall = db.prepare(`
+      INSERT OR IGNORE INTO tool_calls (
+        session_id, message_seq, tool_call_id, tool_name, args_json, status, started_at, completed_at
+      ) VALUES (?, ?, ?, ?, ?, ?, ?, ?)
     `);
-        const updateToolCall = db.prepare(`
-      UPDATE tool_calls
-      SET tool_name = ?, args_json = ?, status = ?, completed_at = ?
-      WHERE session_id = ? AND tool_call_id = ?
+        const updateToolCall = db.prepare(`
+      UPDATE tool_calls
+      SET tool_name = ?, args_json = ?, status = ?, completed_at = ?
+      WHERE session_id = ? AND tool_call_id = ?
     `);
-        const selectToolCall = db.prepare(`
-      SELECT id, tool_call_id, tool_name, args_json
-      FROM tool_calls
-      WHERE session_id = ? AND tool_call_id = ?
+        const selectToolCall = db.prepare(`
+      SELECT id, tool_call_id, tool_name, args_json
+      FROM tool_calls
+      WHERE session_id = ? AND tool_call_id = ?
     `);
-        const insertToolResult = db.prepare(`
-      INSERT INTO tool_results (tool_call_row_id, output_kind, output_json, success, created_at)
-      VALUES (?, ?, ?, ?, ?)
+        const insertToolResult = db.prepare(`
+      INSERT INTO tool_results (tool_call_row_id, output_kind, output_json, success, created_at)
+      VALUES (?, ?, ?, ?, ?)
     `);
-        const updateSession = db.prepare(`
-      UPDATE sessions
-      SET updated_at = ?
-      WHERE id = ?
+        const updateSession = db.prepare(`
+      UPDATE sessions
+      SET updated_at = ?
+      WHERE id = ?
     `);
         messages.forEach((message, index) => {
             const seq = nextSeq + index;
@@ -204,10 +205,10 @@ export function persistToolCallWriteAhead(sessionId, messageSeq, toolCallId, too
     const now = new Date().toISOString();
     try {
         getDatabase()
-            .prepare(`
-        INSERT OR IGNORE INTO tool_calls (
-          session_id, message_seq, tool_call_id, tool_name, args_json, status, started_at, completed_at
-        ) VALUES (?, ?, ?, ?, ?, 'pending', ?, NULL)
+            .prepare(`
+        INSERT OR IGNORE INTO tool_calls (
+          session_id, message_seq, tool_call_id, tool_name, args_json, status, started_at, completed_at
+        ) VALUES (?, ?, ?, ?, ?, 'pending', ?, NULL)
       `)
             .run(sessionId, messageSeq, toolCallId, toolName, argsJson, now);
     }
@@ -253,10 +254,10 @@ export function persistMessageWriteAhead(sessionId, seq, role, messageJson) {
     const now = new Date().toISOString();
     try {
         getDatabase()
-            .prepare(`
-        INSERT OR IGNORE INTO messages (
-          session_id, seq, role, message_json, created_at, status
-        ) VALUES (?, ?, ?, ?, ?, 'pending')
+            .prepare(`
+        INSERT OR IGNORE INTO messages (
+          session_id, seq, role, message_json, created_at, status
+        ) VALUES (?, ?, ?, ?, ?, 'pending')
       `)
             .run(sessionId, seq, role, messageJson, now);
     }
@@ -274,10 +275,10 @@ export function persistMessageWriteAhead(sessionId, seq, role, messageJson) {
 export function markMessageCompleted(sessionId, seq) {
     try {
         getDatabase()
-            .prepare(`
-        UPDATE messages
-        SET status = 'completed'
-        WHERE session_id = ? AND seq = ? AND status = 'pending'
+            .prepare(`
+        UPDATE messages
+        SET status = 'completed'
+        WHERE session_id = ? AND seq = ? AND status = 'pending'
       `)
             .run(sessionId, seq);
     }
@@ -296,10 +297,10 @@ export function markMessageCompleted(sessionId, seq) {
 export function markMessageErrored(sessionId, seq) {
     try {
         getDatabase()
-            .prepare(`
-        UPDATE messages
-        SET status = 'errored'
-        WHERE session_id = ? AND seq = ? AND status = 'pending'
+            .prepare(`
+        UPDATE messages
+        SET status = 'errored'
+        WHERE session_id = ? AND seq = ? AND status = 'pending'
       `)
             .run(sessionId, seq);
     }
@@ -327,21 +328,20 @@ export function markMessageErrored(sessionId, seq) {
  */
 export function sweepStalePendingRows(staleAfterMs = 5 * 60 * 1000) {
     const cutoff = new Date(Date.now() - staleAfterMs).toISOString();
-    const now = new Date().toISOString();
     try {
         const db = getDatabase();
         const toolCalls = db
-            .prepare(`
-        UPDATE tool_calls
-        SET status = 'aborted', completed_at = ?
-        WHERE status = 'pending' AND started_at < ?
+            .prepare(`
+        UPDATE tool_calls
+        SET status = 'aborted', completed_at = started_at
+        WHERE status = 'pending' AND started_at < ?
       `)
-            .run(now, cutoff);
+            .run(cutoff);
         const messages = db
-            .prepare(`
-        UPDATE messages
-        SET status = 'aborted'
-        WHERE status = 'pending' AND created_at < ?
+            .prepare(`
+        UPDATE messages
+        SET status = 'aborted'
+        WHERE status = 'pending' AND created_at < ?
       `)
             .run(cutoff);
         return { toolCalls: toolCalls.changes, messages: messages.changes };
@@ -355,10 +355,10 @@ export function markToolCallErrored(sessionId, toolCallId, errorMessage) {
     const now = new Date().toISOString();
     try {
         getDatabase()
-            .prepare(`
-        UPDATE tool_calls
-        SET status = 'errored', completed_at = ?, args_json = COALESCE(args_json, ?)
-        WHERE session_id = ? AND tool_call_id = ?
+            .prepare(`
+        UPDATE tool_calls
+        SET status = 'errored', completed_at = ?, args_json = COALESCE(args_json, ?)
+        WHERE session_id = ? AND tool_call_id = ?
       `)
             .run(now, JSON.stringify({ error: errorMessage.slice(0, 500) }), sessionId, toolCallId);
     }
@@ -368,109 +368,215 @@ export function markToolCallErrored(sessionId, toolCallId, errorMessage) {
 }
 export function appendCompaction(sessionId, firstKeptSeq, summary, tokensBefore) {
     withTransaction((db) => {
-        db.prepare(`
-      INSERT INTO compactions (session_id, first_kept_seq, summary, tokens_before, created_at)
-      VALUES (?, ?, ?, ?, ?)
+        db.prepare(`
+      INSERT INTO compactions (session_id, first_kept_seq, summary, tokens_before, created_at)
+      VALUES (?, ?, ?, ?, ?)
     `).run(sessionId, firstKeptSeq, summary, tokensBefore, new Date().toISOString());
-        db.prepare(`
-      UPDATE sessions
-      SET updated_at = ?
-      WHERE id = ?
+        db.prepare(`
+      UPDATE sessions
+      SET updated_at = ?
+      WHERE id = ?
     `).run(new Date().toISOString(), sessionId);
     });
 }
+export function revertLatestCompaction(sessionId) {
+    withTransaction((db) => {
+        db.prepare(`
+      DELETE FROM compactions 
+      WHERE session_id = ? 
+      AND id = (SELECT MAX(id) FROM compactions WHERE session_id = ?)
+    `).run(sessionId, sessionId);
+        db.prepare(`
+      UPDATE sessions
+      SET updated_at = ?
+      WHERE id = ?
+    `).run(new Date().toISOString(), sessionId);
+    });
+}
+export function getSessionChain(sessionId) {
+    logger.debug("storage", "Resolving session chain", { sessionId });
+    const db = getDatabase();
+    // 1. Walk up to the root parent
+    let rootId = sessionId;
+    const visitedUp = new Set();
+    while (rootId && !visitedUp.has(rootId)) {
+        visitedUp.add(rootId);
+        try {
+            const row = db.prepare("SELECT parent_session_id FROM sessions WHERE id = ?").get(rootId);
+            if (row?.parent_session_id) {
+                rootId = row.parent_session_id;
+            }
+            else {
+                break;
+            }
+        }
+        catch {
+            break;
+        }
+    }
+    // 2. Find all descendants of the root session using a BFS queue to handle multiple child branches
+    const allIds = new Set([rootId]);
+    const queue = [rootId];
+    while (queue.length > 0) {
+        const currentId = queue.shift();
+        try {
+            const children = db.prepare("SELECT id FROM sessions WHERE parent_session_id = ?").all(currentId);
+            for (const child of children) {
+                if (!allIds.has(child.id)) {
+                    allIds.add(child.id);
+                    queue.push(child.id);
+                }
+            }
+        }
+        catch (err) {
+            logger.error("storage", `Failed to query child sessions for ${currentId} during walk`, { error: err });
+        }
+    }
+    // 3. Sort chronologically by querying the database for all matched sessions ordered by created_at
+    let chain = [];
+    try {
+        const placeholders = Array.from(allIds)
+            .map(() => "?")
+            .join(",");
+        const sortedRows = db
+            .prepare(`
+      SELECT id FROM sessions 
+      WHERE id IN (${placeholders}) 
+      ORDER BY created_at ASC
+    `)
+            .all(...Array.from(allIds));
+        chain = sortedRows.map((r) => r.id);
+    }
+    catch (err) {
+        logger.error("storage", `Failed to sort session chain for root ${rootId}`, { error: err });
+        // Fallback to unsorted array if query fails
+        chain = Array.from(allIds);
+    }
+    logger.debug("storage", "Resolved session chain successfully", { sessionId, chainLength: chain.length, chain });
+    return chain;
+}
 export function buildChatEntries(sessionId) {
-    const toolResults = loadStoredToolResults(sessionId);
+    const chain = getSessionChain(sessionId);
+    logger.debug("storage", "Building chat entries from session chain", { sessionId, chainLength: chain.length });
+    const toolResults = new Map();
+    for (const sid of chain) {
+        const results = loadStoredToolResults(sid);
+        for (const [k, v] of results) {
+            toolResults.set(k, v);
+        }
+    }
+    // Load parent/child session metadata to reconcile user prompt timestamps across rotation/sub-session boundaries
+    const db = getDatabase();
+    const sessionMeta = new Map();
+    try {
+        if (chain.length > 1) {
+            const placeholders = chain.map(() => "?").join(",");
+            const rows = db
+                .prepare(`SELECT id, parent_session_id, created_at FROM sessions WHERE id IN (${placeholders})`)
+                .all(...chain);
+            for (const r of rows) {
+                sessionMeta.set(r.id, { parent_session_id: r.parent_session_id, created_at: r.created_at });
+            }
+        }
+    }
+    catch (err) {
+        logger.error("storage", "Failed to load session metadata for chain in buildChatEntries", { error: err });
+    }
+    // Map each child sub-session's initial user prompt to its creation timestamp
+    const childPrompts = new Map();
+    for (let i = 1; i < chain.length; i++) {
+        const sid = chain[i];
+        const meta = sessionMeta.get(sid);
+        if (meta && meta.parent_session_id) {
+            const records = buildEffectiveMessageRecords(sid);
+            const firstUser = records.find((r) => r.message.role === "user");
+            if (firstUser) {
+                const content = renderUserContent(firstUser.message.content);
+                if (content) {
+                    childPrompts.set(content, new Date(meta.created_at));
+                }
+            }
+        }
+    }
     const callMap = new Map();
     const entries = [];
-    // Response-tool callIds already rendered as a structured_response from a
-    // persisted tool-RESULT row (success path). After the loop we recover any
-    // response tool that was CALLED but produced no result row — i.e. respond_*
-    // whose execution ERRORED (the AI SDK does not persist tool-error parts as
-    // tool-result rows). Its answer lives in the call args; without this it is
-    // dropped on the turn-finalize rebuild and the user sees an empty turn or
-    // only the last council debate round.
     const renderedResponseCallIds = new Set();
     let lastTimestamp;
-    for (const row of buildEffectiveMessageRecords(sessionId)) {
-        const { message, timestamp } = row;
-        lastTimestamp = timestamp;
-        if (message.role === "user") {
-            const content = renderUserContent(message.content);
-            if (content) {
-                entries.push({ type: "user", content, timestamp });
+    for (let i = 0; i < chain.length; i++) {
+        const sid = chain[i];
+        const records = buildEffectiveMessageRecords(sid);
+        const isChildSession = i > 0;
+        for (let j = 0; j < records.length; j++) {
+            const { message, timestamp } = records[j];
+            lastTimestamp = timestamp;
+            if (message.role === "user") {
+                const content = renderUserContent(message.content);
+                if (isChildSession && childPrompts.has(content)) {
+                    // Skip B's user message because it's a duplicate of A's user message
+                    continue;
+                }
+                if (content) {
+                    entries.push({ type: "user", content, timestamp });
+                }
+                continue;
             }
-            continue;
-        }
-        if (message.role === "system") {
-            const summaryText = getCompactionSummaryText(message);
-            const content = summaryText ?? (typeof message.content === "string" ? message.content.trim() : "");
-            if (content && !isInternalCouncilMarker(content)) {
-                // A compaction checkpoint is internal context-management, NOT the
-                // assistant's answer. Tag it with a source label so the UI renders it
-                // as a clearly-marked, collapsible checkpoint instead of dumping the
-                // raw "## Goal / ## Context For Suffix" scaffolding as a chat reply.
-                entries.push(summaryText !== null
-                    ? { type: "assistant", content, timestamp, sourceLabel: "⋯ context checkpoint (auto-compacted)" }
-                    : { type: "assistant", content, timestamp });
+            if (message.role === "system") {
+                const summaryText = getCompactionSummaryText(message);
+                if (isChildSession && j === 0 && summaryText !== null) {
+                    // Skip the first message in a child session if it's the parent compaction summary
+                    // to keep the visual timeline clean, since the user already sees the parent history.
+                    continue;
+                }
+                const content = summaryText ?? (typeof message.content === "string" ? message.content.trim() : "");
+                if (content && !isInternalCouncilMarker(content)) {
+                    entries.push(summaryText !== null
+                        ? { type: "assistant", content, timestamp, sourceLabel: "⋯ context checkpoint (auto-compacted)" }
+                        : { type: "assistant", content, timestamp });
+                }
+                continue;
             }
-            continue;
-        }
-        if (message.role === "assistant") {
-            const text = renderAssistantContent(message.content, callMap);
-            if (text) {
-                entries.push({ type: "assistant", content: text, timestamp });
+            if (message.role === "assistant") {
+                const text = renderAssistantContent(message.content, callMap);
+                if (text) {
+                    entries.push({ type: "assistant", content: text, timestamp });
+                }
+                continue;
             }
-            continue;
-        }
-        if (message.role === "tool" && Array.isArray(message.content)) {
-            for (const part of message.content) {
-                if (part.type !== "tool-result")
-                    continue;
-                // Response tools (respond_general/analyze/plan/...) are the model's
-                // TERMINAL structured answer (identity execute — the payload lives in
-                // the call args and is echoed in the result). The live stream yields
-                // them as a `structured_response` chunk (message-processor.ts:2733),
-                // which the UI renders as the answer block. Persisted history MUST
-                // rebuild the SAME entry: finalizeActiveTurn (app.tsx) replaces the
-                // live message list with getChatEntries() on every normal turn, so if
-                // we rebuilt a respond_* result as a bare tool_result the rendered
-                // answer would be silently dropped the instant streaming ended (live
-                // repro session 9d3d371ca1bd: grok investigated, emitted a 10K-char
-                // respond_general, the answer flashed then vanished on turn finalize —
-                // the user saw only the "→ respond_general" indicator). AI SDK v5/v6
-                // wraps tool outputs as `{type:"json", value:{...}}`; unwrap to expose
-                // the schema-shaped payload to the renderer.
-                if (isResponseTool(part.toolName)) {
-                    renderedResponseCallIds.add(part.toolCallId);
-                    const rawOutput = part.output;
-                    const unwrapped = rawOutput && typeof rawOutput === "object" && rawOutput.type === "json"
-                        ? (rawOutput.value ?? {})
-                        : (rawOutput ?? {});
+            if (message.role === "tool" && Array.isArray(message.content)) {
+                for (const part of message.content) {
+                    if (part.type !== "tool-result")
+                        continue;
+                    if (isResponseTool(part.toolName)) {
+                        renderedResponseCallIds.add(part.toolCallId);
+                        const rawOutput = part.output;
+                        const unwrapped = rawOutput && typeof rawOutput === "object" && rawOutput.type === "json"
+                            ? (rawOutput.value ?? {})
+                            : (rawOutput ?? {});
+                        entries.push({
+                            type: "structured_response",
+                            content: "",
+                            timestamp,
+                            structuredResponse: {
+                                taskType: getResponseTaskType(part.toolName) ?? part.toolName,
+                                data: unwrapped,
+                            },
+                        });
+                        continue;
+                    }
+                    const toolCall = callMap.get(part.toolCallId) ?? toFallbackToolCall(part.toolCallId, part.toolName);
+                    const toolResult = toolResults.get(part.toolCallId) ??
+                        extractToolResultFromOutput(part.output) ?? {
+                        success: isOutputSuccess(part.output),
+                        output: JSON.stringify(part.output),
+                    };
                     entries.push({
-                        type: "structured_response",
-                        content: "",
+                        type: "tool_result",
+                        content: toolResult.success ? toolResult.output || "Success" : toolResult.error || "Error",
                         timestamp,
-                        structuredResponse: {
-                            taskType: getResponseTaskType(part.toolName) ?? part.toolName,
-                            data: unwrapped,
-                        },
+                        toolCall,
+                        toolResult,
                     });
-                    continue;
                 }
-                const toolCall = callMap.get(part.toolCallId) ?? toFallbackToolCall(part.toolCallId, part.toolName);
-                const toolResult = toolResults.get(part.toolCallId) ??
-                    extractToolResultFromOutput(part.output) ?? {
-                    success: isOutputSuccess(part.output),
-                    output: JSON.stringify(part.output),
-                };
-                entries.push({
-                    type: "tool_result",
-                    content: toolResult.success ? toolResult.output || "Success" : toolResult.error || "Error",
-                    timestamp,
-                    toolCall,
-                    toolResult,
-                });
             }
         }
     }
@@ -501,14 +607,27 @@ export function buildChatEntries(sessionId) {
             },
         });
     }
+    // Adjust parent user message timestamps to match their corresponding child session creation time
+    for (const entry of entries) {
+        if (entry.type === "user") {
+            const content = entry.content.trim();
+            const childTime = childPrompts.get(content);
+            if (childTime) {
+                entry.timestamp = childTime;
+            }
+        }
+    }
+    // Sort entries chronologically by timestamp (stable sort)
+    entries.sort((a, b) => a.timestamp.getTime() - b.timestamp.getTime());
+    logger.debug("storage", "Built chat entries successfully", { sessionId, count: entries.length });
     return entries;
 }
 function getNextSequence(db, sessionId) {
     const row = db
-        .prepare(`
-    SELECT COALESCE(MAX(seq), 0) AS max_seq
-    FROM messages
-    WHERE session_id = ?
+        .prepare(`
+    SELECT COALESCE(MAX(seq), 0) AS max_seq
+    FROM messages
+    WHERE session_id = ?
   `)
         .get(sessionId);
     return (row?.max_seq ?? 0) + 1;
@@ -558,12 +677,12 @@ function renderAssistantContent(content, callMap) {
 }
 function loadStoredToolResults(sessionId) {
     const rows = getDatabase()
-        .prepare(`
-    SELECT tc.tool_call_id, tr.output_json
-    FROM tool_results tr
-    JOIN tool_calls tc ON tc.id = tr.tool_call_row_id
-    WHERE tc.session_id = ?
-    ORDER BY tr.id ASC
+        .prepare(`
+    SELECT tc.tool_call_id, tr.output_json
+    FROM tool_results tr
+    JOIN tool_calls tc ON tc.id = tr.tool_call_row_id
+    WHERE tc.session_id = ?
+    ORDER BY tr.id ASC
   `)
         .all(sessionId);
     return new Map(rows.map((row) => [row.tool_call_id, JSON.parse(row.output_json)]));
@@ -579,20 +698,27 @@ function toFallbackToolCall(toolCallId, toolName) {
     };
 }
 export function getLastTodoWriteArgs(sessionId) {
-    try {
-        const row = getDatabase()
-            .prepare(`
-        SELECT args_json
-        FROM tool_calls
-        WHERE session_id = ? AND tool_name = 'todo_write'
-        ORDER BY id DESC
-        LIMIT 1
-      `)
-            .get(sessionId);
-        return row?.args_json ?? null;
-    }
-    catch {
-        return null;
+    const chain = getSessionChain(sessionId);
+    const db = getDatabase();
+    for (let i = chain.length - 1; i >= 0; i--) {
+        try {
+            const row = db
+                .prepare(`
+          SELECT args_json
+          FROM tool_calls
+          WHERE session_id = ? AND tool_name = 'todo_write'
+          ORDER BY id DESC
+          LIMIT 1
+        `)
+                .get(chain[i]);
+            if (row?.args_json) {
+                return row.args_json;
+            }
+        }
+        catch {
+            // ignore and check parent
+        }
     }
+    return null;
 }
 //# sourceMappingURL=transcript.js.map

package/dist/src/storage/usage.js

@@ -1,5 +1,7 @@
 import { getModelInfo } from "../models/registry.js";
+import { logger } from "../utils/logger.js";
 import { getDatabase } from "./db.js";
+import { getSessionChain } from "./transcript.js";
 export function recordUsageEvent(sessionId, source, model, usage, messageSeq, pilActive = false, enrichmentDelta = 0, providerOptionsShape = null) {
     if (!usage)
         return;
@@ -21,24 +23,39 @@ export function recordUsageEvent(sessionId, source, model, usage, messageSeq, pi
         .run(sessionId, messageSeq ?? null, source, model, inputTokens, outputTokens, totalTokens, costMicros, new Date().toISOString(), pilActive ? 1 : 0, enrichmentDelta, cacheReadTokens, cacheCreationTokens, providerOptionsShape);
 }
 export function getSessionTotalTokens(sessionId) {
+    const chain = getSessionChain(sessionId);
+    const placeholders = chain.map(() => "?").join(",");
     const row = getDatabase()
         .prepare(`
     SELECT COALESCE(SUM(total_tokens), 0) AS total_tokens
     FROM usage_events
-    WHERE session_id = ?
+    WHERE session_id IN (${placeholders})
   `)
-        .get(sessionId);
-    return row?.total_tokens ?? 0;
+        .get(...chain);
+    const total = row?.total_tokens ?? 0;
+    logger.debug("storage", "Calculated session total tokens from chain", {
+        sessionId,
+        chainLength: chain.length,
+        totalTokens: total,
+    });
+    return total;
 }
 export function listSessionUsage(sessionId) {
+    const chain = getSessionChain(sessionId);
+    const placeholders = chain.map(() => "?").join(",");
     const rows = getDatabase()
         .prepare(`
     SELECT id, session_id, message_seq, source, model, input_tokens, output_tokens, total_tokens, cost_micros, created_at, cache_read_tokens, cache_creation_tokens
     FROM usage_events
-    WHERE session_id = ?
+    WHERE session_id IN (${placeholders})
     ORDER BY id ASC
   `)
-        .all(sessionId);
+        .all(...chain);
+    logger.debug("storage", "Retrieved session usage events from chain", {
+        sessionId,
+        chainLength: chain.length,
+        eventCount: rows.length,
+    });
     return rows.map((row) => ({
         id: row.id,
         sessionId: row.session_id,