From 7bdeeb8f9358624dc76652cd701acd5ece65e64c Mon Sep 17 00:00:00 2001
From: David Brailovsky <davidbrai@gmail.com>
Date: Wed, 5 Aug 2026 09:57:30 +0200
Subject: [PATCH] fix(agent): count copied messages in fork stats

---
 packages/agent/src/harness/session/memory.ts                | 1 +
 packages/agent/src/harness/session/testing/conformance.ts   | 5 ++++-
 packages/session-backends/sqlite-node/src/sqlite/repo.ts    | 2 +-
 .../sqlite-node/src/sqlite/storage/session-stats.ts         | 6 +++---
 4 files changed, 9 insertions(+), 5 deletions(-)

diff --git a/packages/agent/src/harness/session/memory.ts b/packages/agent/src/harness/session/memory.ts
index 2c212dc0a..6e2525140 100644
--- a/packages/agent/src/harness/session/memory.ts
+++ b/packages/agent/src/harness/session/memory.ts
@@ -96,6 +96,7 @@ export class InMemorySessionStorage implements SessionStorage {
 			storage.entriesById.set(entry.id, entry);
 			storage.usedIds.add(entry.id);
 			storage.log.push({ kind: "entry", seq: entry.seq, entry });
+			if (entry.type === "message") storage.stats.messageCount += 1;
 		}
 		for (const [lane, leafId] of options.scope === "tree" ? this.lanes : []) {
 			storage.log.push({ kind: "lane", seq: storage.nextSequence(), lane, leafId });
diff --git a/packages/agent/src/harness/session/testing/conformance.ts b/packages/agent/src/harness/session/testing/conformance.ts
index 84edad47b..f18ed25a3 100644
--- a/packages/agent/src/harness/session/testing/conformance.ts
+++ b/packages/agent/src/harness/session/testing/conformance.ts
@@ -750,12 +750,14 @@ export function createSessionBackendConformance(
 				strictEqual(await fork.getLabel(threadChild), undefined);
 				deepStrictEqual(await fork.findRecords(), []);
 				deepStrictEqual(await fork.getStats(), {
-					messageCount: 0,
+					messageCount: 3,
 					cachedTokens: 0,
 					uncachedTokens: 0,
 					totalTokens: 0,
 					costTotal: 0,
 				});
+				await fork.appendMessage(createUserMessage("after fork"));
+				strictEqual((await fork.getStats()).messageCount, 4);
 				const metadata = await fork.getMetadata();
 				deepStrictEqual(
 					{ id: metadata.id, parentSessionId: metadata.parentSessionId },
@@ -779,6 +781,7 @@ export function createSessionBackendConformance(
 				{ lane: "thread", leafId: threadChild },
 			]);
 			strictEqual(await fork.getLabel(threadChild), "thread-tip");
+			strictEqual((await fork.getStats()).messageCount, 3);
 			deepStrictEqual(
 				(await fork.getLog()).filter((item) => item.kind === "lane"),
 				[
diff --git a/packages/session-backends/sqlite-node/src/sqlite/repo.ts b/packages/session-backends/sqlite-node/src/sqlite/repo.ts
index 8701ade7d..8afabf003 100644
--- a/packages/session-backends/sqlite-node/src/sqlite/repo.ts
+++ b/packages/session-backends/sqlite-node/src/sqlite/repo.ts
@@ -825,7 +825,7 @@ export class SqliteSessionRepository
 						metadata,
 					});
 					createSequence(db, id);
-					createStats(db, id);
+					createStats(db, id, entries.filter((entry) => entry.type === "message").length);
 
 					let nextSeq = 1;
 					const allocateSeq = () => nextSeq++;
diff --git a/packages/session-backends/sqlite-node/src/sqlite/storage/session-stats.ts b/packages/session-backends/sqlite-node/src/sqlite/storage/session-stats.ts
index 7aa6a90ef..370cfb947 100644
--- a/packages/session-backends/sqlite-node/src/sqlite/storage/session-stats.ts
+++ b/packages/session-backends/sqlite-node/src/sqlite/storage/session-stats.ts
@@ -11,12 +11,12 @@ export interface SessionStatsRow {
 	cost_total: number;
 }
 
-export function createStats(db: SqliteDatabase, sessionId: string): void {
+export function createStats(db: SqliteDatabase, sessionId: string, messageCount = 0): void {
 	db.prepare(
 		`INSERT INTO session_stats
 			(session_id, message_count, cached_tokens, uncached_tokens, total_tokens, cost_total)
-			VALUES (?, 0, 0, 0, 0, 0)`,
-	).run(sessionId);
+			VALUES (?, ?, 0, 0, 0, 0)`,
+	).run(sessionId, messageCount);
 }
 
 export function readStats(db: SqliteDatabase, sessionId: string): SessionStats {