diff --git a/apps/server/src/orchestration-v2/ProjectionStore.test.ts b/apps/server/src/orchestration-v2/ProjectionStore.test.ts index 580cdc24cac46..9dddd56908fe8 100644 --- a/apps/server/src/orchestration-v2/ProjectionStore.test.ts +++ b/apps/server/src/orchestration-v2/ProjectionStore.test.ts @@ -5498,6 +5498,7 @@ it.layer(layerTest)("ProjectionStoreV2", (it) => { }, }); } + let runWrites = 0; const applyRun = ( threadId: ThreadId, suffix: string, @@ -5506,9 +5507,10 @@ it.layer(layerTest)("ProjectionStoreV2", (it) => { ) => { const runId = RunId.make(`run:shell-item-counts:${suffix}`); const rootNodeId = NodeId.make(`node:shell-item-counts:${suffix}`); + runWrites += 1; return projectionStore .apply({ - id: EventId.make(`event:shell-item-counts:${suffix}:run`), + id: EventId.make(`event:shell-item-counts:${suffix}:run:${runWrites}`), type: "run.created", threadId, runId, @@ -5535,11 +5537,14 @@ it.layer(layerTest)("ProjectionStoreV2", (it) => { }) .pipe(Effect.as(runId)); }; - let ordinal = 0; - const applyItem = (threadId: ThreadId, runId: RunId | null) => { - ordinal += 1; + let items = 0; + let writes = 0; + // Passing `rewrite` writes that earlier item again, possibly under another run. + const applyItem = (threadId: ThreadId, runId: RunId | null, rewrite?: number) => { + const ordinal = rewrite ?? (items += 1); + writes += 1; return projectionStore.apply({ - id: EventId.make(`event:shell-item-counts:item:${ordinal}`), + id: EventId.make(`event:shell-item-counts:write:${writes}`), type: "turn-item.updated", threadId, ...(runId === null ? {} : { runId }), @@ -5592,6 +5597,28 @@ it.layer(layerTest)("ProjectionStoreV2", (it) => { ); assert.equal(shell?.itemCount, itemCount); } + + // Writing an item again keeps its count; moving it to another run moves the count. + const itemCountOf = (threadId: ThreadId) => + projectionStore.getThreadShell(threadId).pipe(Effect.map((shell) => shell?.itemCount)); + yield* applyItem(busyThreadId, completedRunId, 3); + assert.equal(yield* itemCountOf(busyThreadId), 4); + yield* applyItem(busyThreadId, rolledBackRunId, 3); + assert.equal(yield* itemCountOf(busyThreadId), 3); + yield* applyItem(busyThreadId, null, 3); + assert.equal(yield* itemCountOf(busyThreadId), 4); + + // A run's rollback uncounts its items, and writing the run again counts them back. + yield* applyRun(busyThreadId, "busy-completed", 1, true); + assert.equal(yield* itemCountOf(busyThreadId), 3); + yield* applyRun(busyThreadId, "busy-completed", 1, false); + assert.equal(yield* itemCountOf(busyThreadId), 4); + // Items written before their run exist count once it does. + const lateRunId = RunId.make("run:shell-item-counts:late"); + yield* applyItem(otherThreadId, lateRunId); + assert.equal(yield* itemCountOf(otherThreadId), 1); + yield* applyRun(otherThreadId, "late", 2, false); + assert.equal(yield* itemCountOf(otherThreadId), 2); }), ); diff --git a/apps/server/src/orchestration-v2/ProjectionStore.ts b/apps/server/src/orchestration-v2/ProjectionStore.ts index 2a0fdc7518d10..42be5f61827d4 100644 --- a/apps/server/src/orchestration-v2/ProjectionStore.ts +++ b/apps/server/src/orchestration-v2/ProjectionStore.ts @@ -2081,6 +2081,24 @@ export const layer: Layer.Layer = `(${DEFERRED_WORK_ITEM_TYPES.map((type) => `'${type}'`).join(", ")})`, ); + // The shell counts a thread's items from these rows, so it never reads the thread's items. + // Items without a run count under run_id ''. `counted` mirrors whether the run exists and was + // not rolled back; run events keep it current. + const bumpRunItemCount = (threadId: ThreadId, runId: RunId | null, delta: 1 | -1) => sql` + INSERT INTO orchestration_v2_projection_run_item_counts (thread_id, run_id, item_count, counted) + VALUES ( + ${threadId}, ${runId ?? ""}, ${delta}, + ${ + runId === null + ? sql`1` + : sql`COALESCE( + (SELECT status <> 'rolled_back' FROM orchestration_v2_projection_runs + WHERE run_id = ${runId}), 0)` + } + ) + ON CONFLICT(thread_id, run_id) DO UPDATE SET item_count = item_count + excluded.item_count + `; + // An ended task whose completion was acknowledged, delivered or disposed, outside an active run. const settledSubagent = (row: Statement.Fragment) => sql`( ${row}.status IN ('completed', 'failed', 'cancelled', 'interrupted') @@ -2294,6 +2312,12 @@ export const layer: Layer.Layer = keepRecordedRunField(sql`excluded.payload_json`, "$.delegatedCompletion"), "$.restartCancelledBackgroundWork", )} + `; + yield* sql` + UPDATE orchestration_v2_projection_run_item_counts + SET counted = ${event.payload.status === "rolled_back" ? 0 : 1} + WHERE thread_id = ${event.payload.threadId} AND run_id = ${event.payload.id} + AND counted <> ${event.payload.status === "rolled_back" ? 0 : 1} `; break; } @@ -2788,6 +2812,13 @@ export const layer: Layer.Layer = case "turn-item.updated": { const payloadJson = yield* encodeTurnItemPayload(event.payload); const updatedAt = yield* encodePayloadTimestamp(event.payload.updatedAt); + const [previous] = yield* sql<{ + readonly thread_id: ThreadId; + readonly run_id: RunId | null; + }>` + SELECT thread_id, run_id FROM orchestration_v2_projection_turn_items + WHERE turn_item_id = ${event.payload.id} + `; yield* sql` INSERT INTO orchestration_v2_projection_turn_items ( turn_item_id, @@ -2831,6 +2862,15 @@ export const layer: Layer.Layer = updated_at = excluded.updated_at, payload_json = excluded.payload_json `; + if ( + previous?.thread_id !== event.payload.threadId || + previous.run_id !== event.payload.runId + ) { + if (previous !== undefined) { + yield* bumpRunItemCount(previous.thread_id, previous.run_id, -1); + } + yield* bumpRunItemCount(event.payload.threadId, event.payload.runId, 1); + } break; } case "checkpoint-scope.created": { @@ -6880,28 +6920,21 @@ export const layer: Layer.Layer = const selectShellRunItemCounts = (threadIds: ReadonlyArray) => sql` - SELECT thread_id, run_id, COUNT(*) AS item_count - FROM orchestration_v2_projection_turn_items - WHERE run_id IS NOT NULL AND thread_id IN (SELECT value FROM json_each(${encodeIdList(threadIds)})) - GROUP BY thread_id, run_id + SELECT thread_id, run_id, item_count + FROM orchestration_v2_projection_run_item_counts + WHERE run_id <> '' AND item_count > 0 + AND thread_id IN (SELECT value FROM json_each(${encodeIdList(threadIds)})) `; - // Count all selected threads in one index walk; fork sources still need per-run counts. + // Fork sources still need per-run counts. const selectShellItemCounts = (threadIds: ReadonlyArray) => sql` - SELECT counts.thread_id, - SUM(CASE WHEN counts.run_id IS NULL OR r.status <> 'rolled_back' - THEN counts.item_count ELSE 0 END) AS item_count, - SUM(CASE WHEN counts.run_id IS NULL THEN counts.item_count ELSE 0 END) - AS runless_item_count - FROM ( - SELECT thread_id, run_id, COUNT(*) AS item_count - FROM orchestration_v2_projection_turn_items - WHERE thread_id IN (SELECT value FROM json_each(${encodeIdList(threadIds)})) - GROUP BY thread_id, run_id - ) counts - LEFT JOIN orchestration_v2_projection_runs r ON r.run_id = counts.run_id - GROUP BY counts.thread_id + SELECT thread_id, + SUM(CASE WHEN counted = 1 THEN item_count ELSE 0 END) AS item_count, + SUM(CASE WHEN run_id = '' THEN item_count ELSE 0 END) AS runless_item_count + FROM orchestration_v2_projection_run_item_counts + WHERE thread_id IN (SELECT value FROM json_each(${encodeIdList(threadIds)})) + GROUP BY thread_id `; const selectShellProviderThreadRows = (threadIds?: ReadonlyArray) => diff --git a/apps/server/src/orchestration-v2/testkit/ProjectionReplay.ts b/apps/server/src/orchestration-v2/testkit/ProjectionReplay.ts index 437db15c375f8..b80e7d59b996d 100644 --- a/apps/server/src/orchestration-v2/testkit/ProjectionReplay.ts +++ b/apps/server/src/orchestration-v2/testkit/ProjectionReplay.ts @@ -108,6 +108,7 @@ export const layer: Layer.Layer< yield* sql`DELETE FROM orchestration_v2_projection_checkpoints`; yield* sql`DELETE FROM orchestration_v2_projection_checkpoint_scopes`; yield* sql`DELETE FROM orchestration_v2_projection_turn_items`; + yield* sql`DELETE FROM orchestration_v2_projection_run_item_counts`; yield* sql`DELETE FROM orchestration_v2_projection_plans`; yield* sql`DELETE FROM orchestration_v2_projection_messages`; yield* sql`DELETE FROM orchestration_v2_projection_runtime_requests`; diff --git a/apps/server/src/persistence/Migrations.ts b/apps/server/src/persistence/Migrations.ts index 422d2122d03bd..a5bb39ca55d15 100644 --- a/apps/server/src/persistence/Migrations.ts +++ b/apps/server/src/persistence/Migrations.ts @@ -86,6 +86,7 @@ import Migration0080 from "./Migrations/080_ScheduledTaskThreadMode.ts"; import Migration0081 from "./Migrations/081_ProjectionMessageSearchText.ts"; import Migration0082 from "./Migrations/082_TurnItemConversationIndex.ts"; import Migration0083 from "./Migrations/083_CommandReadIndexes.ts"; +import Migration0084 from "./Migrations/084_RunItemCounts.ts"; const migrationEntries = [ [1, "OrchestrationEvents", Migration0001], @@ -172,6 +173,7 @@ const migrationEntries = [ [81, "ProjectionMessageSearchText", Migration0081], [82, "TurnItemConversationIndex", Migration0082], [83, "CommandReadIndexes", Migration0083], + [84, "RunItemCounts", Migration0084], ] as const; export const migrationManifest = migrationEntries.map(([id, name]) => [id, name] as const); diff --git a/apps/server/src/persistence/Migrations/064_CanonicalizeThreadPullRequestLinks.test.ts b/apps/server/src/persistence/Migrations/064_CanonicalizeThreadPullRequestLinks.test.ts index b53309c187ed3..59dce03b79b8a 100644 --- a/apps/server/src/persistence/Migrations/064_CanonicalizeThreadPullRequestLinks.test.ts +++ b/apps/server/src/persistence/Migrations/064_CanonicalizeThreadPullRequestLinks.test.ts @@ -156,6 +156,8 @@ layer("064_CanonicalizeThreadPullRequestLinks", (it) => { `; assert.equal(v1?.payload_json, v1Payload); + // The rebuild writes through the current projection tables. + yield* runMigrations(); const maintenance = yield* ProjectionReplay.ProjectionReplay; assert.isTrue((yield* maintenance.rebuild).valid); for (const [host, repository, id] of [ diff --git a/apps/server/src/persistence/Migrations/066_RepairClientSettlementTimestamps.test.ts b/apps/server/src/persistence/Migrations/066_RepairClientSettlementTimestamps.test.ts index d62654b37ece6..02d09087e397e 100644 --- a/apps/server/src/persistence/Migrations/066_RepairClientSettlementTimestamps.test.ts +++ b/apps/server/src/persistence/Migrations/066_RepairClientSettlementTimestamps.test.ts @@ -363,6 +363,8 @@ layer("066_RepairClientSettlementTimestamps", (it) => { yield* Migration0066; assert.deepEqual(yield* snapshot, repaired); + // The rebuild writes through the current projection tables. + yield* runMigrations(); const maintenance = yield* ProjectionReplay.ProjectionReplay; assert.isTrue((yield* maintenance.rebuild).valid); assert.deepEqual(yield* imported, expectedImported); diff --git a/apps/server/src/persistence/Migrations/084_RunItemCounts.test.ts b/apps/server/src/persistence/Migrations/084_RunItemCounts.test.ts new file mode 100644 index 0000000000000..ccf5bef88cd95 --- /dev/null +++ b/apps/server/src/persistence/Migrations/084_RunItemCounts.test.ts @@ -0,0 +1,53 @@ +import { assert, it } from "@effect/vitest"; +import * as NodeSqliteClient from "@t3tools/shared/nodeSqliteClient"; +import * as Effect from "effect/Effect"; +import * as SqlClient from "effect/sql/SqlClient"; + +import { runMigrations } from "../Migrations.ts"; + +const at = "2026-10-01T00:00:00.000Z"; + +it.layer(NodeSqliteClient.layer({ filename: ":memory:" }))("084_RunItemCounts", (it) => { + it.effect( + "counts the existing items per run, marking rolled-back and missing runs uncounted", + () => + Effect.gen(function* () { + const sql = yield* SqlClient.SqlClient; + yield* runMigrations({ toMigrationInclusive: 83 }); + const insertItem = (id: string, threadId: string, runId: string | null) => sql` + INSERT INTO orchestration_v2_projection_turn_items ( + turn_item_id, thread_id, run_id, ordinal, type, status, updated_at, payload_json + ) VALUES (${id}, ${threadId}, ${runId}, 1, 'user_message', 'completed', ${at}, '{}') + `; + const insertRun = (runId: string, ordinal: number, status: string) => sql` + INSERT INTO orchestration_v2_projection_runs ( + run_id, thread_id, ordinal, provider, status, requested_at, payload_json + ) VALUES (${runId}, 'thread:a', ${ordinal}, 'codex', ${status}, ${at}, '{}') + `; + yield* insertRun("run:1", 1, "completed"); + yield* insertRun("run:2", 2, "rolled_back"); + yield* insertItem("a1", "thread:a", "run:1"); + yield* insertItem("a2", "thread:a", "run:1"); + yield* insertItem("a3", "thread:a", "run:2"); + yield* insertItem("a4", "thread:a", null); + yield* insertItem("b1", "thread:b", "run:3"); + + assert.deepEqual(yield* runMigrations({ toMigrationInclusive: 84 }), [ + [84, "RunItemCounts"], + ]); + assert.deepEqual( + yield* sql` + SELECT thread_id, run_id, item_count, counted + FROM orchestration_v2_projection_run_item_counts + ORDER BY thread_id, run_id + `, + [ + { thread_id: "thread:a", run_id: "", item_count: 1, counted: 1 }, + { thread_id: "thread:a", run_id: "run:1", item_count: 2, counted: 1 }, + { thread_id: "thread:a", run_id: "run:2", item_count: 1, counted: 0 }, + { thread_id: "thread:b", run_id: "run:3", item_count: 1, counted: 0 }, + ], + ); + }), + ); +}); diff --git a/apps/server/src/persistence/Migrations/084_RunItemCounts.ts b/apps/server/src/persistence/Migrations/084_RunItemCounts.ts new file mode 100644 index 0000000000000..a1256ff00663d --- /dev/null +++ b/apps/server/src/persistence/Migrations/084_RunItemCounts.ts @@ -0,0 +1,31 @@ +import * as Effect from "effect/Effect"; +import * as SqlClient from "effect/sql/SqlClient"; + +export default Effect.gen(function* () { + const sql = yield* SqlClient.SqlClient; + + // A thread's shell shows how many items it holds. ProjectionStore keeps these counts as items + // and runs are written, so the shell no longer counts every item of a long thread on each update. + // Items without a run count under run_id ''. `counted` is 1 for a run that exists and was not + // rolled back, and for items without a run. + yield* sql` + CREATE TABLE orchestration_v2_projection_run_item_counts ( + thread_id TEXT NOT NULL, + run_id TEXT NOT NULL, + item_count INTEGER NOT NULL, + counted INTEGER NOT NULL, + PRIMARY KEY (thread_id, run_id) + ) WITHOUT ROWID + `; + yield* sql` + INSERT INTO orchestration_v2_projection_run_item_counts (thread_id, run_id, item_count, counted) + SELECT counts.thread_id, counts.run_id, counts.item_count, + CASE WHEN counts.run_id = '' OR run.status <> 'rolled_back' THEN 1 ELSE 0 END + FROM ( + SELECT thread_id, COALESCE(run_id, '') AS run_id, COUNT(*) AS item_count + FROM orchestration_v2_projection_turn_items + GROUP BY thread_id, COALESCE(run_id, '') + ) AS counts + LEFT JOIN orchestration_v2_projection_runs AS run ON run.run_id = counts.run_id + `; +}); diff --git a/apps/server/src/persistence/reconcileV2PreviewMigration.test.ts b/apps/server/src/persistence/reconcileV2PreviewMigration.test.ts index 05f09f8e16850..f2c17477ca1a7 100644 --- a/apps/server/src/persistence/reconcileV2PreviewMigration.test.ts +++ b/apps/server/src/persistence/reconcileV2PreviewMigration.test.ts @@ -63,6 +63,7 @@ describe("V2 preview upgrade", () => { [81, "ProjectionMessageSearchText"], [82, "TurnItemConversationIndex"], [83, "CommandReadIndexes"], + [84, "RunItemCounts"], ]); assert.deepStrictEqual(yield* runMigrations(), []); assert.deepStrictEqual(yield* sql`SELECT * FROM orchestration_v2_legacy_imports`, imports); @@ -169,6 +170,7 @@ describe("V2 preview upgrade", () => { [81, "ProjectionMessageSearchText"], [82, "TurnItemConversationIndex"], [83, "CommandReadIndexes"], + [84, "RunItemCounts"], ]); }).pipe(Effect.provide(NodeSqliteClient.layer({ filename: ":memory:" }))), );