diff --git a/apps/server/src/orchestration-v2/ProjectionStore.test.ts b/apps/server/src/orchestration-v2/ProjectionStore.test.ts index 8b967f648165f..79727a1b46607 100644 --- a/apps/server/src/orchestration-v2/ProjectionStore.test.ts +++ b/apps/server/src/orchestration-v2/ProjectionStore.test.ts @@ -10562,6 +10562,31 @@ it.layer(layerTest)("ProjectionStoreV2", (it) => { ["local", targetThreadId, "command_execution"], ], ); + + // Forks retain a deleted source's history, so their shell counts must keep it too. + const source = yield* projectionStore.getThreadProjection(sourceThreadId); + yield* projectionStore.apply({ + id: EventId.make("event:projection-fork-source-rollback:source-deleted"), + type: "thread.deleted", + threadId: sourceThreadId, + occurredAt: now, + payload: { ...source.thread, deletedAt: now }, + }); + const afterDeletion = yield* projectionStore.getShellSnapshot(); + assert.isNull(yield* projectionStore.getThreadShell(sourceThreadId)); + assert.notInclude( + afterDeletion.threads.map((thread) => thread.id), + sourceThreadId, + ); + for (const threadId of [targetThreadId, nestedThreadId]) { + const projection = yield* projectionStore.getThreadProjection(threadId); + for (const shell of [ + yield* projectionStore.getThreadShell(threadId), + afterDeletion.threads.find((thread) => thread.id === threadId), + ]) { + assert.equal(shell?.visibleItemCount, projection.visibleTurnItems.length, threadId); + } + } }), ); diff --git a/apps/server/src/orchestration-v2/ProjectionStore.ts b/apps/server/src/orchestration-v2/ProjectionStore.ts index 6d845450c79da..10e803ce279e2 100644 --- a/apps/server/src/orchestration-v2/ProjectionStore.ts +++ b/apps/server/src/orchestration-v2/ProjectionStore.ts @@ -7660,7 +7660,11 @@ export const layer: Layer.Layer = worktreePath, threadIds, listedOnly = false, - }: ShellSnapshotOptions & { readonly listedOnly?: boolean } = {}, + includeDeleted = false, + }: ShellSnapshotOptions & { + readonly listedOnly?: boolean; + readonly includeDeleted?: boolean; + } = {}, ) => sql` SELECT @@ -7839,7 +7843,7 @@ export const layer: Layer.Layer = candidate.ordinal DESC, candidate.run_id DESC LIMIT 1 ) AND blocked.status = 'failed' - WHERE t.deleted_at IS NULL${ + WHERE ${includeDeleted ? sql`1 = 1` : sql`t.deleted_at IS NULL`}${ listedOnly ? sql` AND json_extract(t.payload_json, '$.lineage.relationshipToParent') IS NOT 'subagent' AND json_extract(t.payload_json, '$.lineage.sideChat') IS NOT 1` @@ -8480,29 +8484,34 @@ export const layer: Layer.Layer = const shellSnapshotReadError = (cause: unknown) => new ProjectionStoreReadError({ threadId: ThreadId.make("thread:shell"), cause }); + // Adds every fork ancestor of `rows`. Ancestors may be deleted: a retained fork still inherits + // their history, so its visibleItemCount must too. + const withShellForkAncestorRows = Effect.fn("ProjectionStore.withShellForkAncestorRows")( + function* (rows: ReadonlyArray) { + const rowsByThreadId = new Map( + rows.map((row) => [ThreadId.make(row.thread_id), row] as const), + ); + const pendingSourceIds = [...shellForkSourceIds(rows)]; + while (pendingSourceIds.length > 0) { + const sourceId = pendingSourceIds.pop(); + if (sourceId === undefined || rowsByThreadId.has(sourceId)) continue; + const source = (yield* selectShellThreadRows(sourceId, { includeDeleted: true }))[0]; + if (source === undefined) continue; + rowsByThreadId.set(sourceId, source); + if (source.forked_from_run_source_thread_id !== null) { + pendingSourceIds.push(ThreadId.make(source.forked_from_run_source_thread_id)); + } + } + return rowsByThreadId; + }, + ); + const readShellSnapshotRows = Effect.fn("ProjectionStore.readShellSnapshotRows")(function* ( options: ShellSnapshotOptions | undefined, ) { const targetThreadRows = yield* selectShellThreadRows(undefined, options); const targetThreadIds = new Set(targetThreadRows.map((row) => ThreadId.make(row.thread_id))); - const rowsByThreadId = new Map( - targetThreadRows.map((row) => [ThreadId.make(row.thread_id), row] as const), - ); - const pendingSourceIds = targetThreadRows.flatMap((row) => - row.forked_from_run_source_thread_id === null - ? [] - : [ThreadId.make(row.forked_from_run_source_thread_id)], - ); - while (pendingSourceIds.length > 0) { - const sourceId = pendingSourceIds.pop(); - if (sourceId === undefined || rowsByThreadId.has(sourceId)) continue; - const source = (yield* selectShellThreadRows(sourceId))[0]; - if (source === undefined) continue; - rowsByThreadId.set(sourceId, source); - if (source.forked_from_run_source_thread_id !== null) { - pendingSourceIds.push(ThreadId.make(source.forked_from_run_source_thread_id)); - } - } + const rowsByThreadId = yield* withShellForkAncestorRows(targetThreadRows); const threadRows = [...rowsByThreadId.values()]; const threadIds = [...rowsByThreadId.keys()]; const forkSourceIds = shellForkSourceIds(threadRows); @@ -8969,26 +8978,11 @@ export const layer: Layer.Layer = sql .withTransaction( Effect.gen(function* () { - const rowsByThreadId = new Map(); - const pending: Array = [threadId]; - while (pending.length > 0) { - const nextId = pending.pop(); - if (nextId === undefined || rowsByThreadId.has(nextId)) { - continue; - } - const rows = yield* selectShellThreadRows(nextId); - const row = rows[0]; - if (row === undefined) { - continue; - } - rowsByThreadId.set(nextId, row); - if (row.forked_from_run_source_thread_id !== null) { - pending.push(ThreadId.make(row.forked_from_run_source_thread_id)); - } - } - if (!rowsByThreadId.has(threadId)) { + const targetRows = yield* selectShellThreadRows(threadId); + if (targetRows.length === 0) { return null; } + const rowsByThreadId = yield* withShellForkAncestorRows(targetRows); const threadIds = [...rowsByThreadId.keys()]; const forkSourceIds = shellForkSourceIds([...rowsByThreadId.values()]);