diff --git a/src/vs/workbench/api/browser/mainThreadChatSessions.ts b/src/vs/workbench/api/browser/mainThreadChatSessions.ts index 47bff4e1d83..d2d31a6a5da 100644 --- a/src/vs/workbench/api/browser/mainThreadChatSessions.ts +++ b/src/vs/workbench/api/browser/mainThreadChatSessions.ts @@ -33,7 +33,7 @@ import { IEditorGroupsService } from '../../services/editor/common/editorGroupsS import { IEditorService } from '../../services/editor/common/editorService.js'; import { extHostNamedCustomer, IExtHostContext } from '../../services/extensions/common/extHostCustomers.js'; import { Dto } from '../../services/extensions/common/proxyIdentifier.js'; -import { ExtHostChatSessionsShape, ExtHostContext, IChatProgressDto, IChatSessionHistoryItemDto, MainContext, MainThreadChatSessionsShape } from '../common/extHost.protocol.js'; +import { ExtHostChatSessionsShape, ExtHostContext, IChatProgressDto, IChatSessionHistoryItemDto, IChatSessionItemsChange, MainContext, MainThreadChatSessionsShape } from '../common/extHost.protocol.js'; export class ObservableChatSession extends Disposable implements IChatSession { @@ -346,11 +346,13 @@ class MainThreadChatSessionItemController extends Disposable implements IChatSes return this._proxy.$refreshChatSessionItems(this._handle, token); } - setItems(items: readonly IChatSessionItem[]): void { - this._items.clear(); - for (const item of items) { + acceptChange(change: { readonly addedOrUpdated: readonly IChatSessionItem[]; readonly removed: readonly URI[] }): void { + for (const item of change.addedOrUpdated) { this._items.set(item.resource, item); } + for (const uri of change.removed) { + this._items.delete(uri); + } this._onDidChangeChatSessionItems.fire(); } @@ -479,10 +481,13 @@ export class MainThreadChatSessions extends Disposable implements MainThreadChat }; } - async $setChatSessionItems(controllerHandle: number, items: Dto[]): Promise { + async $updateChatSessionItems(controllerHandle: number, change: IChatSessionItemsChange): Promise { const controller = this.getController(controllerHandle); - const resolvedItems = await Promise.all(items.map(item => this._resolveSessionItem(item))); - controller.setItems(resolvedItems); + const resolvedItems = await Promise.all(change.addedOrUpdated.map(item => this._resolveSessionItem(item))); + controller.acceptChange({ + addedOrUpdated: resolvedItems, + removed: change.removed.map(uri => URI.revive(uri)) + }); } async $addOrUpdateChatSessionItem(controllerHandle: number, item: Dto): Promise { diff --git a/src/vs/workbench/api/common/extHost.protocol.ts b/src/vs/workbench/api/common/extHost.protocol.ts index d12f99c0051..566fcf1c240 100644 --- a/src/vs/workbench/api/common/extHost.protocol.ts +++ b/src/vs/workbench/api/common/extHost.protocol.ts @@ -3414,13 +3414,18 @@ export interface IChatSessionProviderOptions { optionGroups?: IChatSessionProviderOptionGroup[]; } +export interface IChatSessionItemsChange { + readonly addedOrUpdated: readonly Dto[]; + readonly removed: readonly UriComponents[]; +} + export interface MainThreadChatSessionsShape extends IDisposable { - $registerChatSessionItemController(handle: number, chatSessionType: string): void; - $unregisterChatSessionItemController(handle: number): void; - $setChatSessionItems(handle: number, items: Dto[]): Promise; - $addOrUpdateChatSessionItem(handle: number, item: Dto): Promise; - $onDidChangeChatSessionItems(handle: number): void; - $onDidCommitChatSessionItem(handle: number, original: UriComponents, modified: UriComponents): void; + $registerChatSessionItemController(controllerHandle: number, chatSessionType: string): void; + $unregisterChatSessionItemController(controllerHandle: number): void; + $updateChatSessionItems(controllerHandle: number, change: IChatSessionItemsChange): Promise; + $addOrUpdateChatSessionItem(controllerHandle: number, item: Dto): Promise; + $onDidChangeChatSessionItems(controllerHandle: number): void; + $onDidCommitChatSessionItem(controllerHandle: number, original: UriComponents, modified: UriComponents): void; $registerChatSessionContentProvider(handle: number, chatSessionScheme: string): void; $unregisterChatSessionContentProvider(handle: number): void; $onDidChangeChatSessionOptions(handle: number, sessionResource: UriComponents, updates: ReadonlyArray): void; diff --git a/src/vs/workbench/api/common/extHostChatSessions.ts b/src/vs/workbench/api/common/extHostChatSessions.ts index 8868ad66110..d867f092b9d 100644 --- a/src/vs/workbench/api/common/extHostChatSessions.ts +++ b/src/vs/workbench/api/common/extHostChatSessions.ts @@ -11,18 +11,19 @@ import { CancellationToken, CancellationTokenSource } from '../../../base/common import { CancellationError } from '../../../base/common/errors.js'; import { Emitter } from '../../../base/common/event.js'; import { Disposable, DisposableStore, toDisposable } from '../../../base/common/lifecycle.js'; -import { ResourceMap } from '../../../base/common/map.js'; +import { ResourceMap, ResourceSet } from '../../../base/common/map.js'; import { MarshalledId } from '../../../base/common/marshallingIds.js'; +import * as objects from '../../../base/common/objects.js'; import { basename } from '../../../base/common/resources.js'; import { URI, UriComponents } from '../../../base/common/uri.js'; import { SymbolKind, SymbolKinds } from '../../../editor/common/languages.js'; import { IExtensionDescription } from '../../../platform/extensions/common/extensions.js'; import { ILogService } from '../../../platform/log/common/log.js'; import { IChatRequestVariableEntry, IDiagnosticVariableEntryFilterData, IPromptFileVariableEntry, ISymbolVariableEntry, PromptFileVariableKind } from '../../contrib/chat/common/attachments/chatVariableEntries.js'; -import { IChatSessionItem, IChatSessionProviderOptionItem } from '../../contrib/chat/common/chatSessionsService.js'; +import { IChatSessionProviderOptionItem } from '../../contrib/chat/common/chatSessionsService.js'; import { ChatAgentLocation } from '../../contrib/chat/common/constants.js'; import { IChatAgentRequest, IChatAgentResult } from '../../contrib/chat/common/participants/chatAgents.js'; -import { Dto, Proxied } from '../../services/extensions/common/proxyIdentifier.js'; +import { Proxied } from '../../services/extensions/common/proxyIdentifier.js'; import { ChatSessionDto, ExtHostChatSessionsShape, IChatAgentProgressShape, IChatSessionProviderOptions, MainContext, MainThreadChatSessionsShape } from './extHost.protocol.js'; import { ChatAgentResponseStream } from './extHostChatAgents2.js'; import { CommandsConverter, ExtHostCommands } from './extHostCommands.js'; @@ -31,7 +32,6 @@ import { IExtHostRpcService } from './extHostRpcService.js'; import * as typeConvert from './extHostTypeConverters.js'; import { Diagnostic } from './extHostTypeConverters.js'; import * as extHostTypes from './extHostTypes.js'; -import * as objects from '../../../base/common/objects.js'; type ChatSessionTiming = vscode.ChatSessionItem['timing']; @@ -176,33 +176,70 @@ class ChatSessionItemImpl implements vscode.ChatSessionItem { } } -interface SessionCollectionListeners { - onItemsChanged(): void; - onItemAddedOrUpdated(item: vscode.ChatSessionItem): void; +interface ChatSessionDelta { + readonly addedOrUpdated?: ResourceMap; + readonly removed?: ResourceSet; +} + +function computeItemsDelta(oldItems: ResourceMap, newItems: ResourceMap): ChatSessionDelta { + const delta = { + addedOrUpdated: new ResourceMap(), + removed: new ResourceSet(), + } satisfies ChatSessionDelta; + + for (const [newResource, newItem] of newItems) { + const oldItem = oldItems.get(newResource); + if (oldItem !== newItem) { + delta.addedOrUpdated.set(newResource, newItem); + } + } + + for (const oldResource of oldItems.keys()) { + if (!newItems.has(oldResource)) { + delta.removed.add(oldResource); + } + } + + return delta; +} + +function convertChatSessionDeltaToDto(delta: ChatSessionDelta): { addedOrUpdated: ReturnType[]; removed: URI[] } { + return { + addedOrUpdated: delta.addedOrUpdated ? Array.from(delta.addedOrUpdated.values(), typeConvert.ChatSessionItem.from) : [], + removed: delta.removed ? Array.from(delta.removed.keys()) : [] + }; } class ChatSessionItemCollectionImpl implements vscode.ChatSessionItemCollection { - readonly #items = new ResourceMap(); - readonly #callbacks: SessionCollectionListeners; + #items = new ResourceMap(); + readonly #proxy: Proxied; + readonly #controllerHandle: number; - constructor(callbacks: SessionCollectionListeners) { - this.#callbacks = callbacks; + constructor(controllerHandle: number, proxy: Proxied) { + this.#proxy = proxy; + this.#controllerHandle = controllerHandle; } get size(): number { return this.#items.size; } - replace(items: readonly vscode.ChatSessionItem[]): void { - if (items.length === 0 && this.#items.size === 0) { + replace(newItems: readonly vscode.ChatSessionItem[]): void { + if (!newItems.length && !this.#items.size) { + // No change return; } - this.#items.clear(); - for (const item of items) { - this.#items.set(item.resource, item); + const newItemsMap = new ResourceMap(newItems.map(item => [item.resource, item] as const)); + + const delta = computeItemsDelta(this.#items, newItemsMap); + if (!delta.addedOrUpdated?.size && !delta.removed?.size) { + // No change + return; } - this.#callbacks.onItemsChanged(); + + this.#items = newItemsMap; + void this.#proxy.$updateChatSessionItems(this.#controllerHandle, convertChatSessionDeltaToDto(delta)); } forEach(callback: (item: vscode.ChatSessionItem, collection: vscode.ChatSessionItemCollection) => unknown, thisArg?: any): void { @@ -219,12 +256,15 @@ class ChatSessionItemCollectionImpl implements vscode.ChatSessionItemCollection } this.#items.set(item.resource, item); - this.#callbacks.onItemAddedOrUpdated(item); + void this.#proxy.$addOrUpdateChatSessionItem(this.#controllerHandle, typeConvert.ChatSessionItem.from(item)); } delete(resource: vscode.Uri): void { if (this.#items.delete(resource)) { - this.#callbacks.onItemsChanged(); + void this.#proxy.$updateChatSessionItems(this.#controllerHandle, { + addedOrUpdated: [], + removed: [resource] + }); } } @@ -286,13 +326,6 @@ export class ExtHostChatSessions extends Disposable implements ExtHostChatSessio readonly disposable: DisposableStore; }>(); - /** - * Map of uri -> chat session items - * - * TODO: this isn't cleared/updated properly - */ - private readonly _sessionItems = new ResourceMap(); - /** * Map of uri -> chat sessions infos */ @@ -314,14 +347,16 @@ export class ExtHostChatSessions extends Disposable implements ExtHostChatSessio commands.registerArgumentProcessor({ processArgument: (arg) => { if (arg && arg.$mid === MarshalledId.AgentSessionContext) { - const id = arg.session.resource || arg.sessionId; - const sessionContent = this._sessionItems.get(id); - if (sessionContent) { - return sessionContent; - } else { - this._logService.warn(`No chat session found for ID: ${id}`); - return arg; + const resource = arg.session.resource; + for (const { controller } of this._chatSessionItemControllers.values()) { + const item = controller.items.get(resource); + if (item) { + return item; + } } + + this._logService.warn(`No chat session found with uri: ${resource}`); + return arg; } return arg; @@ -331,27 +366,12 @@ export class ExtHostChatSessions extends Disposable implements ExtHostChatSessio registerChatSessionItemProvider(extension: IExtensionDescription, chatSessionType: string, provider: vscode.ChatSessionItemProvider): vscode.Disposable { // The legacy provider api is implemented using the new controller API on the backend - const handle = this._itemControllerHandlePool++; + const controllerHandle = this._itemControllerHandlePool++; const disposables = new DisposableStore(); const onDidChangeChatSessionItemStateEmitter = disposables.add(new Emitter()); - const collection = new ChatSessionItemCollectionImpl({ - // Noop for providers - onItemsChanged: () => { }, - onItemAddedOrUpdated: () => { } - }); - - // Helper to push items to main thread - const updateItems = async (items: readonly vscode.ChatSessionItem[]) => { - collection.replace(items); - const convertedItems: Array> = []; - for (const sessionContent of items) { - this._sessionItems.set(sessionContent.resource, sessionContent); - convertedItems.push(typeConvert.ChatSessionItem.from(sessionContent)); - } - void this._proxy.$setChatSessionItems(handle, convertedItems); - }; + const collection = new ChatSessionItemCollectionImpl(controllerHandle, this._proxy); const controller: vscode.ChatSessionItemController = { id: chatSessionType, @@ -365,12 +385,12 @@ export class ExtHostChatSessions extends Disposable implements ExtHostChatSessio }, refreshHandler: async (token: vscode.CancellationToken) => { const items = await provider.provideChatSessionItems(token) ?? []; - updateItems(items); + collection.replace(items); }, }; - this._chatSessionItemControllers.set(handle, { chatSessionType: chatSessionType, controller, extension, disposable: disposables, onDidChangeChatSessionItemStateEmitter }); - this._proxy.$registerChatSessionItemController(handle, chatSessionType); + this._chatSessionItemControllers.set(controllerHandle, { chatSessionType: chatSessionType, controller, extension, disposable: disposables, onDidChangeChatSessionItemStateEmitter }); + this._proxy.$registerChatSessionItemController(controllerHandle, chatSessionType); if (provider.onDidChangeChatSessionItems) { disposables.add(provider.onDidChangeChatSessionItems(() => { @@ -385,15 +405,15 @@ export class ExtHostChatSessions extends Disposable implements ExtHostChatSessio if (provider.onDidCommitChatSessionItem) { disposables.add(provider.onDidCommitChatSessionItem((e) => { const { original, modified } = e; - this._proxy.$onDidCommitChatSessionItem(handle, original.resource, modified.resource); + this._proxy.$onDidCommitChatSessionItem(controllerHandle, original.resource, modified.resource); })); } return { dispose: () => { - this._chatSessionItemControllers.delete(handle); + this._chatSessionItemControllers.delete(controllerHandle); disposables.dispose(); - this._proxy.$unregisterChatSessionItemController(handle); + this._proxy.$unregisterChatSessionItemController(controllerHandle); } }; } @@ -405,21 +425,7 @@ export class ExtHostChatSessions extends Disposable implements ExtHostChatSessio let isDisposed = false; const onDidChangeChatSessionItemStateEmitter = disposables.add(new Emitter()); - const onItemsChanged = () => { - const items: Array> = []; - for (const [_, item] of collection) { - this._sessionItems.set(item.resource, item); - items.push(typeConvert.ChatSessionItem.from(item)); - } - void this._proxy.$setChatSessionItems(controllerHandle, items); - }; - - const collection = new ChatSessionItemCollectionImpl({ - onItemsChanged, - onItemAddedOrUpdated: (item: vscode.ChatSessionItem) => { - void this._proxy.$addOrUpdateChatSessionItem(controllerHandle, typeConvert.ChatSessionItem.from(item)); - } - }); + const collection = new ChatSessionItemCollectionImpl(controllerHandle, this._proxy); const controller = Object.freeze({ id,