@raquezha/notrace 0.1.1 → 0.1.2

package/CHANGELOG.md

@@ -1,5 +1,13 @@
 # @raquezha/notrace
 
+## 0.1.2
+
+### Patch Changes
+
+- a0f076e: Fix notrace index lock behavior under contention and add regression coverage for lock handling and usage normalization.
+- 894da0a: Add Vitest scaffolding and extract `handleSessionShutdown` for testability with no intended runtime behavior change.
+- 9d6c8b2: Skip writing per-session notrace artifacts for ghost sessions and add regression coverage for ghost vs non-ghost shutdown behavior.
+
 ## 0.1.1
 
 ### Patch Changes

package/dist/notrace/index.d.ts

@@ -1,2 +1,36 @@
 import type { ExtensionAPI } from "@earendil-works/pi-coding-agent";
+import type { NotraceCaptureMode, NotraceEvent, NotraceExtensionTelemetry } from "./types.js";
+import { type WorkflowAdapter } from "./adapters.js";
+type UsageLike = {
+    input?: number;
+    output?: number;
+    inputTokens?: number;
+    outputTokens?: number;
+    cacheRead?: number;
+    cacheWrite?: number;
+    cacheReadTokens?: number;
+    cacheWriteTokens?: number;
+    totalTokens?: number;
+    cost?: {
+        input?: number;
+        output?: number;
+        cacheRead?: number;
+        cacheWrite?: number;
+        total?: number;
+    };
+};
+export declare function normalizeUsage(raw: unknown): Required<Pick<UsageLike, "inputTokens" | "outputTokens" | "cacheReadTokens" | "cacheWriteTokens" | "totalTokens">> & {
+    totalCostUsd: number;
+};
+export type SessionShutdownDeps = {
+    events: NotraceEvent[];
+    startTime: number;
+    traceId: string;
+    extensionTelemetry: Map<string, NotraceExtensionTelemetry>;
+    captureMode: NotraceCaptureMode;
+    notraceDir: string;
+    adapter: WorkflowAdapter;
+};
+export declare function handleSessionShutdown(e: any, ctx: any, deps: SessionShutdownDeps): Promise<void>;
 export default function (pi: ExtensionAPI): void;
+export {};

package/dist/notrace/index.js

@@ -1,4 +1,4 @@
-import { readFileSync, writeFileSync, mkdirSync, existsSync, renameSync, chmodSync } from "node:fs";
+import { readFileSync, writeFileSync, mkdirSync, existsSync, renameSync, chmodSync, rmSync } from "node:fs";
 import * as path from "node:path";
 import * as os from "node:os";
 import { execSync } from "node:child_process";
@@ -40,14 +40,18 @@ function sanitizeTraceValue(value) {
 function asNumber(value) {
     return typeof value === "number" && Number.isFinite(value) ? value : 0;
 }
-function normalizeUsage(raw) {
+export function normalizeUsage(raw) {
     const usage = (raw && typeof raw === "object" ? raw : {});
+    const inputTokens = asNumber(usage.inputTokens ?? usage.input);
+    const outputTokens = asNumber(usage.outputTokens ?? usage.output);
+    const cacheReadTokens = asNumber(usage.cacheReadTokens ?? usage.cacheRead);
+    const cacheWriteTokens = asNumber(usage.cacheWriteTokens ?? usage.cacheWrite);
     return {
-        inputTokens: asNumber(usage.inputTokens ?? usage.input),
-        outputTokens: asNumber(usage.outputTokens ?? usage.output),
-        cacheReadTokens: asNumber(usage.cacheReadTokens ?? usage.cacheRead),
-        cacheWriteTokens: asNumber(usage.cacheWriteTokens ?? usage.cacheWrite),
-        totalTokens: asNumber(usage.totalTokens),
+        inputTokens,
+        outputTokens,
+        cacheReadTokens,
+        cacheWriteTokens,
+        totalTokens: usage.totalTokens == null ? inputTokens + outputTokens + cacheReadTokens + cacheWriteTokens : asNumber(usage.totalTokens),
         totalCostUsd: asNumber(usage.cost?.total),
     };
 }
@@ -164,6 +168,129 @@ function createIndexEntry(record, htmlPath, recordPath) {
         },
     };
 }
