fixs bugs, review, make it work

Author: connor4312

src/vs/platform/agentHost/node/agentService.ts

@@ -190,7 +190,8 @@ export class AgentService extends Disposable implements IAgentService {
 			const sourceState = this._stateManager.getSessionState(config.fork.session.toString());
 			let sourceTurns: ITurn[] = [];
 			if (sourceState) {
-				sourceTurns = sourceState.turns.slice(0, config.fork.turnIndex + 1);
+				sourceTurns = sourceState.turns.slice(0, config.fork.turnIndex + 1)
+					.map(t => ({ ...t, id: generateUuid() }));
 			}
 
 			const summary: ISessionSummary = {
@@ -202,7 +203,8 @@ export class AgentService extends Disposable implements IAgentService {
 				modifiedAt: Date.now(),
 				workingDirectory: config.workingDirectory?.toString(),
 			};
-			this._stateManager.restoreSession(summary, sourceTurns);
+			const state = this._stateManager.createSession(summary);
+			state.turns = sourceTurns;
 		} else {
 			// Create empty state for new sessions
 			const summary: ISessionSummary = {

src/vs/platform/agentHost/node/agentSideEffects.ts

@@ -318,16 +318,15 @@ export class AgentSideEffects extends Disposable {
 			}
 			case ActionType.SessionTruncated: {
 				const agent = this._options.getAgent(action.session);
-				// Resolve the protocol turnId to a 0-based index using the
-				// state manager's turn list. The reducer has already applied
-				// the truncation, so we look at the pre-truncation state via
-				// the turnId position.
 				let turnIndex: number | undefined;
 				if (action.turnId !== undefined) {
 					const state = this._stateManager.getSessionState(action.session);
-					// After the reducer, the turns array is already truncated.
-					// The kept turns include the target, so its index = length - 1.
-					turnIndex = state ? state.turns.length - 1 : undefined;
+					if (state) {
+						const idx = state.turns.findIndex(t => t.id === action.turnId);
+						if (idx >= 0) {
+							turnIndex = idx;
+						}
+					}
 				}
 				agent?.truncateSession?.(URI.parse(action.session), turnIndex).catch(err => {
 					this._logService.error('[AgentSideEffects] truncateSession failed', err);

src/vs/platform/agentHost/node/copilot/copilotAgent.ts

@@ -337,10 +337,8 @@ export class CopilotAgent extends Disposable implements IAgent {
 			}
 			this._sessions.deleteAndDispose(sessionId);
 
-			if (keepUpToTurnIndex >= 0) {
-				const copilotDataDir = getCopilotDataDir();
-				await truncateCopilotSessionOnDisk(copilotDataDir, sessionId, keepUpToTurnIndex);
-			}
+			const copilotDataDir = getCopilotDataDir();
+			await truncateCopilotSessionOnDisk(copilotDataDir, sessionId, keepUpToTurnIndex);
 
 			// Resume the session from the modified on-disk data
 			await this._resumeSession(sessionId);

src/vs/platform/agentHost/node/copilot/copilotAgentForking.ts

@@ -169,7 +169,7 @@ export function buildForkedEventLog(
 		}
 
 		const newParentId = entry.parentId !== null
-			? (idMap.get(entry.parentId) ?? null)
+			? (idMap.get(entry.parentId) ?? result[result.length - 1]?.id ?? null)
 			: null;
 
 		result.push({
@@ -541,7 +541,24 @@ export async function truncateCopilotSessionOnDisk(
 	const eventsPath = path.join(sessionDir, 'events.jsonl');
 	const content = await fs.promises.readFile(eventsPath, 'utf-8');
 	const entries = parseEventLog(content);
-	const truncatedEntries = buildTruncatedEventLog(entries, keepUpToTurnIndex);
+
+	let truncatedEntries: ICopilotEventLogEntry[];
+	if (keepUpToTurnIndex < 0) {
+		// Truncate all turns: keep only a fresh session.start event
+		const originalStart = entries.find(e => e.type === 'session.start');
+		if (!originalStart) {
+			throw new Error('No session.start event found in event log');
+		}
+		truncatedEntries = [{
+			type: 'session.start',
+			data: { ...originalStart.data, startTime: new Date().toISOString() },
+			id: generateUuid(),
+			timestamp: new Date().toISOString(),
+			parentId: null,
+		}];
+	} else {
+		truncatedEntries = buildTruncatedEventLog(entries, keepUpToTurnIndex);
+	}
 
 	// Overwrite events.jsonl
 	await fs.promises.writeFile(eventsPath, serializeEventLog(truncatedEntries), 'utf-8');

src/vs/platform/agentHost/node/protocolServerHandler.ts

@@ -336,9 +336,12 @@ export class ProtocolServerHandler extends Disposable {
 			let fork: { session: URI; turnIndex: number } | undefined;
 			if (params.fork) {
 				const sourceState = this._stateManager.getSessionState(params.fork.session);
-				const turnIndex = sourceState?.turns.findIndex(t => t.id === params.fork!.turnId) ?? -1;
+				if (!sourceState) {
+					throw new ProtocolError(AHP_SESSION_NOT_FOUND, `Fork source session not found: ${params.fork.session}`);
+				}
+				const turnIndex = sourceState.turns.findIndex(t => t.id === params.fork!.turnId);
 				if (turnIndex < 0) {
-					throw new ProtocolError(AHP_PROVIDER_NOT_FOUND, `Fork turn ID ${params.fork.turnId} not found in session ${params.fork.session}`);
+					throw new ProtocolError(AHP_SESSION_NOT_FOUND, `Fork turn ID ${params.fork.turnId} not found in session ${params.fork.session}`);
 				}
 				fork = { session: URI.parse(params.fork.session), turnIndex };
 			}

src/vs/platform/agentHost/test/node/copilotAgentForking.test.ts

@@ -246,6 +246,30 @@ suite('CopilotAgentForking', () => {
 			const entries = buildTestEventLog(1);
 			assert.throws(() => buildForkedEventLog(entries, 5, 'new-session-id'));
 		});
+
+		test('falls back to last kept event when lifecycle event parent is stripped', () => {
+			const entries = buildTestEventLog(2);
+			// Insert shutdown event between turns, then make the next turn's
+			// user.message point to the shutdown event (which will be stripped)
+			const shutdownEntry = makeEntry('session.shutdown', {
+				id: 'shutdown-1',
+				parentId: entries[4].id, // turn-end-0
+			});
+			entries.splice(5, 0, shutdownEntry);
+			// Next entry (user-msg-1) now points to the shutdown event
+			entries[6] = { ...entries[6], parentId: 'shutdown-1' };
+
+			const forked = buildForkedEventLog(entries, 1, 'new-session-id');
+
+			// All parentIds should be valid
+			const idSet = new Set(forked.map(e => e.id));
+			for (let i = 1; i < forked.length; i++) {
+				assert.ok(
+					forked[i].parentId !== null && idSet.has(forked[i].parentId!),
+					`Event ${i} (${forked[i].type}) should have a valid parentId, got ${forked[i].parentId}`,
+				);
+			}
+		});
 	});
 
 	// ---- buildTruncatedEventLog -----------------------------------------
@@ -646,5 +670,36 @@ suite('CopilotAgentForking', () => {
 				await close(db);
 			}
 		});
+
+		test('removes all turns when keepUpToTurnIndex is -1', async () => {
+			const db = await openTestDb();
+			try {
+				await setupSchema(db);
+				await exec(db, `
+					INSERT INTO sessions (id, cwd, summary, created_at, updated_at)
+					VALUES ('sess', '/test', 'Test', '2026-01-01', '2026-01-01');
+				`);
+				for (let i = 0; i < 3; i++) {
+					await exec(db, `
+						INSERT INTO turns (session_id, turn_index, user_message, timestamp)
+						VALUES ('sess', ${i}, 'msg ${i}', '2026-01-01');
+					`);
+					await exec(db, `
+						INSERT INTO search_index (content, session_id, source_type, source_id)
+						VALUES ('content ${i}', 'sess', 'turn', 'sess:turn:${i}');
+					`);
+				}
+
+				await truncateSessionInDb(db, 'sess', -1);
+
+				const turns = await all(db, 'SELECT * FROM turns WHERE session_id = ?', ['sess']);
+				assert.strictEqual(turns.length, 0, 'all turns should be removed');
+
+				const searchEntries = await all(db, 'SELECT * FROM search_index WHERE session_id = ?', ['sess']);
+				assert.strictEqual(searchEntries.length, 0, 'all search entries should be removed');
+			} finally {
+				await close(db);
+			}
+		});
 	});
 });

src/vs/platform/agentHost/test/node/mockAgent.ts

@@ -418,6 +418,16 @@ export class ScriptedMockAgent implements IAgent {
 		// Mock agent doesn't track model state
 	}
 
+	async truncateSession(_session: URI, _turnIndex?: number): Promise<void> {
+		// Mock agent accepts truncation without side effects
+	}
+
+	async forkSession(_sourceSession: URI, newSessionId: string, _turnIndex: number): Promise<void> {
+		// Create the forked session so it can be resumed
+		const session = AgentSession.uri('mock', newSessionId);
+		this._sessions.set(newSessionId, session);
+	}
+
 	respondToPermissionRequest(toolCallId: string, approved: boolean): void {
 		const callback = this._pendingPermissions.get(toolCallId);
 		if (callback) {

src/vs/platform/agentHost/test/node/protocolWebSocket.integrationTest.ts

@@ -1270,4 +1270,175 @@ suite('Protocol WebSocket E2E', function () {
 
 		await client.waitForNotification(n => isActionNotification(n, 'session/turnComplete'));
 	});
+
+	// ---- Truncation tests ---------------------------------------------------
+
+	test('truncate session removes turns after specified turn', async function () {
+		this.timeout(15_000);
+
+		const sessionUri = await createAndSubscribeSession(client, 'test-truncate');
+
+		// Create two turns
+		dispatchTurnStarted(client, sessionUri, 'turn-t1', 'hello', 1);
+		await client.waitForNotification(n => isActionNotification(n, 'session/turnComplete') && (getActionEnvelope(n).action as { turnId: string }).turnId === 'turn-t1');
+
+		client.clearReceived();
+		dispatchTurnStarted(client, sessionUri, 'turn-t2', 'hello', 2);
+		await client.waitForNotification(n => isActionNotification(n, 'session/turnComplete') && (getActionEnvelope(n).action as { turnId: string }).turnId === 'turn-t2');
+
+		// Verify 2 turns exist
+		let snapshot = await client.call<ISubscribeResult>('subscribe', { resource: sessionUri });
+		let state = snapshot.snapshot.state as ISessionState;
+		assert.strictEqual(state.turns.length, 2);
+
+		client.clearReceived();
+
+		// Truncate: keep only turn-t1
+		client.notify('dispatchAction', {
+			clientSeq: 3,
+			action: { type: 'session/truncated', session: sessionUri, turnId: 'turn-t1' },
+		});
+
+		await client.waitForNotification(n => isActionNotification(n, 'session/truncated'));
+
+		snapshot = await client.call<ISubscribeResult>('subscribe', { resource: sessionUri });
+		state = snapshot.snapshot.state as ISessionState;
+		assert.strictEqual(state.turns.length, 1);
+		assert.strictEqual(state.turns[0].id, 'turn-t1');
+	});
+
+	test('truncate all turns clears session history', async function () {
+		this.timeout(15_000);
+
+		const sessionUri = await createAndSubscribeSession(client, 'test-truncate-all');
+
+		dispatchTurnStarted(client, sessionUri, 'turn-ta1', 'hello', 1);
+		await client.waitForNotification(n => isActionNotification(n, 'session/turnComplete'));
+
+		client.clearReceived();
+
+		// Truncate all (no turnId)
+		client.notify('dispatchAction', {
+			clientSeq: 2,
+			action: { type: 'session/truncated', session: sessionUri },
+		});
+
+		await client.waitForNotification(n => isActionNotification(n, 'session/truncated'));
+
+		const snapshot = await client.call<ISubscribeResult>('subscribe', { resource: sessionUri });
+		const state = snapshot.snapshot.state as ISessionState;
+		assert.strictEqual(state.turns.length, 0);
+	});
+
+	test('new turn after truncation works correctly', async function () {
+		this.timeout(15_000);
+
+		const sessionUri = await createAndSubscribeSession(client, 'test-truncate-resume');
+
+		dispatchTurnStarted(client, sessionUri, 'turn-tr1', 'hello', 1);
+		await client.waitForNotification(n => isActionNotification(n, 'session/turnComplete') && (getActionEnvelope(n).action as { turnId: string }).turnId === 'turn-tr1');
+
+		client.clearReceived();
+		dispatchTurnStarted(client, sessionUri, 'turn-tr2', 'hello', 2);
+		await client.waitForNotification(n => isActionNotification(n, 'session/turnComplete') && (getActionEnvelope(n).action as { turnId: string }).turnId === 'turn-tr2');
+
+		client.clearReceived();
+
+		// Truncate to turn-tr1
+		client.notify('dispatchAction', {
+			clientSeq: 3,
+			action: { type: 'session/truncated', session: sessionUri, turnId: 'turn-tr1' },
+		});
+
+		await client.waitForNotification(n => isActionNotification(n, 'session/truncated'));
+
+		// Send a new turn after truncation
+		dispatchTurnStarted(client, sessionUri, 'turn-tr3', 'hello', 4);
+		await client.waitForNotification(n => isActionNotification(n, 'session/turnComplete'));
+
+		const snapshot = await client.call<ISubscribeResult>('subscribe', { resource: sessionUri });
+		const state = snapshot.snapshot.state as ISessionState;
+		assert.strictEqual(state.turns.length, 2);
+		assert.strictEqual(state.turns[0].id, 'turn-tr1');
+		assert.strictEqual(state.turns[1].id, 'turn-tr3');
+	});
+
+	// ---- Fork tests ---------------------------------------------------------
+
+	test('fork creates a new session with source history', async function () {
+		this.timeout(15_000);
+
+		const sessionUri = await createAndSubscribeSession(client, 'test-fork');
+
+		// Create two turns
+		dispatchTurnStarted(client, sessionUri, 'turn-f1', 'hello', 1);
+		await client.waitForNotification(n => isActionNotification(n, 'session/turnComplete') && (getActionEnvelope(n).action as { turnId: string }).turnId === 'turn-f1');
+
+		client.clearReceived();
+		dispatchTurnStarted(client, sessionUri, 'turn-f2', 'hello', 2);
+		await client.waitForNotification(n => isActionNotification(n, 'session/turnComplete') && (getActionEnvelope(n).action as { turnId: string }).turnId === 'turn-f2');
+
+		client.clearReceived();
+
+		// Fork at turn-f1 (keep turns up to and including turn-f1)
+		const forkedSessionUri = nextSessionUri();
+		await client.call('createSession', {
+			session: forkedSessionUri,
+			provider: 'mock',
+			fork: { session: sessionUri, turnId: 'turn-f1' },
+		});
+
+		const addedNotif = await client.waitForNotification(n =>
+			n.method === 'notification' && (n.params as INotificationBroadcastParams).notification.type === 'notify/sessionAdded'
+		);
+		const addedSession = (addedNotif.params as INotificationBroadcastParams).notification as ISessionAddedNotification;
+
+		// Subscribe — forked session should have 1 turn (from the protocol state
+		// populated during createSession with fork params).
+		const snapshot = await client.call<ISubscribeResult>('subscribe', { resource: addedSession.summary.resource });
+		const state = snapshot.snapshot.state as ISessionState;
+		assert.strictEqual(state.lifecycle, 'ready');
+		assert.strictEqual(state.turns.length, 1, 'forked session should have 1 turn');
+
+		// Source session should be unaffected
+		const sourceSnapshot = await client.call<ISubscribeResult>('subscribe', { resource: sessionUri });
+		const sourceState = sourceSnapshot.snapshot.state as ISessionState;
+		assert.strictEqual(sourceState.turns.length, 2);
+	});
+
+	test('fork with invalid turn ID returns error', async function () {
+		this.timeout(10_000);
+
+		const sessionUri = await createAndSubscribeSession(client, 'test-fork-invalid');
+
+		let gotError = false;
+		try {
+			await client.call('createSession', {
+				session: nextSessionUri(),
+				provider: 'mock',
+				fork: { session: sessionUri, turnId: 'nonexistent-turn' },
+			});
+		} catch {
+			gotError = true;
+		}
+		assert.ok(gotError, 'should get error for invalid fork turn ID');
+	});
+
+	test('fork with invalid source session returns error', async function () {
+		this.timeout(10_000);
+
+		await client.call('initialize', { protocolVersion: PROTOCOL_VERSION, clientId: 'test-fork-no-source' });
+
+		let gotError = false;
+		try {
+			await client.call('createSession', {
+				session: nextSessionUri(),
+				provider: 'mock',
+				fork: { session: 'mock://nonexistent-session', turnId: 'turn-1' },
+			});
+		} catch {
+			gotError = true;
+		}
+		assert.ok(gotError, 'should get error for invalid fork source session');
+	});
 });