make forking work

This commit is contained in:
Connor Peet
2026-03-31 20:16:07 -07:00
parent 8fe8366dac
commit ff9044ba14
3 changed files with 41 additions and 20 deletions

View File

@@ -179,10 +179,7 @@ export class AgentService extends Disposable implements IAgentService {
const sourceState = this._stateManager.getSessionState(config.fork.session.toString());
let sourceTurns: ITurn[] = [];
if (sourceState) {
const forkIdx = sourceState.turns.findIndex(t => t.id === config.fork!.turnId);
if (forkIdx >= 0) {
sourceTurns = sourceState.turns.slice(0, forkIdx + 1);
}
sourceTurns = sourceState.turns.slice(0, config.fork.turnIndex + 1);
}
const summary: ISessionSummary = {

View File

@@ -310,24 +310,41 @@ export async function forkSessionInDb(
[newSessionId, sourceSessionId, forkTurnIndex],
);
// Copy search index entries for kept turns
await dbRun(db,
`INSERT INTO search_index (content, session_id, source_type, source_id)
SELECT content, ?, source_type,
REPLACE(source_id, ?, ?)
// Copy search index entries for kept turns only.
// source_id format is "<session_id>:turn:<turn_index>"; filter by
// parsing the turn index so we don't leak content from later turns.
await dbAll(db,
`SELECT content, source_type, source_id
FROM search_index
WHERE session_id = ?
AND source_type = 'turn'`,
[newSessionId, sourceSessionId, newSessionId, sourceSessionId],
);
WHERE session_id = ? AND source_type = 'turn'`,
[sourceSessionId],
).then(async rows => {
const prefix = `${sourceSessionId}:turn:`;
for (const row of rows) {
const sourceId = row.source_id as string;
if (sourceId.startsWith(prefix)) {
const turnIdx = parseInt(sourceId.substring(prefix.length), 10);
if (!isNaN(turnIdx) && turnIdx <= forkTurnIndex) {
const newSourceId = sourceId.replace(sourceSessionId, newSessionId);
await dbRun(db,
`INSERT INTO search_index (content, session_id, source_type, source_id)
VALUES (?, ?, ?, ?)`,
[row.content, newSessionId, row.source_type, newSourceId],
);
}
}
}
});
// Copy checkpoints at or before the fork point
// Copy checkpoints at or before the fork point.
// checkpoint_number is 1-based and correlates to turns, so we keep
// only those where checkpoint_number <= forkTurnIndex + 1.
await dbRun(db,
`INSERT INTO checkpoints (session_id, checkpoint_number, title, overview, history, work_done, technical_details, important_files, next_steps, created_at)
SELECT ?, checkpoint_number, title, overview, history, work_done, technical_details, important_files, next_steps, created_at
FROM checkpoints
WHERE session_id = ?`,
[newSessionId, sourceSessionId],
WHERE session_id = ? AND checkpoint_number <= ?`,
[newSessionId, sourceSessionId, forkTurnIndex + 1],
);
await dbExec(db, 'COMMIT');

View File

@@ -1270,15 +1270,22 @@ export class AgentHostSessionHandler extends Disposable implements IChatSessionC
}
// Determine the turn index to fork at. If a specific request is
// provided, find its position in the protocol state's turn list.
// Otherwise fork the entire session.
// provided, fork BEFORE it (keeping turns up to the previous one).
// This matches the non-contributed path in ForkConversationAction
// which uses `requestIndex - 1`. If no request is provided, fork
// the entire session.
const protocolState = this._clientState.getSessionState(backendSession.toString());
let turnIndex: number | undefined;
if (request) {
turnIndex = protocolState?.turns.findIndex(t => t.id === request.id);
if (turnIndex === undefined || turnIndex < 0) {
const requestIdx = protocolState?.turns.findIndex(t => t.id === request.id);
if (requestIdx === undefined || requestIdx < 0) {
throw new Error(`Cannot fork: turn for request ${request.id} not found in protocol state`);
}
// Fork before this request — keep turns [0..requestIdx-1]
turnIndex = requestIdx - 1;
if (turnIndex < 0) {
throw new Error('Cannot fork: cannot fork before the first request');
}
} else if (protocolState && protocolState.turns.length > 0) {
turnIndex = protocolState.turns.length - 1;
}