+export async function handleSessionShutdown(e, ctx, deps) {
+    const shutdownReason = typeof e?.reason === "string" ? e.reason : null;
+    const endedAt = Date.now();
+    const context = deps.adapter.getContext(ctx.cwd);
+    const finalTraceId = ctx.sessionManager?.getSessionId?.() || deps.traceId;
+    const outputDir = path.join(deps.notraceDir, "sessions", finalTraceId.replace(/[^a-z0-9]/gi, "-"));
+    const repositoryName = path.basename(ctx.cwd);
+    let branchName = null;
+    try {
+        branchName = execSync("git branch --show-current", { cwd: ctx.cwd, stdio: ["ignore", "pipe", "ignore"], encoding: "utf8", timeout: 1000 }).trim() || null;
+    }
+    catch {
+        // not a git repo or no commits yet
+    }
+    const recordPath = path.join(outputDir, "notrace.json");
+    let mergedEvents = deps.events;
+    let originalStartedAt = deps.startTime;
+    let originalTask = null;
+    if (existsSync(recordPath)) {
+        try {
+            const oldRecord = readJsonFile(recordPath, null);
+            if (Array.isArray(oldRecord.events)) {
+                mergedEvents = [...oldRecord.events, ...deps.events];
+            }
+            if (oldRecord.session?.startedAt) {
+                originalStartedAt = new Date(oldRecord.session.startedAt).getTime();
+            }
+            if (oldRecord.task) {
+                originalTask = oldRecord.task;
+            }
+        }
+        catch (err) {
+            // ignore parse errors
+        }
+    }
+    const activity = collectActivity(mergedEvents, originalStartedAt, endedAt);
+    // Do not index purely empty ghost sessions
+    const isGhostSession = activity.llmCallCount === 0 && activity.toolCallCount === 0 && activity.totals.totalTokens === 0;
+    const telemetry = Object.fromEntries([...deps.extensionTelemetry.entries()].sort(([a], [b]) => a.localeCompare(b)));
+    const record = {
+        kind: "notrace-run",
+        schemaVersion: SCHEMA_VERSION,
+        traceId: finalTraceId,
+        repository: {
+            name: repositoryName,
+            cwd: ctx.cwd,
+            branch: branchName,
+        },
+        session: {
+            id: finalTraceId,
+            startedAt: new Date(originalStartedAt).toISOString(),
+            endedAt: new Date(endedAt).toISOString(),
+            durationMs: activity.durationMs,
+            shutdownReason,
+        },
+        task: toTaskInfo(context) || originalTask,
+        captureMode: deps.captureMode,
+        conditions: buildConditions(mergedEvents, telemetry),
+        activity,
+        telemetry: { extensions: telemetry },
+        events: mergedEvents,
+    };
+    validateRunRecord(record);
+    const htmlPath = path.join(outputDir, "notrace.html");
+    if (!isGhostSession) {
+        const html = generateHtmlReport(record);
+        mkdirSync(outputDir, { recursive: true });
+        writePrivateFileAtomic(htmlPath, html);
+        writePrivateFileAtomic(recordPath, `${JSON.stringify(record, null, 2)}\n`);
+    }
+    const indexPath = path.join(deps.notraceDir, "index.json");
+    const lockPath = `${indexPath}.lock`;
+    let lockAcquired = false;
+    for (let i = 0; i < 20; i++) {
+        try {
+            writeFileSync(lockPath, `${process.pid}`, { flag: "wx" });
+            lockAcquired = true;
+            break;
+        }
+        catch {
+            const t = Date.now();
+            while (Date.now() - t < 50) { } // busy wait 50ms
+        }
+    }
+    if (!lockAcquired) {
+        // Could not get exclusive access to the index after retrying. Skip the
+        // index/dashboard update rather than racing another process's
+        // read-modify-write on index.json. The per-session record and HTML
+        // report were already written above and are not affected.
+        console.warn(`[notrace] Could not acquire index lock, skipping index update for ${finalTraceId}`);
+    }
+    else {
+        try {
+            const existing = readJsonFile(indexPath, { sessions: [] });
+            let sessions = Array.isArray(existing.sessions) ? existing.sessions.filter((s) => s.sessionId !== finalTraceId) : [];
+            if (!isGhostSession) {
+                sessions.push(createIndexEntry(record, htmlPath, recordPath));
+            }
+            writePrivateFileAtomic(indexPath, `${JSON.stringify({ sessions }, null, 2)}\n`);
+            writePrivateFileAtomic(path.join(deps.notraceDir, "index.html"), generateDashboardHtml(sessions, {}));
+        }
+        finally {
+            if (existsSync(lockPath)) {
+                try {
+                    rmSync(lockPath);
+                }
+                catch { }
+            }
+        }
+    }
+    if (context && !isGhostSession) {
+        const displayPath = htmlPath.startsWith(os.homedir())
+            ? `~${htmlPath.slice(os.homedir().length)}`
+            : htmlPath;
+        deps.adapter.attach(context, {
+            html: displayPath,
+            record: recordPath
+        });
+    }
+    if (!isGhostSession) {
+        console.log(`\n\x1b[1m\x1b[38;5;208m[notrace] Session Retrospective: file://${htmlPath}\x1b[0m\n`);
+    }
+}
 function normalizeTelemetryPayload(raw) {
     if (!raw || typeof raw !== "object")
         return null;
@@ -194,7 +321,6 @@ export default function (pi) {
     const startTime = Date.now();
     let traceId = "";
     let activeLlmPayload = null;
-    let shutdownReason = null;
     const extensionTelemetry = new Map();
     currentMode = getInitialMode();
     if (typeof pi.events?.on === "function") {
@@ -249,115 +375,14 @@ export default function (pi) {
         }
     });
     pi.on("session_shutdown", async (e, ctx) => {
-        shutdownReason = typeof e?.reason === "string" ? e.reason : null;
-        const endedAt = Date.now();
-        const adapter = getActiveAdapter(ctx.cwd);
-        const context = adapter.getContext(ctx.cwd);
-        const notraceDir = process.env.NOTRACE_DIR || path.join(os.homedir(), ".notrace");
-        const finalTraceId = ctx.sessionManager?.getSessionId?.() || traceId;
-        const outputDir = path.join(notraceDir, "sessions", finalTraceId.replace(/[^a-z0-9]/gi, "-"));
-        const repositoryName = path.basename(ctx.cwd);
-        let branchName = null;
-        try {
-            branchName = execSync("git branch --show-current", { cwd: ctx.cwd, stdio: ["ignore", "pipe", "ignore"], encoding: "utf8", timeout: 1000 }).trim() || null;
-        }
-        catch {
-            // not a git repo or no commits yet
-        }
-        const recordPath = path.join(outputDir, "notrace.json");
-        let mergedEvents = events;
-        let originalStartedAt = startTime;
-        let originalTask = null;
-        if (existsSync(recordPath)) {
-            try {
-                const oldRecord = readJsonFile(recordPath, null);
-                if (Array.isArray(oldRecord.events)) {
-                    mergedEvents = [...oldRecord.events, ...events];
-                }
-                if (oldRecord.session?.startedAt) {
-                    originalStartedAt = new Date(oldRecord.session.startedAt).getTime();
-                }
-                if (oldRecord.task) {
-                    originalTask = oldRecord.task;
-                }
-            }
-            catch (err) {
-                // ignore parse errors
-            }
-        }
-        const activity = collectActivity(mergedEvents, originalStartedAt, endedAt);
-        // Do not index purely empty ghost sessions
-        const isGhostSession = activity.llmCallCount === 0 && activity.toolCallCount === 0 && activity.totals.totalTokens === 0;
-        const telemetry = Object.fromEntries([...extensionTelemetry.entries()].sort(([a], [b]) => a.localeCompare(b)));
-        const record = {
-            kind: "notrace-run",
-            schemaVersion: SCHEMA_VERSION,
-            traceId: finalTraceId,
-            repository: {
-                name: repositoryName,
-                cwd: ctx.cwd,
-                branch: branchName,
-            },
-            session: {
-                id: finalTraceId,
-                startedAt: new Date(originalStartedAt).toISOString(),
-                endedAt: new Date(endedAt).toISOString(),
-                durationMs: activity.durationMs,
-                shutdownReason,
-            },
-            task: toTaskInfo(context) || originalTask,
+        await handleSessionShutdown(e, ctx, {
+            events,
+            startTime,
+            traceId,
+            extensionTelemetry,
             captureMode: currentMode,
-            conditions: buildConditions(mergedEvents, telemetry),
-            activity,
-            telemetry: { extensions: telemetry },
-            events: mergedEvents,
-        };
-        validateRunRecord(record);
-        const html = generateHtmlReport(record);
-        mkdirSync(outputDir, { recursive: true });
-        const htmlPath = path.join(outputDir, "notrace.html");
-        writePrivateFileAtomic(htmlPath, html);
-        writePrivateFileAtomic(recordPath, `${JSON.stringify(record, null, 2)}\n`);
-        const indexPath = path.join(notraceDir, "index.json");
-        const lockPath = `${indexPath}.lock`;
-        let lockAcquired = false;
-        for (let i = 0; i < 20; i++) {
-            try {
-                writeFileSync(lockPath, `${process.pid}`, { flag: "wx" });
-                lockAcquired = true;
-                break;
-            }
-            catch {
-                const t = Date.now();
-                while (Date.now() - t < 50) { } // busy wait 50ms
-            }
-        }
-        try {
-            const existing = readJsonFile(indexPath, { sessions: [] });
-            let sessions = Array.isArray(existing.sessions) ? existing.sessions.filter((s) => s.sessionId !== finalTraceId) : [];
-            if (!isGhostSession) {
-                sessions.push(createIndexEntry(record, htmlPath, recordPath));
-            }
-            writePrivateFileAtomic(indexPath, `${JSON.stringify({ sessions }, null, 2)}\n`);
-            writePrivateFileAtomic(path.join(notraceDir, "index.html"), generateDashboardHtml(sessions, {}));
-        }
-        finally {
-            if (lockAcquired && existsSync(lockPath)) {
-                try {
-                    import("node:fs").then(fs => fs.rmSync ? fs.rmSync(lockPath) : fs.unlinkSync(lockPath));
-                }
-                catch { }
-            }
-        }
-        if (context) {
-            const displayPath = htmlPath.startsWith(os.homedir())
-                ? `~${htmlPath.slice(os.homedir().length)}`
-                : htmlPath;
-            adapter.attach(context, {
-                html: displayPath,
-                record: recordPath
-            });
-        }
-        console.log(`\n\x1b[1m\x1b[38;5;208m[notrace] Session Retrospective: file://${htmlPath}\x1b[0m\n`);
+            notraceDir: process.env.NOTRACE_DIR || path.join(os.homedir(), ".notrace"),
+            adapter: getActiveAdapter(ctx.cwd),
+        });
     });
 }

