agentHost: Fix Codex recovery for paginated threads

src/vs/platform/agentHost/node/codex/codexAgent.ts

@@ -5126,11 +5126,10 @@ export class CodexAgent extends Disposable implements IAgent {
 	 * through a host-supplied session hint or chat-URI shape. An unbound
 	 * source therefore fails fast rather than guessing its owning session.
 	 *
-	 * We `thread/fork` the source thread — which copies its full history — then
-	 * `thread/rollback` the trailing turns so the fork retains only the turns up
-	 * to and including `fork.turnId`. The forked thread already exists on the
-	 * app-server, so the runtime is registered as resumable (its first send
-	 * issues a `thread/resume`).
+	 * We `thread/fork` the source thread through the requested source turn, so
+	 * trailing history is never copied into the fork. The forked thread already
+	 * exists on the app-server, so the runtime is registered as resumable (its
+	 * first send issues a `thread/resume`).
 	 *
 	 * `adoptedSessionId`, when set, is the owning session's identity this
 	 * backing adopts (the session's runtime is stood up by this fork); otherwise
@@ -5166,10 +5165,10 @@ export class CodexAgent extends Disposable implements IAgent {
 			? distinctAbsolutePaths(inheritedWorkingDirectories.map(directory => directory.fsPath))
 			: undefined;
 
-		// Resolve how many trailing turns to drop so the fork keeps turns up to
-		// and including `fork.turnId`. A live source maps host turn ids to codex
-		// turn ids; a restored source already uses codex ids. Fall back to the
-		// caller-supplied `turnIndex` when the id can't be resolved.
+		// Resolve the last source turn the fork should include. A live source maps
+		// host turn ids to codex turn ids; a restored source already uses codex
+		// ids. Fall back to the caller-supplied `turnIndex` when the id can't be
+		// resolved.
 		const codexTurnId = sourceSession?.codexTurnIdByHostTurnId.get(fork.turnId) ?? fork.turnId;
 		// Reject an unresolvable fork boundary rather than silently keeping the
 		// full history: if neither the mapped codex turn id nor the caller's
@@ -5183,6 +5182,7 @@ export class CodexAgent extends Disposable implements IAgent {
 			throw new Error(`Cannot fork codex session ${sourceThreadId}: unable to resolve fork boundary for turn ${fork.turnId} (turnIndex=${fallbackTurnIndex}, turns=${sourceTurns.length})`);
 		}
 		const { keepThroughIndex, numTurnsToDrop } = boundary;
+		const lastTurnId = keepThroughIndex >= 0 ? sourceTurns[keepThroughIndex].id : undefined;
 
 		const inheritedModel = sourceSession?.model
 			?? (sourceRead.persistedModelId ? { id: sourceRead.persistedModelId } : undefined)
@@ -5240,6 +5240,7 @@ export class CodexAgent extends Disposable implements IAgent {
 			this._assertCurrentConnection(forkConnection);
 			forkResult = await forkConnection.client.request<'thread/fork', ThreadForkResponse>('thread/fork', {
 				threadId: sourceThreadId,
+				...(lastTurnId !== undefined ? { lastTurnId } : {}),
 				...(forkManagedWorkingDirectory ? {
 					cwd: forkManagedWorkingDirectory.fsPath,
 				} : runtimeWorkspaceRoots?.length ? {
@@ -5260,25 +5261,6 @@ export class CodexAgent extends Disposable implements IAgent {
 		}
 		const newThreadId = forkResult.thread.id;
 
-		// The fork copies the full source history; drop the trailing turns so
-		// the new thread ends at the requested fork point. A failed rollback
-		// would leave the fork carrying the very turns the user asked to branch
-		// away from, so treat it as a hard failure: archive the orphaned fork
-		// and reject rather than returning a session with the wrong history.
-		if (numTurnsToDrop > 0) {
-			try {
-				await forkConnection.client.request<'thread/rollback'>('thread/rollback', { threadId: newThreadId, numTurns: numTurnsToDrop });
-			} catch (err) {
-				const message = err instanceof Error ? err.message : String(err);
-				this._logService.warn(`[Codex:${newThreadId}] fork rollback failed (numTurns=${numTurnsToDrop}); discarding fork: ${message}`);
-				await this._archiveThreadBestEffort(newThreadId, 'fork rollback failed', forkConnection);
-				if (forkManagedWorkingDirectory) {
-					await this._removeManagedWorkingDirectory(forkManagedWorkingDirectory);
-				}
-				throw new Error(`Failed to fork codex session ${sourceThreadId}: could not roll back forked thread ${newThreadId} to the requested turn (${message})`);
-			}
-		}
-
 		// The runtime's durable id: the owning session's when this fork stands
 		// that session up (so every session-addressed call keeps resolving), and
 		// otherwise the forked thread id — the Codex convention that a
@@ -5321,7 +5303,7 @@ export class CodexAgent extends Disposable implements IAgent {
 		// Seed the host→codex turn-id map for the copied turns so a later
 		// edit/truncate of an inherited turn can resolve its app-server turn id.
 		// Without this, `truncateChat` can't map the host id and skips the
-		// rollback. `thread/fork` may regenerate turn ids, so read the forked
+		// truncation. `thread/fork` may regenerate turn ids, so read the forked
 		// thread's authoritative kept turns and pair them, in order, with the new
 		// host turn ids from `fork.turnIdMapping`. Best-effort: a failed read just
 		// leaves the map unseeded (same as before), never blocking the fork.
@@ -6349,10 +6331,10 @@ export class CodexAgent extends Disposable implements IAgent {
 	 * runtime bound to `chat` — resolved through the recorded binding, never by
 	 * re-deriving membership from its configuration scope or URI shape.
 	 *
-	 * Codex rolls back by a count of trailing turns. Resolve how many turns
-	 * follow `turnId` (or all of them when omitted) from the persisted thread,
-	 * whose turn ids match the workbench's restored turn ids (see
-	 * {@link replayThreadToTurns}). Unknown ids no-op to avoid data loss.
+	 * Resolve the first turn to remove from the persisted thread, whose turn ids
+	 * match the workbench's restored turn ids (see {@link replayThreadToTurns}).
+	 * Paginated threads revert before that turn; legacy threads roll back by the
+	 * equivalent trailing-turn count. Unknown ids no-op to avoid data loss.
 	 */
 	async truncateChat(chat: URI, turnId?: string, context?: URI | IAgentChatContext): Promise<void> {
 		const targetUri = this._resolveConversationSession(chat, context);
@@ -6371,9 +6353,9 @@ export class CodexAgent extends Disposable implements IAgent {
 		if (turns.length === 0) {
 			return;
 		}
-		let numTurns: number;
+		let firstTurnToRemove: number;
 		if (turnId === undefined) {
-			numTurns = turns.length;
+			firstTurnToRemove = 0;
 		} else {
 			// A live session's workbench turn id maps to a codex turn id; a
 			// restored session already uses codex turn ids, so fall back to the
@@ -6384,18 +6366,28 @@ export class CodexAgent extends Disposable implements IAgent {
 				this._logService.warn(`[Codex] truncateChat: turnId ${turnId} not found in thread ${read.thread.id}; skipping`);
 				return;
 			}
-			numTurns = turns.length - (index + 1);
+			firstTurnToRemove = index + 1;
 		}
-		if (numTurns <= 0) {
+		if (firstTurnToRemove >= turns.length) {
 			return;
 		}
 		try {
 			const conn = targetSession
 				? (await this._ensureThreadConnection(targetSession)).connection
 				: await this._ensureConnection();
-			await conn.client.request<'thread/rollback'>('thread/rollback', { threadId: read.thread.id, numTurns });
+			if (read.thread.historyMode === 'paginated') {
+				await conn.client.request<'thread/revert'>('thread/revert', {
+					threadId: read.thread.id,
+					beforeTurnId: turns[firstTurnToRemove].id,
+				});
+			} else {
+				await conn.client.request<'thread/rollback'>('thread/rollback', {
+					threadId: read.thread.id,
+					numTurns: turns.length - firstTurnToRemove,
+				});
+			}
 		} catch (err) {
-			this._logService.warn(`[Codex:${read.thread.id}] thread/rollback failed: ${err instanceof Error ? err.message : String(err)}`);
+			this._logService.warn(`[Codex:${read.thread.id}] thread/${read.thread.historyMode === 'paginated' ? 'revert' : 'rollback'} failed: ${err instanceof Error ? err.message : String(err)}`);
 		}
 	}
 

src/vs/platform/agentHost/node/codex/codexForkPlan.ts

@@ -22,9 +22,8 @@ export type ForkBoundaryResolution =
 /**
  * Resolve the fork boundary from the source thread's ordered turn ids.
  *
- * `thread/fork` copies the full source history; the returned `numTurnsToDrop`
- * is how many trailing turns must be rolled back so the fork ends at (and
- * includes) the requested turn.
+ * The returned boundary identifies the last turn passed to `thread/fork` and
+ * records how many source turns are omitted from the bounded fork.
  *
  * @param sourceTurnIds Codex turn ids of the source thread, in order.
  * @param codexTurnId The resolved codex turn id of the requested fork point.
@@ -54,12 +53,12 @@ export type ForkedTurnIdMapEntry = readonly [hostTurnId: string, codexTurnId: st
  * A later edit/truncate of a copied turn needs to map the workbench's (new)
  * host turn id back to the forked thread's app-server turn id. The kept turns
  * line up by index between the source and the forked thread (the fork copies
- * them in order, then trailing turns are rolled back), so for each kept turn we
+ * them in order through the requested boundary), so for each kept turn we
  * derive its new host id via `turnIdMapping` and pair it with the forked
  * thread's authoritative codex id (which `thread/fork` may have regenerated).
  *
  * @param sourceTurnIds Codex turn ids of the source thread, in order.
- * @param forkedTurnIds Codex turn ids of the forked thread (post-rollback), in order.
+ * @param forkedTurnIds Codex turn ids of the bounded fork, in order.
  * @param keepThroughIndex Index of the last kept turn (inclusive).
  * @param hostTurnIdBySourceCodexId Source session's codex→host turn id map (live sessions only).
  * @param turnIdMapping Old→new host turn id remapping supplied by the fork caller.

src/vs/platform/agentHost/test/node/codex/codexCreateChat.test.ts

@@ -67,6 +67,8 @@ interface ITestWireRequest {
 		readonly config?: Readonly<Record<string, JsonValue>>;
 		readonly includeTurns?: boolean;
 		readonly numTurns?: number;
+		readonly lastTurnId?: string;
+		readonly beforeTurnId?: string;
 		readonly input?: readonly { readonly type: string; readonly text?: string; readonly text_elements?: readonly object[] }[];
 		readonly additionalContext?: Readonly<Record<string, { readonly kind: string; readonly value: string }>>;
 		readonly dynamicTools?: readonly { readonly name: string }[];
@@ -1471,6 +1473,60 @@ suite('CodexAgent createChat', () => {
 		}
 	});
 
+	test('fork: bounds paginated history in the fork request', async () => {
+		const agent = await createAgent(disposables, { sdkResolvableWithoutDownload: true });
+		const peer = disposables.add(createTestPeer());
+		connectPeer(agent, peer);
+
+		try {
+			const sourceSession = AgentSession.uri('codex', 'paginated-fork-source');
+			const sourceChat = URI.parse(buildDefaultChatUri(sourceSession));
+			const folder = URI.file('/repo/paginated-fork');
+			await createSessionBackedChat(agent, sourceChat, { configurationResource: sourceSession, resource: sourceChat }, {
+				workingDirectories: [folder],
+				model: { id: COPILOT_TEST_MODEL },
+			});
+			const sourceStart = await readNextRequest(peer.outbound);
+			peer.push({ id: sourceStart.id, result: { thread: { id: 'paginated-source-thread', cwd: folder.fsPath } } });
+			await agent['_sessions'].get('paginated-fork-source')!.materializePromise;
+
+			const forkSession = AgentSession.uri('codex', 'paginated-fork-target');
+			const forkChat = URI.parse(buildDefaultChatUri(forkSession));
+			const forking = createSessionBackedChat(agent, forkChat, { configurationResource: forkSession, resource: forkChat }, {
+				fork: { source: sourceChat, turnId: 'source-turn-1', turnIndex: 0 },
+			});
+			const read = await readNextRequest(peer.outbound);
+			peer.push({
+				id: read.id,
+				result: { thread: { id: 'paginated-source-thread', cwd: folder.fsPath, historyMode: 'paginated', turns: [] } },
+			});
+			const turns = await readNextRequest(peer.outbound);
+			peer.push({
+				id: turns.id,
+				result: { data: [{ id: 'source-turn-1' }, { id: 'source-turn-2' }], nextCursor: null },
+			});
+
+			const fork = await readNextRequest(peer.outbound);
+			assert.deepStrictEqual({
+				method: fork.method,
+				threadId: fork.params.threadId,
+				lastTurnId: fork.params.lastTurnId,
+			}, {
+				method: 'thread/fork',
+				threadId: 'paginated-source-thread',
+				lastTurnId: 'source-turn-1',
+			});
+			peer.push({ id: fork.id, result: { thread: { id: 'paginated-forked-thread', cwd: folder.fsPath }, cwd: folder.fsPath } });
+			await forking;
+
+			const inventory = await readNextRequest(peer.outbound);
+			assert.strictEqual(inventory.method, 'mcpServerStatus/list');
+			peer.push({ id: inventory.id, result: { data: [], nextCursor: null } });
+		} finally {
+			peer.dispose();
+		}
+	});
+
 	test('fork resumes a source from a replacement app-server before reading or forking it', async () => {
 		const agent = await createAgent(disposables);
 		const peer = disposables.add(createTestPeer());
@@ -2644,6 +2700,52 @@ suite('CodexAgent exact chat routing', () => {
 		}
 	});
 
+	test('truncateChat reverts paginated history before the first removed turn', async () => {
+		const agent = await createAgent(disposables, { sdkResolvableWithoutDownload: true });
+		const peer = disposables.add(createTestPeer());
+		connectPeer(agent, peer);
+
+		try {
+			const session = AgentSession.uri('codex', 'paginated-truncate');
+			const chat = URI.parse(buildDefaultChatUri(session));
+			const folder = URI.file('/repo/paginated-truncate');
+			await createSessionBackedChat(agent, chat, { configurationResource: session, resource: chat }, {
+				workingDirectories: [folder],
+				model: { id: COPILOT_TEST_MODEL },
+			});
+			const start = await readNextRequest(peer.outbound);
+			peer.push({ id: start.id, result: { thread: { id: 'paginated-truncate-thread', cwd: folder.fsPath } } });
+			await agent['_sessions'].get('paginated-truncate')!.materializePromise;
+
+			const truncating = agent.truncateChat(chat, 'turn-2', { configurationResource: session, resource: chat });
+			const read = await readNextRequest(peer.outbound);
+			peer.push({
+				id: read.id,
+				result: { thread: { id: 'paginated-truncate-thread', cwd: folder.fsPath, historyMode: 'paginated', turns: [] } },
+			});
+			const turns = await readNextRequest(peer.outbound);
+			peer.push({
+				id: turns.id,
+				result: { data: [{ id: 'turn-1' }, { id: 'turn-2' }, { id: 'turn-3' }], nextCursor: null },
+			});
+			const revert = await readNextRequest(peer.outbound);
+			peer.push({ id: revert.id, result: {} });
+			await truncating;
+
+			assert.deepStrictEqual({
+				method: revert.method,
+				threadId: revert.params.threadId,
+				beforeTurnId: revert.params.beforeTurnId,
+			}, {
+				method: 'thread/revert',
+				threadId: 'paginated-truncate-thread',
+				beforeTurnId: 'turn-3',
+			});
+		} finally {
+			peer.dispose();
+		}
+	});
+
 	test('truncateChat resumes a replacement app-server before reading or rolling back', async () => {
 		const agent = await createAgent(disposables);
 		const peer = disposables.add(createTestPeer());