Add rollback support for Claude agent SDK file edits

src/extension/agents/claude/AGENTS.md

@@ -67,6 +67,7 @@ All interactions are displayed through VS Code's native chat UI, providing a sea
 - Starts and manages the language model server (`LanguageModelServer`)
 - Creates and caches `ClaudeCodeSession` instances by session ID
 - Resolves prompts by replacing VS Code references (files, locations) with actual paths
+- Provides `rewindFiles` method to revert file changes to a previous message state
 
 **ClaudeCodeSession**
 - Represents a single Claude Code conversation session
@@ -79,6 +80,8 @@ All interactions are displayed through VS Code's native chat UI, providing a sea
 - Handles tool confirmation dialogs via VS Code's chat API
 - Auto-approves safe operations (file edits in workspace)
 - Tracks external edits to show proper diffs
+- Enables file checkpointing for rollback functionality
+- Provides `rewindFiles` method to restore files to their state at a specific user message
 
 ### `node/claudeCodeSdkService.ts`
 
@@ -136,6 +139,27 @@ Claude Code sessions are persisted to `~/.claude/projects/<workspace-slug>/` as
 - Resume a previous session by ID
 - Cache sessions with mtime-based invalidation
 
+## File Checkpointing and Rollback
+
+The integration enables file checkpointing to support rolling back edits:
+
+**How it works:**
+- File checkpointing is enabled automatically via the `enableFileCheckpointing: true` option when starting a Claude session
+- The Claude SDK creates backups of files before they are modified during the session
+- Each user message in the conversation can serve as a restore point
+
+**API Methods:**
+- `ClaudeCodeSession.rewindFiles(userMessageId, dryRun?)`: Rewinds tracked files to their state at a specific user message
+  - `userMessageId`: UUID of the user message to rewind to
+  - `dryRun` (optional): When `true`, previews changes without modifying files
+  - Returns: `RewindFilesResult` with information about which files were changed and statistics
+- `ClaudeAgentManager.rewindFiles(sessionId, userMessageId, dryRun?)`: Manager-level method to rewind files in a specific session
+
+**Use Cases:**
+- Revert unwanted edits made by Claude
+- Experiment with different approaches by rolling back to earlier states
+- Preview changes before committing to them using dry-run mode
+
 ## Testing
 
 Unit tests are located in `node/test/`:

src/extension/agents/claude/node/claudeCodeAgent.ts

@@ -3,7 +3,7 @@
  *  Licensed under the MIT License. See License.txt in the project root for license information.
  *--------------------------------------------------------------------------------------------*/
 
-import { HookCallbackMatcher, HookEvent, HookInput, HookJSONOutput, Options, PermissionMode, PreToolUseHookInput, Query, SDKAssistantMessage, SDKResultMessage, SDKUserMessage } from '@anthropic-ai/claude-agent-sdk';
+import { HookCallbackMatcher, HookEvent, HookInput, HookJSONOutput, Options, PermissionMode, PreToolUseHookInput, Query, RewindFilesResult, SDKAssistantMessage, SDKResultMessage, SDKUserMessage } from '@anthropic-ai/claude-agent-sdk';
 import { TodoWriteInput } from '@anthropic-ai/claude-agent-sdk/sdk-tools';
 import Anthropic from '@anthropic-ai/sdk';
 import * as l10n from '@vscode/l10n';
@@ -140,6 +140,31 @@ export class ClaudeAgentManager extends Disposable {
 
 		return prompt;
 	}
+
+	/**
+	 * Rewinds tracked files in a Claude session to their state at a specific user message.
+	 * 
+	 * This is a manager-level method that delegates to the session's rewindFiles method.
+	 * Use this when you have a session ID but not direct access to the session instance.
+	 * 
+	 * @param claudeSessionId - The session ID to rewind files in
+	 * @param userMessageId - UUID of the user message to rewind to
+	 * @param dryRun - If true, preview changes without modifying files
+	 * @returns Result containing whether rewind is possible, any errors, and file change statistics
+	 */
+	public async rewindFiles(claudeSessionId: string, userMessageId: string, dryRun?: boolean): Promise<RewindFilesResult> {
+		const session = this._sessions.get(claudeSessionId);
+		if (!session) {
+			this.logService.warn(`[ClaudeAgentManager] Cannot rewind files: session ${claudeSessionId} not found`);
+			return {
+				canRewind: false,
+				error: 'Session not found'
+			};
+		}
+
+		this.logService.trace(`[ClaudeAgentManager] Rewinding files in session ${claudeSessionId} to message ${userMessageId} (dryRun=${dryRun})`);
+		return session.rewindFiles(userMessageId, dryRun);
+	}
 }
 
 class KnownClaudeError extends Error { }