package/extensions/notrace/__tests__/ghost-session.test.ts

@@ -0,0 +1,103 @@
+import { existsSync, readFileSync } from "node:fs";
+import * as path from "node:path";
+import { afterEach, describe, expect, it, vi } from "vitest";
+import { handleSessionShutdown, type SessionShutdownDeps } from "../index.js";
+import { cleanupTempNotraceDir, makeTempNotraceDir } from "./helpers.js";
+
+const tempDirs: string[] = [];
+
+function makeCtx(sessionId: string) {
+  return {
+    cwd: process.cwd(),
+    sessionManager: {
+      getSessionId: () => sessionId,
+    },
+  };
+}
+
+function makeDeps(notraceDir: string, traceId: string, events: SessionShutdownDeps["events"], attach = vi.fn()): SessionShutdownDeps {
+  return {
+    events,
+    startTime: Date.now() - 1000,
+    traceId,
+    extensionTelemetry: new Map(),
+    captureMode: "full",
+    notraceDir,
+    adapter: {
+      name: "test",
+      detect: () => true,
+      getContext: () => ({ workflow: "test", taskId: "task-1", taskPath: null, taskDir: null }),
+      attach,
+    },
+  };
+}
+
+function sessionDir(notraceDir: string, sessionId: string): string {
+  return path.join(notraceDir, "sessions", sessionId.replace(/[^a-z0-9]/gi, "-"));
+}
+
+function readJson(filePath: string) {
+  return JSON.parse(readFileSync(filePath, "utf-8"));
+}
+
+afterEach(() => {
+  vi.restoreAllMocks();
+  while (tempDirs.length) cleanupTempNotraceDir(tempDirs.pop()!);
+});
+
+describe("handleSessionShutdown ghost sessions", () => {
+  it("skips session artifact writes, index entry creation, attach, and console log for ghost sessions", async () => {
+    const notraceDir = makeTempNotraceDir();
+    tempDirs.push(notraceDir);
+    const sessionId = "ghost-session";
+    const attach = vi.fn();
+    const log = vi.spyOn(console, "log").mockImplementation(() => {});
+
+    await handleSessionShutdown(
+      { reason: "ghost" },
+      makeCtx(sessionId),
+      makeDeps(notraceDir, sessionId, [], attach),
+    );
+
+    const dir = sessionDir(notraceDir, sessionId);
+    expect(existsSync(path.join(dir, "notrace.html"))).toBe(false);
+    expect(existsSync(path.join(dir, "notrace.json"))).toBe(false);
+
+    const indexPath = path.join(notraceDir, "index.json");
+    if (existsSync(indexPath)) {
+      const index = readJson(indexPath);
+      expect(index.sessions.filter((session: { sessionId: string }) => session.sessionId === sessionId)).toHaveLength(0);
+    } else {
+      expect(existsSync(indexPath)).toBe(false);
+    }
+
+    expect(attach).not.toHaveBeenCalled();
+    expect(log).not.toHaveBeenCalled();
+  });
+
+  it("still writes artifacts and indexes non-ghost sessions", async () => {
+    const notraceDir = makeTempNotraceDir();
+    tempDirs.push(notraceDir);
+    const sessionId = "real-session";
+    const attach = vi.fn();
+    vi.spyOn(console, "log").mockImplementation(() => {});
+
+    await handleSessionShutdown(
+      { reason: "normal" },
+      makeCtx(sessionId),
+      makeDeps(notraceDir, sessionId, [
+        { type: "tool_start", toolName: "test-tool", args: {}, timestamp: Date.now() },
+      ], attach),
+    );
+
+    const dir = sessionDir(notraceDir, sessionId);
+    const htmlPath = path.join(dir, "notrace.html");
+    const recordPath = path.join(dir, "notrace.json");
+    expect(existsSync(htmlPath)).toBe(true);
+    expect(existsSync(recordPath)).toBe(true);
+
+    const index = readJson(path.join(notraceDir, "index.json"));
+    expect(index.sessions.some((session: { sessionId: string }) => session.sessionId === sessionId)).toBe(true);
+    expect(attach).toHaveBeenCalledOnce();
+  });
+});

