Skip to content

Commit cd6804d

Browse files
TylerLeonhardtCopilot
TylerLeonhardt
and
Copilot
authored
agentHost: implement Claude truncateSession (Phase 6.7 — Restore Checkpoint + Start Over) (#323197)
* agentHost: implement Claude truncateSession (Phase 6.7 — Restore Checkpoint + Start Over) Implements IAgent.truncateSession for the Claude agent host so the workbench "Restore Checkpoint" and "Start Over" actions work end-to-end: - Point-restore (turnId): truncates the conversation in place on the same SDK session id / protocol URI via the SDK's `resumeSessionAt` option. The protocol turn is resolved to its SDK assistant-envelope uuid and staged as a one-shot anchor the next turn's rebuild applies (lazy — full history is preserved if the user restores then walks away). - Remove-all (no turnId, "Start Over"): tears down the live subprocess, deletes the on-disk transcript, and recreates a fresh provisional under the same id, preserving the model/agent/permissionMode overlay. Teardown awaits the subprocess's actual exit (Query.return() -> the SDK's memoized cleanup -> transport.waitForExit()) before deleting + respawning the same `--session-id`, fixing a "Session ID ... is already in use" race. Also fixes a pre-existing rebind consumer-loop handoff race surfaced by truncation (post-restore turn could hang), and simplifies the pipeline's query handles (single `_query` mirroring the warm subprocess; `_needsRebind` as the sole health signal). Unit + integration green (427 Claude tests); both flows live-E2E verified against real Claude. See node/claude/phase6.7-plan.md for the implementation log and node/claude/roadmap.md (Phase 6.7) marked done. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> * agentHost: address Phase 6.7 PR review - claudeAgentSession: retain the pending truncation anchor until a rebuild succeeds (read it without clearing, clear only after materialize / rebuild installs the pipeline). A throw/cancel after reading no longer silently drops `resumeSessionAt`, so the next send retries the checkpoint restore. + regression test. - claudeSdkPipeline: `_ensureQueryBound` now honors `_needsRebind` (rebuilds via the rematerializer like `send()` does) so pre-flight helpers (reloadPlugins / snapshotResolvedCustomizations) never operate on a dead stream after an abort/crash. - claudeAgent: cold remove-all fails fast when no working directory is available (SDK cwd absent and no live session), mirroring `_resumeSession` and the fork path, instead of recreating a provisional that fails later. - claudeSdkPipeline: pass the error object to the logger in shutdownAndWait instead of stringifying it, preserving stack traces. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --------- Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
1 parent a646e32 commit cd6804d

13 files changed

Lines changed: 1242 additions & 40 deletions

‎src/vs/platform/agentHost/node/claude/claudeAgent.ts‎

Lines changed: 90 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -578,6 +578,96 @@ export class ClaudeAgent extends Disposable implements IAgent {
578578
};
579579
}
580580