@@ -334,6 +359,7 @@ export class ClaudeCodeSession extends Disposable {
 				preset: 'claude_code'
 			},
 			settingSources: ['user', 'project', 'local'],
+			enableFileCheckpointing: true,
 			...(isDebugEnabled && {
 				stderr: data => {
 					this.logService.trace(`claude-agent-sdk stderr: ${data}`);
@@ -612,6 +638,36 @@ export class ClaudeCodeSession extends Disposable {
 		}
 	}
 
+	/**
+	 * Rewinds tracked files to their state at a specific user message.
+	 * This allows reverting edits made after a particular point in the conversation.
+	 * 
+	 * Note: Requires an active session created via invoke(). File checkpointing must be
+	 * enabled in the session options for this to work (which is done automatically).
+	 * 
+	 * @param userMessageId - UUID of the user message to rewind to
+	 * @param dryRun - If true, preview changes without modifying files
+	 * @returns Result containing whether rewind is possible, any errors, and file change statistics
+	 */
+	public async rewindFiles(userMessageId: string, dryRun?: boolean): Promise<RewindFilesResult> {
+		if (!this._queryGenerator) {
+			return {
+				canRewind: false,
+				error: 'No active session'
+			};
+		}
+
+		try {
+			return await this._queryGenerator.rewindFiles(userMessageId, { dryRun });
+		} catch (error) {
+			this.logService.error('[ClaudeCodeSession] Error rewinding files', error);
+			return {
+				canRewind: false,
+				error: error instanceof Error ? error.message : String(error)
+			};
+		}
+	}
+
 }
 
 interface IManageTodoListToolInputParams {

src/extension/agents/claude/node/test/claudeCodeAgent.spec.ts

@@ -61,6 +61,34 @@ describe('ClaudeAgentManager', () => {
 		// Verify that the service's query method was called only once (proving session reuse)
 		expect(mockService.queryCallCount).toBe(1);
 	});
+
+	it('rewinds files through the manager', async () => {
+		const manager = instantiationService.createInstance(ClaudeAgentManager);
+
+		// Create a session first
+		const stream = new MockChatResponseStream();
+		const req = new TestChatRequest('Hi');
+		const res = await manager.handleRequest(undefined, req, { history: [] } as any, stream, CancellationToken.None);
+
+		expect(res.claudeSessionId).toBe('sess-1');
+
+		// Now call rewindFiles through the manager
+		const result = await manager.rewindFiles(res.claudeSessionId!, 'test-message-id', false);
+
+		expect(mockService.rewindFilesCallCount).toBe(1);
+		expect(mockService.lastRewindMessageId).toBe('test-message-id');
+		expect(result.canRewind).toBe(true);
+	});
+
+	it('returns error when rewinding non-existent session', async () => {
+		const manager = instantiationService.createInstance(ClaudeAgentManager);
+
+		// Try to rewind a session that doesn't exist
+		const result = await manager.rewindFiles('non-existent-session', 'test-message-id', false);
+
+		expect(result.canRewind).toBe(false);
+		expect(result.error).toBe('Session not found');
+	});
 });
 
 describe('ClaudeCodeSession', () => {
@@ -199,4 +227,51 @@ describe('ClaudeCodeSession', () => {
 		await session.invoke('Hello again', {} as vscode.ChatParticipantToolToken, stream2, CancellationToken.None, 'claude-3-sonnet');
 		expect(mockService.queryCallCount).toBe(1); // Same query reused
 	});
+
+	it('rewinds files when rewindFiles is called', async () => {
+		const serverConfig = { port: 8080, nonce: 'test-nonce' };
+		const mockService = instantiationService.invokeFunction(accessor => accessor.get(IClaudeCodeSdkService)) as MockClaudeCodeSdkService;
+		const session = store.add(instantiationService.createInstance(ClaudeCodeSession, serverConfig, 'test-session', undefined, undefined));
+
+		// First invoke to create the query generator
+		const stream = new MockChatResponseStream();
+		await session.invoke('Hello', {} as vscode.ChatParticipantToolToken, stream, CancellationToken.None);
+
+		// Now call rewindFiles
+		const result = await session.rewindFiles('test-message-id', false);
+
+		expect(mockService.rewindFilesCallCount).toBe(1);
+		expect(mockService.lastRewindMessageId).toBe('test-message-id');
+		expect(mockService.lastRewindDryRun).toBe(false);
+		expect(result.canRewind).toBe(true);
+		expect(result.filesChanged).toEqual(['file1.ts', 'file2.ts']);
+	});
+
+	it('returns error when rewindFiles is called without active session', async () => {
+		const serverConfig = { port: 8080, nonce: 'test-nonce' };
+		const session = store.add(instantiationService.createInstance(ClaudeCodeSession, serverConfig, 'test-session', undefined, undefined));
+
+		// Call rewindFiles without invoking first
+		const result = await session.rewindFiles('test-message-id', false);
+
+		expect(result.canRewind).toBe(false);
+		expect(result.error).toBe('No active session');
+	});
+
+	it('supports dry run mode for rewindFiles', async () => {
+		const serverConfig = { port: 8080, nonce: 'test-nonce' };
+		const mockService = instantiationService.invokeFunction(accessor => accessor.get(IClaudeCodeSdkService)) as MockClaudeCodeSdkService;
+		const session = store.add(instantiationService.createInstance(ClaudeCodeSession, serverConfig, 'test-session', undefined, undefined));
+
+		// First invoke to create the query generator
+		const stream = new MockChatResponseStream();
+		await session.invoke('Hello', {} as vscode.ChatParticipantToolToken, stream, CancellationToken.None);
+
+		// Now call rewindFiles with dry run
+		const result = await session.rewindFiles('test-message-id', true);
+
+		expect(mockService.rewindFilesCallCount).toBe(1);
+		expect(mockService.lastRewindDryRun).toBe(true);
+		expect(result.canRewind).toBe(true);
+	});
 });