package/extensions/notrace/__tests__/helpers.ts

@@ -0,0 +1,11 @@
+import { mkdtempSync, rmSync } from "node:fs";
+import * as os from "node:os";
+import * as path from "node:path";
+
+export function makeTempNotraceDir(): string {
+  return mkdtempSync(path.join(os.tmpdir(), "notrace-test-"));
+}
+
+export function cleanupTempNotraceDir(dir: string): void {
+  rmSync(dir, { recursive: true, force: true });
+}

package/extensions/notrace/__tests__/lock-race.test.ts

@@ -0,0 +1,176 @@
+import { existsSync, mkdirSync, readFileSync, writeFileSync } from "node:fs";
+import * as path from "node:path";
+import { execSync } from "node:child_process";
+import { Worker, isMainThread, parentPort, workerData } from "node:worker_threads";
+import { afterEach, describe, expect, it, vi } from "vitest";
+import { handleSessionShutdown, type SessionShutdownDeps } from "../index.js";
+import { cleanupTempNotraceDir, makeTempNotraceDir } from "./helpers.js";
+
+const tempDirs: string[] = [];
+
+function makeDeps(notraceDir: string, traceId: string): SessionShutdownDeps {
+  return {
+    events: [{ type: "tool_start", toolName: "test-tool", args: {}, timestamp: Date.now() }],
+    startTime: Date.now() - 1000,
+    traceId,
+    extensionTelemetry: new Map(),
+    captureMode: "full",
+    notraceDir,
+    adapter: {
+      name: "test",
+      detect: () => true,
+      getContext: () => null,
+      attach: () => {},
+    },
+  };
+}
+
+function makeCtx(sessionId: string) {
+  return {
+    cwd: process.cwd(),
+    sessionManager: {
+      getSessionId: () => sessionId,
+    },
+  };
+}
+
+function readJson(filePath: string) {
+  return JSON.parse(readFileSync(filePath, "utf-8"));
+}
+
+function sessionDir(notraceDir: string, sessionId: string): string {
+  return path.join(notraceDir, "sessions", sessionId.replace(/[^a-z0-9]/gi, "-"));
+}
+
+async function runWorker(sessionId: string, notraceDir: string): Promise<void> {
+  const distIndexPath = path.resolve(process.cwd(), "dist/notrace/index.js");
+  const workerScript = `
+    import { parentPort, workerData } from "node:worker_threads";
+    const { handleSessionShutdown } = await import(workerData.distIndexPath);
+    const deps = {
+      events: [{ type: "tool_start", toolName: "test-tool", args: {}, timestamp: Date.now() }],
+      startTime: Date.now() - 1000,
+      traceId: workerData.sessionId,
+      extensionTelemetry: new Map(),
+      captureMode: "full",
+      notraceDir: workerData.notraceDir,
+      adapter: { name: "test", detect: () => true, getContext: () => null, attach: () => {} },
+    };
+    const ctx = {
+      cwd: workerData.cwd,
+      sessionManager: { getSessionId: () => workerData.sessionId },
+    };
+    await handleSessionShutdown({ reason: "worker-test" }, ctx, deps);
+    parentPort.postMessage("done");
+  `;
+
+  await new Promise<void>((resolve, reject) => {
+    const worker = new Worker(workerScript, {
+      eval: true,
+      type: "module",
+      workerData: { sessionId, notraceDir, cwd: process.cwd(), distIndexPath },
+    });
+    worker.once("message", () => resolve());
+    worker.once("error", reject);
+    worker.once("exit", (code) => {
+      if (code !== 0) reject(new Error(`worker exited with code ${code}`));
+    });
+  });
+}
+
+if (!isMainThread) {
+  const { sessionId, notraceDir } = workerData as { sessionId: string; notraceDir: string };
+  await handleSessionShutdown({ reason: "worker-test" }, makeCtx(sessionId), makeDeps(notraceDir, sessionId));
+  parentPort?.postMessage("done");
+}
+
+if (isMainThread) {
+  afterEach(() => {
+    vi.restoreAllMocks();
+    while (tempDirs.length) cleanupTempNotraceDir(tempDirs.pop()!);
+  });
+
+  describe("handleSessionShutdown lock behavior", () => {
+    it("keeps both index entries from two concurrent shutdowns", async () => {
+      const notraceDir = makeTempNotraceDir();
+      tempDirs.push(notraceDir);
+      execSync("npm run build", { cwd: process.cwd(), stdio: "ignore" });
+
+      await Promise.all([
+        runWorker("session-a", notraceDir),
+        runWorker("session-b", notraceDir),
+      ]);
+
+      const index = readJson(path.join(notraceDir, "index.json"));
+      const sessionIds = index.sessions.map((session: { sessionId: string }) => session.sessionId).sort();
+
+      expect(sessionIds).toEqual(["session-a", "session-b"]);
+    });
+
+    it("warns and skips index update when lock acquisition fails", async () => {
+      const notraceDir = makeTempNotraceDir();
+      tempDirs.push(notraceDir);
+
+      const indexPath = path.join(notraceDir, "index.json");
+      const lockPath = `${indexPath}.lock`;
+      const seed = { sessions: [{ sessionId: "existing-session" }] };
+      writeFileSync(indexPath, `${JSON.stringify(seed, null, 2)}\n`);
+      writeFileSync(lockPath, "held");
+
+      const warn = vi.spyOn(console, "warn").mockImplementation(() => {});
+      const log = vi.spyOn(console, "log").mockImplementation(() => {});
+
+      await expect(handleSessionShutdown(
+        { reason: "lock-held" },
+        makeCtx("skipped-session"),
+        makeDeps(notraceDir, "skipped-session"),
+      )).resolves.toBeUndefined();
+
+      expect(warn).toHaveBeenCalledWith(expect.stringContaining("Could not acquire index lock"));
+      expect(readJson(indexPath)).toEqual(seed);
+      expect(existsSync(path.join(sessionDir(notraceDir, "skipped-session"), "notrace.json"))).toBe(true);
+      expect(existsSync(path.join(sessionDir(notraceDir, "skipped-session"), "notrace.html"))).toBe(true);
+      expect(existsSync(lockPath)).toBe(true);
+      expect(log).toHaveBeenCalled();
+    });
+
+    it("removes the lock file after a successful shutdown", async () => {
+      const notraceDir = makeTempNotraceDir();
+      tempDirs.push(notraceDir);
+
+      const lockPath = path.join(notraceDir, "index.json.lock");
+      const log = vi.spyOn(console, "log").mockImplementation(() => {});
+
+      await handleSessionShutdown(
+        { reason: "success" },
+        makeCtx("lock-cleanup-session"),
+        makeDeps(notraceDir, "lock-cleanup-session"),
+      );
+
+      expect(existsSync(lockPath)).toBe(false);
+      expect(log).toHaveBeenCalled();
+    });
+
+    it("writes a correct index entry for a normal shutdown", async () => {
+      const notraceDir = makeTempNotraceDir();
+      tempDirs.push(notraceDir);
+      const sessionId = "single-session";
+
+      mkdirSync(notraceDir, { recursive: true });
+      vi.spyOn(console, "log").mockImplementation(() => {});
+
+      await handleSessionShutdown(
+        { reason: "normal" },
+        makeCtx(sessionId),
+        makeDeps(notraceDir, sessionId),
+      );
+
+      const index = readJson(path.join(notraceDir, "index.json"));
+      expect(index.sessions).toHaveLength(1);
+      expect(index.sessions[0].sessionId).toBe(sessionId);
+      expect(index.sessions[0].repositoryName).toBe(path.basename(process.cwd()));
+      expect(existsSync(index.sessions[0].artifacts.html)).toBe(true);
+      expect(existsSync(index.sessions[0].artifacts.record)).toBe(true);
+    });
+  });
+}

