From 7dfddb09d1274e5c0b0133b08b02f1b6900aa9db Mon Sep 17 00:00:00 2001 From: Giuseppe Cianci Date: Wed, 16 Sep 2026 12:09:42 +0200 Subject: [PATCH] agentHost: fix Codex recovery for paginated threads Bound forks at creation with lastTurnId and use thread/revert for paginated truncation. This avoids deprecated rollback and prevents oversized history from surviving restore and fork operations. --- .../agentHost/node/codex/codexAgent.ts | 70 ++++++------ .../agentHost/node/codex/codexForkPlan.ts | 9 +- .../test/node/codex/codexCreateChat.test.ts | 102 ++++++++++++++++++ 3 files changed, 137 insertions(+), 44 deletions(-) diff --git a/src/vs/platform/agentHost/node/codex/codexAgent.ts b/src/vs/platform/agentHost/node/codex/codexAgent.ts index 4f7a5cba80070..8e1186b62169c 100644 --- a/src/vs/platform/agentHost/node/codex/codexAgent.ts +++ b/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 { 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)}`); } } diff --git a/src/vs/platform/agentHost/node/codex/codexForkPlan.ts b/src/vs/platform/agentHost/node/codex/codexForkPlan.ts index 58b8c9eed96aa..cfbfe15223ac4 100644 --- a/src/vs/platform/agentHost/node/codex/codexForkPlan.ts +++ b/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. diff --git a/src/vs/platform/agentHost/test/node/codex/codexCreateChat.test.ts b/src/vs/platform/agentHost/test/node/codex/codexCreateChat.test.ts index 7a56480230a5f..da99a253d28fb 100644 --- a/src/vs/platform/agentHost/test/node/codex/codexCreateChat.test.ts +++ b/src/vs/platform/agentHost/test/node/codex/codexCreateChat.test.ts @@ -67,6 +67,8 @@ interface ITestWireRequest { readonly config?: Readonly>; 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>; 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());