diff --git a/apps/server/src/orchestration-v2/ProjectionStore.test.ts b/apps/server/src/orchestration-v2/ProjectionStore.test.ts index 041d4d0946b12..06fc3ca868ae1 100644 --- a/apps/server/src/orchestration-v2/ProjectionStore.test.ts +++ b/apps/server/src/orchestration-v2/ProjectionStore.test.ts @@ -9006,6 +9006,193 @@ it.layer(layerTest)("ProjectionStoreV2", (it) => { }), ); + it.effect("hides source turns rewound before the fork from a fork's visible items", () => + Effect.gen(function* () { + const projectionStore = yield* ProjectionStore.ProjectionStoreV2; + const now = yield* DateTime.now; + const rewoundAt = DateTime.add(now, { minutes: 1 }); + const forkedAt = DateTime.add(now, { minutes: 2 }); + const prefix = "projection-fork-after-rewind"; + const projectId = ProjectId.make(`project:${prefix}`); + const sourceThreadId = ThreadId.make(`thread:${prefix}:source`); + const targetThreadId = ThreadId.make(`thread:${prefix}:target`); + const sourceProviderThreadId = ProviderThreadId.make(`provider-thread:${prefix}:source`); + const runIds = [1, 2, 3].map((ordinal) => RunId.make(`run:${prefix}:${ordinal}`)); + const nodeIds = [1, 2, 3].map((ordinal) => NodeId.make(`node:${prefix}:${ordinal}`)); + const threadPayload = (input: { + readonly id: ThreadId; + readonly forkedFrom: { type: "run"; threadId: ThreadId; runId: RunId } | null; + readonly createdAt: DateTime.Utc; + }) => ({ + createdBy: "user" as const, + creationSource: "web" as const, + id: input.id, + projectId, + title: input.id, + providerInstanceId, + modelSelection, + runtimeMode: "full-access" as const, + interactionMode: "default" as const, + branch: null, + worktreePath: null, + activeProviderThreadId: sourceProviderThreadId, + lineage: + input.forkedFrom === null + ? { parentThreadId: null, relationshipToParent: null, rootThreadId: sourceThreadId } + : { + parentThreadId: sourceThreadId, + relationshipToParent: "fork" as const, + rootThreadId: sourceThreadId, + }, + forkedFrom: input.forkedFrom, + createdAt: input.createdAt, + updatedAt: input.createdAt, + archivedAt: null, + settledOverride: null, + settledAt: null, + lastVisitedAt: null, + deletedAt: null, + }); + const runPayload = (index: number, status: "completed" | "rolled_back") => ({ + id: runIds[index]!, + threadId: sourceThreadId, + ordinal: index + 1, + providerInstanceId, + modelSelection, + providerThreadId: sourceProviderThreadId, + userMessageId: MessageId.make(`message:${prefix}:user:${index + 1}`), + rootNodeId: nodeIds[index]!, + activeAttemptId: null, + status, + requestedAt: now, + startedAt: now, + completedAt: status === "rolled_back" ? rewoundAt : now, + checkpointId: null, + contextHandoffId: null, + }); + + yield* projectionStore.apply({ + id: EventId.make(`event:${prefix}:source-thread`), + type: "thread.created", + threadId: sourceThreadId, + occurredAt: now, + payload: threadPayload({ id: sourceThreadId, forkedFrom: null, createdAt: now }), + }); + for (const [index, text] of ["one", "two", "three"].entries()) { + const runId = runIds[index]!; + const nodeId = nodeIds[index]!; + const ordinal = index + 1; + yield* projectionStore.apply({ + id: EventId.make(`event:${prefix}:run-${ordinal}`), + type: "run.created", + threadId: sourceThreadId, + runId, + nodeId, + driver, + occurredAt: now, + payload: runPayload(index, "completed"), + }); + for (const [offset, role] of (["user", "assistant"] as const).entries()) { + const itemBase = { + id: TurnItemId.make(`turn-item:${prefix}:${role}:${ordinal}`), + threadId: sourceThreadId, + runId, + nodeId, + providerThreadId: sourceProviderThreadId, + providerTurnId: null, + nativeItemRef: null, + parentItemId: null, + ordinal: ordinal * 100 + offset, + status: "completed" as const, + title: null, + startedAt: now, + completedAt: now, + updatedAt: now, + messageId: MessageId.make(`message:${prefix}:${role}:${ordinal}`), + }; + yield* projectionStore.apply({ + id: EventId.make(`event:${prefix}:${role}-item-${ordinal}`), + type: "turn-item.updated", + threadId: sourceThreadId, + runId, + nodeId, + driver, + occurredAt: now, + payload: + role === "user" + ? { + ...itemBase, + createdBy: "user", + creationSource: "web", + type: "user_message", + inputIntent: "turn_start", + text: `prompt ${text}`, + attachments: [], + } + : { + ...itemBase, + type: "assistant_message", + text: `reply ${text}`, + streaming: false, + }, + }); + } + // Edit from here on prompt two rewinds run 2 before prompt three is sent. + if (ordinal === 2) { + yield* projectionStore.apply({ + id: EventId.make(`event:${prefix}:run-2-rolled-back`), + type: "run.updated", + threadId: sourceThreadId, + runId, + nodeId, + driver, + occurredAt: rewoundAt, + payload: runPayload(index, "rolled_back"), + }); + } + } + + const forkedFrom = { type: "run" as const, threadId: sourceThreadId, runId: runIds[2]! }; + yield* projectionStore.apply({ + id: EventId.make(`event:${prefix}:target-thread`), + type: "thread.created", + threadId: targetThreadId, + occurredAt: forkedAt, + payload: threadPayload({ id: targetThreadId, forkedFrom, createdAt: forkedAt }), + }); + + const visibleTexts = (rows: ReadonlyArray<{ readonly item: OrchestrationV2TurnItem }>) => + rows.map((row) => + row.item.type === "user_message" || row.item.type === "assistant_message" + ? row.item.text + : row.item.type, + ); + const source = yield* projectionStore.getThreadProjection(sourceThreadId); + assert.deepEqual(visibleTexts(source.visibleTurnItems), [ + "prompt one", + "reply one", + "prompt three", + "reply three", + ]); + const expected = ["prompt one", "reply one", "prompt three", "reply three", "fork"]; + const target = yield* projectionStore.getThreadProjection(targetThreadId); + assert.deepEqual(visibleTexts(target.visibleTurnItems), expected); + const page = yield* projectionStore.getTimelinePage(targetThreadId, { + view: "activity", + limit: 10, + }); + assert.deepEqual(visibleTexts(page.items), expected); + const window = yield* projectionStore.getThreadSnapshotWindow(targetThreadId, { + rowLimit: 10, + }); + assert.deepEqual(visibleTexts(window.projection.visibleTurnItems), expected); + assert.equal( + (yield* projectionStore.getThreadShell(targetThreadId))?.visibleItemCount, + expected.length, + ); + }), + ); + it.effect("a pull request watch keeps a finished thread working until it ends", () => Effect.gen(function* () { const store = yield* ProjectionStore.ProjectionStoreV2; diff --git a/apps/server/src/orchestration-v2/ProjectionStore.ts b/apps/server/src/orchestration-v2/ProjectionStore.ts index d869eb9890c56..9a058f2cddb88 100644 --- a/apps/server/src/orchestration-v2/ProjectionStore.ts +++ b/apps/server/src/orchestration-v2/ProjectionStore.ts @@ -741,6 +741,7 @@ type ProjectionWindow = Omit< "anchorThreadId" > & { readonly suppressLocal?: boolean | undefined; + readonly forkCreatedAt?: DateTime.Utc | undefined; readonly historyAnchor?: { readonly threadId: ThreadId; readonly itemId: TurnItemId } | undefined; }; @@ -1071,6 +1072,12 @@ type ShellRunRow = { readonly thread_id: string; readonly run_id: string; readonly ordinal: number; + readonly rolled_back_at: string | null; +}; + +type ShellRun = { + readonly ordinal: number; + readonly rolledBackAt: DateTime.Utc | null; }; type SettlementThreadRow = Pick< @@ -1415,16 +1422,40 @@ function isTurnItemAtOrBeforeRun(input: { return ordinal !== undefined && ordinal <= input.sourceRunOrdinal; } +// Rollback stamps completedAt; a fork made after the rewind never received that run. +function runRolledBackAt(run: Pick) { + return run.status === "rolled_back" ? run.completedAt : null; +} + +function isRolledBackBeforeFork( + rolledBackAt: DateTime.Utc | null, + forkCreatedAt: DateTime.Utc, +): boolean { + return rolledBackAt !== null && DateTime.isLessThan(rolledBackAt, forkCreatedAt); +} + +function forkInheritedRunOrdinals( + runs: ReadonlyArray, + forkCreatedAt: DateTime.Utc, +): Map { + return new Map( + runs + .filter((run) => !isRolledBackBeforeFork(runRolledBackAt(run), forkCreatedAt)) + .map((run) => [run.id, run.ordinal]), + ); +} + function visibleTurnItemsThroughRun(input: { readonly sourceProjection: OrchestrationV2ThreadProjection; readonly sourceRunId: NonNullable; + readonly forkCreatedAt: DateTime.Utc; }): Array> { const sourceRun = input.sourceProjection.runs.find((run) => run.id === input.sourceRunId); if (sourceRun === undefined) { return []; } - const runOrdinalById = new Map(input.sourceProjection.runs.map((run) => [run.id, run.ordinal])); + const runOrdinalById = forkInheritedRunOrdinals(input.sourceProjection.runs, input.forkCreatedAt); const inheritedPrefix = input.sourceProjection.visibleTurnItems .filter( (row) => row.item.threadId !== input.sourceProjection.thread.id || row.item.type === "fork", @@ -1484,6 +1515,7 @@ function buildVisibleTurnItems(input: { : visibleTurnItemsThroughRun({ sourceProjection: input.sourceProjection, sourceRunId: forkedFrom.runId, + forkCreatedAt: input.projection.thread.createdAt, }); const markerItem = makeForkMarkerTurnItem({ targetProjection: input.projection, @@ -1832,7 +1864,7 @@ type ShellThreadState = { readonly itemCount: number; readonly runlessItemCount: number; readonly updatedAt: OrchestrationV2ThreadProjection["updatedAt"]; - readonly runOrdinalById: ReadonlyMap; + readonly runsById: ReadonlyMap; readonly itemCountByRunId: ReadonlyMap; }; @@ -1859,16 +1891,21 @@ function shellStatusFromStoredRunStatus(status: string | null): OrchestrationV2S function itemCountThroughRun(input: { readonly state: ShellThreadState; readonly runId: RunId; + readonly forkCreatedAt: DateTime.Utc; }): number { - const runOrdinal = input.state.runOrdinalById.get(input.runId); + const runOrdinal = input.state.runsById.get(input.runId)?.ordinal; if (runOrdinal === undefined) { return 0; } let count = input.state.thread.historyOrigin === "v1_import" ? input.state.runlessItemCount : 0; for (const [runId, itemCount] of input.state.itemCountByRunId) { - const itemRunOrdinal = input.state.runOrdinalById.get(runId); - if (itemRunOrdinal !== undefined && itemRunOrdinal <= runOrdinal) { + const itemRun = input.state.runsById.get(runId); + if ( + itemRun !== undefined && + itemRun.ordinal <= runOrdinal && + !isRolledBackBeforeFork(itemRun.rolledBackAt, input.forkCreatedAt) + ) { count += itemCount; } } @@ -1925,7 +1962,11 @@ function visibleItemCountForShell(input: { return ( inheritedPrefixCount + - itemCountThroughRun({ state: sourceState, runId: forkedFrom.runId }) + + itemCountThroughRun({ + state: sourceState, + runId: forkedFrom.runId, + forkCreatedAt: state.thread.createdAt, + }) + 1 + state.itemCount ); @@ -3035,12 +3076,23 @@ export const layer: Layer.Layer = AND json_extract(source_thread.payload_json, '$.historyOrigin') = 'v1_import' ) ) - OR run.ordinal <= ( - SELECT required_run.ordinal - FROM orchestration_v2_projection_runs AS required_run - WHERE required_run.thread_id = item.thread_id - AND required_run.run_id = ${window.requiredRunId ?? null} - LIMIT 1 + OR ( + run.ordinal <= ( + SELECT required_run.ordinal + FROM orchestration_v2_projection_runs AS required_run + WHERE required_run.thread_id = item.thread_id + AND required_run.run_id = ${window.requiredRunId ?? null} + LIMIT 1 + ) + AND NOT COALESCE( + run.status = 'rolled_back' + AND julianday(run.completed_at) < julianday(${ + window.forkCreatedAt === undefined + ? null + : DateTime.formatIso(window.forkCreatedAt) + }), + 0 + ) ) ) AND NOT ( @@ -4154,6 +4206,7 @@ export const layer: Layer.Layer = userTurnLimit: window.userTurnLimit, deferWork: window.deferWork, requiredRunId: forkedFrom.runId, + forkCreatedAt: projection.thread.createdAt, suppressLocal: anchor === undefined || anchorIsInDescendant, ...(anchor === undefined || anchorIsInDescendant ? {} @@ -6045,7 +6098,7 @@ export const layer: Layer.Layer = const sourceRun = source.records.runs.find((run) => run.id === fork.runId); let inherited: Array = []; if (sourceRun !== undefined) { - const ordinals = new Map(source.records.runs.map((run) => [run.id, run.ordinal])); + const ordinals = forkInheritedRunOrdinals(source.records.runs, records.thread.createdAt); const sourceItems = source.local.map((row) => row.item); inherited = [ ...source.visible.filter( @@ -6725,11 +6778,13 @@ export const layer: Layer.Layer = const selectShellRunRows = (threadIds?: ReadonlyArray) => threadIds === undefined ? sql` - SELECT thread_id, run_id, ordinal + SELECT thread_id, run_id, ordinal, + CASE WHEN status = 'rolled_back' THEN completed_at END AS rolled_back_at FROM orchestration_v2_projection_runs ` : sql` - SELECT thread_id, run_id, ordinal + SELECT thread_id, run_id, ordinal, + CASE WHEN status = 'rolled_back' THEN completed_at END AS rolled_back_at FROM orchestration_v2_projection_runs WHERE thread_id IN (SELECT value FROM json_each(${encodeIdList(threadIds)})) `; @@ -6806,13 +6861,17 @@ export const layer: Layer.Layer = readonly runRows: ReadonlyArray; readonly itemCountRows: ReadonlyArray; }) => { - const runOrdinalsByThreadId = new Map>(); + const runsByThreadId = new Map>(); for (const row of input.runRows) { const threadId = ThreadId.make(row.thread_id); const runId = RunId.make(row.run_id); - const existing = runOrdinalsByThreadId.get(threadId) ?? new Map(); - existing.set(runId, row.ordinal); - runOrdinalsByThreadId.set(threadId, existing); + const existing = runsByThreadId.get(threadId) ?? new Map(); + existing.set(runId, { + ordinal: row.ordinal, + rolledBackAt: + row.rolled_back_at === null ? null : DateTime.makeUnsafe(row.rolled_back_at), + }); + runsByThreadId.set(threadId, existing); } const itemCountsByThreadId = new Map>(); for (const row of input.itemCountRows) { @@ -6822,7 +6881,7 @@ export const layer: Layer.Layer = existing.set(runId, row.item_count); itemCountsByThreadId.set(threadId, existing); } - return { runOrdinalsByThreadId, itemCountsByThreadId }; + return { runsByThreadId, itemCountsByThreadId }; }; const shellItemCountMap = (rows: ReadonlyArray) => @@ -7093,7 +7152,7 @@ export const layer: Layer.Layer = /* Runs once per shell row; the enclosing snapshot or thread read owns the span. */ Effect.fnUntraced(function* (input: { readonly row: ShellThreadRow; - readonly runOrdinalsByThreadId: ReadonlyMap>; + readonly runsByThreadId: ReadonlyMap>; readonly itemCountsByThreadId: ReadonlyMap>; readonly shellItemCountsByThreadId: ReadonlyMap< ThreadId, @@ -7110,7 +7169,7 @@ export const layer: Layer.Layer = }) { const { row, - runOrdinalsByThreadId, + runsByThreadId, itemCountsByThreadId, shellItemCountsByThreadId, providerThreadsByThreadId, @@ -7276,7 +7335,7 @@ export const layer: Layer.Layer = itemCount: shellItemCountsByThreadId.get(thread.id)?.itemCount ?? 0, runlessItemCount: shellItemCountsByThreadId.get(thread.id)?.runlessItemCount ?? 0, updatedAt: thread.updatedAt, - runOrdinalById: runOrdinalsByThreadId.get(ThreadId.make(row.thread_id)) ?? new Map(), + runsById: runsByThreadId.get(ThreadId.make(row.thread_id)) ?? new Map(), itemCountByRunId: itemCountsByThreadId.get(ThreadId.make(row.thread_id)) ?? new Map(), } satisfies ShellThreadState; }); @@ -7365,7 +7424,7 @@ export const layer: Layer.Layer = providerThreadRows, pendingTurnItemRows, }: Effect.Success>) { - const { runOrdinalsByThreadId, itemCountsByThreadId } = runMapsByThreadId({ + const { runsByThreadId, itemCountsByThreadId } = runMapsByThreadId({ runRows, itemCountRows, }); @@ -7375,7 +7434,7 @@ export const layer: Layer.Layer = const states = yield* Effect.forEach(threadRows, (row) => shellThreadStateFromRow({ row, - runOrdinalsByThreadId, + runsByThreadId, itemCountsByThreadId, shellItemCountsByThreadId, providerThreadsByThreadId, @@ -7481,7 +7540,7 @@ export const layer: Layer.Layer = shellThreadStateFromRow({ row, // Awareness never consumes timeline counts or inherited fork history. - runOrdinalsByThreadId: new Map(), + runsByThreadId: new Map(), itemCountsByThreadId: new Map(), shellItemCountsByThreadId: new Map(), providerThreadsByThreadId, @@ -7772,7 +7831,7 @@ export const layer: Layer.Layer = ), ), ]); - const { runOrdinalsByThreadId, itemCountsByThreadId } = runMapsByThreadId({ + const { runsByThreadId, itemCountsByThreadId } = runMapsByThreadId({ runRows, itemCountRows, }); @@ -7783,7 +7842,7 @@ export const layer: Layer.Layer = const states = yield* Effect.forEach([...rowsByThreadId.values()], (row) => shellThreadStateFromRow({ row, - runOrdinalsByThreadId, + runsByThreadId, itemCountsByThreadId, shellItemCountsByThreadId, providerThreadsByThreadId,