Optimize ChatSessionController.replace

Follow up on #294798

Do a simple diff instead of re-transmitting everything
This commit is contained in:
Matt Bierner
2026-02-12 01:36:45 -08:00
parent 37bde10eee
commit ac73dc720b
3 changed files with 100 additions and 84 deletions

View File

@@ -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<IChatSessionItem>[]): Promise<void> {
async $updateChatSessionItems(controllerHandle: number, change: IChatSessionItemsChange): Promise<void> {
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<IChatSessionItem>): Promise<void> {

View File

@@ -3414,13 +3414,18 @@ export interface IChatSessionProviderOptions {
optionGroups?: IChatSessionProviderOptionGroup[];
}
export interface IChatSessionItemsChange {
readonly addedOrUpdated: readonly Dto<IChatSessionItem>[];
readonly removed: readonly UriComponents[];
}
export interface MainThreadChatSessionsShape extends IDisposable {
$registerChatSessionItemController(handle: number, chatSessionType: string): void;
$unregisterChatSessionItemController(handle: number): void;
$setChatSessionItems(handle: number, items: Dto<IChatSessionItem>[]): Promise<void>;
$addOrUpdateChatSessionItem(handle: number, item: Dto<IChatSessionItem>): Promise<void>;
$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<void>;
$addOrUpdateChatSessionItem(controllerHandle: number, item: Dto<IChatSessionItem>): Promise<void>;
$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<ChatSessionOptionUpdateDto2>): void;

View File

@@ -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<vscode.ChatSessionItem>;
readonly removed?: ResourceSet;
}
function computeItemsDelta(oldItems: ResourceMap<vscode.ChatSessionItem>, newItems: ResourceMap<vscode.ChatSessionItem>): ChatSessionDelta {
const delta = {
addedOrUpdated: new ResourceMap<vscode.ChatSessionItem>(),
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<typeof typeConvert.ChatSessionItem.from>[]; 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<vscode.ChatSessionItem>();
readonly #callbacks: SessionCollectionListeners;
#items = new ResourceMap<vscode.ChatSessionItem>();
readonly #proxy: Proxied<MainThreadChatSessionsShape>;
readonly #controllerHandle: number;
constructor(callbacks: SessionCollectionListeners) {
this.#callbacks = callbacks;
constructor(controllerHandle: number, proxy: Proxied<MainThreadChatSessionsShape>) {
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<vscode.ChatSessionItem>();
/**
* 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<vscode.ChatSessionItem>());
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<Dto<IChatSessionItem>> = [];
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<vscode.ChatSessionItem>());
const onItemsChanged = () => {
const items: Array<Dto<IChatSessionItem>> = [];
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<vscode.ChatSessionItemController>({
id,