package/extensions/notrace/__tests__/usage-normalization.test.ts

@@ -0,0 +1,80 @@
+import { describe, expect, it } from "vitest";
+import { normalizeUsage } from "../index.js";
+
+describe("normalizeUsage", () => {
+  it("uses explicit totalTokens as-is", () => {
+    expect(normalizeUsage({
+      inputTokens: 10,
+      outputTokens: 20,
+      cacheReadTokens: 30,
+      cacheWriteTokens: 40,
+      totalTokens: 7,
+    })).toMatchObject({
+      inputTokens: 10,
+      outputTokens: 20,
+      cacheReadTokens: 30,
+      cacheWriteTokens: 40,
+      totalTokens: 7,
+    });
+  });
+
+  it("sums component fields when totalTokens is missing", () => {
+    expect(normalizeUsage({
+      inputTokens: 10,
+      outputTokens: 20,
+      cacheReadTokens: 3,
+      cacheWriteTokens: 4,
+    })).toMatchObject({
+      inputTokens: 10,
+      outputTokens: 20,
+      cacheReadTokens: 3,
+      cacheWriteTokens: 4,
+      totalTokens: 37,
+    });
+  });
+
+  it("normalizes empty or null usage to zero", () => {
+    expect(normalizeUsage({})).toMatchObject({
+      inputTokens: 0,
+      outputTokens: 0,
+      cacheReadTokens: 0,
+      cacheWriteTokens: 0,
+      totalTokens: 0,
+      totalCostUsd: 0,
+    });
+    expect(normalizeUsage(null)).toMatchObject({
+      inputTokens: 0,
+      outputTokens: 0,
+      cacheReadTokens: 0,
+      cacheWriteTokens: 0,
+      totalTokens: 0,
+      totalCostUsd: 0,
+    });
+  });
+
+  it("supports mixed field name variants when totalTokens is absent", () => {
+    expect(normalizeUsage({
+      input: 5,
+      outputTokens: 6,
+      cacheRead: 7,
+      cacheWriteTokens: 8,
+    })).toMatchObject({
+      inputTokens: 5,
+      outputTokens: 6,
+      cacheReadTokens: 7,
+      cacheWriteTokens: 8,
+      totalTokens: 26,
+    });
+  });
+
+  it("keeps totalCostUsd normalization unchanged", () => {
+    expect(normalizeUsage({
+      inputTokens: 1,
+      outputTokens: 2,
+      cost: { total: 0.123 },
+    })).toMatchObject({
+      totalTokens: 3,
+      totalCostUsd: 0.123,
+    });
+  });
+});