src/extension/agents/claude/node/test/mockClaudeCodeSdkService.ts

@@ -3,17 +3,20 @@
  *  Licensed under the MIT License. See License.txt in the project root for license information.
  *--------------------------------------------------------------------------------------------*/
 
-import { Options, Query, SDKAssistantMessage, SDKResultMessage, SDKUserMessage } from '@anthropic-ai/claude-agent-sdk';
+import { Options, Query, RewindFilesResult, SDKAssistantMessage, SDKResultMessage, SDKUserMessage } from '@anthropic-ai/claude-agent-sdk';
 import { IClaudeCodeSdkService } from '../claudeCodeSdkService';
 
 /**
- * Mock implementation of IClaudeCodeService for testing
+ * Mock implementation of IClaudeCodeSdkService for testing
  */
 export class MockClaudeCodeSdkService implements IClaudeCodeSdkService {
 	readonly _serviceBrand: undefined;
 	public queryCallCount = 0;
 	public setModelCallCount = 0;
 	public lastSetModel: string | undefined;
+	public rewindFilesCallCount = 0;
+	public lastRewindMessageId: string | undefined;
+	public lastRewindDryRun: boolean | undefined;
 
 	public async query(options: {
 		prompt: AsyncIterable<SDKUserMessage>;
@@ -33,6 +36,17 @@ export class MockClaudeCodeSdkService implements IClaudeCodeSdkService {
 			},
 			setPermissionMode: async (_mode: string) => { /* no-op for mock */ },
 			abort: () => { /* no-op for mock */ },
+			rewindFiles: async (userMessageId: string, options?: { dryRun?: boolean }): Promise<RewindFilesResult> => {
+				this.rewindFilesCallCount++;
+				this.lastRewindMessageId = userMessageId;
+				this.lastRewindDryRun = options?.dryRun;
+				return {
+					canRewind: true,
+					filesChanged: ['file1.ts', 'file2.ts'],
+					insertions: 10,
+					deletions: 5
+				};
+			},
 		} as unknown as Query;
 	}