581+
/**
582+
* In-place "Restore Checkpoint" truncation. Keeps turns
583+
* `[0..turnId]` INCLUSIVE (or removes all turns when `turnId` is
584+
* omitted) on the **same** session id / URI — unlike fork, which mints a
585+
* new id. The `turnId` path resolves the protocol turn to its SDK
586+
* assistant-envelope uuid ({@link resolveForkAnchorUuid}) and stages it
587+
* as a one-shot `resumeSessionAt` anchor that the next turn's rebuild
588+
* applies (the truncation finalizes when the next turn writes the
589+
* branch). Serialized on {@link _sessionSequencer} (same key as
590+
* `sendMessage`) so the `ChatTruncated` → `ChatTurnStarted` dispatch pair
591+
* stays ordered. Provisional sessions short-circuit.
592+
*/
593+
async truncateSession(session: URI, turnId?: string): Promise<void> {
594+
const sessionId = AgentSession.id(session);
595+
await this._sessionSequencer.queue(sessionId, async () => {
596+
const existing = this._findAnySession(sessionId);
597+
if (existing && !existing.isPipelineReady) {
598+
this._logService.info(`[Claude:${sessionId}] truncateSession on a provisional session — nothing to truncate`);
599+
return;
600+
}
601+
602+
if (turnId === undefined) {
603+
await this._removeAllTurns(session, sessionId, existing);
604+
return;
605+
}
606+
607+
const messages = await this._sdkService.getSessionMessages(sessionId, { includeSystemMessages: true });
608+
const anchor = resolveForkAnchorUuid(messages, turnId);
609+
if (anchor === undefined) {
610+
throw new Error(`Cannot truncate session ${sessionId}: turn ${turnId} not found in transcript`);
611+
}
612+
613+
// Operate on a live session; cold-resume an unloaded one first so
614+
// there is a single code path that sets the anchor on a live
615+
// pipeline (the next send applies it).
616+
const live = existing ?? await this._resumeSession(sessionId, session);
617+
await live.truncateToTurn(turnId, anchor);
618+
this._logService.info(`[Claude:${sessionId}] truncateSession kept [0..${turnId}] (anchor=${anchor})`);
619+
});
620+
}
621+
622+
/**
623+
* Remove-all ("start over") branch of {@link truncateSession}: there is no
624+
* anchor to resume at, so tear down the live Query, delete the on-disk
625+
* transcript via the SDK, then recreate a fresh provisional under the SAME
626+
* id/URI so the next `sendMessage` materializes non-resume `{ sessionId }`
627+
* on a clean transcript (keeps the id stable). `deleteSession` is eagerly
628+
* durable (unlike the lazy `turnId` path), matching its "clear / start
629+
* over" semantic. `existing` is the live session, or `undefined` on the
630+
* cold path (unloaded session). Caller serializes on {@link _sessionSequencer}.
631+
*/
632+
private async _removeAllTurns(session: URI, sessionId: string, existing: ClaudeAgentSession | undefined): Promise<void> {
633+
const info = existing ? undefined : await this._sdkService.getSessionInfo(sessionId);
634+
const workingDirectory = existing?.workingDirectory ?? (info?.cwd ? URI.file(info.cwd) : undefined);
635+
if (!workingDirectory) {
636+
// Mirror `_resumeSession` / fork: fail fast rather than recreate a
637+
// provisional with no cwd that would only fail later at materialize.
638+
throw new Error(`Cannot clear session ${sessionId}: workingDirectory missing (SDK cwd absent and no live session)`);
639+
}
640+
let overlay: IClaudeSessionOverlay = {};
641+
try {
642+
overlay = await this._metadataStore.read(session);
643+
} catch (err) {
644+
this._logService.warn(`[Claude:${sessionId}] overlay read failed during remove-all; continuing with defaults`, err);
645+
}
646+
647+
// `shutdownLiveQuery` awaits the subprocess's actual exit (and its final
648+
// transcript flush), so the on-disk `<id>.jsonl` is now stable and safe
649+
// to delete: no live writer can recreate it before the next turn
650+
// respawns a fresh `--session-id <id>`.
651+
await existing?.shutdownLiveQuery();
652+
this._sessions.deleteAndDispose(sessionId);
653+
await this._sdkService.deleteSession(sessionId);
654+
655+
await this.createSession({
656+
session,
657+
workingDirectory,
658+
...(overlay.model ? { model: overlay.model } : {}),
659+
...(overlay.agent ? { agent: overlay.agent } : {}),
660+
...(overlay.permissionMode ? { config: { [ClaudeSessionConfigKey.PermissionMode]: overlay.permissionMode } } : {}),
661+
});
662+
// Re-fetch (not reuse `existing`): `existing` is the OLD session, already
663+
// torn down by `deleteAndDispose` above, and is `undefined` entirely on
664+
// the cold path. `createSession` registered a fresh instance under the
665+
// same id — prune through that live session so a single path covers both
666+
// warm and cold remove-all.
667+
await this._findAnySession(sessionId)?.pruneAllTurns();
668+
this._logService.info(`[Claude:${sessionId}] truncateSession removed all turns (deleteSession + fresh same-id)`);
669+
}
670+
581671
/**
582672
* Fork an existing session at a protocol `turnId` (keep `[0..N]`
583673
* INCLUSIVE) into a new, non-provisional session. The SDK `Query` is

‎src/vs/platform/agentHost/node/claude/claudeAgentSdkService.ts‎

Lines changed: 8 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,7 @@
33
* Licensed under the MIT License. See License.txt in the project root for license information.
44
*--------------------------------------------------------------------------------------------*/
55