package/extensions/notrace/index.ts

@@ -1,5 +1,5 @@
 import type { ExtensionAPI } from "@earendil-works/pi-coding-agent";
-import { readFileSync, writeFileSync, mkdirSync, existsSync, renameSync, chmodSync } from "node:fs";
+import { readFileSync, writeFileSync, mkdirSync, existsSync, renameSync, chmodSync, rmSync } from "node:fs";
 import * as path from "node:path";
 import * as os from "node:os";
 import { execSync } from "node:child_process";
@@ -12,7 +12,7 @@ import type {
   NotraceRunRecord,
   WorkflowContext,
 } from "./types.js";
-import { getActiveAdapter } from "./adapters.js";
+import { getActiveAdapter, type WorkflowAdapter } from "./adapters.js";
 import { generateHtmlReport, generateDashboardHtml } from "./renderer.js";
 
 const REDACTED = "[REDACTED by notrace]";
@@ -81,14 +81,18 @@ function asNumber(value: unknown): number {
   return typeof value === "number" && Number.isFinite(value) ? value : 0;
 }
 
-function normalizeUsage(raw: unknown): Required<Pick<UsageLike, "inputTokens" | "outputTokens" | "cacheReadTokens" | "cacheWriteTokens" | "totalTokens">> & { totalCostUsd: number } {
+export function normalizeUsage(raw: unknown): Required<Pick<UsageLike, "inputTokens" | "outputTokens" | "cacheReadTokens" | "cacheWriteTokens" | "totalTokens">> & { totalCostUsd: number } {
   const usage = (raw && typeof raw === "object" ? raw : {}) as UsageLike;
+  const inputTokens = asNumber(usage.inputTokens ?? usage.input);
+  const outputTokens = asNumber(usage.outputTokens ?? usage.output);
+  const cacheReadTokens = asNumber(usage.cacheReadTokens ?? usage.cacheRead);
+  const cacheWriteTokens = asNumber(usage.cacheWriteTokens ?? usage.cacheWrite);
   return {
-    inputTokens: asNumber(usage.inputTokens ?? usage.input),
-    outputTokens: asNumber(usage.outputTokens ?? usage.output),
-    cacheReadTokens: asNumber(usage.cacheReadTokens ?? usage.cacheRead),
-    cacheWriteTokens: asNumber(usage.cacheWriteTokens ?? usage.cacheWrite),
-    totalTokens: asNumber(usage.totalTokens),
+    inputTokens,
+    outputTokens,
+    cacheReadTokens,
+    cacheWriteTokens,
+    totalTokens: usage.totalTokens == null ? inputTokens + outputTokens + cacheReadTokens + cacheWriteTokens : asNumber(usage.totalTokens),
     totalCostUsd: asNumber(usage.cost?.total),
   };
 }
@@ -203,6 +207,144 @@ function createIndexEntry(record: NotraceRunRecord, htmlPath: string, recordPath
   };
 }
 
