Skip to content

fix(server): forks no longer show turns the source rewound before the fork #29695

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
saphid wants to merge 1 commit into from
+73 −25
Closed

fix(server): forks no longer show turns the source rewound before the fork #29695

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
24 changes: 20 additions & 4 deletions apps/server/src/orchestration-v2/ProjectionStore.test.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -6622,14 +6622,30 @@ it.layer(layerTest)("ProjectionStoreV2", (it) => {
const emptyBoundaryFork = yield* projectionStore.getThreadSnapshotWindow(targetThreadId, {
rowLimit: 2,
});
assert.deepEqual(
emptyBoundaryFork.projection.visibleTurnItems.map((row) =>
// Run 2 was rolled back before this later fork run existed, so the fork never had it.
const forkRowLabels = (rows: ReadonlyArray<{ readonly item: OrchestrationV2TurnItem }>) =>
rows.map((row) =>
row.item.type === "user_message" || row.item.type === "assistant_message"
? row.item.text
: row.item.type,
),
["source two", "two", "fork"],
);
assert.deepEqual(forkRowLabels(emptyBoundaryFork.projection.visibleTurnItems), [
"source one",
"one",
"fork",
]);
const emptyBoundaryProjection = yield* projectionStore.getThreadProjection(targetThreadId);
assert.deepEqual(forkRowLabels(emptyBoundaryProjection.visibleTurnItems), [
"source one",
"one",
"fork",
]);
assert.deepEqual(
(yield* projectionStore.getTimelinePage(targetThreadId, { view: "activity", limit: 10 }))
.items,
emptyBoundaryProjection.visibleTurnItems,
);
assert.equal((yield* projectionStore.getThreadShell(targetThreadId))?.visibleItemCount, 3);
yield* sql`
UPDATE orchestration_v2_projection_threads
SET payload_json = json_set(payload_json, '$.forkedFrom.runId', ${sourceRun2Id})
Expand Down
74 changes: 53 additions & 21 deletions apps/server/src/orchestration-v2/ProjectionStore.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -989,8 +989,11 @@ type ShellRunRow = {
readonly thread_id: string;
readonly run_id: string;
readonly ordinal: number;
readonly status: string;
};

type ForkSourceRun = { readonly id: RunId; readonly ordinal: number; readonly status: string };

type SettlementThreadRow = Pick<
ShellThreadRow,
| "thread_id"
Expand Down Expand Up @@ -1298,6 +1301,22 @@ function isTurnItemAtOrBeforeRun(input: {
return ordinal !== undefined && ordinal <= input.sourceRunOrdinal;
}

/**
* Source runs a fork inherits from. Runs rolled back before the fork had already left the
* source's history; once the fork run itself is rolled back, the fork keeps every row it inherited.
*/
function forkSourceRunOrdinals(
runs: ReadonlyArray<ForkSourceRun>,
forkRunId: RunId,
): Map<RunId, number> {
const keepRolledBack = runs.some((run) => run.id === forkRunId && run.status === "rolled_back");
Comment thread
coderabbitai[bot] marked this conversation as resolved.
Resolved
return new Map(
runs.flatMap((run) =>
run.status === "rolled_back" && !keepRolledBack ? [] : [[run.id, run.ordinal] as const],
),
);
}

function visibleTurnItemsThroughRun(input: {
readonly sourceProjection: OrchestrationV2ThreadProjection;
readonly sourceRunId: NonNullable<OrchestrationV2TurnItem["runId"]>;
Expand All @@ -1307,7 +1326,7 @@ function visibleTurnItemsThroughRun(input: {
return [];
}

const runOrdinalById = new Map(input.sourceProjection.runs.map((run) => [run.id, run.ordinal]));
const runOrdinalById = forkSourceRunOrdinals(input.sourceProjection.runs, sourceRun.id);
const inheritedPrefix = input.sourceProjection.visibleTurnItems
.filter(
(row) => row.item.threadId !== input.sourceProjection.thread.id || row.item.type === "fork",
Expand Down Expand Up @@ -1685,7 +1704,7 @@ type ShellThreadState = {
readonly itemCount: number;
readonly runlessItemCount: number;
readonly updatedAt: OrchestrationV2ThreadProjection["updatedAt"];
readonly runOrdinalById: ReadonlyMap<RunId, number>;
readonly runs: ReadonlyArray<ForkSourceRun>;
readonly itemCountByRunId: ReadonlyMap<RunId, number>;
};

Expand Down Expand Up @@ -1713,14 +1732,15 @@ function itemCountThroughRun(input: {
readonly state: ShellThreadState;
readonly runId: RunId;
}): number {
const runOrdinal = input.state.runOrdinalById.get(input.runId);
const runOrdinalById = forkSourceRunOrdinals(input.state.runs, input.runId);
const runOrdinal = runOrdinalById.get(input.runId);
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);
const itemRunOrdinal = runOrdinalById.get(runId);
if (itemRunOrdinal !== undefined && itemRunOrdinal <= runOrdinal) {
count += itemCount;
}
Expand Down Expand Up @@ -2798,7 +2818,20 @@ export const layer: Layer.Layer<ProjectionStoreV2, never, SqlClient.SqlClient> =
9223372036854775807
)
AND (
${window.requiredRunId ?? null} IS NOT NULL
(
${window.requiredRunId ?? null} IS NOT NULL
-- Matches forkSourceRunOrdinals.
AND (
run.status IS NULL
OR run.status <> 'rolled_back'
OR EXISTS (
SELECT 1 FROM orchestration_v2_projection_runs AS fork_run
WHERE fork_run.thread_id = ${threadId}
AND fork_run.run_id = ${window.requiredRunId ?? null}
AND fork_run.status = 'rolled_back'
)
)
)
OR (
(run.status IS NULL OR run.status <> 'rolled_back')
AND NOT (
Expand Down Expand Up @@ -5580,7 +5613,7 @@ export const layer: Layer.Layer<ProjectionStoreV2, never, SqlClient.SqlClient> =
const sourceRun = source.records.runs.find((run) => run.id === fork.runId);
let inherited: Array<TimelineIndexRow> = [];
if (sourceRun !== undefined) {
const ordinals = new Map(source.records.runs.map((run) => [run.id, run.ordinal]));
const ordinals = forkSourceRunOrdinals(source.records.runs, sourceRun.id);
const sourceItems = source.local.map((row) => row.item);
inherited = [
...source.visible.filter(
Expand Down Expand Up @@ -6152,11 +6185,11 @@ export const layer: Layer.Layer<ProjectionStoreV2, never, SqlClient.SqlClient> =
const selectShellRunRows = (threadIds?: ReadonlyArray<ThreadId>) =>
threadIds === undefined
? sql<ShellRunRow>`
SELECT thread_id, run_id, ordinal
SELECT thread_id, run_id, ordinal, status
FROM orchestration_v2_projection_runs
`
: sql<ShellRunRow>`
SELECT thread_id, run_id, ordinal
SELECT thread_id, run_id, ordinal, status
FROM orchestration_v2_projection_runs
WHERE thread_id IN (SELECT value FROM json_each(${encodeIdList(threadIds)}))
`;
Expand Down Expand Up @@ -6233,13 +6266,12 @@ export const layer: Layer.Layer<ProjectionStoreV2, never, SqlClient.SqlClient> =
readonly runRows: ReadonlyArray<ShellRunRow>;
readonly itemCountRows: ReadonlyArray<ShellRunItemCountRow>;
}) => {
const runOrdinalsByThreadId = new Map<ThreadId, Map<RunId, number>>();
const runsByThreadId = new Map<ThreadId, Array<ForkSourceRun>>();
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<RunId, number>();
existing.set(runId, row.ordinal);
runOrdinalsByThreadId.set(threadId, existing);
const existing = runsByThreadId.get(threadId) ?? [];
existing.push({ id: RunId.make(row.run_id), ordinal: row.ordinal, status: row.status });
runsByThreadId.set(threadId, existing);
}
const itemCountsByThreadId = new Map<ThreadId, Map<RunId, number>>();
for (const row of input.itemCountRows) {
Expand All @@ -6249,7 +6281,7 @@ export const layer: Layer.Layer<ProjectionStoreV2, never, SqlClient.SqlClient> =
existing.set(runId, row.item_count);
itemCountsByThreadId.set(threadId, existing);
}
return { runOrdinalsByThreadId, itemCountsByThreadId };
return { runsByThreadId, itemCountsByThreadId };
};

const shellItemCountMap = (rows: ReadonlyArray<ShellItemCountRow>) =>
Expand Down Expand Up @@ -6500,7 +6532,7 @@ export const layer: Layer.Layer<ProjectionStoreV2, never, SqlClient.SqlClient> =
/* Runs once per shell row; the enclosing snapshot or thread read owns the span. */
Effect.fnUntraced(function* (input: {
readonly row: ShellThreadRow;
readonly runOrdinalsByThreadId: ReadonlyMap<ThreadId, Map<RunId, number>>;
readonly runsByThreadId: ReadonlyMap<ThreadId, ReadonlyArray<ForkSourceRun>>;
readonly itemCountsByThreadId: ReadonlyMap<ThreadId, Map<RunId, number>>;
readonly shellItemCountsByThreadId: ReadonlyMap<
ThreadId,
Expand All @@ -6517,7 +6549,7 @@ export const layer: Layer.Layer<ProjectionStoreV2, never, SqlClient.SqlClient> =
}) {
const {
row,
runOrdinalsByThreadId,
runsByThreadId,
itemCountsByThreadId,
shellItemCountsByThreadId,
providerThreadsByThreadId,
Expand Down Expand Up @@ -6673,7 +6705,7 @@ export const layer: Layer.Layer<ProjectionStoreV2, never, SqlClient.SqlClient> =
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(),
runs: runsByThreadId.get(ThreadId.make(row.thread_id)) ?? [],
itemCountByRunId: itemCountsByThreadId.get(ThreadId.make(row.thread_id)) ?? new Map(),
} satisfies ShellThreadState;
});
Expand Down Expand Up @@ -6762,7 +6794,7 @@ export const layer: Layer.Layer<ProjectionStoreV2, never, SqlClient.SqlClient> =
providerThreadRows,
pendingTurnItemRows,
}: Effect.Success<ReturnType<typeof readShellSnapshotRows>>) {
const { runOrdinalsByThreadId, itemCountsByThreadId } = runMapsByThreadId({
const { runsByThreadId, itemCountsByThreadId } = runMapsByThreadId({
runRows,
itemCountRows,
});
Expand All @@ -6772,7 +6804,7 @@ export const layer: Layer.Layer<ProjectionStoreV2, never, SqlClient.SqlClient> =
const states = yield* Effect.forEach(threadRows, (row) =>
shellThreadStateFromRow({
row,
runOrdinalsByThreadId,
runsByThreadId,
itemCountsByThreadId,
shellItemCountsByThreadId,
providerThreadsByThreadId,
Expand Down Expand Up @@ -7034,7 +7066,7 @@ export const layer: Layer.Layer<ProjectionStoreV2, never, SqlClient.SqlClient> =
),
),
]);
const { runOrdinalsByThreadId, itemCountsByThreadId } = runMapsByThreadId({
const { runsByThreadId, itemCountsByThreadId } = runMapsByThreadId({
runRows,
itemCountRows,
});
Expand All @@ -7045,7 +7077,7 @@ export const layer: Layer.Layer<ProjectionStoreV2, never, SqlClient.SqlClient> =
const states = yield* Effect.forEach([...rowsByThreadId.values()], (row) =>
shellThreadStateFromRow({
row,
runOrdinalsByThreadId,
runsByThreadId,
itemCountsByThreadId,
shellItemCountsByThreadId,
providerThreadsByThreadId,
Expand Down
Loading