6-
import type { AnyZodRawShape, ForkSessionOptions, ForkSessionResult, GetSessionMessagesOptions, GetSubagentMessagesOptions, InferShape, ListSessionsOptions, ListSubagentsOptions, McpSdkServerConfigWithInstance, Options, Query, SDKSessionInfo, SDKUserMessage, SdkMcpToolDefinition, SessionMessage, WarmQuery } from '@anthropic-ai/claude-agent-sdk';
6+
import type { AnyZodRawShape, ForkSessionOptions, ForkSessionResult, GetSessionMessagesOptions, GetSubagentMessagesOptions, InferShape, ListSessionsOptions, ListSubagentsOptions, McpSdkServerConfigWithInstance, Options, Query, SDKSessionInfo, SDKUserMessage, SdkMcpToolDefinition, SessionMessage, SessionMutationOptions, WarmQuery } from '@anthropic-ai/claude-agent-sdk';
77
import type { CallToolResult } from '@modelcontextprotocol/sdk/types.js';
88
import { pathToFileURL } from 'url';
99
import { CancellationToken } from '../../../../base/common/cancellation.js';
@@ -55,6 +55,7 @@ export interface IClaudeAgentSdkService {
5555
listSubagents(sessionId: string, options?: ListSubagentsOptions): Promise<readonly string[]>;
5656
getSubagentMessages(sessionId: string, agentId: string, options?: GetSubagentMessagesOptions): Promise<readonly SessionMessage[]>;
5757
forkSession(sessionId: string, options?: ForkSessionOptions): Promise<ForkSessionResult>;
58+
deleteSession(sessionId: string, options?: SessionMutationOptions): Promise<void>;
5859
createSdkMcpServer(options: {
5960
name: string;
6061
version?: string;
@@ -92,6 +93,7 @@ export interface IClaudeSdkBindings {
9293
listSubagents(sessionId: string, options?: ListSubagentsOptions): Promise<string[]>;
9394
getSubagentMessages(sessionId: string, agentId: string, options?: GetSubagentMessagesOptions): Promise<SessionMessage[]>;
9495
forkSession(sessionId: string, options?: ForkSessionOptions): Promise<ForkSessionResult>;
96+
deleteSession(sessionId: string, options?: SessionMutationOptions): Promise<void>;
9597
createSdkMcpServer(options: {
9698
name: string;
9799
version?: string;
@@ -167,6 +169,11 @@ export class ClaudeAgentSdkService implements IClaudeAgentSdkService {
167169
return sdk.forkSession(sessionId, options);
168170
}
169171

172+
async deleteSession(sessionId: string, options?: SessionMutationOptions): Promise<void> {
173+
const sdk = await this._getSdk();
174+
return sdk.deleteSession(sessionId, options);
175+
}
176+
170177
async createSdkMcpServer(options: {
171178
name: string;
172179
version?: string;

‎src/vs/platform/agentHost/node/claude/claudeAgentSession.ts‎

Lines changed: 71 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -20,7 +20,7 @@ import { ClaudeRuntimeEffortLevel, clampEffortForRuntime, resolveClaudeEffort }
2020
import { AgentSignal, IAgentSessionProjectInfo } from '../../common/agentService.js';
2121
import type { IAgentServerToolHost } from '../../common/agentServerTools.js';
2222
import { PendingRequestRegistry } from '../../common/pendingRequestRegistry.js';
23-
import { ISessionDataService } from '../../common/sessionDataService.js';
23+
import { ISessionDatabase, ISessionDataService } from '../../common/sessionDataService.js';
2424
import { ActionType } from '../../common/state/sessionActions.js';
2525
import { PendingMessage, ChatInputAnswer, ChatInputRequest, ChatInputResponseKind, ToolCallPendingConfirmationState, type AgentSelection, type ModelSelection, type ToolDefinition } from '../../common/state/protocol/state.js';
2626
import type { Customization, ToolCallResult } from '../../common/state/sessionState.js';
@@ -295,6 +295,53 @@ export class ClaudeAgentSession extends Disposable {
295295
this._register(customizationWatcher.onDidChange(() => this._onDidCustomizationsChange.fire()));
296296
}
297297

298+
/**
299+
* One-shot SDK assistant-message uuid that the next materialize / rebuild
300+
* resumes *up to and including* (the SDK's `Options.resumeSessionAt`).
301+
* Staged by {@link truncateToTurn}; read by the next build and cleared
302+
* only once that build *succeeds* (so a thrown / cancelled rebuild keeps
303+
* the anchor staged and the next send retries the truncation rather than
304+
* silently proceeding without it and undoing the checkpoint restore).
305+
*/
306+
private _pendingResumeSessionAt: string | undefined;
307+
308+
/**
309+
* In-place truncation to `turnId` ("Restore Checkpoint"): prune the
310+
* per-turn DB rows (file edits, checkpoint refs) past the boundary AND
311+
* stage the SDK resume anchor that the next rebuild applies via
312+
* `Options.resumeSessionAt`. These two halves are one invariant — pruning
313+
* without staging the anchor would drop DB rows while the SDK still
314+
* replays the truncated turns; staging without pruning would leave stale
315+
* rows — so they live behind a single call rather than two the caller
316+
* could half-invoke. The prune runs first because it is the fallible half:
317+
* a DB failure then rejects without leaving an anchor staged for the next
318+
* turn. `turnId` is the protocol turn id (DB key); `resumeAnchorUuid` is
319+
* the SDK assistant-message uuid the agent resolved for it.
320+
*/
321+
async truncateToTurn(turnId: string, resumeAnchorUuid: string): Promise<void> {
322+
await this._withDatabase(db => db.deleteTurnsAfter(turnId));
323+
this._pendingResumeSessionAt = resumeAnchorUuid;
324+
}
325+
326+
/** Prunes all per-turn DB rows (remove-all truncation). */
327+
async pruneAllTurns(): Promise<void> {
328+
await this._withDatabase(db => db.deleteAllTurns());
329+
}
330+
331+
/**
332+
* Runs `fn` against a short-lived, ref-counted session DB handle so the
333+
* write is safe regardless of the pipeline's own dbRef lifecycle (the
334+
* ref-count keeps the shared DB alive; disposing only decrements).
335+
*/
336+
private async _withDatabase(fn: (db: ISessionDatabase) => Promise<void>): Promise<void> {
337+
const ref = this._sessionDataService.openDatabase(this.sessionUri);
338+
try {
339+
await fn(ref.object);
340+
} finally {
341+
ref.dispose();
342+
}
343+
}
344+
298345
/**
299346
* Bring the session up: build SDK `Options`, start the SDK, open the
300347
* session-scoped DB ref, construct the pipeline, and attach the
@@ -330,6 +377,7 @@ export class ClaudeAgentSession extends Disposable {
330377
permissionMode,
331378
canUseTool: ctx.canUseTool,
332379
isResume: ctx.isResume,
380+
resumeSessionAt: this._pendingResumeSessionAt,
333381
mcpServers,
334382
allowedTools,
335383
plugins: this.clientCustomizationsDiff.consume(),
@@ -369,6 +417,10 @@ export class ClaudeAgentSession extends Disposable {
369417
}
370418
this._register(pipeline.onDidProduceSignal(s => this._onDidSessionProgress.fire(this._enrichSignalWithCredits(s))));
371419
this._pipeline = pipeline;
420+
// The materialize succeeded with the staged anchor applied to `Options`
421+
// — clear it now so it isn't re-applied. A throw before this point (e.g.
422+
// `startup` / pipeline-create) leaves it staged for the next retry.
423+
this._pendingResumeSessionAt = undefined;
372424

373425
// Seed the pipeline's bijective config cache so a rebuild re-applies
374426
// the user's last-chosen model / effort without losing the picker
@@ -421,6 +473,7 @@ export class ClaudeAgentSession extends Disposable {
421473
permissionMode: liveMode,
422474
canUseTool: ctx.canUseTool,
423475
isResume: true,
476+
resumeSessionAt: this._pendingResumeSessionAt,
424477
mcpServers: rebuildMcp,
425478
allowedTools: rebuildAllowedTools,
426479
plugins: this.clientCustomizationsDiff.consume(),
@@ -432,6 +485,11 @@ export class ClaudeAgentSession extends Disposable {
432485
);
433486
this._logService.info(`[Claude] session ${this.sessionId}: resume rebuild agent=${rebuildOptions.agent ?? '(none)'}`);
434487
const rebuildWarm = await this._sdkService.startup({ options: rebuildOptions });
488+
// Rebuild succeeded with the anchor applied — clear it so it
489+
// isn't re-applied. A throw above keeps it staged (handled in the
490+
// catch alongside the tool/customization diffs) so the next send
491+
// retries the truncation instead of dropping the restore.
492+
this._pendingResumeSessionAt = undefined;
435493
return { warm: rebuildWarm, abortController: rebuildAbort };
436494
} catch (err) {
437495
this.toolDiff.markDirty();
@@ -506,6 +564,17 @@ export class ClaudeAgentSession extends Disposable {
506564

507565
get isResumed(): boolean { return this._requirePipeline().isResumed; }
508566

567+
/**
568+
* Abort the live SDK subprocess and await its full teardown so the
569+
* session id is released. No-op when the session was never materialized
570+
* (no subprocess to stop). Used by remove-all truncation before it
571+
* recreates a fresh session under the same id — the CLI keeps the id
572+
* locked until the old subprocess exits.
573+
*/
574+
async shutdownLiveQuery(): Promise<void> {
575+
await this._pipeline?.shutdownAndWait();
576+
}
577+
509578
/**
510579
* Seed the pipeline's current + applied config cache from
511580
* materialize-time `Options`. The SDK already starts with these
@@ -545,7 +614,7 @@ export class ClaudeAgentSession extends Disposable {
545614
// New turn: reset the per-turn credit accumulator so proxy reports
546615
// for this turn's `/v1/messages` calls sum from zero.
547616
this._currentTurnNanoAiu = 0;
548-
if (this.toolDiff.hasDifference || this.clientCustomizationsDiff.hasDifference) {
617+
if (this.toolDiff.hasDifference || this.clientCustomizationsDiff.hasDifference || this._pendingResumeSessionAt !== undefined) {
549618
await this._rebindForSyncedState();
550619
} else {
551620
await pipeline.setPermissionMode(resolveCurrentPermissionMode(this._configurationService, this.sessionUri, this._permissionModeFallback));

‎src/vs/platform/agentHost/node/claude/claudeSdkOptions.ts‎

Lines changed: 11 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -33,6 +33,16 @@ export interface IBuildOptionsInput {
3333
readonly permissionMode: ClaudePermissionMode;
3434
readonly canUseTool: NonNullable<Options['canUseTool']>;
3535
readonly isResume: boolean;
36+
/**
37+
* One-shot SDK assistant-message uuid to resume *up to and including*
38+
* (the SDK's `Options.resumeSessionAt`). Only meaningful with
39+
* {@link isResume}; truncates the loaded transcript to this anchor so
40+
* the next turn continues from the restored point on the same session
41+
* id. Omitted in the non-resume (`sessionId`) branch and on ordinary
42+
* resumes. Set by `truncateSession` for the rebuild that immediately
43+
* precedes the post-restore turn.
44+
*/
45+
readonly resumeSessionAt?: string;
3646
readonly mcpServers: Record<string, McpSdkServerConfigWithInstance> | undefined;
3747
/**
3848
* SDK-prefixed tool names to auto-approve without prompting (projected
@@ -122,7 +132,7 @@ export async function buildOptions(
122132
effort: resolveClaudeEffort(input.model),
123133
permissionMode: input.permissionMode,
124134
...(input.isResume
125-
? { resume: input.sessionId }
135+
? { resume: input.sessionId, ...(input.resumeSessionAt ? { resumeSessionAt: input.resumeSessionAt } : {}) }
126136
: { sessionId: input.sessionId }),
127137
...(input.mcpServers ? { mcpServers: input.mcpServers } : {}),
128138
...(input.allowedTools && input.allowedTools.length > 0 ? { allowedTools: [...input.allowedTools] } : {}),

0 commit comments

Comments
 (0)