+export type SessionShutdownDeps = {
+  events: NotraceEvent[];
+  startTime: number;
+  traceId: string;
+  extensionTelemetry: Map<string, NotraceExtensionTelemetry>;
+  captureMode: NotraceCaptureMode;
+  notraceDir: string;
+  adapter: WorkflowAdapter;
+};
+
+export async function handleSessionShutdown(e: any, ctx: any, deps: SessionShutdownDeps): Promise<void> {
+  const shutdownReason = typeof e?.reason === "string" ? e.reason : null;
+  const endedAt = Date.now();
+  const context = deps.adapter.getContext(ctx.cwd);
+  const finalTraceId = ctx.sessionManager?.getSessionId?.() || deps.traceId;
+  const outputDir = path.join(deps.notraceDir, "sessions", finalTraceId.replace(/[^a-z0-9]/gi, "-"));
+  const repositoryName = path.basename(ctx.cwd);
+  let branchName: string | null = null;
+  try {
+    branchName = execSync("git branch --show-current", { cwd: ctx.cwd, stdio: ["ignore", "pipe", "ignore"], encoding: "utf8", timeout: 1000 }).trim() || null;
+  } catch {
+    // not a git repo or no commits yet
+  }
+  const recordPath = path.join(outputDir, "notrace.json");
+
+  let mergedEvents = deps.events;
+  let originalStartedAt = deps.startTime;
+  let originalTask: any = null;
+  if (existsSync(recordPath)) {
+    try {
+      const oldRecord = readJsonFile<any>(recordPath, null);
+      if (Array.isArray(oldRecord.events)) {
+        mergedEvents = [...oldRecord.events, ...deps.events];
+      }
+      if (oldRecord.session?.startedAt) {
+        originalStartedAt = new Date(oldRecord.session.startedAt).getTime();
+      }
+      if (oldRecord.task) {
+        originalTask = oldRecord.task;
+      }
+    } catch (err) {
+      // ignore parse errors
+    }
+  }
+
+  const activity = collectActivity(mergedEvents, originalStartedAt, endedAt);
+  
+  // Do not index purely empty ghost sessions
+  const isGhostSession = activity.llmCallCount === 0 && activity.toolCallCount === 0 && activity.totals.totalTokens === 0;
+
+  const telemetry = Object.fromEntries([...deps.extensionTelemetry.entries()].sort(([a], [b]) => a.localeCompare(b)));
+
+  const record: NotraceRunRecord = {
+    kind: "notrace-run",
+    schemaVersion: SCHEMA_VERSION,
+    traceId: finalTraceId,
+    repository: {
+      name: repositoryName,
+      cwd: ctx.cwd,
+      branch: branchName,
+    },
+    session: {
+      id: finalTraceId,
+      startedAt: new Date(originalStartedAt).toISOString(),
+      endedAt: new Date(endedAt).toISOString(),
+      durationMs: activity.durationMs,
+      shutdownReason,
+    },
+    task: toTaskInfo(context) || originalTask,
+    captureMode: deps.captureMode,
+    conditions: buildConditions(mergedEvents, telemetry),
+    activity,
+    telemetry: { extensions: telemetry },
+    events: mergedEvents,
+  };
+
+  validateRunRecord(record);
+  const htmlPath = path.join(outputDir, "notrace.html");
+
+  if (!isGhostSession) {
+    const html = generateHtmlReport(record);
+    mkdirSync(outputDir, { recursive: true });
+    writePrivateFileAtomic(htmlPath, html);
+    writePrivateFileAtomic(recordPath, `${JSON.stringify(record, null, 2)}\n`);
+  }
+
+  const indexPath = path.join(deps.notraceDir, "index.json");
+  const lockPath = `${indexPath}.lock`;
+  let lockAcquired = false;
+  for (let i = 0; i < 20; i++) {
+    try {
+      writeFileSync(lockPath, `${process.pid}`, { flag: "wx" });
+      lockAcquired = true;
+      break;
+    } catch {
+      const t = Date.now(); while (Date.now() - t < 50) {} // busy wait 50ms
+    }
+  }
+
+  if (!lockAcquired) {
+    // Could not get exclusive access to the index after retrying. Skip the
+    // index/dashboard update rather than racing another process's
+    // read-modify-write on index.json. The per-session record and HTML
+    // report were already written above and are not affected.
+    console.warn(`[notrace] Could not acquire index lock, skipping index update for ${finalTraceId}`);
+  } else {
+    try {
+      const existing = readJsonFile<any>(indexPath, { sessions: [] });
+      let sessions = Array.isArray(existing.sessions) ? existing.sessions.filter((s: any) => s.sessionId !== finalTraceId) : [];
+
+      if (!isGhostSession) {
+        sessions.push(createIndexEntry(record, htmlPath, recordPath));
+      }
+
+      writePrivateFileAtomic(indexPath, `${JSON.stringify({ sessions }, null, 2)}\n`);
+      writePrivateFileAtomic(path.join(deps.notraceDir, "index.html"), generateDashboardHtml(sessions, {}));
+    } finally {
+      if (existsSync(lockPath)) {
+        try { rmSync(lockPath); } catch {}
+      }
+    }
+  }
+
+  if (context && !isGhostSession) {
+    const displayPath = htmlPath.startsWith(os.homedir()) 
+      ? `~${htmlPath.slice(os.homedir().length)}` 
+      : htmlPath;
+    deps.adapter.attach(context, {
+      html: displayPath,
+      record: recordPath
+    });
+  }
+
+  if (!isGhostSession) {
+    console.log(`\n\x1b[1m\x1b[38;5;208m[notrace] Session Retrospective: file://${htmlPath}\x1b[0m\n`);
+  }
+}
+
 function normalizeTelemetryPayload(raw: unknown): { extension: string; telemetry: NotraceExtensionTelemetry } | null {
   if (!raw || typeof raw !== "object") return null;
   const payload = raw as ExtensionTelemetryPayload;
@@ -235,7 +377,6 @@ export default function (pi: ExtensionAPI) {
   const startTime = Date.now();
   let traceId = "";
   let activeLlmPayload: unknown = null;
-  let shutdownReason: string | null = null;
   const extensionTelemetry = new Map<string, NotraceExtensionTelemetry>();
   currentMode = getInitialMode();
 
@@ -297,120 +438,14 @@ export default function (pi: ExtensionAPI) {
   });
 
   pi.on("session_shutdown" as any, async (e: any, ctx: any) => {
-    shutdownReason = typeof e?.reason === "string" ? e.reason : null;
-    const endedAt = Date.now();
-    const adapter = getActiveAdapter(ctx.cwd);
-    const context = adapter.getContext(ctx.cwd);
-    const notraceDir = process.env.NOTRACE_DIR || path.join(os.homedir(), ".notrace");
-    const finalTraceId = ctx.sessionManager?.getSessionId?.() || traceId;
-    const outputDir = path.join(notraceDir, "sessions", finalTraceId.replace(/[^a-z0-9]/gi, "-"));
-    const repositoryName = path.basename(ctx.cwd);
-    let branchName: string | null = null;
-    try {
-      branchName = execSync("git branch --show-current", { cwd: ctx.cwd, stdio: ["ignore", "pipe", "ignore"], encoding: "utf8", timeout: 1000 }).trim() || null;
-    } catch {
-      // not a git repo or no commits yet
-    }
-    const recordPath = path.join(outputDir, "notrace.json");
-
-    let mergedEvents = events;
-    let originalStartedAt = startTime;
-    let originalTask: any = null;
-    if (existsSync(recordPath)) {
-      try {
-        const oldRecord = readJsonFile<any>(recordPath, null);
-        if (Array.isArray(oldRecord.events)) {
-          mergedEvents = [...oldRecord.events, ...events];
-        }
-        if (oldRecord.session?.startedAt) {
-          originalStartedAt = new Date(oldRecord.session.startedAt).getTime();
-        }
-        if (oldRecord.task) {
-          originalTask = oldRecord.task;
-        }
-      } catch (err) {
-        // ignore parse errors
-      }
-    }
-
-    const activity = collectActivity(mergedEvents, originalStartedAt, endedAt);
-    
-    // Do not index purely empty ghost sessions
-    const isGhostSession = activity.llmCallCount === 0 && activity.toolCallCount === 0 && activity.totals.totalTokens === 0;
-
-    const telemetry = Object.fromEntries([...extensionTelemetry.entries()].sort(([a], [b]) => a.localeCompare(b)));
-
-    const record: NotraceRunRecord = {
-      kind: "notrace-run",
-      schemaVersion: SCHEMA_VERSION,
-      traceId: finalTraceId,
-      repository: {
-        name: repositoryName,
-        cwd: ctx.cwd,
-        branch: branchName,
-      },
-      session: {
-        id: finalTraceId,
-        startedAt: new Date(originalStartedAt).toISOString(),
-        endedAt: new Date(endedAt).toISOString(),
-        durationMs: activity.durationMs,
-        shutdownReason,
-      },
-      task: toTaskInfo(context) || originalTask,
+    await handleSessionShutdown(e, ctx, {
+      events,
+      startTime,
+      traceId,
+      extensionTelemetry,
       captureMode: currentMode,
-      conditions: buildConditions(mergedEvents, telemetry),
-      activity,
-      telemetry: { extensions: telemetry },
-      events: mergedEvents,
-    };
-
-    validateRunRecord(record);
-    const html = generateHtmlReport(record);
-
-    mkdirSync(outputDir, { recursive: true });
-    const htmlPath = path.join(outputDir, "notrace.html");
-    writePrivateFileAtomic(htmlPath, html);
-    writePrivateFileAtomic(recordPath, `${JSON.stringify(record, null, 2)}\n`);
-
-    const indexPath = path.join(notraceDir, "index.json");
-    const lockPath = `${indexPath}.lock`;
-    let lockAcquired = false;
-    for (let i = 0; i < 20; i++) {
-      try {
-        writeFileSync(lockPath, `${process.pid}`, { flag: "wx" });
-        lockAcquired = true;
-        break;
-      } catch {
-        const t = Date.now(); while (Date.now() - t < 50) {} // busy wait 50ms
-      }
-    }
-
-    try {
-      const existing = readJsonFile<any>(indexPath, { sessions: [] });
-      let sessions = Array.isArray(existing.sessions) ? existing.sessions.filter((s: any) => s.sessionId !== finalTraceId) : [];
-      
-      if (!isGhostSession) {
-        sessions.push(createIndexEntry(record, htmlPath, recordPath));
-      }
-      
-      writePrivateFileAtomic(indexPath, `${JSON.stringify({ sessions }, null, 2)}\n`);
-      writePrivateFileAtomic(path.join(notraceDir, "index.html"), generateDashboardHtml(sessions, {}));
-    } finally {
-      if (lockAcquired && existsSync(lockPath)) {
-        try { import("node:fs").then(fs => fs.rmSync ? fs.rmSync(lockPath) : fs.unlinkSync(lockPath)); } catch {}
-      }
-    }
-
-    if (context) {
-      const displayPath = htmlPath.startsWith(os.homedir()) 
-        ? `~${htmlPath.slice(os.homedir().length)}` 
-        : htmlPath;
-      adapter.attach(context, {
-        html: displayPath,
-        record: recordPath
-      });
-    }
-
-    console.log(`\n\x1b[1m\x1b[38;5;208m[notrace] Session Retrospective: file://${htmlPath}\x1b[0m\n`);
+      notraceDir: process.env.NOTRACE_DIR || path.join(os.homedir(), ".notrace"),
+      adapter: getActiveAdapter(ctx.cwd),
+    });
   });
 }

package/package.json

@@ -1,12 +1,13 @@
 {
   "name": "@raquezha/notrace",
-  "version": "0.1.1",
+  "version": "0.1.2",
   "description": "Zero-dependency, local-first interactive HTML Trace Viewer for the Pi Coding Agent",
   "main": "dist/notrace/index.js",
   "types": "dist/notrace/index.d.ts",
   "type": "module",
   "scripts": {
     "build": "tsc",
+    "test": "vitest run --passWithNoTests",
     "render:samples": "npm run build && node ./templates/render-samples.mjs",
     "review": "node ./bin/notrace-review.mjs",
     "compare": "node ./bin/notrace-compare.mjs",
@@ -16,7 +17,8 @@
     "@earendil-works/pi-coding-agent": ">=0.74.0"
   },
   "devDependencies": {
-    "typescript": "^5.0.0"
+    "typescript": "^5.0.0",
+    "vitest": "^3.2.4"
   },
   "keywords": [
     "pi-agent",