diff --git a/src/extension/agents/claude/AGENTS.md b/src/extension/agents/claude/AGENTS.md index f8c1f5f6fb..871160f21d 100644 --- a/src/extension/agents/claude/AGENTS.md +++ b/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//` 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/`: diff --git a/src/extension/agents/claude/node/claudeCodeAgent.ts b/src/extension/agents/claude/node/claudeCodeAgent.ts index 80e9052a84..e3fdfc302c 100644 --- a/src/extension/agents/claude/node/claudeCodeAgent.ts +++ b/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 { + 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 { + 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 { diff --git a/src/extension/agents/claude/node/test/claudeCodeAgent.spec.ts b/src/extension/agents/claude/node/test/claudeCodeAgent.spec.ts index ea9f7fb8db..235b88f19c 100644 --- a/src/extension/agents/claude/node/test/claudeCodeAgent.spec.ts +++ b/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); + }); }); diff --git a/src/extension/agents/claude/node/test/mockClaudeCodeSdkService.ts b/src/extension/agents/claude/node/test/mockClaudeCodeSdkService.ts index e75ad98269..48ed97c71c 100644 --- a/src/extension/agents/claude/node/test/mockClaudeCodeSdkService.ts +++ b/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; @@ -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 => { + 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; }