edits: integrate claude code edits with our edit session model (#1699)

The extension side of microsoft/vscode#274025
(lockstep w/ main API is not required)

Author: connor4312

src/extension/agents/claude/common/claudeTools.ts

@@ -3,6 +3,9 @@
  *  Licensed under the MIT License. See License.txt in the project root for license information.
  *--------------------------------------------------------------------------------------------*/
 
+import { PreToolUseHookInput } from '@anthropic-ai/claude-code';
+import { URI } from '../../../../util/vs/base/common/uri';
+
 export enum ClaudeToolNames {
 	Task = 'Task',
 	Bash = 'Bash',
@@ -39,3 +42,19 @@ export interface ITaskToolInput {
 	readonly subagent_type: string;
 	readonly prompt: string;
 }
+
+export const claudeEditTools: readonly string[] = [ClaudeToolNames.Edit, ClaudeToolNames.MultiEdit, ClaudeToolNames.Write, ClaudeToolNames.NotebookEdit];
+
+export function getAffectedUrisForEditTool(input: PreToolUseHookInput): URI[] {
+	switch (input.tool_name) {
+		case ClaudeToolNames.Edit:
+		case ClaudeToolNames.MultiEdit:
+			return [URI.file((input.tool_input as any).file_path)];
+		case ClaudeToolNames.Write:
+			return [URI.file((input.tool_input as any).file_path)];
+		case ClaudeToolNames.NotebookEdit:
+			return [URI.file((input.tool_input as any).notebook_path)];
+		default:
+			return [];
+	}
+}

src/extension/agents/claude/common/toolInvocationFormatter.ts

@@ -35,9 +35,9 @@ export function createFormattedToolInvocation(
 	} else if (toolUse.name === ClaudeToolNames.LS) {
 		formatLSInvocation(invocation, toolUse);
 	} else if (toolUse.name === ClaudeToolNames.Edit || toolUse.name === ClaudeToolNames.MultiEdit) {
-		formatEditInvocation(invocation, toolUse);
+		return; // edit diff is shown
 	} else if (toolUse.name === ClaudeToolNames.Write) {
-		formatWriteInvocation(invocation, toolUse);
+		return; // edit diff is shown
 	} else if (toolUse.name === ClaudeToolNames.ExitPlanMode) {
 		formatExitPlanModeInvocation(invocation, toolUse);
 	} else if (toolUse.name === ClaudeToolNames.Task) {
@@ -84,18 +84,6 @@ function formatLSInvocation(invocation: ChatToolInvocationPart, toolUse: Anthrop
 	invocation.invocationMessage = new MarkdownString(l10n.t("Read {0}", display));
 }
 
-function formatEditInvocation(invocation: ChatToolInvocationPart, toolUse: Anthropic.ToolUseBlock): void {
-	const filePath: string = (toolUse.input as any)?.file_path ?? '';
-	const display = filePath ? formatUriForMessage(filePath) : '';
-	invocation.invocationMessage = new MarkdownString(l10n.t("Edited {0}", display));
-}
-
-function formatWriteInvocation(invocation: ChatToolInvocationPart, toolUse: Anthropic.ToolUseBlock): void {
-	const filePath: string = (toolUse.input as any)?.file_path ?? '';
-	const display = filePath ? formatUriForMessage(filePath) : '';
-	invocation.invocationMessage = new MarkdownString(l10n.t("Wrote {0}", display));
-}
-
 function formatExitPlanModeInvocation(invocation: ChatToolInvocationPart, toolUse: Anthropic.ToolUseBlock): void {
 	invocation.invocationMessage = `Here is Claude's plan:\n\n${(toolUse.input as IExitPlanModeInput)?.plan}`;
 }

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 { Options, Query, SDKAssistantMessage, SDKResultMessage, SDKUserMessage } from '@anthropic-ai/claude-code';
+import { HookInput, HookJSONOutput, Options, PreToolUseHookInput, Query, SDKAssistantMessage, SDKResultMessage, SDKUserMessage } from '@anthropic-ai/claude-code';
 import Anthropic from '@anthropic-ai/sdk';
 import type * as vscode from 'vscode';
 import { ConfigKey, IConfigurationService } from '../../../../platform/configuration/common/configurationService';
@@ -22,7 +22,7 @@ import { ToolName } from '../../../tools/common/toolNames';
 import { IToolsService } from '../../../tools/common/toolsService';
 import { isFileOkForTool } from '../../../tools/node/toolUtils';
 import { ILanguageModelServerConfig, LanguageModelServer } from '../../node/langModelServer';
-import { ClaudeToolNames, IExitPlanModeInput, ITodoWriteInput } from '../common/claudeTools';
+import { claudeEditTools, ClaudeToolNames, getAffectedUrisForEditTool, IExitPlanModeInput, ITodoWriteInput } from '../common/claudeTools';
 import { createFormattedToolInvocation } from '../common/toolInvocationFormatter';
 import { IClaudeCodeSdkService } from './claudeCodeSdkService';
 
@@ -153,6 +153,7 @@ export class ClaudeCodeSession extends Disposable {
 	private _currentRequest: CurrentRequest | undefined;
 	private _pendingPrompt: DeferredPromise<QueuedRequest> | undefined;
 	private _abortController = new AbortController();
+	private _ongoingEdits = new Map<string | undefined, { complete: () => void; onDidComplete: Thenable<void> }>();
 
 	constructor(
 		private readonly serverConfig: ILanguageModelServerConfig,
@@ -163,7 +164,8 @@ export class ClaudeCodeSession extends Disposable {
 		@IEnvService private readonly envService: IEnvService,
 		@IInstantiationService private readonly instantiationService: IInstantiationService,
 		@IToolsService private readonly toolsService: IToolsService,
-		@IClaudeCodeSdkService private readonly claudeCodeService: IClaudeCodeSdkService
+		@IClaudeCodeSdkService private readonly claudeCodeService: IClaudeCodeSdkService,
+		@ILogService private readonly _log: ILogService,
 	) {
 		super();
 	}
@@ -195,7 +197,7 @@ export class ClaudeCodeSession extends Disposable {
 		}
 
 		if (!this._queryGenerator) {
-			await this._startSession();
+			await this._startSession(token);
 		}
 
 		// Add this request to the queue and wait for completion
@@ -232,7 +234,7 @@ export class ClaudeCodeSession extends Disposable {
 	/**
 	 * Starts a new Claude Code session with the configured options
 	 */
-	private async _startSession(): Promise<void> {
+	private async _startSession(token: vscode.CancellationToken): Promise<void> {
 		// Build options for the Claude Code SDK
 		// process.env.DEBUG = '1'; // debug messages from sdk.mjs
 		const isDebugEnabled = this.configService.getConfig(ConfigKey.Internal.ClaudeCodeDebugEnabled);
@@ -252,6 +254,20 @@ export class ClaudeCodeSession extends Disposable {
 				PATH: `${this.envService.appRoot}/node_modules/@vscode/ripgrep/bin${pathSep}${process.env.PATH}`
 			},
 			resume: this.sessionId,
+			hooks: {
+				PreToolUse: [
+					{
+						matcher: claudeEditTools.join('|'),
+						hooks: [(input, toolID) => this._onWillEditTool(input, toolID, token)]
+					}
+				],
+				PostToolUse: [
+					{
+						matcher: claudeEditTools.join('|'),
+						hooks: [(input, toolID) => this._onDidEditTool(input, toolID)]
+					}
+				],
+			},
 			canUseTool: async (name, input) => {
 				return this._currentRequest ?
 					this.canUseTool(name, input, this._currentRequest.toolInvocationToken) :
@@ -270,6 +286,47 @@ export class ClaudeCodeSession extends Disposable {
 		this._processMessages();
 	}
 
+	private async _onWillEditTool(input: HookInput, toolUseID: string | undefined, token: CancellationToken): Promise<HookJSONOutput> {
+		let uris: URI[] = [];
+		try {
+			uris = getAffectedUrisForEditTool(input as PreToolUseHookInput);
+		} catch (error) {
+			this._log.error('Error getting affected URIs for edit tool', error);
+		}
+		if (!uris.length || !this._currentRequest || token.isCancellationRequested) {
+			return {};
+		}
+
+		if (!this._currentRequest!.stream.externalEdit) {
+			return {}; // back-compat during 1.106 insiders
+		}
+
+		return new Promise(proceedWithEdit => {
+			const deferred = new DeferredPromise<void>();
+			const cancelListen = token.onCancellationRequested(() => {
+				this._ongoingEdits.delete(toolUseID);
+				deferred.complete();
+			});
+			const onDidComplete = this._currentRequest!.stream.externalEdit(uris, async () => {
+				proceedWithEdit({});
+				await deferred.p;
+				cancelListen.dispose();
+			});
+
+			this._ongoingEdits.set(toolUseID, { onDidComplete, complete: () => deferred.complete() });
+		});
+	}
+
+	private async _onDidEditTool(input: HookInput, toolUseID: string | undefined) {
+		const ongoingEdit = this._ongoingEdits.get(toolUseID);
+		if (ongoingEdit) {
+			this._ongoingEdits.delete(toolUseID);
+			ongoingEdit.complete();
+			await ongoingEdit.onDidComplete;
+		}
+		return {};
+	}
+
 	private async *_createPromptIterable(): AsyncIterable<SDKUserMessage> {
 		while (true) {
 			// Wait for a request to be available

src/extension/chatSessions/vscode-node/test/__snapshots__/claudeChatSessionContentProvider.spec.ts.snap

@@ -70,15 +70,7 @@ exports[`ChatSessionContentProvider > loads real fixture file with tool invocati
     "type": "response",
   },
   {
-    "parts": [
-      {
-        "invocationMessage": "Edited [](file:///Users/roblou/code/vscode-copilot-chat/src/extension/agents/claude/vscode-node/claudeCodeAgent.ts)",
-        "isError": undefined,
-        "toolCallId": "toolu_01NXDY5nya4UzHwUxPnhmQDX",
-        "toolName": "Edit",
-        "type": "tool",
-      },
-    ],
+    "parts": [],
     "type": "response",
   },
   {

src/extension/codeBlocks/node/codeBlockProcessor.ts

@@ -120,6 +120,7 @@ export class CodeBlockTrackingChatResponseStream implements ChatResponseStream {
 	reference2 = this.forward(this._wrapped.reference2.bind(this._wrapped));
 	codeCitation = this.forward(this._wrapped.codeCitation.bind(this._wrapped));
 	anchor = this.forward(this._wrapped.anchor.bind(this._wrapped));
+	externalEdit = this.forward(this._wrapped.externalEdit.bind(this._wrapped));
 	prepareToolInvocation = this.forward(this._wrapped.prepareToolInvocation.bind(this._wrapped));
 }
 

src/extension/linkify/common/responseStreamWithLinkification.ts

@@ -93,6 +93,10 @@ export class ResponseStreamWithLinkification implements FinalizableChatResponseS
 		return this;
 	}
 
+	externalEdit<T>(target: Uri | Uri[], callback: () => Thenable<T>): Thenable<T> {
+		return this.enqueue(() => this._progress.externalEdit(target, callback), true);
+	}
+
 	push(part: ChatResponsePart): ChatResponseStream {
 		if (part instanceof ChatResponseMarkdownPart) {
 			this.appendMarkdown(part.value);
@@ -157,14 +161,14 @@ export class ResponseStreamWithLinkification implements FinalizableChatResponseS
 
 	//#endregion
 
-	private sequencer: Promise<void> = Promise.resolve();
+	private sequencer: Promise<unknown> = Promise.resolve();
 
-	private enqueue(f: () => any | Promise<any>, flush: boolean) {
+	private enqueue<T>(f: () => T | Thenable<T>, flush: boolean) {
 		if (flush) {
 			this.sequencer = this.sequencer.then(() => this.doFinalize());
 		}
 		this.sequencer = this.sequencer.then(f);
-		return this.sequencer;
+		return this.sequencer as Promise<T>;
 	}
 
 	private async appendMarkdown(md: MarkdownString): Promise<void> {

src/extension/tools/node/test/__snapshots__/getErrorsResult.spec.tsx.snap

@@ -73,77 +73,3 @@ error 2
 </errors>
 "
 `;
-
-exports[`GetErrorsTool > diagnostics with max 1`] = `
-"Showing first 1 results out of 2
-<errors path="/test/workspace/file.ts">
-This code at line 1
-\`\`\`
-line 1
-\`\`\`
-has the problem reported:
-<compileError>
-error
-</compileError>
-
-</errors>
-"
-`;
-
-exports[`GetErrorsTool > diagnostics with more complex max 1`] = `
-"Showing first 3 results out of 4
-<errors path="/test/workspace/file.ts">
-This code at line 1
-\`\`\`
-line 1
-\`\`\`
-has the problem reported:
-<compileError>
-error
-</compileError>
-This code at line 2
-\`\`\`
-line 2
-\`\`\`
-has the problem reported:
-<compileError>
-error 2
-</compileError>
-
-</errors>
-<errors path="/test/workspace/file2.ts">
-This code at line 1
-\`\`\`
-line 1
-\`\`\`
-has the problem reported:
-<compileError>
-error
-</compileError>
-
-</errors>
-"
-`;
-
-exports[`GetErrorsTool > simple diagnostics 1`] = `
-"<errors path="/test/workspace/file.ts">
-This code at line 1
-\`\`\`
-line 1
-\`\`\`
-has the problem reported:
-<compileError>
-error
-</compileError>
-This code at line 2
-\`\`\`
-line 2
-\`\`\`
-has the problem reported:
-<compileError>
-error 2
-</compileError>
-
-</errors>
-"
-`;

src/extension/vscode.proposed.chatParticipantAdditions.d.ts

@@ -167,7 +167,14 @@ declare module 'vscode' {
 		constructor(value: ChatResponseDiffEntry[], title: string, readOnly?: boolean);
 	}
 
-	export type ExtendedChatResponsePart = ChatResponsePart | ChatResponseTextEditPart | ChatResponseNotebookEditPart | ChatResponseConfirmationPart | ChatResponseCodeCitationPart | ChatResponseReferencePart2 | ChatResponseMovePart | ChatResponseExtensionsPart | ChatResponsePullRequestPart | ChatPrepareToolInvocationPart | ChatToolInvocationPart | ChatResponseMultiDiffPart | ChatResponseThinkingProgressPart;
+	export class ChatResponseExternalEditPart {
+		uris: Uri[];
+		callback: () => Thenable<unknown>;
+		applied: Thenable<void>;
+		constructor(uris: Uri[], callback: () => Thenable<unknown>);
+	}
+
+	export type ExtendedChatResponsePart = ChatResponsePart | ChatResponseTextEditPart | ChatResponseNotebookEditPart | ChatResponseConfirmationPart | ChatResponseCodeCitationPart | ChatResponseReferencePart2 | ChatResponseMovePart | ChatResponseExtensionsPart | ChatResponsePullRequestPart | ChatPrepareToolInvocationPart | ChatToolInvocationPart | ChatResponseMultiDiffPart | ChatResponseThinkingProgressPart | ChatResponseExternalEditPart;
 	export class ChatResponseWarningPart {
 		value: MarkdownString;
 		constructor(value: string | MarkdownString);
@@ -301,6 +308,14 @@ declare module 'vscode' {
 
 		notebookEdit(target: Uri, isDone: true): void;
 
+		/**
+		 * Makes an external edit to one or more resources. Changes to the
+		 * resources made within the `callback` and before it resolves will be
+		 * tracked as agent edits. This can be used to track edits made from
+		 * external tools that don't generate simple {@link textEdit textEdits}.
+		 */
+		externalEdit<T>(target: Uri | Uri[], callback: () => Thenable<T>): Thenable<T>;
+
 		markdownWithVulnerabilities(value: string | MarkdownString, vulnerabilities: ChatVulnerability[]): void;
 		codeblockUri(uri: Uri, isEdit?: boolean): void;
 		push(part: ChatResponsePart | ChatResponseTextEditPart | ChatResponseWarningPart | ChatResponseProgressPart2): void;

src/util/common/chatResponseStreamImpl.ts

@@ -5,7 +5,7 @@
 
 import { ChatResponseReferencePartStatusKind } from '@vscode/prompt-tsx';
 import type { ChatResponseFileTree, ChatResponseStream, ChatVulnerability, Command, ExtendedChatResponsePart, Location, NotebookEdit, Progress, ThinkingDelta, Uri } from 'vscode';
-import { ChatPrepareToolInvocationPart, ChatResponseAnchorPart, ChatResponseClearToPreviousToolInvocationReason, ChatResponseCodeblockUriPart, ChatResponseCodeCitationPart, ChatResponseCommandButtonPart, ChatResponseConfirmationPart, ChatResponseFileTreePart, ChatResponseMarkdownPart, ChatResponseMarkdownWithVulnerabilitiesPart, ChatResponseNotebookEditPart, ChatResponseProgressPart, ChatResponseProgressPart2, ChatResponseReferencePart, ChatResponseReferencePart2, ChatResponseTextEditPart, ChatResponseThinkingProgressPart, ChatResponseWarningPart, MarkdownString, TextEdit } from '../../vscodeTypes';
+import { ChatPrepareToolInvocationPart, ChatResponseAnchorPart, ChatResponseClearToPreviousToolInvocationReason, ChatResponseCodeblockUriPart, ChatResponseCodeCitationPart, ChatResponseCommandButtonPart, ChatResponseConfirmationPart, ChatResponseExternalEditPart, ChatResponseFileTreePart, ChatResponseMarkdownPart, ChatResponseMarkdownWithVulnerabilitiesPart, ChatResponseNotebookEditPart, ChatResponseProgressPart, ChatResponseProgressPart2, ChatResponseReferencePart, ChatResponseReferencePart2, ChatResponseTextEditPart, ChatResponseThinkingProgressPart, ChatResponseWarningPart, MarkdownString, TextEdit } from '../../vscodeTypes';
 import type { ThemeIcon } from '../vs/base/common/themables';
 
 
@@ -99,6 +99,12 @@ export class ChatResponseStreamImpl implements FinalizableChatResponseStream {
 		this._push(new ChatResponseFileTreePart(value, baseUri));
 	}
 
+	externalEdit<T>(target: Uri | Uri[], callback: () => Thenable<T>): Thenable<T> {
+		const part = new ChatResponseExternalEditPart(target instanceof Array ? target : [target], callback);
+		this._push(part);
+		return part.applied as Thenable<T>;
+	}
+
 	progress(value: string, task?: (progress: Progress<ChatResponseWarningPart | ChatResponseReferencePart>) => Thenable<string | void>): void {
 		if (typeof task === 'undefined') {
 			this._push(new ChatResponseProgressPart(value));

src/util/common/test/shims/chatTypes.ts

@@ -57,6 +57,20 @@ export class ChatResponseThinkingProgressPart {
 	}
 }
 
+export class ChatResponseExternalEditPart {
+	applied: Thenable<void>;
+	didGetApplied!: () => void;
+
+	constructor(
+		public uris: vscode.Uri[],
+		public callback: () => Thenable<unknown>,
+	) {
+		this.applied = new Promise<void>((resolve) => {
+			this.didGetApplied = resolve;
+		});
+	}
+}
+
 export class ChatResponseProgressPart2 {
 	value: string;
 	task?: (progress: vscode.Progress<vscode.ChatResponseWarningPart>) => Thenable<string | void>;