From fe781b70ab4d2fd10378367815598cbc2c5968b9 Mon Sep 17 00:00:00 2001 From: Connor Peet Date: Fri, 27 Mar 2026 13:53:33 -0700 Subject: [PATCH 1/3] agentHost: actually really track (and restore) file edits Historically we've tracked edits in the ChatEditingSession, which is owned by the editor and _very_ internal and _very_ complex, for Reasons. In the agent host world, the agent host now owns edits. This is a minimal implementation of an IChatEditingSession that is used for the agent host. It does not have keep/undo (the writing has been on the wall for that for a while) which removes a large chunk of complexity. Nevertheless, it can deliver diffs, undo/redo and restore state. Diffs still happen client-side, but this could be optimized in the future. Closes #305332 --- .../platform/agentHost/common/agentService.ts | 8 +- .../common/state/protocol/commands.ts | 53 ++ .../common/state/protocol/messages.ts | 3 +- .../common/state/protocol/version/registry.ts | 107 --- .../agentHost/common/state/sessionProtocol.ts | 2 + .../electron-browser/agentHostService.ts | 5 +- .../remoteAgentHostProtocolClient.ts | 4 + .../platform/agentHost/node/agentService.ts | 6 +- .../agentHost/node/agentSideEffects.ts | 29 +- .../node/copilot/mapSessionEvents.ts | 3 +- .../agentHost/node/protocolServerHandler.ts | 6 + .../agentHost/agentHostEditingSession.ts | 688 ++++++++++++++++++ .../agentHost/agentHostSessionHandler.ts | 197 +++-- .../agentHost/loggingAgentConnection.ts | 6 +- .../agentHost/stateToProgressAdapter.ts | 6 +- .../chatEditing/chatEditingServiceImpl.ts | 25 +- .../browser/chatEditing/chatEditingSession.ts | 1 + .../chat/browser/widget/chatListRenderer.ts | 4 +- .../chat/common/editing/chatEditingService.ts | 13 + .../agentHost/agentHostEditingSession.test.ts | 678 +++++++++++++++++ .../agentHostChatContribution.test.ts | 4 + .../test/common/chatDebugServiceImpl.test.ts | 11 +- 22 files changed, 1663 insertions(+), 196 deletions(-) delete mode 100644 src/vs/platform/agentHost/common/state/protocol/version/registry.ts create mode 100644 src/vs/workbench/contrib/chat/browser/agentSessions/agentHost/agentHostEditingSession.ts create mode 100644 src/vs/workbench/contrib/chat/test/browser/agentHost/agentHostEditingSession.test.ts diff --git a/src/vs/platform/agentHost/common/agentService.ts b/src/vs/platform/agentHost/common/agentService.ts index 0c99ed3addc..f16c1c76a18 100644 --- a/src/vs/platform/agentHost/common/agentService.ts +++ b/src/vs/platform/agentHost/common/agentService.ts @@ -8,7 +8,7 @@ import { IAuthorizationProtectedResourceMetadata } from '../../../base/common/oa import { URI } from '../../../base/common/uri.js'; import { createDecorator } from '../../instantiation/common/instantiation.js'; import type { IActionEnvelope, INotification, ISessionAction } from './state/sessionActions.js'; -import type { IBrowseDirectoryResult, IFetchContentResult, IStateSnapshot } from './state/sessionProtocol.js'; +import type { IBrowseDirectoryResult, IFetchContentResult, IStateSnapshot, IWriteFileParams, IWriteFileResult } from './state/sessionProtocol.js'; import { AttachmentType, type IPendingMessage, type IToolCallResult, type PolicyState, type StringOrMarkdown } from './state/sessionState.js'; // IPC contract between the renderer and the agent host utility process. @@ -452,6 +452,12 @@ export interface IAgentService { * or reading files from the remote filesystem). */ fetchContent(uri: URI): Promise; + + /** + * Write content to a file on the agent host's filesystem. + * Used for undo/redo operations on file edits. + */ + writeFile(params: IWriteFileParams): Promise; } /** diff --git a/src/vs/platform/agentHost/common/state/protocol/commands.ts b/src/vs/platform/agentHost/common/state/protocol/commands.ts index 4c06a7cc45a..2f1ea9241f4 100644 --- a/src/vs/platform/agentHost/common/state/protocol/commands.ts +++ b/src/vs/platform/agentHost/common/state/protocol/commands.ts @@ -283,6 +283,59 @@ export interface IFetchContentResult { contentType?: string; } +// ─── writeFile ─────────────────────────────────────────────────────────────── + +/** + * Writes content to a file on the server's filesystem. + * + * Binary content (images, etc.) MUST use `base64` encoding. Text content MAY + * use `utf-8` encoding. + * + * If the file does not exist, it is created. If the file already exists, it is + * overwritten unless `createOnly` is set. + * + * @category Commands + * @method writeFile + * @direction Client → Server + * @messageType Request + * @version 1 + * @throws `NotFound` (`-32008`) if the parent directory does not exist. + * @throws `PermissionDenied` (`-32009`) if the client is not permitted to write to the path. + * @example + * ```jsonc + * // Client → Server + * { "jsonrpc": "2.0", "id": 11, "method": "writeFile", + * "params": { "uri": "file:///workspace/hello.txt", "data": "SGVsbG8=", + * "encoding": "base64", "contentType": "text/plain" } } + * + * // Server → Client + * { "jsonrpc": "2.0", "id": 11, "result": {} } + * ``` + */ +export interface IWriteFileParams { + /** Target file URI on the server filesystem */ + uri: URI; + /** Content encoded as a string */ + data: string; + /** How `data` is encoded */ + encoding: ContentEncoding; + /** Content type (e.g. `"text/plain"`, `"image/png"`) */ + contentType?: string; + /** + * If `true`, the server MUST fail if the file already exists instead of + * overwriting it. Useful for safe creation of new files. + */ + createOnly?: boolean; +} + +/** + * Result of the `writeFile` command. + * + * An empty object on success. + */ +export interface IWriteFileResult { +} + // ─── browseDirectory ──────────────────────────────────────────────────────── /** diff --git a/src/vs/platform/agentHost/common/state/protocol/messages.ts b/src/vs/platform/agentHost/common/state/protocol/messages.ts index 4e61208780d..829945a73c5 100644 --- a/src/vs/platform/agentHost/common/state/protocol/messages.ts +++ b/src/vs/platform/agentHost/common/state/protocol/messages.ts @@ -6,7 +6,7 @@ // allow-any-unicode-comment-file // DO NOT EDIT -- auto-generated by scripts/sync-agent-host-protocol.ts -import type { IInitializeParams, IInitializeResult, IReconnectParams, IReconnectResult, ISubscribeParams, ISubscribeResult, ICreateSessionParams, IDisposeSessionParams, IListSessionsParams, IListSessionsResult, IFetchContentParams, IFetchContentResult, IBrowseDirectoryParams, IBrowseDirectoryResult, IFetchTurnsParams, IFetchTurnsResult, IUnsubscribeParams, IDispatchActionParams, IAuthenticateParams, IAuthenticateResult } from './commands.js'; +import type { IInitializeParams, IInitializeResult, IReconnectParams, IReconnectResult, ISubscribeParams, ISubscribeResult, ICreateSessionParams, IDisposeSessionParams, IListSessionsParams, IListSessionsResult, IFetchContentParams, IFetchContentResult, IWriteFileParams, IWriteFileResult, IBrowseDirectoryParams, IBrowseDirectoryResult, IFetchTurnsParams, IFetchTurnsResult, IUnsubscribeParams, IDispatchActionParams, IAuthenticateParams, IAuthenticateResult } from './commands.js'; import type { IActionEnvelope } from './actions.js'; import type { IProtocolNotification } from './notifications.js'; @@ -64,6 +64,7 @@ export interface ICommandMap { 'disposeSession': { params: IDisposeSessionParams; result: null }; 'listSessions': { params: IListSessionsParams; result: IListSessionsResult }; 'fetchContent': { params: IFetchContentParams; result: IFetchContentResult }; + 'writeFile': { params: IWriteFileParams; result: IWriteFileResult }; 'browseDirectory': { params: IBrowseDirectoryParams; result: IBrowseDirectoryResult }; 'fetchTurns': { params: IFetchTurnsParams; result: IFetchTurnsResult }; 'authenticate': { params: IAuthenticateParams; result: IAuthenticateResult }; diff --git a/src/vs/platform/agentHost/common/state/protocol/version/registry.ts b/src/vs/platform/agentHost/common/state/protocol/version/registry.ts deleted file mode 100644 index c3d45ae3dfa..00000000000 --- a/src/vs/platform/agentHost/common/state/protocol/version/registry.ts +++ /dev/null @@ -1,107 +0,0 @@ -/*--------------------------------------------------------------------------------------------- - * Copyright (c) Microsoft Corporation. All rights reserved. - * Licensed under the MIT License. See License.txt in the project root for license information. - *--------------------------------------------------------------------------------------------*/ - -// allow-any-unicode-comment-file -// DO NOT EDIT -- auto-generated by scripts/sync-agent-host-protocol.ts - -import { ActionType, type IStateAction } from '../actions.js'; -import { NotificationType, type IProtocolNotification } from '../notifications.js'; - -// ─── Protocol Version Constants ────────────────────────────────────────────── - -/** The current protocol version that new code speaks. */ -export const PROTOCOL_VERSION = 1; - -/** The oldest protocol version the implementation maintains compatibility with. */ -export const MIN_PROTOCOL_VERSION = 1; - -// ─── Exhaustive Action → Version Map ───────────────────────────────────────── - -/** - * Maps every action type to the protocol version that introduced it. - * Adding a new action to `IStateAction` without adding it here is a compile error. - */ -export const ACTION_INTRODUCED_IN: { readonly [K in IStateAction['type']]: number } = { - [ActionType.RootAgentsChanged]: 1, - [ActionType.RootActiveSessionsChanged]: 1, - [ActionType.SessionReady]: 1, - [ActionType.SessionCreationFailed]: 1, - [ActionType.SessionTurnStarted]: 1, - [ActionType.SessionDelta]: 1, - [ActionType.SessionResponsePart]: 1, - [ActionType.SessionToolCallStart]: 1, - [ActionType.SessionToolCallDelta]: 1, - [ActionType.SessionToolCallReady]: 1, - [ActionType.SessionToolCallConfirmed]: 1, - [ActionType.SessionToolCallComplete]: 1, - [ActionType.SessionToolCallResultConfirmed]: 1, - [ActionType.SessionTurnComplete]: 1, - [ActionType.SessionTurnCancelled]: 1, - [ActionType.SessionError]: 1, - [ActionType.SessionTitleChanged]: 1, - [ActionType.SessionUsage]: 1, - [ActionType.SessionReasoning]: 1, - [ActionType.SessionModelChanged]: 1, - [ActionType.SessionServerToolsChanged]: 1, - [ActionType.SessionActiveClientChanged]: 1, - [ActionType.SessionActiveClientToolsChanged]: 1, - [ActionType.SessionPendingMessageSet]: 1, - [ActionType.SessionPendingMessageRemoved]: 1, - [ActionType.SessionQueuedMessagesReordered]: 1, - [ActionType.SessionCustomizationsChanged]: 1, - [ActionType.SessionCustomizationToggled]: 1, -}; - -/** - * Returns whether the given action type is known to the specified protocol version. - */ -export function isActionKnownToVersion(action: IStateAction, clientVersion: number): boolean { - return ACTION_INTRODUCED_IN[action.type] <= clientVersion; -} - -// ─── Exhaustive Notification → Version Map ───────────────────────────────── - -/** - * Maps every notification type to the protocol version that introduced it. - * Adding a new notification to `IProtocolNotification` without adding it here - * is a compile error. - */ -export const NOTIFICATION_INTRODUCED_IN: { readonly [K in IProtocolNotification['type']]: number } = { - [NotificationType.SessionAdded]: 1, - [NotificationType.SessionRemoved]: 1, - [NotificationType.AuthRequired]: 1, -}; - -/** - * Returns whether the given notification type is known to the specified protocol version. - */ -export function isNotificationKnownToVersion(notification: IProtocolNotification, clientVersion: number): boolean { - return NOTIFICATION_INTRODUCED_IN[notification.type] <= clientVersion; -} - -// ─── Capabilities ──────────────────────────────────────────────────────────── - -/** - * Feature capabilities gated by protocol version. - */ -export interface ProtocolCapabilities { - /** v1 — always present */ - readonly sessions: true; - /** v1 — always present */ - readonly tools: true; - /** v1 — always present */ - readonly permissions: true; -} - -/** - * Derives capabilities from a protocol version number. - */ -export function capabilitiesForVersion(_version: number): ProtocolCapabilities { - return { - sessions: true, - tools: true, - permissions: true, - }; -} diff --git a/src/vs/platform/agentHost/common/state/sessionProtocol.ts b/src/vs/platform/agentHost/common/state/sessionProtocol.ts index deca30524f4..ba3a8c33ecb 100644 --- a/src/vs/platform/agentHost/common/state/sessionProtocol.ts +++ b/src/vs/platform/agentHost/common/state/sessionProtocol.ts @@ -59,6 +59,8 @@ export type { IReconnectSnapshotResult, ISubscribeParams, IUnsubscribeParams, + IWriteFileParams, + IWriteFileResult, } from './protocol/commands.js'; export { ContentEncoding, ReconnectResultType } from './protocol/commands.js'; diff --git a/src/vs/platform/agentHost/electron-browser/agentHostService.ts b/src/vs/platform/agentHost/electron-browser/agentHostService.ts index b414308ebe8..3e887b8bcc8 100644 --- a/src/vs/platform/agentHost/electron-browser/agentHostService.ts +++ b/src/vs/platform/agentHost/electron-browser/agentHostService.ts @@ -15,7 +15,7 @@ import { IConfigurationService } from '../../configuration/common/configuration. import { ILogService } from '../../log/common/log.js'; import { AgentHostEnabledSettingId, AgentHostIpcChannels, IAgentCreateSessionConfig, IAgentDescriptor, IAgentHostService, IAgentService, IAgentSessionMetadata, IAuthenticateParams, IAuthenticateResult, IResourceMetadata } from '../common/agentService.js'; import type { IActionEnvelope, INotification, ISessionAction } from '../common/state/sessionActions.js'; -import type { IBrowseDirectoryResult, IFetchContentResult, IStateSnapshot } from '../common/state/sessionProtocol.js'; +import type { IBrowseDirectoryResult, IFetchContentResult, IStateSnapshot, IWriteFileParams, IWriteFileResult } from '../common/state/sessionProtocol.js'; import { revive } from '../../../base/common/marshalling.js'; import { URI } from '../../../base/common/uri.js'; @@ -122,6 +122,9 @@ class AgentHostServiceClient extends Disposable implements IAgentHostService { fetchContent(uri: URI): Promise { return this._proxy.fetchContent(uri); } + writeFile(params: IWriteFileParams): Promise { + return this._proxy.writeFile(params); + } async restartAgentHost(): Promise { // Restart is handled by the main process side } diff --git a/src/vs/platform/agentHost/electron-browser/remoteAgentHostProtocolClient.ts b/src/vs/platform/agentHost/electron-browser/remoteAgentHostProtocolClient.ts index de400d71531..4a7ebaab2cb 100644 --- a/src/vs/platform/agentHost/electron-browser/remoteAgentHostProtocolClient.ts +++ b/src/vs/platform/agentHost/electron-browser/remoteAgentHostProtocolClient.ts @@ -207,6 +207,10 @@ export class RemoteAgentHostProtocolClient extends Disposable implements IAgentC return this._sendRequest('fetchContent', { uri: uri.toString() }); } + async writeFile(params: ICommandMap['writeFile']['params']): Promise { + return this._sendRequest('writeFile', params); + } + private _handleMessage(msg: IProtocolMessage): void { if (isJsonRpcResponse(msg)) { const pending = this._pendingRequests.get(msg.id); diff --git a/src/vs/platform/agentHost/node/agentService.ts b/src/vs/platform/agentHost/node/agentService.ts index b90fa91c0c0..f40a238b58a 100644 --- a/src/vs/platform/agentHost/node/agentService.ts +++ b/src/vs/platform/agentHost/node/agentService.ts @@ -12,7 +12,7 @@ import { ILogService } from '../../log/common/log.js'; import { AgentProvider, AgentSession, IAgent, IAgentCreateSessionConfig, IAgentDescriptor, IAgentService, IAgentSessionMetadata, IAuthenticateParams, IAuthenticateResult, IResourceMetadata } from '../common/agentService.js'; import { ISessionDataService } from '../common/sessionDataService.js'; import { ActionType, IActionEnvelope, INotification, ISessionAction } from '../common/state/sessionActions.js'; -import type { IBrowseDirectoryResult, IFetchContentResult, IStateSnapshot } from '../common/state/sessionProtocol.js'; +import type { IBrowseDirectoryResult, IFetchContentResult, IStateSnapshot, IWriteFileParams, IWriteFileResult } from '../common/state/sessionProtocol.js'; import { SessionStatus, type ISessionSummary } from '../common/state/sessionState.js'; import { AgentSideEffects } from './agentSideEffects.js'; import { SessionStateManager } from './sessionStateManager.js'; @@ -217,6 +217,10 @@ export class AgentService extends Disposable implements IAgentService { return this._sideEffects.handleFetchContent(uri.toString()); } + async writeFile(params: IWriteFileParams): Promise { + return this._sideEffects.handleWriteFile(params); + } + async shutdown(): Promise { this._logService.info('AgentService: shutting down all providers...'); const promises: Promise[] = []; diff --git a/src/vs/platform/agentHost/node/agentSideEffects.ts b/src/vs/platform/agentHost/node/agentSideEffects.ts index fa90488d644..c2cff8c8f25 100644 --- a/src/vs/platform/agentHost/node/agentSideEffects.ts +++ b/src/vs/platform/agentHost/node/agentSideEffects.ts @@ -4,6 +4,7 @@ *--------------------------------------------------------------------------------------------*/ import * as os from 'os'; +import { decodeBase64, VSBuffer } from '../../../base/common/buffer.js'; import { Disposable, DisposableStore, IDisposable } from '../../../base/common/lifecycle.js'; import { autorun, IObservable } from '../../../base/common/observable.js'; import { URI } from '../../../base/common/uri.js'; @@ -13,7 +14,7 @@ import { ILogService } from '../../log/common/log.js'; import { IAgent, IAgentAttachment, IAgentMessageEvent, IAgentToolCompleteEvent, IAgentToolStartEvent, IAuthenticateParams, IAuthenticateResult, IResourceMetadata } from '../common/agentService.js'; import { ISessionDataService } from '../common/sessionDataService.js'; import { ActionType, ISessionAction } from '../common/state/sessionActions.js'; -import { AhpErrorCodes, AHP_PROVIDER_NOT_FOUND, AHP_SESSION_NOT_FOUND, ContentEncoding, IBrowseDirectoryResult, ICreateSessionParams, IDirectoryEntry, IFetchContentResult, JSON_RPC_INTERNAL_ERROR, ProtocolError } from '../common/state/sessionProtocol.js'; +import { AhpErrorCodes, AHP_PROVIDER_NOT_FOUND, AHP_SESSION_NOT_FOUND, ContentEncoding, IBrowseDirectoryResult, ICreateSessionParams, IDirectoryEntry, IFetchContentResult, IWriteFileParams, IWriteFileResult, JSON_RPC_INTERNAL_ERROR, ProtocolError } from '../common/state/sessionProtocol.js'; import { PendingMessageKind, ResponsePartKind, @@ -442,8 +443,6 @@ export class AgentSideEffects extends Disposable implements IProtocolSideEffectH pendingTools: Map; } | undefined; - let turnCounter = 0; - const finalizeTurn = (turn: NonNullable, state: TurnState): void => { turns.push({ id: turn.id, @@ -454,8 +453,8 @@ export class AgentSideEffects extends Disposable implements IProtocolSideEffectH }); }; - const startTurn = (text: string): NonNullable => ({ - id: `restored-${turnCounter++}`, + const startTurn = (id: string, text: string): NonNullable => ({ + id, userMessage: { text }, responseParts: [], pendingTools: new Map(), @@ -468,10 +467,10 @@ export class AgentSideEffects extends Disposable implements IProtocolSideEffectH if (currentTurn) { finalizeTurn(currentTurn, TurnState.Cancelled); } - currentTurn = startTurn(msg.content); + currentTurn = startTurn(msg.messageId, msg.content); } else if (msg.type === 'message' && msg.role === 'assistant') { if (!currentTurn) { - currentTurn = startTurn(''); + currentTurn = startTurn(msg.messageId, ''); } if (msg.content) { @@ -586,6 +585,22 @@ export class AgentSideEffects extends Disposable implements IProtocolSideEffectH } } + async handleWriteFile(params: IWriteFileParams): Promise { + const fileUri = typeof params.uri === 'string' ? URI.parse(params.uri) : URI.revive(params.uri); + let content: VSBuffer; + if (params.encoding === ContentEncoding.Base64) { + content = decodeBase64(params.data); + } else { + content = VSBuffer.fromString(params.data); + } + try { + await this._fileService.writeFile(fileUri, content); + return {}; + } catch (_e) { + throw new ProtocolError(AhpErrorCodes.NotFound, `Failed to write file: ${fileUri.toString()}`); + } + } + private async _fetchSessionDbContent(fields: ISessionDbUriFields): Promise { const sessionUri = URI.parse(fields.sessionUri); const ref = this._options.sessionDataService.openDatabase(sessionUri); diff --git a/src/vs/platform/agentHost/node/copilot/mapSessionEvents.ts b/src/vs/platform/agentHost/node/copilot/mapSessionEvents.ts index 6f697f02cae..f9bfbf1470d 100644 --- a/src/vs/platform/agentHost/node/copilot/mapSessionEvents.ts +++ b/src/vs/platform/agentHost/node/copilot/mapSessionEvents.ts @@ -50,6 +50,7 @@ export interface ISessionEventMessage { type: 'assistant.message' | 'user.message'; data?: { messageId?: string; + interactionId?: string; content?: string; toolRequests?: readonly { toolCallId: string; name: string; arguments?: unknown; type?: 'function' | 'custom' }[]; reasoningOpaque?: string; @@ -130,7 +131,7 @@ export async function mapSessionEvents( session, type: 'message', role: e.type === 'user.message' ? 'user' : 'assistant', - messageId: d?.messageId ?? '', + messageId: d?.messageId ?? d?.interactionId ?? '', content: d?.content ?? '', toolRequests: d?.toolRequests?.map((tr) => ({ toolCallId: tr.toolCallId, diff --git a/src/vs/platform/agentHost/node/protocolServerHandler.ts b/src/vs/platform/agentHost/node/protocolServerHandler.ts index 50442444a36..232eb8c372c 100644 --- a/src/vs/platform/agentHost/node/protocolServerHandler.ts +++ b/src/vs/platform/agentHost/node/protocolServerHandler.ts @@ -25,6 +25,8 @@ import { type IJsonRpcResponse, type IReconnectParams, type IStateSnapshot, + type IWriteFileParams, + type IWriteFileResult, } from '../common/state/sessionProtocol.js'; import { ROOT_STATE_URI, type ISessionSummary, type URI } from '../common/state/sessionState.js'; import type { IProtocolServer, IProtocolTransport } from '../common/state/sessionTransport.js'; @@ -313,6 +315,9 @@ export class ProtocolServerHandler extends Disposable { this._sideEffectHandler.handleDisposeSession(params.session); return null; }, + writeFile: async (_client, params) => { + return this._sideEffectHandler.handleWriteFile(params); + }, listSessions: async () => { const items = await this._sideEffectHandler.handleListSessions(); return { items }; @@ -452,6 +457,7 @@ export interface IProtocolSideEffectHandler { handleAuthenticate(params: IAuthenticateParams): Promise; handleBrowseDirectory(uri: URI): Promise; handleFetchContent(uri: URI): Promise; + handleWriteFile(params: IWriteFileParams): Promise; /** Returns the server's default browsing directory, if available. */ getDefaultDirectory?(): URI; /** Refresh models from all providers (VS Code extension method). */ diff --git a/src/vs/workbench/contrib/chat/browser/agentSessions/agentHost/agentHostEditingSession.ts b/src/vs/workbench/contrib/chat/browser/agentSessions/agentHost/agentHostEditingSession.ts new file mode 100644 index 00000000000..46075aef564 --- /dev/null +++ b/src/vs/workbench/contrib/chat/browser/agentSessions/agentHost/agentHostEditingSession.ts @@ -0,0 +1,688 @@ +/*--------------------------------------------------------------------------------------------- + * Copyright (c) Microsoft Corporation. All rights reserved. + * Licensed under the MIT License. See License.txt in the project root for license information. + *--------------------------------------------------------------------------------------------*/ + +import { Sequencer } from '../../../../../../base/common/async.js'; +import { VSBuffer } from '../../../../../../base/common/buffer.js'; +import { Emitter, Event } from '../../../../../../base/common/event.js'; +import { MarkdownString } from '../../../../../../base/common/htmlContent.js'; +import { Disposable, DisposableStore } from '../../../../../../base/common/lifecycle.js'; +import { Schemas } from '../../../../../../base/common/network.js'; +import { constObservable, derived, derivedOpts, IObservable, IReader, ObservablePromise, observableValue, transaction } from '../../../../../../base/common/observable.js'; +import { isEqual } from '../../../../../../base/common/resources.js'; +import { isDefined } from '../../../../../../base/common/types.js'; +import { URI } from '../../../../../../base/common/uri.js'; +import { IDocumentDiff } from '../../../../../../editor/common/diff/documentDiffProvider.js'; +import { ITextModel } from '../../../../../../editor/common/model.js'; +import { IEditorWorkerService } from '../../../../../../editor/common/services/editorWorker.js'; +import { ITextModelService } from '../../../../../../editor/common/services/resolverService.js'; +import { localize } from '../../../../../../nls.js'; +import { toAgentHostUri } from '../../../../../../platform/agentHost/common/agentHostUri.js'; +import { ContentEncoding, IWriteFileParams } from '../../../../../../platform/agentHost/common/state/sessionProtocol.js'; +import { getToolFileEdits, ToolCallStatus, type IToolCallState } from '../../../../../../platform/agentHost/common/state/sessionState.js'; +import { EditorActivation } from '../../../../../../platform/editor/common/editor.js'; +import { IFileService } from '../../../../../../platform/files/common/files.js'; +import { IInstantiationService } from '../../../../../../platform/instantiation/common/instantiation.js'; +import { ILogService } from '../../../../../../platform/log/common/log.js'; +import { IEditorPane } from '../../../../../common/editor.js'; +import { IEditorService } from '../../../../../services/editor/common/editorService.js'; +import { MultiDiffEditor } from '../../../../multiDiffEditor/browser/multiDiffEditor.js'; +import { MultiDiffEditorInput } from '../../../../multiDiffEditor/browser/multiDiffEditorInput.js'; +import { IChatProgress, IChatWorkspaceEdit } from '../../../common/chatService/chatService.js'; +import { ChatEditingSessionState, emptySessionEntryDiff, getMultiDiffSourceUri, IChatEditingSession, IEditSessionDiffStats, IEditSessionEntryDiff, IModifiedFileEntry, IModifiedFileEntryChangeHunk, IModifiedFileEntryEditorIntegration, IStreamingEdits, ModifiedFileEntryState } from '../../../common/editing/chatEditingService.js'; +import { IChatRequestDisablement, IChatResponseModel } from '../../../common/model/chatModel.js'; +import { fileEditsToExternalEdits, type IToolCallFileEdit } from './stateToProgressAdapter.js'; + +// ---- Internal data model ---------------------------------------------------- + +interface IAgentHostFileEdit { + readonly resource: URI; + readonly beforeContentUri: URI; + readonly afterContentUri: URI; + readonly undoStopId: string; + readonly diff?: { added?: number; removed?: number }; +} + +interface IAgentHostCheckpoint { + readonly requestId: string; + readonly undoStopId: string; + readonly edits: IAgentHostFileEdit[]; +} + +// ---- Modified file entry ---------------------------------------------------- + +class AgentHostModifiedFileEntry implements IModifiedFileEntry { + + readonly entryId: string; + readonly originalURI: URI; + readonly modifiedURI: URI; + readonly lastModifyingRequestId: string; + + readonly state = constObservable(ModifiedFileEntryState.Accepted); + readonly isCurrentlyBeingModifiedBy = constObservable<{ responseModel: IChatResponseModel; undoStopId: string | undefined } | undefined>(undefined); + readonly lastModifyingResponse = constObservable(undefined); + readonly rewriteRatio = constObservable(1); + readonly waitsForLastEdits = constObservable(false); + readonly reviewMode = constObservable(false); + readonly autoAcceptController = constObservable<{ total: number; remaining: number; cancel(): void } | undefined>(undefined); + readonly changesCount = constObservable(0); + readonly diffInfo?: IObservable; + readonly linesAdded?: IObservable; + readonly linesRemoved?: IObservable; + + constructor( + resource: URI, + beforeContentUri: URI, + lastModifyingRequestId: string, + added: number, + removed: number, + ) { + this.entryId = `agenthost-${resource.toString()}`; + this.modifiedURI = resource; + this.originalURI = beforeContentUri; + this.lastModifyingRequestId = lastModifyingRequestId; + if (added > 0 || removed > 0) { + this.linesAdded = constObservable(added); + this.linesRemoved = constObservable(removed); + } + } + + async accept(): Promise { /* no-op */ } + async reject(): Promise { /* no-op */ } + enableReviewModeUntilSettled(): void { /* no-op */ } + + getEditorIntegration(_editor: IEditorPane): IModifiedFileEntryEditorIntegration { + return { + currentIndex: observableValue('currentIndex', 0), + reveal(): void { /* no-op */ }, + next(): boolean { return false; }, + previous(): boolean { return false; }, + enableAccessibleDiffView(): void { /* no-op */ }, + async acceptNearestChange(_change?: IModifiedFileEntryChangeHunk): Promise { /* no-op */ }, + async rejectNearestChange(_change?: IModifiedFileEntryChangeHunk): Promise { /* no-op */ }, + async toggleDiff(_change: IModifiedFileEntryChangeHunk | undefined, _show?: boolean): Promise { /* no-op */ }, + dispose(): void { /* no-op */ }, + }; + } +} + +// ---- Editing session -------------------------------------------------------- + +export class AgentHostEditingSession extends Disposable implements IChatEditingSession { + + readonly supportsKeepUndo = true; + readonly isGlobalEditingSession = false; + + private readonly _state = observableValue(this, ChatEditingSessionState.Idle); + readonly state: IObservable = this._state; + + private readonly _entriesObs = observableValue(this, []); + readonly entries: IObservable = this._entriesObs; + + readonly requestDisablement: IObservable = constObservable([]); + + private readonly _onDidDispose = this._register(new Emitter()); + readonly onDidDispose: Event = this._onDidDispose.event; + + private readonly _onDidRequestFileWrite = this._register(new Emitter()); + readonly onDidRequestFileWrite: Event = this._onDidRequestFileWrite.event; + + private readonly _checkpoints: IAgentHostCheckpoint[] = []; + private readonly _currentCheckpointIndex = observableValue(this, -1); + private readonly _diffCache = new Map(); + private readonly _undoRedoSequencer = new Sequencer(); + + private _editorPane: MultiDiffEditor | undefined; + private _hasExplanations = false; + + readonly canUndo: IObservable = derived(this, r => this._currentCheckpointIndex.read(r) >= 0); + readonly canRedo: IObservable = derived(this, r => this._currentCheckpointIndex.read(r) < this._checkpoints.length - 1); + + constructor( + readonly chatSessionResource: URI, + private readonly _connectionAuthority: string, + @IEditorService private readonly _editorService: IEditorService, + @IInstantiationService private readonly _instantiationService: IInstantiationService, + @ILogService private readonly _logService: ILogService, + @IFileService private readonly _fileService: IFileService, + @ITextModelService private readonly _textModelService: ITextModelService, + @IEditorWorkerService private readonly _editorWorkerService: IEditorWorkerService, + ) { + super(); + } + + // ---- Hydration from protocol state -------------------------------------- + + addToolCallEdits(requestId: string, tc: IToolCallState): IChatProgress[] { + if (tc.status !== ToolCallStatus.Completed) { + return []; + } + + // Deduplicate: ignore if this tool call was already added + if (this._checkpoints.some(cp => cp.undoStopId === tc.toolCallId)) { + return []; + } + + const fileEdits = fileEditsToExternalEdits(tc); + if (fileEdits.length === 0) { + return []; + } + + const authority = this._connectionAuthority; + const protocolEdits = getToolFileEdits(tc); + + const edits: IAgentHostFileEdit[] = fileEdits.map((edit: IToolCallFileEdit, i: number) => ({ + resource: toAgentHostUri(edit.resource, authority), + beforeContentUri: toAgentHostUri(edit.beforeContentUri, authority), + afterContentUri: toAgentHostUri(edit.afterContentUri, authority), + undoStopId: edit.undoStopId, + diff: protocolEdits[i]?.diff, + })); + + const checkpoint: IAgentHostCheckpoint = { + requestId, + undoStopId: tc.toolCallId, + edits, + }; + + this._checkpoints.push(checkpoint); + + transaction(tx => { + this._currentCheckpointIndex.set(this._checkpoints.length - 1, tx); + if (this._state.get() === ChatEditingSessionState.Initial) { + this._state.set(ChatEditingSessionState.Idle, tx); + } + }); + + this._rebuildEntries(); + + // Build progress parts for the file edit pills in the chat response + const progressParts: IChatProgress[] = []; + for (const edit of edits) { + progressParts.push({ kind: 'markdownContent', content: new MarkdownString('\n````\n') }); + progressParts.push({ kind: 'codeblockUri', uri: edit.resource, isEdit: true, undoStopId: tc.toolCallId }); + progressParts.push({ kind: 'textEdit', uri: edit.resource, edits: [], done: false, isExternalEdit: true }); + progressParts.push({ kind: 'textEdit', uri: edit.resource, edits: [], done: true, isExternalEdit: true }); + progressParts.push({ kind: 'markdownContent', content: new MarkdownString('\n````\n') }); + } + return progressParts; + } + + // ---- Show diff editor --------------------------------------------------- + + async show(previousChanges?: boolean): Promise { + if (this._editorPane?.isVisible()) { + return; + } + + if (this._editorPane?.input) { + await this._editorService.openEditor(this._editorPane.input, { pinned: true, activation: EditorActivation.ACTIVATE }); + return; + } + + const input = MultiDiffEditorInput.fromResourceMultiDiffEditorInput({ + multiDiffSource: getMultiDiffSourceUri(this, previousChanges), + label: localize('multiDiffEditorInput.name', "Suggested Edits") + }, this._instantiationService); + + this._editorPane = await this._editorService.openEditor(input, { pinned: true, activation: EditorActivation.ACTIVATE }) as MultiDiffEditor | undefined; + } + + // ---- Entry lookups ------------------------------------------------------ + + getEntry(uri: URI): IModifiedFileEntry | undefined { + return this._entriesObs.get().find(e => isEqual(e.modifiedURI, uri)); + } + + readEntry(uri: URI, reader: IReader): IModifiedFileEntry | undefined { + return this._entriesObs.read(reader).find(e => isEqual(e.modifiedURI, uri)); + } + + // ---- Accept / Reject (no-op) -------------------------------------------- + + async accept(..._uris: URI[]): Promise { /* no-op */ } + async reject(..._uris: URI[]): Promise { /* no-op */ } + + // ---- Snapshots ---------------------------------------------------------- + + async restoreSnapshot(requestId: string, _stopId: string | undefined): Promise { + const idx = this._checkpoints.findIndex(cp => cp.requestId === requestId); + if (idx < 0) { + this._logService.warn(`[AgentHostEditingSession] No checkpoint found for requestId=${requestId}`); + return; + } + + // Navigate to the target checkpoint + const currentIdx = this._currentCheckpointIndex.get(); + if (idx < currentIdx) { + // Undo forward checkpoints + for (let i = currentIdx; i > idx; i--) { + await this._writeCheckpointContent(this._checkpoints[i], 'before'); + } + } else if (idx > currentIdx) { + // Redo to reach the target + for (let i = currentIdx + 1; i <= idx; i++) { + await this._writeCheckpointContent(this._checkpoints[i], 'after'); + } + } + + transaction(tx => { + this._currentCheckpointIndex.set(idx, tx); + }); + this._rebuildEntries(); + } + + getSnapshotUri(requestId: string, uri: URI, stopId: string | undefined): URI | undefined { + const cp = this._checkpoints.find(c => c.requestId === requestId); + if (!cp) { + return undefined; + } + + const uriStr = uri.toString(); + const edit = cp.edits.find(e => e.resource.toString() === uriStr); + if (!edit) { + return undefined; + } + + return URI.from({ + scheme: Schemas.chatEditingSnapshotScheme, + path: uri.path, + query: JSON.stringify({ session: this.chatSessionResource.toString(), requestId, undoStop: stopId ?? '' }), + }); + } + + async getSnapshotContents(requestId: string, uri: URI, _stopId: string | undefined): Promise { + const cp = this._checkpoints.find(c => c.requestId === requestId); + if (!cp) { + return undefined; + } + + const uriStr = uri.toString(); + const edit = cp.edits.find(e => e.resource.toString() === uriStr); + if (!edit) { + return undefined; + } + + try { + const content = await this._fileService.readFile(edit.afterContentUri); + return content.value; + } catch (err) { + this._logService.warn(`[AgentHostEditingSession] Failed to fetch snapshot content`, err); + return undefined; + } + } + + async getSnapshotModel(_requestId: string, _undoStop: string | undefined, _snapshotUri: URI): Promise { + return null; + } + + // ---- Diffs -------------------------------------------------------------- + + getEntryDiffBetweenStops(uri: URI, requestId: string | undefined, stopId: string | undefined): IObservable | undefined { + // Find the checkpoint for this stop + const startIdx = requestId !== undefined + ? this._checkpoints.findIndex(cp => cp.requestId === requestId && (stopId === undefined || cp.undoStopId === stopId)) + : -1; + if (startIdx < 0 && requestId !== undefined) { + return undefined; + } + + // fromIdx is the boundary *before* the range, toIdx is the last + // checkpoint in the range. For a single stop, the stop checkpoint + // should be the range, so fromIdx is one before it. + const fromIdx = requestId !== undefined ? startIdx - 1 : -1; + const toIdx = requestId !== undefined ? startIdx : (this._checkpoints.length > 0 ? 0 : -1); + if (toIdx < 0) { + return undefined; + } + + return this._getFileDiffObservable(uri, fromIdx, toIdx); + } + + getEntryDiffBetweenRequests(uri: URI, startRequestId: string, stopRequestId: string): IObservable { + const startIndices = this._checkpoints + .map((cp, i) => cp.requestId === startRequestId ? i : -1) + .filter(i => i >= 0); + const stopIndices = this._checkpoints + .map((cp, i) => cp.requestId === stopRequestId ? i : -1) + .filter(i => i >= 0); + + if (startIndices.length === 0 || stopIndices.length === 0) { + return constObservable(undefined); + } + + const fromIdx = startIndices[0]; + const toIdx = stopIndices[stopIndices.length - 1]; + + return this._getFileDiffObservable(uri, fromIdx - 1, toIdx); + } + + /** + * Returns an observable diff for a single file between two checkpoint + * boundary positions. `fromIdx` is the checkpoint *before* the range + * (use -1 for "from baseline"), `toIdx` is the last checkpoint in the range. + */ + private _getFileDiffObservable(uri: URI, fromIdx: number, toIdx: number): IObservable { + const uriStr = uri.toString(); + + // Determine the "before" content URI: the state of the file at the + // fromIdx boundary. If fromIdx >= 0, this is the afterContentUri of + // the last edit at or before that checkpoint. If fromIdx is -1 + // (baseline), it's the first edit's beforeContentUri. + let beforeContentUri: URI | undefined; + if (fromIdx >= 0) { + for (let i = fromIdx; i >= 0; i--) { + for (const edit of this._checkpoints[i].edits) { + if (edit.resource.toString() === uriStr) { + beforeContentUri = edit.afterContentUri; + break; + } + } + if (beforeContentUri) { + break; + } + } + } + + // Determine the "after" content URI: the state after the last edit + // in the range. Also pick up the first beforeContentUri if we didn't + // find one above (file wasn't edited before fromIdx). + let afterContentUri: URI | undefined; + for (let i = Math.max(0, fromIdx); i <= toIdx && i < this._checkpoints.length; i++) { + for (const edit of this._checkpoints[i].edits) { + if (edit.resource.toString() === uriStr) { + if (!beforeContentUri) { + beforeContentUri = edit.beforeContentUri; + } + if (i > fromIdx) { + afterContentUri = edit.afterContentUri; + } + } + } + } + + if (!beforeContentUri || !afterContentUri) { + return constObservable(undefined); + } + + return this._computeFileDiffObservable(beforeContentUri, afterContentUri, uri); + } + + /** + * Returns a cached observable that computes the diff between two content URIs. + * The result is cached by the URI pair since content is immutable. + */ + private _computeFileDiffObservable(beforeUri: URI, afterUri: URI, fileUri: URI): IObservable { + const cacheKey = `${beforeUri.toString()}\0${afterUri.toString()}`; + const cached = this._diffCache.get(cacheKey); + if (cached) { + return constObservable(cached); + } + + const promise = new ObservablePromise(this._computeFileDiff(beforeUri, afterUri, fileUri)); + + return derivedOpts({ owner: this }, reader => { + const result = promise.promiseResult.read(reader); + if (!result) { + return { ...emptySessionEntryDiff(beforeUri, afterUri), isBusy: true }; + } + if (result.data) { + this._diffCache.set(cacheKey, result.data); + return result.data; + } + return emptySessionEntryDiff(beforeUri, afterUri); + }); + } + + /** + * Fetches before/after content, creates temporary text models, computes + * the diff via {@link IEditorWorkerService}, and returns the result. + */ + private async _computeFileDiff(beforeUri: URI, afterUri: URI, fileUri: URI): Promise { + const refs = new DisposableStore(); + try { + const beforeRef = await this._textModelService.createModelReference(beforeUri); + refs.add(beforeRef); + const afterRef = await this._textModelService.createModelReference(afterUri); + refs.add(afterRef); + + const diff = await this._editorWorkerService.computeDiff( + beforeRef.object.textEditorModel.uri, + afterRef.object.textEditorModel.uri, + { ignoreTrimWhitespace: false, computeMoves: false, maxComputationTimeMs: 3000 }, + 'advanced', + ); + + const entryDiff: IEditSessionEntryDiff = { + originalURI: beforeUri, + modifiedURI: fileUri, + identical: !!diff?.identical, + isFinal: true, + quitEarly: !diff || diff.quitEarly, + added: 0, + removed: 0, + isBusy: false, + }; + + if (diff) { + for (const change of diff.changes) { + entryDiff.removed += change.original.endLineNumberExclusive - change.original.startLineNumber; + entryDiff.added += change.modified.endLineNumberExclusive - change.modified.startLineNumber; + } + } + + return entryDiff; + } catch (err) { + this._logService.warn('[AgentHostEditingSession] diff computation failed', err); + return { ...emptySessionEntryDiff(beforeUri, afterUri), isFinal: true }; + } finally { + refs.dispose(); + } + } + + getDiffsForFilesInSession(): IObservable { + return derived(this, r => { + const currentIdx = this._currentCheckpointIndex.read(r); + return this._readDiffsFromCheckpoints(-1, currentIdx); + }).map((diffs, r) => diffs.read(r)); + } + + getDiffsForFilesInRequest(requestId: string): IObservable { + return derived(this, r => { + const currentIdx = this._currentCheckpointIndex.read(r); + const filteredCheckpoints: number[] = []; + for (let i = 0; i <= currentIdx && i < this._checkpoints.length; i++) { + if (this._checkpoints[i].requestId === requestId) { + filteredCheckpoints.push(i); + } + } + if (filteredCheckpoints.length === 0) { + return undefined; + } + return this._readDiffsFromCheckpoints(filteredCheckpoints[0] - 1, filteredCheckpoints[filteredCheckpoints.length - 1]); + }).map((diffs, r) => diffs?.read(r) || []); + } + + hasEditsInRequest(requestId: string, _reader?: IReader): boolean { + return this._checkpoints.some(cp => cp.requestId === requestId); + } + + getDiffForSession(): IObservable { + const sessionDiffs = this.getDiffsForFilesInSession(); + return derived(this, r => { + const diffs = sessionDiffs.read(r); + let added = 0; + let removed = 0; + for (const diff of diffs) { + added += diff.added; + removed += diff.removed; + } + return { added, removed }; + }); + } + + // ---- Undo / Redo -------------------------------------------------------- + + async undoInteraction(): Promise { + return this._undoRedoSequencer.queue(() => this._undoInteractionImpl()); + } + + async redoInteraction(): Promise { + return this._undoRedoSequencer.queue(() => this._redoInteractionImpl()); + } + + private async _undoInteractionImpl(): Promise { + const idx = this._currentCheckpointIndex.get(); + if (idx < 0) { + return; + } + + await this._writeCheckpointContent(this._checkpoints[idx], 'before'); + + transaction(tx => { + this._currentCheckpointIndex.set(idx - 1, tx); + }); + this._rebuildEntries(); + } + + private async _redoInteractionImpl(): Promise { + const idx = this._currentCheckpointIndex.get(); + if (idx >= this._checkpoints.length - 1) { + return; + } + + const nextIdx = idx + 1; + await this._writeCheckpointContent(this._checkpoints[nextIdx], 'after'); + + transaction(tx => { + this._currentCheckpointIndex.set(nextIdx, tx); + }); + this._rebuildEntries(); + } + + // ---- Explanations (stubs) ----------------------------------------------- + + async triggerExplanationGeneration(): Promise { + this._hasExplanations = true; + } + + clearExplanations(): void { + this._hasExplanations = false; + } + + hasExplanations(): boolean { + return this._hasExplanations; + } + + // ---- Unsupported operations (agent host owns edits server-side) ---------- + + startStreamingEdits(_resource: URI, _responseModel: IChatResponseModel, _inUndoStop: string | undefined): IStreamingEdits { + throw new Error('Not supported for agent host sessions'); + } + + applyWorkspaceEdit(_edit: IChatWorkspaceEdit, _responseModel: IChatResponseModel, _undoStopId: string): void { + throw new Error('Not supported for agent host sessions'); + } + + async startExternalEdits(_responseModel: IChatResponseModel, _operationId: number, _resources: URI[], _undoStopId: string, _contentFor?: URI[]): Promise { + throw new Error('Not supported for agent host sessions'); + } + + async stopExternalEdits(_responseModel: IChatResponseModel, _operationId: number, _contentFor?: URI[]): Promise { + throw new Error('Not supported for agent host sessions'); + } + + // ---- Stop / Dispose ----------------------------------------------------- + + async stop(_clearState?: boolean): Promise { + this.dispose(); + } + + override dispose(): void { + this._state.set(ChatEditingSessionState.Disposed, undefined); + this._onDidDispose.fire(); + this._diffCache.clear(); + super.dispose(); + } + + // ---- Private helpers ---------------------------------------------------- + + private _rebuildEntries(): void { + const currentIdx = this._currentCheckpointIndex.get(); + const resourceMap = new Map(); + + for (let i = 0; i <= currentIdx && i < this._checkpoints.length; i++) { + const cp = this._checkpoints[i]; + for (const edit of cp.edits) { + const key = edit.resource.toString(); + const existing = resourceMap.get(key); + if (existing) { + // Update after-content to the latest, accumulate diff counts + existing.afterContentUri = edit.afterContentUri; + existing.requestId = cp.requestId; + existing.added += edit.diff?.added ?? 0; + existing.removed += edit.diff?.removed ?? 0; + } else { + resourceMap.set(key, { + resource: edit.resource, + beforeContentUri: edit.beforeContentUri, + afterContentUri: edit.afterContentUri, + requestId: cp.requestId, + added: edit.diff?.added ?? 0, + removed: edit.diff?.removed ?? 0, + }); + } + } + } + + const entries = [...resourceMap.values()].map(v => + new AgentHostModifiedFileEntry(v.resource, v.beforeContentUri, v.requestId, v.added, v.removed) + ); + + this._entriesObs.set(entries, undefined); + } + + private async _writeCheckpointContent(checkpoint: IAgentHostCheckpoint, direction: 'before' | 'after'): Promise { + const writes = checkpoint.edits.map(async edit => { + const contentUri = direction === 'before' ? edit.beforeContentUri : edit.afterContentUri; + try { + const file = await this._fileService.readFile(contentUri); + this._onDidRequestFileWrite.fire({ + uri: edit.resource.toString(), + data: file.value.toString(), + encoding: ContentEncoding.Utf8, + }); + } catch (err) { + this._logService.warn(`[AgentHostEditingSession] Failed to fetch content for ${direction}`, contentUri.toString(), err); + } + }); + await Promise.all(writes); + } + + /** + * Collects unique file URIs from checkpoints in the given range and + * computes diffs for each via {@link _getFileDiffObservable}. + */ + private _readDiffsFromCheckpoints( + fromIdx: number, + toIdx: number, + ): IObservable { + // Collect unique resource URIs from checkpoints in the range + const seen = new Set(); + const uris: URI[] = []; + for (let i = Math.max(0, fromIdx + 1); i <= toIdx && i < this._checkpoints.length; i++) { + for (const edit of this._checkpoints[i].edits) { + const key = edit.resource.toString(); + if (!seen.has(key)) { + seen.add(key); + uris.push(edit.resource); + } + } + } + + const observables = uris.map(uri => this._getFileDiffObservable(uri, fromIdx, toIdx)); + + return derived(reader => observables.flatMap(o => o.read(reader)).filter(isDefined)); + } +} diff --git a/src/vs/workbench/contrib/chat/browser/agentSessions/agentHost/agentHostSessionHandler.ts b/src/vs/workbench/contrib/chat/browser/agentSessions/agentHost/agentHostSessionHandler.ts index 89bb620781b..8cb4c25df5f 100644 --- a/src/vs/workbench/contrib/chat/browser/agentSessions/agentHost/agentHostSessionHandler.ts +++ b/src/vs/workbench/contrib/chat/browser/agentSessions/agentHost/agentHostSessionHandler.ts @@ -5,32 +5,35 @@ import { Throttler } from '../../../../../../base/common/async.js'; import { CancellationToken, CancellationTokenSource } from '../../../../../../base/common/cancellation.js'; -import { Emitter } from '../../../../../../base/common/event.js'; +import { Emitter, Event } from '../../../../../../base/common/event.js'; import { MarkdownString } from '../../../../../../base/common/htmlContent.js'; -import { Disposable, DisposableMap, DisposableStore, MutableDisposable, type IDisposable, toDisposable } from '../../../../../../base/common/lifecycle.js'; +import { Disposable, DisposableMap, DisposableResourceMap, DisposableStore, MutableDisposable, toDisposable, type IDisposable } from '../../../../../../base/common/lifecycle.js'; +import { ResourceMap } from '../../../../../../base/common/map.js'; import { observableValue } from '../../../../../../base/common/observable.js'; +import { isEqual } from '../../../../../../base/common/resources.js'; import { URI } from '../../../../../../base/common/uri.js'; import { generateUuid } from '../../../../../../base/common/uuid.js'; import { localize } from '../../../../../../nls.js'; -import { toAgentHostUri } from '../../../../../../platform/agentHost/common/agentHostUri.js'; import { AgentProvider, AgentSession, IAgentAttachment, type IAgentConnection } from '../../../../../../platform/agentHost/common/agentService.js'; import { ActionType, isSessionAction, type ISessionAction } from '../../../../../../platform/agentHost/common/state/sessionActions.js'; import { SessionClientState } from '../../../../../../platform/agentHost/common/state/sessionClientState.js'; import { AHP_AUTH_REQUIRED, ProtocolError } from '../../../../../../platform/agentHost/common/state/sessionProtocol.js'; import { getToolKind, getToolLanguage } from '../../../../../../platform/agentHost/common/state/sessionReducers.js'; -import { AttachmentType, PendingMessageKind, ResponsePartKind, ToolCallCancellationReason, ToolCallConfirmationReason, ToolCallStatus, TurnState, type IMessageAttachment, type ISessionState } from '../../../../../../platform/agentHost/common/state/sessionState.js'; +import { AttachmentType, getToolFileEdits, PendingMessageKind, ResponsePartKind, ToolCallCancellationReason, ToolCallConfirmationReason, ToolCallStatus, TurnState, type IMessageAttachment, type ISessionState, type IToolCallState, type ITurn } from '../../../../../../platform/agentHost/common/state/sessionState.js'; import { ExtensionIdentifier } from '../../../../../../platform/extensions/common/extensions.js'; import { IInstantiationService } from '../../../../../../platform/instantiation/common/instantiation.js'; import { ILogService } from '../../../../../../platform/log/common/log.js'; import { IProductService } from '../../../../../../platform/product/common/productService.js'; import { IWorkspaceContextService } from '../../../../../../platform/workspace/common/workspace.js'; -import { IChatProgress, IChatService, IChatToolInvocation, ToolConfirmKind, ChatRequestQueueKind } from '../../../common/chatService/chatService.js'; +import { ChatRequestQueueKind, IChatProgress, IChatService, IChatToolInvocation, ToolConfirmKind } from '../../../common/chatService/chatService.js'; import { IChatSession, IChatSessionContentProvider, IChatSessionHistoryItem } from '../../../common/chatSessionsService.js'; import { ChatAgentLocation, ChatModeKind } from '../../../common/constants.js'; +import { IChatEditingService } from '../../../common/editing/chatEditingService.js'; import { ChatToolInvocation } from '../../../common/model/chatProgressTypes/chatToolInvocation.js'; import { IChatAgentData, IChatAgentImplementation, IChatAgentRequest, IChatAgentResult, IChatAgentService } from '../../../common/participants/chatAgents.js'; import { getAgentHostIcon } from '../agentSessions.js'; -import { activeTurnToProgress, finalizeToolInvocation, toolCallStateToInvocation, turnsToHistory, type IToolCallFileEdit } from './stateToProgressAdapter.js'; +import { AgentHostEditingSession } from './agentHostEditingSession.js'; +import { activeTurnToProgress, finalizeToolInvocation, toolCallStateToInvocation, turnsToHistory } from './stateToProgressAdapter.js'; // ============================================================================= // AgentHostSessionHandler - renderer-side handler for a single agent host @@ -159,13 +162,17 @@ export interface IAgentHostSessionHandlerConfig { export class AgentHostSessionHandler extends Disposable implements IChatSessionContentProvider { - private readonly _activeSessions = new Map(); + private readonly _activeSessions = new ResourceMap(); /** Maps UI resource keys to resolved backend session URIs. */ - private readonly _sessionToBackend = new Map(); + private readonly _sessionToBackend = new ResourceMap(); /** Per-session subscription to chat model pending request changes. */ - private readonly _pendingMessageSubscriptions = this._register(new DisposableMap()); + private readonly _pendingMessageSubscriptions = this._register(new DisposableResourceMap()); /** Per-session subscription watching for server-initiated turns. */ - private readonly _serverTurnWatchers = this._register(new DisposableMap()); + private readonly _serverTurnWatchers = this._register(new DisposableMap()); + /** Per-session writeFile listeners for agent host editing sessions. */ + private readonly _editingSessionListeners = this._register(new DisposableResourceMap()); + /** Historical turns with file edits, pending hydration into the editing session. */ + private readonly _pendingHistoryTurns = new ResourceMap(); /** Turn IDs dispatched by this client, used to distinguish server-originated turns. */ private readonly _clientDispatchedTurnIds = new Set(); private readonly _config: IAgentHostSessionHandlerConfig; @@ -177,6 +184,7 @@ export class AgentHostSessionHandler extends Disposable implements IChatSessionC config: IAgentHostSessionHandlerConfig, @IChatAgentService private readonly _chatAgentService: IChatAgentService, @IChatService private readonly _chatService: IChatService, + @IChatEditingService private readonly _chatEditingService: IChatEditingService, @ILogService private readonly _logService: ILogService, @IProductService private readonly _productService: IProductService, @IWorkspaceContextService private readonly _workspaceContextService: IWorkspaceContextService, @@ -188,6 +196,20 @@ export class AgentHostSessionHandler extends Disposable implements IChatSessionC // Create shared client state manager for this handler instance this._clientState = this._register(new SessionClientState(config.connection.clientId, this._logService)); + // Register an editing session provider for this handler's session type + this._register(this._chatEditingService.registerEditingSessionProvider( + config.sessionType, + { + createEditingSession: (chatSessionResource: URI) => { + return this._instantiationService.createInstance( + AgentHostEditingSession, + chatSessionResource, + config.connectionAuthority, + ); + }, + }, + )); + // Forward action envelopes from IPC to client state this._register(config.connection.onDidAction(envelope => { if (isSessionAction(envelope.action)) { @@ -199,19 +221,18 @@ export class AgentHostSessionHandler extends Disposable implements IChatSessionC } async provideChatSessionContent(sessionResource: URI, _token: CancellationToken): Promise { - const resourceKey = sessionResource.path.substring(1); // For untitled (new) sessions, defer backend session creation until the // first request arrives so the user-selected model is available. // For existing sessions we resolve immediately to load history. let resolvedSession: URI | undefined; - const isUntitled = resourceKey.startsWith('untitled-'); + const isUntitled = sessionResource.path.substring(1).startsWith('untitled-'); const history: IChatSessionHistoryItem[] = []; let initialProgress: IChatProgress[] | undefined; let activeTurnId: string | undefined; if (!isUntitled) { resolvedSession = this._resolveSessionUri(sessionResource); - this._sessionToBackend.set(resourceKey, resolvedSession); + this._sessionToBackend.set(sessionResource, resolvedSession); try { const snapshot = await this._config.connection.subscribe(resolvedSession); if (snapshot?.state) { @@ -220,6 +241,16 @@ export class AgentHostSessionHandler extends Disposable implements IChatSessionC if (sessionState) { history.push(...turnsToHistory(sessionState.turns, this._config.agentId)); + // Store turns with file edits so the editing session + // can be hydrated when it's created lazily. + const hasTurnsWithEdits = sessionState.turns.some(t => + t.responseParts.some(rp => rp.kind === ResponsePartKind.ToolCall + && rp.toolCall.status === ToolCallStatus.Completed + && getToolFileEdits(rp.toolCall).length > 0)); + if (hasTurnsWithEdits) { + this._pendingHistoryTurns.set(sessionResource, sessionState.turns); + } + // If there's an active turn, include its request in history // with an empty response so the chat service creates a // pending request, then provide accumulated progress via @@ -253,19 +284,21 @@ export class AgentHostSessionHandler extends Disposable implements IChatSessionC const backendSession = resolvedSession ?? await this._createAndSubscribe(sessionResource, request.userSelectedModelId); if (!resolvedSession) { resolvedSession = backendSession; - this._sessionToBackend.set(resourceKey, backendSession); + this._sessionToBackend.set(sessionResource, backendSession); } // For existing sessions, set up pending message sync on the first turn // (after the ChatModel becomes available in the ChatService). - this._ensurePendingMessageSubscription(resourceKey, sessionResource, backendSession); + this._ensurePendingMessageSubscription(sessionResource, backendSession); return this._handleTurn(backendSession, request, progress, token); }, initialProgress, () => { - this._activeSessions.delete(resourceKey); - this._sessionToBackend.delete(resourceKey); - this._pendingMessageSubscriptions.deleteAndDispose(resourceKey); - this._serverTurnWatchers.deleteAndDispose(resourceKey); + this._activeSessions.delete(sessionResource); + this._sessionToBackend.delete(sessionResource); + this._pendingMessageSubscriptions.deleteAndDispose(sessionResource); + this._serverTurnWatchers.deleteAndDispose(sessionResource); + this._editingSessionListeners.deleteAndDispose(sessionResource); + this._pendingHistoryTurns.delete(sessionResource); if (resolvedSession) { this._clientState.unsubscribe(resolvedSession.toString()); this._config.connection.unsubscribe(resolvedSession); @@ -273,9 +306,20 @@ export class AgentHostSessionHandler extends Disposable implements IChatSessionC } }, ); - this._activeSessions.set(resourceKey, session); + this._activeSessions.set(sessionResource, session); if (resolvedSession) { + // If there are historical turns with file edits, eagerly create + // the editing session once the ChatModel is available so that + // edit pills render with diff info on session restore. + if (this._pendingHistoryTurns.has(sessionResource)) { + session.registerDisposable(Event.once(this._chatService.onDidCreateModel)(model => { + if (isEqual(model.sessionResource, sessionResource)) { + this._ensureEditingSession(sessionResource); + } + })); + } + // If reconnecting to an active turn, wire up an ongoing state listener // to stream new progress into the session's progressObs. if (activeTurnId && initialProgress !== undefined) { @@ -329,16 +373,15 @@ export class AgentHostSessionHandler extends Disposable implements IChatSessionC this._logService.info(`[AgentHost] _invokeAgent called for resource: ${request.sessionResource.toString()}`); // Resolve or create backend session - const resourceKey = request.sessionResource.path.substring(1); - let resolvedSession = this._sessionToBackend.get(resourceKey); + let resolvedSession = this._sessionToBackend.get(request.sessionResource); if (!resolvedSession) { resolvedSession = await this._createAndSubscribe(request.sessionResource, request.userSelectedModelId); - this._sessionToBackend.set(resourceKey, resolvedSession); + this._sessionToBackend.set(request.sessionResource, resolvedSession); } await this._handleTurn(resolvedSession, request, progress, cancellationToken); - const activeSession = this._activeSessions.get(resourceKey); + const activeSession = this._activeSessions.get(request.sessionResource); if (activeSession) { activeSession.isCompleteObs.set(true, undefined); } @@ -490,7 +533,7 @@ export class AgentHostSessionHandler extends Disposable implements IChatSessionC return; } - const chatSession = this._activeSessions.get(resourceKey); + const chatSession = this._activeSessions.get(sessionResource); if (!chatSession) { previousQueuedIds = currentQueuedIds; return; @@ -811,7 +854,10 @@ export class AgentHostSessionHandler extends Disposable implements IChatSessionC if (existing && (tc.status === ToolCallStatus.Completed || tc.status === ToolCallStatus.Cancelled) && !IChatToolInvocation.isComplete(existing)) { const fileEdits = finalizeToolInvocation(existing, tc); if (fileEdits.length > 0) { - await this._applyFileEdits(request.sessionResource, request, fileEdits, progress); + const editParts = this._hydrateFileEdits(request.sessionResource, request.requestId, tc); + if (editParts.length > 0) { + progress(editParts); + } } } break; @@ -1068,47 +1114,68 @@ export class AgentHostSessionHandler extends Disposable implements IChatSessionC // ---- File edit routing --------------------------------------------------- /** - * Routes file edits from completed tool calls through the editing session's - * external edits pipeline. Calls start/stop in sequence since the edit has - * already happened on the remote by the time we receive the tool completion. + * Ensures the chat model has an editing session and returns it if it's an + * {@link AgentHostEditingSession}. The editing session is created via the + * provider registered in the constructor if one doesn't exist yet. */ - private async _applyFileEdits( - sessionResource: URI, - request: IChatAgentRequest, - fileEdits: IToolCallFileEdit[], - progress: (parts: IChatProgress[]) => void, - ): Promise { - const chatSession = this._chatService.getSession(sessionResource); - const editingSession = chatSession?.editingSession; - const response = chatSession?.getRequests().find(req => req.id === request.requestId)?.response; - if (!editingSession || !response) { - return; + private _ensureEditingSession(sessionResource: URI): AgentHostEditingSession | undefined { + const chatModel = this._chatService.getSession(sessionResource); + if (!chatModel) { + return undefined; } - const authority = this._config.connectionAuthority; - const wrapUri = (uri: URI) => toAgentHostUri(uri, authority); - - for (const edit of fileEdits) { - const operationId = this._nextOperationId++; - const resource = wrapUri(edit.resource); - const beforeUri = wrapUri(edit.beforeContentUri); - const afterUri = wrapUri(edit.afterContentUri); - - const startProgress = await editingSession.startExternalEdits( - response, operationId, [resource], edit.undoStopId, - [beforeUri], - ); - progress(startProgress); - - const stopProgress = await editingSession.stopExternalEdits( - response, operationId, - [afterUri], - ); - progress(stopProgress); + // Start the editing session if not already started — this will use + // our registered provider to create an AgentHostEditingSession. + if (!chatModel.editingSession) { + chatModel.startEditingSession(); } + + const editingSession = chatModel.editingSession; + if (!(editingSession instanceof AgentHostEditingSession)) { + return undefined; + } + + // Wire up the writeFile listener if not already done + if (!this._editingSessionListeners.has(sessionResource)) { + this._editingSessionListeners.set(sessionResource, editingSession.onDidRequestFileWrite(params => { + this._config.connection.writeFile(params).catch(err => { + this._logService.warn('[AgentHost] writeFile failed for undo/redo', err); + }); + })); + + // Hydrate from historical turns if this is the first time + // the editing session is accessed for this chat session. + const pendingTurns = this._pendingHistoryTurns.get(sessionResource); + if (pendingTurns) { + this._pendingHistoryTurns.delete(sessionResource); + for (const turn of pendingTurns) { + for (const rp of turn.responseParts) { + if (rp.kind === ResponsePartKind.ToolCall) { + editingSession.addToolCallEdits(turn.id, rp.toolCall); + } + } + } + } + } + + return editingSession; } - private _nextOperationId = 0; + /** + * Hydrates the editing session with file edits from a completed tool call + * and returns progress parts for the file edit pills. + */ + private _hydrateFileEdits( + sessionResource: URI, + requestId: string, + tc: IToolCallState, + ): IChatProgress[] { + const editingSession = this._ensureEditingSession(sessionResource); + if (editingSession) { + return editingSession.addToolCallEdits(requestId, tc); + } + return []; + } // ---- Session resolution ------------------------------------------------- @@ -1164,7 +1231,7 @@ export class AgentHostSessionHandler extends Disposable implements IChatSessionC } // Start syncing the chat model's pending requests to the protocol - this._ensurePendingMessageSubscription(resourceKey, sessionResource, session); + this._ensurePendingMessageSubscription(sessionResource, session); // Start watching for server-initiated turns on this session this._watchForServerInitiatedTurns(session, sessionResource); @@ -1176,13 +1243,13 @@ export class AgentHostSessionHandler extends Disposable implements IChatSessionC * Ensures that the chat model's pending request changes are synced to the * protocol for a given session. No-ops if already subscribed. */ - private _ensurePendingMessageSubscription(resourceKey: string, sessionResource: URI, backendSession: URI): void { - if (this._pendingMessageSubscriptions.has(resourceKey)) { + private _ensurePendingMessageSubscription(sessionResource: URI, backendSession: URI): void { + if (this._pendingMessageSubscriptions.has(sessionResource)) { return; } const chatModel = this._chatService?.getSession(sessionResource); if (chatModel) { - this._pendingMessageSubscriptions.set(resourceKey, chatModel.onDidChangePendingRequests(() => { + this._pendingMessageSubscriptions.set(sessionResource, chatModel.onDidChangePendingRequests(() => { this._syncPendingMessages(sessionResource, backendSession); })); } diff --git a/src/vs/workbench/contrib/chat/browser/agentSessions/agentHost/loggingAgentConnection.ts b/src/vs/workbench/contrib/chat/browser/agentSessions/agentHost/loggingAgentConnection.ts index 01266552539..da92b77f45f 100644 --- a/src/vs/workbench/contrib/chat/browser/agentSessions/agentHost/loggingAgentConnection.ts +++ b/src/vs/workbench/contrib/chat/browser/agentSessions/agentHost/loggingAgentConnection.ts @@ -9,7 +9,7 @@ import { URI, UriComponents } from '../../../../../../base/common/uri.js'; import { Registry } from '../../../../../../platform/registry/common/platform.js'; import { IAgentConnection, IAgentCreateSessionConfig, IAgentDescriptor, IAgentSessionMetadata, IAuthenticateParams, IAuthenticateResult, IResourceMetadata, AgentHostIpcLoggingSettingId } from '../../../../../../platform/agentHost/common/agentService.js'; import type { IActionEnvelope, INotification, ISessionAction } from '../../../../../../platform/agentHost/common/state/sessionActions.js'; -import type { IBrowseDirectoryResult, IFetchContentResult, IStateSnapshot } from '../../../../../../platform/agentHost/common/state/sessionProtocol.js'; +import type { IBrowseDirectoryResult, IFetchContentResult, IStateSnapshot, IWriteFileParams, IWriteFileResult } from '../../../../../../platform/agentHost/common/state/sessionProtocol.js'; import { Extensions, IOutputChannel, IOutputChannelRegistry, IOutputService } from '../../../../../services/output/common/output.js'; import { IConfigurationService } from '../../../../../../platform/configuration/common/configuration.js'; @@ -155,6 +155,10 @@ export class LoggingAgentConnection extends Disposable implements IAgentConnecti return this._logCall('fetchContent', uri, () => this._inner.fetchContent(uri)); } + async writeFile(params: IWriteFileParams): Promise { + return this._logCall('writeFile', params, () => this._inner.writeFile(params)); + } + // ---- Public logging API for callers' catch blocks ----------------------- /** diff --git a/src/vs/workbench/contrib/chat/browser/agentSessions/agentHost/stateToProgressAdapter.ts b/src/vs/workbench/contrib/chat/browser/agentSessions/agentHost/stateToProgressAdapter.ts index b7bf76c94fc..69a2c1d7109 100644 --- a/src/vs/workbench/contrib/chat/browser/agentSessions/agentHost/stateToProgressAdapter.ts +++ b/src/vs/workbench/contrib/chat/browser/agentSessions/agentHost/stateToProgressAdapter.ts @@ -36,7 +36,7 @@ export function turnsToHistory(turns: readonly ITurn[], participantId: string): const history: IChatSessionHistoryItem[] = []; for (const turn of turns) { // Request - history.push({ type: 'request', prompt: turn.userMessage.text, participant: participantId }); + history.push({ id: turn.id, type: 'request', prompt: turn.userMessage.text, participant: participantId }); // Response parts — iterate the unified responseParts array const parts: IChatProgress[] = []; @@ -339,7 +339,7 @@ export function finalizeToolInvocation(invocation: ChatToolInvocation, tc: ITool * converts them to {@link IToolCallFileEdit} data for routing through * the editing session's external edits pipeline. */ -function fileEditsToExternalEdits(tc: IToolCallState): IToolCallFileEdit[] { +export function fileEditsToExternalEdits(tc: IToolCallState): IToolCallFileEdit[] { if (tc.status !== ToolCallStatus.Completed) { return []; } @@ -366,7 +366,7 @@ function fileEditsToExternalEdits(tc: IToolCallState): IToolCallFileEdit[] { * Extracts the file path from a tool call's input parameters. * Edit tools store the file path in JSON parameters as `path`. */ -function getFilePathFromToolInput(tc: IToolCallState): string | undefined { +export function getFilePathFromToolInput(tc: IToolCallState): string | undefined { if (tc.status !== ToolCallStatus.Completed || !tc.toolInput) { return undefined; } diff --git a/src/vs/workbench/contrib/chat/browser/chatEditing/chatEditingServiceImpl.ts b/src/vs/workbench/contrib/chat/browser/chatEditing/chatEditingServiceImpl.ts index a8d49698b6e..ba31aec8a0f 100644 --- a/src/vs/workbench/contrib/chat/browser/chatEditing/chatEditingServiceImpl.ts +++ b/src/vs/workbench/contrib/chat/browser/chatEditing/chatEditingServiceImpl.ts @@ -10,7 +10,7 @@ import { groupBy } from '../../../../../base/common/collections.js'; import { ErrorNoTelemetry } from '../../../../../base/common/errors.js'; import { Emitter, Event } from '../../../../../base/common/event.js'; import { Iterable } from '../../../../../base/common/iterator.js'; -import { Disposable, DisposableStore, dispose, IDisposable } from '../../../../../base/common/lifecycle.js'; +import { Disposable, DisposableStore, dispose, IDisposable, toDisposable } from '../../../../../base/common/lifecycle.js'; import { LinkedList } from '../../../../../base/common/linkedList.js'; import { ResourceMap } from '../../../../../base/common/map.js'; import { Schemas } from '../../../../../base/common/network.js'; @@ -37,7 +37,7 @@ import { ILifecycleService } from '../../../../services/lifecycle/common/lifecyc import { IMultiDiffSourceResolver, IMultiDiffSourceResolverService, IResolvedMultiDiffSource, MultiDiffEditorItem } from '../../../multiDiffEditor/browser/multiDiffSourceResolverService.js'; import { CellUri, ICellEditOperation } from '../../../notebook/common/notebookCommon.js'; import { INotebookService } from '../../../notebook/common/notebookService.js'; -import { CHAT_EDITING_MULTI_DIFF_SOURCE_RESOLVER_SCHEME, chatEditingAgentSupportsReadonlyReferencesContextKey, chatEditingResourceContextKey, ChatEditingSessionState, IChatEditingService, IChatEditingSession, IModifiedFileEntry, inChatEditingSessionContextKey, IStreamingEdits, ModifiedFileEntryState, parseChatMultiDiffUri } from '../../common/editing/chatEditingService.js'; +import { CHAT_EDITING_MULTI_DIFF_SOURCE_RESOLVER_SCHEME, chatEditingAgentSupportsReadonlyReferencesContextKey, chatEditingResourceContextKey, ChatEditingSessionState, IChatEditingService, IChatEditingSession, IChatEditingSessionProvider, IModifiedFileEntry, inChatEditingSessionContextKey, IStreamingEdits, ModifiedFileEntryState, parseChatMultiDiffUri } from '../../common/editing/chatEditingService.js'; import { ChatModel, ICellTextEditOperation, IChatResponseModel, isCellTextEditOperationArray } from '../../common/model/chatModel.js'; import { IChatService } from '../../common/chatService/chatService.js'; import { ChatEditorInput } from '../widgetHosts/editor/chatEditorInput.js'; @@ -49,8 +49,9 @@ export class ChatEditingService extends Disposable implements IChatEditingServic _serviceBrand: undefined; + private readonly _providers = new Map(); - private readonly _sessionsObs = observableValueOpts>({ equalsFn: (a, b) => false }, new LinkedList()); + private readonly _sessionsObs = observableValueOpts>({ equalsFn: (a, b) => false }, new LinkedList()); readonly editingSessionsObs: IObservable = derived(r => { const result = Array.from(this._sessionsObs.read(r)); @@ -172,7 +173,10 @@ export class ChatEditingService extends Disposable implements IChatEditingServic assertType(this.getEditingSession(chatModel.sessionResource) === undefined, 'CANNOT have more than one editing session per chat session'); - const session = this._instantiationService.createInstance(ChatEditingSession, chatModel.sessionResource, global, this._lookupEntry.bind(this), initFrom); + const provider = this._providers.get(chatModel.sessionResource.scheme); + const session = provider + ? provider.createEditingSession(chatModel.sessionResource) + : this._instantiationService.createInstance(ChatEditingSession, chatModel.sessionResource, global, this._lookupEntry.bind(this), initFrom); const list = this._sessionsObs.get(); const removeSession = list.unshift(session); @@ -180,7 +184,9 @@ export class ChatEditingService extends Disposable implements IChatEditingServic const store = new DisposableStore(); this._store.add(store); - store.add(this.installAutoApplyObserver(session, chatModel)); + if (!provider && session instanceof ChatEditingSession) { + store.add(this.installAutoApplyObserver(session, chatModel)); + } store.add(session.onDidDispose(e => { removeSession(); @@ -193,6 +199,15 @@ export class ChatEditingService extends Disposable implements IChatEditingServic return session; } + registerEditingSessionProvider(scheme: string, provider: IChatEditingSessionProvider): IDisposable { + this._providers.set(scheme, provider); + return toDisposable(() => { + if (this._providers.get(scheme) === provider) { + this._providers.delete(scheme); + } + }); + } + private installAutoApplyObserver(session: ChatEditingSession, chatModel: ChatModel): IDisposable { if (!chatModel) { throw new ErrorNoTelemetry(`Edit session was created for a non-existing chat session: ${session.chatSessionResource}`); diff --git a/src/vs/workbench/contrib/chat/browser/chatEditing/chatEditingSession.ts b/src/vs/workbench/contrib/chat/browser/chatEditing/chatEditingSession.ts index f585c7f0996..7f819104627 100644 --- a/src/vs/workbench/contrib/chat/browser/chatEditing/chatEditingSession.ts +++ b/src/vs/workbench/contrib/chat/browser/chatEditing/chatEditingSession.ts @@ -145,6 +145,7 @@ function createOpeningEditCodeBlock(uri: URI, isNotebook: boolean, undoStopId: s export class ChatEditingSession extends Disposable implements IChatEditingSession { + readonly supportsKeepUndo = false; private readonly _state = observableValue(this, ChatEditingSessionState.Initial); private readonly _timeline: IChatEditingCheckpointTimeline; diff --git a/src/vs/workbench/contrib/chat/browser/widget/chatListRenderer.ts b/src/vs/workbench/contrib/chat/browser/widget/chatListRenderer.ts index d1c8c32b4ec..a33f39ab401 100644 --- a/src/vs/workbench/contrib/chat/browser/widget/chatListRenderer.ts +++ b/src/vs/workbench/contrib/chat/browser/widget/chatListRenderer.ts @@ -107,6 +107,7 @@ import { HookType } from '../../common/promptSyntax/hookTypes.js'; import { IWorkbenchEnvironmentService } from '../../../../services/environment/common/environmentService.js'; import { AccessibilityWorkbenchSettingId } from '../../../accessibility/browser/accessibilityConfiguration.js'; import { isMcpToolInvocation } from './chatContentParts/toolInvocationParts/chatToolPartUtilities.js'; +import { isAgentHostTarget } from '../agentSessions/agentSessions.js'; const $ = dom.$; @@ -1453,7 +1454,8 @@ export class ChatListItemRenderer extends Disposable implements ITreeRenderer('chat.checkpoints.showFileChanges'); } diff --git a/src/vs/workbench/contrib/chat/common/editing/chatEditingService.ts b/src/vs/workbench/contrib/chat/common/editing/chatEditingService.ts index 0572f62a498..2fec427255d 100644 --- a/src/vs/workbench/contrib/chat/common/editing/chatEditingService.ts +++ b/src/vs/workbench/contrib/chat/common/editing/chatEditingService.ts @@ -26,6 +26,10 @@ import { IChatAgentResult } from '../participants/chatAgents.js'; export const IChatEditingService = createDecorator('chatEditingService'); +export interface IChatEditingSessionProvider { + createEditingSession(chatSessionResource: URI): IChatEditingSession; +} + export interface IChatEditingService { _serviceBrand: undefined; @@ -48,6 +52,14 @@ export interface IChatEditingService { * Creates an editing session with state transferred from the provided session. */ transferEditingSession(chatModel: ChatModel, session: IChatEditingSession): IChatEditingSession; + + /** + * Registers a provider that creates editing sessions for chat sessions + * with the given URI scheme. When {@link createEditingSession} is called + * for a chat model whose sessionResource matches the scheme, the provider + * is used instead of the default implementation. + */ + registerEditingSessionProvider(scheme: string, provider: IChatEditingSessionProvider): IDisposable; } export interface WorkingSetDisplayMetadata { @@ -89,6 +101,7 @@ export interface ISnapshotEntry { export interface IChatEditingSession extends IDisposable { readonly isGlobalEditingSession: boolean; + readonly supportsKeepUndo: boolean; readonly chatSessionResource: URI; readonly onDidDispose: Event; readonly state: IObservable; diff --git a/src/vs/workbench/contrib/chat/test/browser/agentHost/agentHostEditingSession.test.ts b/src/vs/workbench/contrib/chat/test/browser/agentHost/agentHostEditingSession.test.ts new file mode 100644 index 00000000000..956bbb28710 --- /dev/null +++ b/src/vs/workbench/contrib/chat/test/browser/agentHost/agentHostEditingSession.test.ts @@ -0,0 +1,678 @@ +/*--------------------------------------------------------------------------------------------- + * Copyright (c) Microsoft Corporation. All rights reserved. + * Licensed under the MIT License. See License.txt in the project root for license information. + *--------------------------------------------------------------------------------------------*/ + +import assert from 'assert'; +import { VSBuffer } from '../../../../../../base/common/buffer.js'; +import { Event } from '../../../../../../base/common/event.js'; +import { DisposableStore } from '../../../../../../base/common/lifecycle.js'; +import { URI } from '../../../../../../base/common/uri.js'; +import { mock } from '../../../../../../base/test/common/mock.js'; +import { ensureNoDisposablesAreLeakedInTestSuite } from '../../../../../../base/test/common/utils.js'; +import { IDocumentDiff } from '../../../../../../editor/common/diff/documentDiffProvider.js'; +import { DetailedLineRangeMapping } from '../../../../../../editor/common/diff/rangeMapping.js'; +import { IEditorWorkerService } from '../../../../../../editor/common/services/editorWorker.js'; +import { IResolvedTextEditorModel, ITextModelService } from '../../../../../../editor/common/services/resolverService.js'; +import { toAgentHostUri } from '../../../../../../platform/agentHost/common/agentHostUri.js'; +import { IToolCallState, ToolCallConfirmationReason, ToolCallStatus, ToolResultContentType } from '../../../../../../platform/agentHost/common/state/protocol/state.js'; +import { ContentEncoding, IWriteFileParams } from '../../../../../../platform/agentHost/common/state/sessionProtocol.js'; +import type { IToolCallCompletedState } from '../../../../../../platform/agentHost/common/state/sessionState.js'; +import { IFileContent, IFileService } from '../../../../../../platform/files/common/files.js'; +import { IInstantiationService } from '../../../../../../platform/instantiation/common/instantiation.js'; +import { NullLogService } from '../../../../../../platform/log/common/log.js'; +import { IEditorService } from '../../../../../services/editor/common/editorService.js'; +import { AgentHostEditingSession } from '../../../browser/agentSessions/agentHost/agentHostEditingSession.js'; +import { ChatEditingSessionState, ModifiedFileEntryState } from '../../../common/editing/chatEditingService.js'; +import { autorun, IObservable } from '../../../../../../base/common/observable.js'; + +// ---- Test helpers ----------------------------------------------------------- + +/** + * Waits for an observable to satisfy a condition by subscribing to it. + */ +function waitForObservable(obs: IObservable, predicate: (v: T) => boolean, timeoutMs = 2000): Promise { + return new Promise((resolve, reject) => { + const timeout = setTimeout(() => { + d.dispose(); + reject(new Error(`waitForObservable timed out after ${timeoutMs}ms`)); + }, timeoutMs); + const d = autorun(reader => { + const value = obs.read(reader); + if (predicate(value)) { + clearTimeout(timeout); + queueMicrotask(() => { + d.dispose(); + resolve(value); + }); + } + }); + }); +} + +function makeToolCall(opts: { + toolCallId: string; + filePath: string; + beforeURI: string; + afterURI: string; + added?: number; + removed?: number; +}): IToolCallCompletedState { + return { + status: ToolCallStatus.Completed, + toolCallId: opts.toolCallId, + toolName: 'codeEdit', + displayName: 'Edit File', + invocationMessage: 'Editing file', + toolInput: JSON.stringify({ path: opts.filePath }), + success: true, + pastTenseMessage: 'Edited file', + confirmed: ToolCallConfirmationReason.NotNeeded, + content: [{ + type: ToolResultContentType.FileEdit, + beforeURI: opts.beforeURI, + afterURI: opts.afterURI, + diff: { + added: opts.added ?? 0, + removed: opts.removed ?? 0, + }, + }], + }; +} + +function makeMockFileService(contentMap: Map): IFileService { + return new class extends mock() { + override async readFile(uri: URI) { + const key = uri.toString(); + const data = contentMap.get(key); + if (data === undefined) { + throw new Error(`Content not found: ${key}`); + } + return { value: VSBuffer.fromString(data) } as IFileContent; + } + }; +} + +function createSession(store: DisposableStore, contentMap: Map, opts?: { computeDiffResult?: IDocumentDiff | null }): AgentHostEditingSession { + const sessionResource = URI.from({ scheme: 'agent-host-copilot', path: '/test-session' }); + const mockEditorService = new class extends mock() { + override readonly onDidActiveEditorChange = Event.None; + }; + const mockInstantiationService = new class extends mock() { }; + const mockFileService = makeMockFileService(contentMap); + const mockTextModelService = new class extends mock() { + override async createModelReference(uri: URI) { + const content = contentMap.get(uri.toString()) ?? ''; + return { + object: { textEditorModel: { uri, getValue: () => content } } as IResolvedTextEditorModel, + dispose: () => { }, + }; + } + }; + const mockEditorWorkerService = new class extends mock() { + override async computeDiff(_original: URI, _modified: URI): Promise { + return opts?.computeDiffResult ?? { identical: false, quitEarly: false, changes: [], moves: [] }; + } + }; + const session = new AgentHostEditingSession( + sessionResource, + 'local', + mockEditorService, + mockInstantiationService, + new NullLogService(), + mockFileService, + mockTextModelService, + mockEditorWorkerService, + ); + store.add(session); + return session; +} + +// ---- Tests ------------------------------------------------------------------ + +suite('AgentHostEditingSession', () => { + + const store = new DisposableStore(); + + teardown(() => store.clear()); + + ensureNoDisposablesAreLeakedInTestSuite(); + + test('initial state', () => { + const session = createSession(store, new Map()); + + assert.strictEqual(session.supportsKeepUndo, true); + assert.strictEqual(session.isGlobalEditingSession, false); + assert.strictEqual(session.state.get(), ChatEditingSessionState.Idle); + assert.deepStrictEqual(session.entries.get(), []); + assert.strictEqual(session.canUndo.get(), false); + assert.strictEqual(session.canRedo.get(), false); + }); + + test('addToolCallEdits hydrates entries and transitions state', () => { + const session = createSession(store, new Map()); + + session.addToolCallEdits('req-1', makeToolCall({ + toolCallId: 'tc-1', + filePath: '/workspace/file.ts', + beforeURI: 'content://before-1', + afterURI: 'content://after-1', + added: 5, + removed: 2, + })); + + assert.strictEqual(session.state.get(), ChatEditingSessionState.Idle); + assert.strictEqual(session.entries.get().length, 1); + + const entry = session.entries.get()[0]; + assert.strictEqual(entry.lastModifyingRequestId, 'req-1'); + assert.strictEqual(entry.state.get(), ModifiedFileEntryState.Accepted); + assert.strictEqual(entry.linesAdded?.get(), 5); + assert.strictEqual(entry.linesRemoved?.get(), 2); + assert.strictEqual(session.canUndo.get(), true); + assert.strictEqual(session.canRedo.get(), false); + }); + + test('addToolCallEdits ignores non-completed tool calls', () => { + const session = createSession(store, new Map()); + + const tc = { ...makeToolCall({ toolCallId: 'tc-1', filePath: '/f.ts', beforeURI: 'b', afterURI: 'a' }), status: ToolCallStatus.Running } as IToolCallState; + session.addToolCallEdits('req-1', tc); + + assert.strictEqual(session.state.get(), ChatEditingSessionState.Idle); + assert.deepStrictEqual(session.entries.get(), []); + }); + + test('addToolCallEdits deduplicates by toolCallId', () => { + const session = createSession(store, new Map()); + + const tc = makeToolCall({ toolCallId: 'tc-1', filePath: '/f.ts', beforeURI: 'b', afterURI: 'a', added: 3 }); + session.addToolCallEdits('req-1', tc); + session.addToolCallEdits('req-1', tc); + + assert.strictEqual(session.entries.get().length, 1); + }); + + test('multiple tool calls to same file accumulate diffs', () => { + const session = createSession(store, new Map()); + + session.addToolCallEdits('req-1', makeToolCall({ + toolCallId: 'tc-1', + filePath: '/workspace/file.ts', + beforeURI: 'content://before-1', + afterURI: 'content://after-1', + added: 5, + removed: 2, + })); + + session.addToolCallEdits('req-2', makeToolCall({ + toolCallId: 'tc-2', + filePath: '/workspace/file.ts', + beforeURI: 'content://after-1', + afterURI: 'content://after-2', + added: 3, + removed: 1, + })); + + // Should merge into one entry with accumulated counts + assert.strictEqual(session.entries.get().length, 1); + const entry = session.entries.get()[0]; + assert.strictEqual(entry.linesAdded?.get(), 8); + assert.strictEqual(entry.linesRemoved?.get(), 3); + assert.strictEqual(entry.lastModifyingRequestId, 'req-2'); + }); + + test('multiple tool calls to different files create separate entries', () => { + const session = createSession(store, new Map()); + + session.addToolCallEdits('req-1', makeToolCall({ + toolCallId: 'tc-1', + filePath: '/workspace/a.ts', + beforeURI: 'content://before-a', + afterURI: 'content://after-a', + added: 2, + })); + + session.addToolCallEdits('req-1', makeToolCall({ + toolCallId: 'tc-2', + filePath: '/workspace/b.ts', + beforeURI: 'content://before-b', + afterURI: 'content://after-b', + added: 4, + })); + + assert.strictEqual(session.entries.get().length, 2); + }); + + test('getEntry finds entry by URI', () => { + const session = createSession(store, new Map()); + + session.addToolCallEdits('req-1', makeToolCall({ + toolCallId: 'tc-1', + filePath: '/workspace/file.ts', + beforeURI: 'content://before-1', + afterURI: 'content://after-1', + })); + + const entry = session.entries.get()[0]; + assert.ok(session.getEntry(entry.modifiedURI)); + assert.strictEqual(session.getEntry(URI.parse('file:///nonexistent')), undefined); + }); + + test('hasEditsInRequest', () => { + const session = createSession(store, new Map()); + + session.addToolCallEdits('req-1', makeToolCall({ + toolCallId: 'tc-1', + filePath: '/workspace/file.ts', + beforeURI: 'b', + afterURI: 'a', + })); + + assert.strictEqual(session.hasEditsInRequest('req-1'), true); + assert.strictEqual(session.hasEditsInRequest('req-other'), false); + }); + + test('getDiffForSession aggregates diff stats', async () => { + const diffResult: IDocumentDiff = { + identical: false, quitEarly: false, moves: [], + changes: [{ + original: { startLineNumber: 1, endLineNumberExclusive: 3 }, + modified: { startLineNumber: 1, endLineNumberExclusive: 6 }, + innerChanges: null, + } as unknown as DetailedLineRangeMapping], + }; + const session = createSession(store, new Map(), { computeDiffResult: diffResult }); + + session.addToolCallEdits('req-1', makeToolCall({ + toolCallId: 'tc-1', + filePath: '/workspace/a.ts', + beforeURI: 'b1', + afterURI: 'a1', + })); + + session.addToolCallEdits('req-2', makeToolCall({ + toolCallId: 'tc-2', + filePath: '/workspace/b.ts', + beforeURI: 'b2', + afterURI: 'a2', + })); + + // Each file produces +5 -2 from the mock diff result, wait for both + const stats = await waitForObservable(session.getDiffForSession(), s => s.added === 10); + assert.deepStrictEqual(stats, { added: 10, removed: 4 }); + }); + + suite('undo/redo', () => { + + test('undo fires writeFile with before-content and updates state', async () => { + const beforeUri = toAgentHostUri(URI.file('/workspace/file.ts'), 'local'); + const contentMap = new Map(); + contentMap.set(beforeUri.toString(), 'before-content'); + const session = createSession(store, contentMap); + + session.addToolCallEdits('req-1', makeToolCall({ + toolCallId: 'tc-1', + filePath: '/workspace/file.ts', + beforeURI: URI.file('/workspace/file.ts').toString(), + afterURI: 'content://after-1', + })); + + const writes: IWriteFileParams[] = []; + store.add(session.onDidRequestFileWrite(p => writes.push(p))); + + await session.undoInteraction(); + + assert.strictEqual(writes.length, 1); + assert.strictEqual(writes[0].data, 'before-content'); + assert.strictEqual(writes[0].encoding, ContentEncoding.Utf8); + assert.strictEqual(session.canUndo.get(), false); + assert.strictEqual(session.canRedo.get(), true); + assert.deepStrictEqual(session.entries.get(), []); + }); + + test('redo fires writeFile with after-content', async () => { + const beforeUri = toAgentHostUri(URI.file('/workspace/file.ts'), 'local'); + const afterUri = toAgentHostUri(URI.parse('content://after-1'), 'local'); + const contentMap = new Map(); + contentMap.set(beforeUri.toString(), 'before-content'); + contentMap.set(afterUri.toString(), 'after-content'); + const session = createSession(store, contentMap); + + session.addToolCallEdits('req-1', makeToolCall({ + toolCallId: 'tc-1', + filePath: '/workspace/file.ts', + beforeURI: URI.file('/workspace/file.ts').toString(), + afterURI: 'content://after-1', + })); + + await session.undoInteraction(); + + const writes: IWriteFileParams[] = []; + store.add(session.onDidRequestFileWrite(p => writes.push(p))); + + await session.redoInteraction(); + + assert.strictEqual(writes.length, 1); + assert.strictEqual(writes[0].data, 'after-content'); + assert.strictEqual(session.canUndo.get(), true); + assert.strictEqual(session.canRedo.get(), false); + assert.strictEqual(session.entries.get().length, 1); + }); + + test('undo when nothing to undo is no-op', async () => { + const session = createSession(store, new Map()); + + const writes: IWriteFileParams[] = []; + store.add(session.onDidRequestFileWrite(p => writes.push(p))); + + await session.undoInteraction(); + + assert.strictEqual(writes.length, 0); + }); + + test('redo when nothing to redo is no-op', async () => { + const session = createSession(store, new Map()); + + session.addToolCallEdits('req-1', makeToolCall({ + toolCallId: 'tc-1', + filePath: '/f.ts', + beforeURI: 'b', + afterURI: 'a', + })); + + const writes: IWriteFileParams[] = []; + store.add(session.onDidRequestFileWrite(p => writes.push(p))); + + await session.redoInteraction(); + + assert.strictEqual(writes.length, 0); + }); + + test('undo after multiple checkpoints removes entries correctly', async () => { + const contentMap = new Map(); + contentMap.set(toAgentHostUri(URI.file('/workspace/a.ts'), 'local').toString(), 'a-before'); + contentMap.set(toAgentHostUri(URI.file('/workspace/b.ts'), 'local').toString(), 'b-before'); + contentMap.set(toAgentHostUri(URI.parse('content://after-a'), 'local').toString(), 'a-after'); + contentMap.set(toAgentHostUri(URI.parse('content://after-b'), 'local').toString(), 'b-after'); + const session = createSession(store, contentMap); + + session.addToolCallEdits('req-1', makeToolCall({ + toolCallId: 'tc-1', + filePath: '/workspace/a.ts', + beforeURI: URI.file('/workspace/a.ts').toString(), + afterURI: 'content://after-a', + added: 5, + })); + + session.addToolCallEdits('req-2', makeToolCall({ + toolCallId: 'tc-2', + filePath: '/workspace/b.ts', + beforeURI: URI.file('/workspace/b.ts').toString(), + afterURI: 'content://after-b', + added: 3, + })); + + assert.strictEqual(session.entries.get().length, 2); + + await session.undoInteraction(); + + // Only a.ts should remain (b.ts was undone) + assert.strictEqual(session.entries.get().length, 1); + assert.ok(session.entries.get()[0].modifiedURI.path.includes('a.ts')); + }); + }); + + suite('getDiffsForFilesInSession', () => { + test('returns diffs for all files', async () => { + const diffResult: IDocumentDiff = { + identical: false, quitEarly: false, moves: [], + changes: [{ + original: { startLineNumber: 1, endLineNumberExclusive: 3 }, + modified: { startLineNumber: 1, endLineNumberExclusive: 11 }, + innerChanges: null, + } as unknown as DetailedLineRangeMapping], + }; + const session = createSession(store, new Map(), { computeDiffResult: diffResult }); + + session.addToolCallEdits('req-1', makeToolCall({ + toolCallId: 'tc-1', + filePath: '/workspace/a.ts', + beforeURI: 'b1', + afterURI: 'a1', + })); + + const diffs = await waitForObservable(session.getDiffsForFilesInSession(), d => d.length > 0 && d[0].added > 0); + assert.strictEqual(diffs.length, 1); + assert.strictEqual(diffs[0].added, 10); + assert.strictEqual(diffs[0].removed, 2); + assert.strictEqual(diffs[0].isFinal, true); + }); + }); + + suite('getDiffsForFilesInRequest', () => { + test('returns diffs scoped to a request', async () => { + const diffResult: IDocumentDiff = { + identical: false, quitEarly: false, moves: [], + changes: [{ + original: { startLineNumber: 1, endLineNumberExclusive: 1 }, + modified: { startLineNumber: 1, endLineNumberExclusive: 6 }, + innerChanges: null, + } as unknown as DetailedLineRangeMapping], + }; + const session = createSession(store, new Map(), { computeDiffResult: diffResult }); + + session.addToolCallEdits('req-1', makeToolCall({ + toolCallId: 'tc-1', + filePath: '/workspace/a.ts', + beforeURI: 'b1', + afterURI: 'a1', + })); + + session.addToolCallEdits('req-2', makeToolCall({ + toolCallId: 'tc-2', + filePath: '/workspace/b.ts', + beforeURI: 'b2', + afterURI: 'a2', + })); + + const req1Diffs = await waitForObservable(session.getDiffsForFilesInRequest('req-1'), d => d.length > 0 && d[0].added > 0); + assert.strictEqual(req1Diffs.length, 1); + assert.strictEqual(req1Diffs[0].added, 5); + + const req2Diffs = await waitForObservable(session.getDiffsForFilesInRequest('req-2'), d => d.length > 0 && d[0].added > 0); + assert.strictEqual(req2Diffs.length, 1); + assert.strictEqual(req2Diffs[0].added, 5); + + const noReqDiffs = session.getDiffsForFilesInRequest('req-none').get(); + assert.strictEqual(noReqDiffs.length, 0); + }); + }); + + suite('snapshots', () => { + test('getSnapshotUri returns URI for valid checkpoint', () => { + const session = createSession(store, new Map()); + + session.addToolCallEdits('req-1', makeToolCall({ + toolCallId: 'tc-1', + filePath: '/workspace/file.ts', + beforeURI: 'b', + afterURI: 'a', + })); + + const entry = session.entries.get()[0]; + const snapshotUri = session.getSnapshotUri('req-1', entry.modifiedURI, undefined); + assert.ok(snapshotUri); + assert.strictEqual(snapshotUri!.scheme, 'chat-editing-snapshot-text-model'); + }); + + test('getSnapshotUri returns undefined for unknown request', () => { + const session = createSession(store, new Map()); + + const result = session.getSnapshotUri('nonexistent', URI.file('/x'), undefined); + assert.strictEqual(result, undefined); + }); + + test('getSnapshotContents fetches content from connection', async () => { + const contentMap = new Map(); + const session = createSession(store, contentMap); + + session.addToolCallEdits('req-1', makeToolCall({ + toolCallId: 'tc-1', + filePath: '/workspace/file.ts', + beforeURI: 'content://before', + afterURI: 'content://after', + })); + + // The afterContentUri is wrapped via toAgentHostUri(URI.parse('content://after'), 'local') + const wrappedAfterUri = toAgentHostUri(URI.parse('content://after'), 'local'); + contentMap.set(wrappedAfterUri.toString(), 'file content here'); + + const entry = session.entries.get()[0]; + const content = await session.getSnapshotContents('req-1', entry.modifiedURI, undefined); + assert.ok(content); + assert.strictEqual(content!.toString(), 'file content here'); + }); + }); + + suite('dispose', () => { + test('sets state to Disposed and fires event', () => { + const session = createSession(store, new Map()); + + let disposed = false; + store.add(session.onDidDispose(() => { disposed = true; })); + + session.dispose(); + + assert.strictEqual(session.state.get(), ChatEditingSessionState.Disposed); + assert.strictEqual(disposed, true); + }); + }); + + suite('explanations', () => { + test('trigger/clear/has cycle', async () => { + const session = createSession(store, new Map()); + + assert.strictEqual(session.hasExplanations(), false); + + await session.triggerExplanationGeneration(); + assert.strictEqual(session.hasExplanations(), true); + + session.clearExplanations(); + assert.strictEqual(session.hasExplanations(), false); + }); + }); + + suite('accept/reject are no-ops', () => { + test('accept does not throw', async () => { + const session = createSession(store, new Map()); + await session.accept(URI.file('/test')); + }); + + test('reject does not throw', async () => { + const session = createSession(store, new Map()); + await session.reject(URI.file('/test')); + }); + }); + + suite('getEntryDiffBetweenStops', () => { + test('returns undefined for unknown requestId', () => { + const session = createSession(store, new Map()); + const result = session.getEntryDiffBetweenStops(URI.file('/f'), 'unknown-req', undefined); + assert.strictEqual(result, undefined); + }); + + test('computes diff for a known stop', async () => { + const contentMap = new Map(); + const beforeUri = toAgentHostUri(URI.parse('content://before-1'), 'local'); + const afterUri = toAgentHostUri(URI.parse('content://after-1'), 'local'); + contentMap.set(beforeUri.toString(), 'line1\nline2\n'); + contentMap.set(afterUri.toString(), 'line1\nline2\nline3\n'); + + const diffResult: IDocumentDiff = { + identical: false, + quitEarly: false, + changes: [{ + original: { startLineNumber: 3, endLineNumberExclusive: 3 }, + modified: { startLineNumber: 3, endLineNumberExclusive: 4 }, + innerChanges: null, + } as unknown as DetailedLineRangeMapping], + moves: [], + }; + const session = createSession(store, contentMap, { computeDiffResult: diffResult }); + + session.addToolCallEdits('req-1', makeToolCall({ + toolCallId: 'tc-1', + filePath: '/workspace/file.ts', + beforeURI: 'content://before-1', + afterURI: 'content://after-1', + })); + + const entry = session.entries.get()[0]; + const diffObs = session.getEntryDiffBetweenStops(entry.modifiedURI, 'req-1', 'tc-1'); + assert.ok(diffObs); + + // Wait for the async diff computation by polling the observable + let diff = diffObs!.get(); + if (diff?.isBusy) { + await new Promise(r => setTimeout(r, 50)); + diff = diffObs!.get(); + } + assert.ok(diff); + assert.strictEqual(diff!.added, 1); + assert.strictEqual(diff!.removed, 0); + assert.strictEqual(diff!.isFinal, true); + }); + }); + + suite('getEntryDiffBetweenRequests', () => { + test('returns undefined for unknown requests', () => { + const session = createSession(store, new Map()); + session.addToolCallEdits('req-1', makeToolCall({ + toolCallId: 'tc-1', + filePath: '/f.ts', + beforeURI: 'b', + afterURI: 'a', + })); + const diffObs = session.getEntryDiffBetweenRequests(URI.file('/f'), 'unknown', 'also-unknown'); + assert.strictEqual(diffObs.get(), undefined); + }); + + test('computes diff spanning multiple requests', async () => { + const contentMap = new Map(); + const beforeUri = toAgentHostUri(URI.parse('content://before'), 'local'); + const afterUri1 = toAgentHostUri(URI.parse('content://after-1'), 'local'); + const afterUri2 = toAgentHostUri(URI.parse('content://after-2'), 'local'); + contentMap.set(beforeUri.toString(), 'original'); + contentMap.set(afterUri1.toString(), 'modified-1'); + contentMap.set(afterUri2.toString(), 'modified-2'); + + const session = createSession(store, contentMap); + + session.addToolCallEdits('req-1', makeToolCall({ + toolCallId: 'tc-1', + filePath: '/workspace/file.ts', + beforeURI: 'content://before', + afterURI: 'content://after-1', + })); + + session.addToolCallEdits('req-2', makeToolCall({ + toolCallId: 'tc-2', + filePath: '/workspace/file.ts', + beforeURI: 'content://after-1', + afterURI: 'content://after-2', + })); + + const entry = session.entries.get()[0]; + const diffObs = session.getEntryDiffBetweenRequests(entry.modifiedURI, 'req-1', 'req-2'); + + let diff = diffObs.get(); + if (diff?.isBusy) { + await new Promise(r => setTimeout(r, 50)); + diff = diffObs.get(); + } + assert.ok(diff); + assert.strictEqual(diff!.isFinal, true); + }); + }); +}); diff --git a/src/vs/workbench/contrib/chat/test/browser/agentSessions/agentHostChatContribution.test.ts b/src/vs/workbench/contrib/chat/test/browser/agentSessions/agentHostChatContribution.test.ts index fa0d6674c79..3ccd69f785d 100644 --- a/src/vs/workbench/contrib/chat/test/browser/agentSessions/agentHostChatContribution.test.ts +++ b/src/vs/workbench/contrib/chat/test/browser/agentSessions/agentHostChatContribution.test.ts @@ -22,6 +22,7 @@ import { IAuthenticationService } from '../../../../../services/authentication/c import { IChatAgentData, IChatAgentImplementation, IChatAgentRequest, IChatAgentService } from '../../../common/participants/chatAgents.js'; import { ChatAgentLocation } from '../../../common/constants.js'; import { IChatMarkdownContent, IChatProgress, IChatTerminalToolInvocationData, IChatToolInvocation, IChatToolInvocationSerialized, ToolConfirmKind } from '../../../common/chatService/chatService.js'; +import { IChatEditingService } from '../../../common/editing/chatEditingService.js'; import { IMarkdownString } from '../../../../../../base/common/htmlContent.js'; import { IChatSessionsService } from '../../../common/chatSessionsService.js'; import { ILanguageModelsService } from '../../../common/languageModels.js'; @@ -175,6 +176,9 @@ function createTestServices(disposables: DisposableStore) { instantiationService.stub(IConfigurationService, { getValue: () => true }); instantiationService.stub(IOutputService, { getChannel: () => undefined }); instantiationService.stub(IWorkspaceContextService, { getWorkspace: () => ({ id: '', folders: [] }), getWorkspaceFolder: () => null }); + instantiationService.stub(IChatEditingService, { + registerEditingSessionProvider: () => toDisposable(() => { }), + }); return { instantiationService, agentHostService, chatAgentService }; } diff --git a/src/vs/workbench/contrib/chat/test/common/chatDebugServiceImpl.test.ts b/src/vs/workbench/contrib/chat/test/common/chatDebugServiceImpl.test.ts index 62e41c7fcb8..55ed1e84adf 100644 --- a/src/vs/workbench/contrib/chat/test/common/chatDebugServiceImpl.test.ts +++ b/src/vs/workbench/contrib/chat/test/common/chatDebugServiceImpl.test.ts @@ -5,6 +5,7 @@ import assert from 'assert'; import { CancellationToken, CancellationTokenSource } from '../../../../../base/common/cancellation.js'; +import { errorHandler } from '../../../../../base/common/errors.js'; import { URI } from '../../../../../base/common/uri.js'; import { ensureNoDisposablesAreLeakedInTestSuite } from '../../../../../base/test/common/utils.js'; import { ChatDebugLogLevel, IChatDebugEvent, IChatDebugGenericEvent, IChatDebugLogProvider, IChatDebugModelTurnEvent, IChatDebugResolvedEventContent, IChatDebugToolCallEvent } from '../../common/chatDebugService.js'; @@ -355,8 +356,14 @@ suite('ChatDebugServiceImpl', () => { }; disposables.add(service.registerProvider(provider)); - // Should not throw - await service.invokeProviders(errorSession); + // Suppress the expected onUnexpectedError from _invokeProvider + const origHandler = errorHandler.getUnexpectedErrorHandler(); + errorHandler.setUnexpectedErrorHandler(() => { }); + try { + await service.invokeProviders(errorSession); + } finally { + errorHandler.setUnexpectedErrorHandler(origHandler); + } assert.strictEqual(service.getEvents(errorSession).length, 0); }); }); From bd2f91618ede8b837a76f0691fa10c043a8b20fd Mon Sep 17 00:00:00 2001 From: Connor Peet Date: Fri, 27 Mar 2026 16:23:07 -0700 Subject: [PATCH 2/3] fix compile --- src/vs/platform/agentHost/node/agentHostMain.ts | 3 +++ .../platform/agentHost/test/node/protocolServerHandler.test.ts | 1 + 2 files changed, 4 insertions(+) diff --git a/src/vs/platform/agentHost/node/agentHostMain.ts b/src/vs/platform/agentHost/node/agentHostMain.ts index 7df7461682c..04a779fc59e 100644 --- a/src/vs/platform/agentHost/node/agentHostMain.ts +++ b/src/vs/platform/agentHost/node/agentHostMain.ts @@ -173,6 +173,9 @@ async function startWebSocketServer(agentService: AgentService, logService: ILog handleBrowseDirectory(uri) { return agentService.browseDirectory(URI.parse(uri)); }, + handleWriteFile(params) { + return agentService.writeFile(params); + }, async handleRestoreSession(session) { return agentService.restoreSession(URI.parse(session)); }, diff --git a/src/vs/platform/agentHost/test/node/protocolServerHandler.test.ts b/src/vs/platform/agentHost/test/node/protocolServerHandler.test.ts index 9a87a5edb02..f4e2fae680e 100644 --- a/src/vs/platform/agentHost/test/node/protocolServerHandler.test.ts +++ b/src/vs/platform/agentHost/test/node/protocolServerHandler.test.ts @@ -77,6 +77,7 @@ class MockSideEffectHandler implements IProtocolSideEffectHandler { async handleListSessions(): Promise { return []; } async handleRestoreSession(_session: string): Promise { } handleGetResourceMetadata() { return { resources: [] }; } + handleWriteFile() { return Promise.resolve({}); } async handleAuthenticate(_params: { resource: string; token: string }) { return { authenticated: true }; } async handleBrowseDirectory(uri: string): Promise<{ entries: { name: string; type: 'file' | 'directory' }[] }> { this.browsedUris.push(URI.parse(uri)); From 36f8813cfa5d9156a71eca30b92f1fafcacffb49 Mon Sep 17 00:00:00 2001 From: Connor Peet Date: Fri, 27 Mar 2026 16:40:44 -0700 Subject: [PATCH 3/3] address PR review comments --- .../common/state/protocol/.ahp-version | 2 +- .../common/state/protocol/commands.ts | 1 + .../agentHost/common/state/protocol/errors.ts | 5 + .../common/state/protocol/version/registry.ts | 107 ++++++++++++++++++ .../agentHost/node/agentSideEffects.ts | 17 ++- .../agentHost/agentHostEditingSession.ts | 30 ++++- .../agentHost/agentHostSessionHandler.ts | 7 +- 7 files changed, 155 insertions(+), 14 deletions(-) create mode 100644 src/vs/platform/agentHost/common/state/protocol/version/registry.ts diff --git a/src/vs/platform/agentHost/common/state/protocol/.ahp-version b/src/vs/platform/agentHost/common/state/protocol/.ahp-version index 1e7e823363b..eda0ecd4b8b 100644 --- a/src/vs/platform/agentHost/common/state/protocol/.ahp-version +++ b/src/vs/platform/agentHost/common/state/protocol/.ahp-version @@ -1 +1 @@ -101c091 +95cbb57 diff --git a/src/vs/platform/agentHost/common/state/protocol/commands.ts b/src/vs/platform/agentHost/common/state/protocol/commands.ts index 2f1ea9241f4..96ea242ae2e 100644 --- a/src/vs/platform/agentHost/common/state/protocol/commands.ts +++ b/src/vs/platform/agentHost/common/state/protocol/commands.ts @@ -301,6 +301,7 @@ export interface IFetchContentResult { * @version 1 * @throws `NotFound` (`-32008`) if the parent directory does not exist. * @throws `PermissionDenied` (`-32009`) if the client is not permitted to write to the path. + * @throws `AlreadyExists` (`-32010`) if `createOnly` is set and the file already exists. * @example * ```jsonc * // Client → Server diff --git a/src/vs/platform/agentHost/common/state/protocol/errors.ts b/src/vs/platform/agentHost/common/state/protocol/errors.ts index f582cdae0b2..d22126e9222 100644 --- a/src/vs/platform/agentHost/common/state/protocol/errors.ts +++ b/src/vs/platform/agentHost/common/state/protocol/errors.ts @@ -66,6 +66,11 @@ export const AhpErrorCodes = { * directory or workspace roots). */ PermissionDenied: -32009, + /** + * The target resource already exists and the operation does not allow + * overwriting (e.g. `writeFile` with `createOnly: true`). + */ + AlreadyExists: -32010, } as const; /** Union type of all AHP application error codes. */ diff --git a/src/vs/platform/agentHost/common/state/protocol/version/registry.ts b/src/vs/platform/agentHost/common/state/protocol/version/registry.ts new file mode 100644 index 00000000000..c3d45ae3dfa --- /dev/null +++ b/src/vs/platform/agentHost/common/state/protocol/version/registry.ts @@ -0,0 +1,107 @@ +/*--------------------------------------------------------------------------------------------- + * Copyright (c) Microsoft Corporation. All rights reserved. + * Licensed under the MIT License. See License.txt in the project root for license information. + *--------------------------------------------------------------------------------------------*/ + +// allow-any-unicode-comment-file +// DO NOT EDIT -- auto-generated by scripts/sync-agent-host-protocol.ts + +import { ActionType, type IStateAction } from '../actions.js'; +import { NotificationType, type IProtocolNotification } from '../notifications.js'; + +// ─── Protocol Version Constants ────────────────────────────────────────────── + +/** The current protocol version that new code speaks. */ +export const PROTOCOL_VERSION = 1; + +/** The oldest protocol version the implementation maintains compatibility with. */ +export const MIN_PROTOCOL_VERSION = 1; + +// ─── Exhaustive Action → Version Map ───────────────────────────────────────── + +/** + * Maps every action type to the protocol version that introduced it. + * Adding a new action to `IStateAction` without adding it here is a compile error. + */ +export const ACTION_INTRODUCED_IN: { readonly [K in IStateAction['type']]: number } = { + [ActionType.RootAgentsChanged]: 1, + [ActionType.RootActiveSessionsChanged]: 1, + [ActionType.SessionReady]: 1, + [ActionType.SessionCreationFailed]: 1, + [ActionType.SessionTurnStarted]: 1, + [ActionType.SessionDelta]: 1, + [ActionType.SessionResponsePart]: 1, + [ActionType.SessionToolCallStart]: 1, + [ActionType.SessionToolCallDelta]: 1, + [ActionType.SessionToolCallReady]: 1, + [ActionType.SessionToolCallConfirmed]: 1, + [ActionType.SessionToolCallComplete]: 1, + [ActionType.SessionToolCallResultConfirmed]: 1, + [ActionType.SessionTurnComplete]: 1, + [ActionType.SessionTurnCancelled]: 1, + [ActionType.SessionError]: 1, + [ActionType.SessionTitleChanged]: 1, + [ActionType.SessionUsage]: 1, + [ActionType.SessionReasoning]: 1, + [ActionType.SessionModelChanged]: 1, + [ActionType.SessionServerToolsChanged]: 1, + [ActionType.SessionActiveClientChanged]: 1, + [ActionType.SessionActiveClientToolsChanged]: 1, + [ActionType.SessionPendingMessageSet]: 1, + [ActionType.SessionPendingMessageRemoved]: 1, + [ActionType.SessionQueuedMessagesReordered]: 1, + [ActionType.SessionCustomizationsChanged]: 1, + [ActionType.SessionCustomizationToggled]: 1, +}; + +/** + * Returns whether the given action type is known to the specified protocol version. + */ +export function isActionKnownToVersion(action: IStateAction, clientVersion: number): boolean { + return ACTION_INTRODUCED_IN[action.type] <= clientVersion; +} + +// ─── Exhaustive Notification → Version Map ───────────────────────────────── + +/** + * Maps every notification type to the protocol version that introduced it. + * Adding a new notification to `IProtocolNotification` without adding it here + * is a compile error. + */ +export const NOTIFICATION_INTRODUCED_IN: { readonly [K in IProtocolNotification['type']]: number } = { + [NotificationType.SessionAdded]: 1, + [NotificationType.SessionRemoved]: 1, + [NotificationType.AuthRequired]: 1, +}; + +/** + * Returns whether the given notification type is known to the specified protocol version. + */ +export function isNotificationKnownToVersion(notification: IProtocolNotification, clientVersion: number): boolean { + return NOTIFICATION_INTRODUCED_IN[notification.type] <= clientVersion; +} + +// ─── Capabilities ──────────────────────────────────────────────────────────── + +/** + * Feature capabilities gated by protocol version. + */ +export interface ProtocolCapabilities { + /** v1 — always present */ + readonly sessions: true; + /** v1 — always present */ + readonly tools: true; + /** v1 — always present */ + readonly permissions: true; +} + +/** + * Derives capabilities from a protocol version number. + */ +export function capabilitiesForVersion(_version: number): ProtocolCapabilities { + return { + sessions: true, + tools: true, + permissions: true, + }; +} diff --git a/src/vs/platform/agentHost/node/agentSideEffects.ts b/src/vs/platform/agentHost/node/agentSideEffects.ts index c2cff8c8f25..51288d5d1e1 100644 --- a/src/vs/platform/agentHost/node/agentSideEffects.ts +++ b/src/vs/platform/agentHost/node/agentSideEffects.ts @@ -9,7 +9,7 @@ import { Disposable, DisposableStore, IDisposable } from '../../../base/common/l import { autorun, IObservable } from '../../../base/common/observable.js'; import { URI } from '../../../base/common/uri.js'; import { generateUuid } from '../../../base/common/uuid.js'; -import { IFileService } from '../../files/common/files.js'; +import { FileSystemProviderErrorCode, IFileService, toFileSystemProviderErrorCode } from '../../files/common/files.js'; import { ILogService } from '../../log/common/log.js'; import { IAgent, IAgentAttachment, IAgentMessageEvent, IAgentToolCompleteEvent, IAgentToolStartEvent, IAuthenticateParams, IAuthenticateResult, IResourceMetadata } from '../common/agentService.js'; import { ISessionDataService } from '../common/sessionDataService.js'; @@ -594,9 +594,20 @@ export class AgentSideEffects extends Disposable implements IProtocolSideEffectH content = VSBuffer.fromString(params.data); } try { - await this._fileService.writeFile(fileUri, content); + if (params.createOnly) { + await this._fileService.createFile(fileUri, content, { overwrite: false }); + } else { + await this._fileService.writeFile(fileUri, content); + } return {}; - } catch (_e) { + } catch (e) { + const code = toFileSystemProviderErrorCode(e as Error); + if (code === FileSystemProviderErrorCode.FileExists) { + throw new ProtocolError(AhpErrorCodes.AlreadyExists, `File already exists: ${fileUri.toString()}`); + } + if (code === FileSystemProviderErrorCode.NoPermissions) { + throw new ProtocolError(AhpErrorCodes.PermissionDenied, `Permission denied: ${fileUri.toString()}`); + } throw new ProtocolError(AhpErrorCodes.NotFound, `Failed to write file: ${fileUri.toString()}`); } } diff --git a/src/vs/workbench/contrib/chat/browser/agentSessions/agentHost/agentHostEditingSession.ts b/src/vs/workbench/contrib/chat/browser/agentSessions/agentHost/agentHostEditingSession.ts index 46075aef564..019e7459e33 100644 --- a/src/vs/workbench/contrib/chat/browser/agentSessions/agentHost/agentHostEditingSession.ts +++ b/src/vs/workbench/contrib/chat/browser/agentSessions/agentHost/agentHostEditingSession.ts @@ -246,10 +246,28 @@ export class AgentHostEditingSession extends Disposable implements IChatEditingS // ---- Snapshots ---------------------------------------------------------- - async restoreSnapshot(requestId: string, _stopId: string | undefined): Promise { - const idx = this._checkpoints.findIndex(cp => cp.requestId === requestId); + private _findCheckpointIndex(requestId: string, stopId: string | undefined): number { + if (stopId !== undefined) { + return this._checkpoints.findIndex(cp => cp.requestId === requestId && cp.undoStopId === stopId); + } + // No specific stop: use the last checkpoint for this request + for (let i = this._checkpoints.length - 1; i >= 0; i--) { + if (this._checkpoints[i].requestId === requestId) { + return i; + } + } + return -1; + } + + private _findCheckpoint(requestId: string, stopId: string | undefined): IAgentHostCheckpoint | undefined { + const idx = this._findCheckpointIndex(requestId, stopId); + return idx >= 0 ? this._checkpoints[idx] : undefined; + } + + async restoreSnapshot(requestId: string, stopId: string | undefined): Promise { + const idx = this._findCheckpointIndex(requestId, stopId); if (idx < 0) { - this._logService.warn(`[AgentHostEditingSession] No checkpoint found for requestId=${requestId}`); + this._logService.warn(`[AgentHostEditingSession] No checkpoint found for requestId=${requestId}${stopId ? `, stopId=${stopId}` : ''}`); return; } @@ -274,7 +292,7 @@ export class AgentHostEditingSession extends Disposable implements IChatEditingS } getSnapshotUri(requestId: string, uri: URI, stopId: string | undefined): URI | undefined { - const cp = this._checkpoints.find(c => c.requestId === requestId); + const cp = this._findCheckpoint(requestId, stopId); if (!cp) { return undefined; } @@ -292,8 +310,8 @@ export class AgentHostEditingSession extends Disposable implements IChatEditingS }); } - async getSnapshotContents(requestId: string, uri: URI, _stopId: string | undefined): Promise { - const cp = this._checkpoints.find(c => c.requestId === requestId); + async getSnapshotContents(requestId: string, uri: URI, stopId: string | undefined): Promise { + const cp = this._findCheckpoint(requestId, stopId); if (!cp) { return undefined; } diff --git a/src/vs/workbench/contrib/chat/browser/agentSessions/agentHost/agentHostSessionHandler.ts b/src/vs/workbench/contrib/chat/browser/agentSessions/agentHost/agentHostSessionHandler.ts index 8cb4c25df5f..80dfe0d46d9 100644 --- a/src/vs/workbench/contrib/chat/browser/agentSessions/agentHost/agentHostSessionHandler.ts +++ b/src/vs/workbench/contrib/chat/browser/agentSessions/agentHost/agentHostSessionHandler.ts @@ -7,7 +7,7 @@ import { Throttler } from '../../../../../../base/common/async.js'; import { CancellationToken, CancellationTokenSource } from '../../../../../../base/common/cancellation.js'; import { Emitter, Event } from '../../../../../../base/common/event.js'; import { MarkdownString } from '../../../../../../base/common/htmlContent.js'; -import { Disposable, DisposableMap, DisposableResourceMap, DisposableStore, MutableDisposable, toDisposable, type IDisposable } from '../../../../../../base/common/lifecycle.js'; +import { Disposable, DisposableResourceMap, DisposableStore, MutableDisposable, toDisposable, type IDisposable } from '../../../../../../base/common/lifecycle.js'; import { ResourceMap } from '../../../../../../base/common/map.js'; import { observableValue } from '../../../../../../base/common/observable.js'; import { isEqual } from '../../../../../../base/common/resources.js'; @@ -168,7 +168,7 @@ export class AgentHostSessionHandler extends Disposable implements IChatSessionC /** Per-session subscription to chat model pending request changes. */ private readonly _pendingMessageSubscriptions = this._register(new DisposableResourceMap()); /** Per-session subscription watching for server-initiated turns. */ - private readonly _serverTurnWatchers = this._register(new DisposableMap()); + private readonly _serverTurnWatchers = this._register(new DisposableResourceMap()); /** Per-session writeFile listeners for agent host editing sessions. */ private readonly _editingSessionListeners = this._register(new DisposableResourceMap()); /** Historical turns with file edits, pending hydration into the editing session. */ @@ -497,7 +497,6 @@ export class AgentHostSessionHandler extends Disposable implements IChatSessionC * if applicable, and pipes turn progress through `progressObs`. */ private _watchForServerInitiatedTurns(backendSession: URI, sessionResource: URI): void { - const resourceKey = sessionResource.path.substring(1); const sessionStr = backendSession.toString(); // Seed from the current state so we don't treat any pre-existing active @@ -561,7 +560,7 @@ export class AgentHostSessionHandler extends Disposable implements IChatSessionC this._trackServerTurnProgress(backendSession, activeTurn.id, chatSession, sessionResource, turnStore); })); - this._serverTurnWatchers.set(resourceKey, disposables); + this._serverTurnWatchers.set(sessionResource, disposables); } /**