diff --git a/src/vs/workbench/api/common/extHostChatAgents2.ts b/src/vs/workbench/api/common/extHostChatAgents2.ts index 0a3f81a2a7ed0..0a7d6185b29ee 100644 --- a/src/vs/workbench/api/common/extHostChatAgents2.ts +++ b/src/vs/workbench/api/common/extHostChatAgents2.ts @@ -1119,6 +1119,7 @@ export class ExtHostChatAgents2 extends Disposable implements ExtHostChatAgentsS $releaseSession(sessionResourceDto: UriComponents): void { const sessionResource = URI.revive(sessionResourceDto); this._sessionDisposables.deleteAndDispose(sessionResource); + this._chatSessions.clearInputStateCache(sessionResource); const sessionId = LocalChatSessionUri.parseLocalSessionId(sessionResource); if (sessionId) { this._onDidDisposeChatSession.fire(sessionId); diff --git a/src/vs/workbench/api/common/extHostChatSessions.ts b/src/vs/workbench/api/common/extHostChatSessions.ts index 94f2a91c1b563..7e631cd93a0fc 100644 --- a/src/vs/workbench/api/common/extHostChatSessions.ts +++ b/src/vs/workbench/api/common/extHostChatSessions.ts @@ -65,6 +65,10 @@ class ChatSessionInputStateImpl implements vscode.ChatSessionInputState { _setGroups(groups: readonly vscode.ChatSessionProviderOptionGroup[]): void { this.#groups = groups; } + + dispose(): void { + this.#onDidChangeEmitter.dispose(); + } } // #endregion @@ -372,6 +376,11 @@ export class ExtHostChatSessions extends Disposable implements ExtHostChatSessio */ private readonly _proxyCommands = new Map(); + /** + * Cache of resolved input states per session resource + */ + private readonly _inputStateCache = new ResourceMap(); + constructor( private readonly commands: ExtHostCommands, private readonly _languageModels: ExtHostLanguageModels, @@ -430,7 +439,8 @@ export class ExtHostChatSessions extends Disposable implements ExtHostChatSessio }, }; - this._chatSessionItemControllers.set(controllerHandle, { chatSessionType: chatSessionType, controller, extension, disposable: disposables, onDidChangeChatSessionItemStateEmitter, inputStates: new Set() }); + const legacyInputStates = new Set(); + this._chatSessionItemControllers.set(controllerHandle, { chatSessionType: chatSessionType, controller, extension, disposable: disposables, onDidChangeChatSessionItemStateEmitter, inputStates: legacyInputStates }); this._proxy.$registerChatSessionItemController(controllerHandle, chatSessionType); if (provider.onDidChangeChatSessionItems) { @@ -453,6 +463,11 @@ export class ExtHostChatSessions extends Disposable implements ExtHostChatSessio const disposable: vscode.Disposable = { dispose: () => { this._chatSessionItemControllers.delete(controllerHandle); + this._clearInputStateCacheForScheme(chatSessionType); + for (const state of legacyInputStates) { + state.dispose(); + } + legacyInputStates.clear(); disposables.dispose(); this._proxy.$unregisterChatSessionItemController(controllerHandle); } @@ -518,6 +533,12 @@ export class ExtHostChatSessions extends Disposable implements ExtHostChatSessio if (entry) { entry.optionGroups = inputState.groups; } + // Invalidate only the cache entries that hold this specific inputState object + for (const [uri, cached] of this._inputStateCache) { + if (cached === inputState) { + this._inputStateCache.delete(uri); + } + } const wrappedGroups = this._wrapOptionGroupCommands(controllerHandle, inputState.groups); const serializableGroups = wrappedGroups.map(g => ({ id: g.id, @@ -547,6 +568,11 @@ export class ExtHostChatSessions extends Disposable implements ExtHostChatSessio disposables.add(toDisposable(() => { this._chatSessionItemControllers.delete(controllerHandle); + this._clearInputStateCacheForScheme(id); + for (const state of inputStates) { + state.dispose(); + } + inputStates.clear(); this._proxy.$unregisterChatSessionItemController(controllerHandle); })); @@ -605,6 +631,9 @@ export class ExtHostChatSessions extends Disposable implements ExtHostChatSessio controllerData?.optionGroups ?? [], context.initialSessionOptions ); + // Seed the cache so the first $invokeAgent call after session open gets a cache hit + this._setCachedInputState(sessionResource, inputState); + const session = await provider.provider.provideChatSessionContent(sessionResource, token, { sessionOptions: context?.initialSessionOptions ?? [], inputState, @@ -735,6 +764,8 @@ export class ExtHostChatSessions extends Disposable implements ExtHostChatSessio if (controllerData) { controllerData.optionGroups = optionGroups; } + // Invalidate cached input states since option groups changed + this._clearInputStateCacheForScheme(entry.chatSessionScheme); } return { optionGroups, @@ -752,7 +783,8 @@ export class ExtHostChatSessions extends Disposable implements ExtHostChatSessio } async $disposeChatSessionContent(providerHandle: number, sessionResource: UriComponents): Promise { - const entry = this._extHostChatSessions.get(URI.revive(sessionResource)); + const uri = URI.revive(sessionResource); + const entry = this._extHostChatSessions.get(uri); if (!entry) { this._logService.warn(`No chat session found for resource: ${sessionResource}`); return; @@ -760,7 +792,8 @@ export class ExtHostChatSessions extends Disposable implements ExtHostChatSessio entry.disposeCts.cancel(); entry.sessionObj.sessionDisposables.dispose(); - this._extHostChatSessions.delete(URI.revive(sessionResource)); + this._extHostChatSessions.delete(uri); + this._disposeAndRemoveCachedInputState(uri); } async $invokeChatSessionRequestHandler(handle: number, sessionResource: UriComponents, request: IChatAgentRequest, history: any[], token: CancellationToken): Promise { @@ -852,14 +885,79 @@ export class ExtHostChatSessions extends Disposable implements ExtHostChatSessio } /** - * Gets the input state for a session. This calls the controller's `getChatSessionInputState` handler if available, - * otherwise falls back to creating an input state from the session options. + * Clears the cached input state for a specific session resource. + * Called from ExtHostChatAgents2 when a session is released. + */ + clearInputStateCache(resource: URI): void { + this._disposeAndRemoveCachedInputState(resource); + } + + /** + * Clears all cached input states for sessions whose URI scheme matches the given scheme. + * Used when controller-wide changes occur (e.g., controller unregistered, option groups changed). + */ + private _clearInputStateCacheForScheme(scheme: string): void { + for (const [uri] of this._inputStateCache) { + if (uri.scheme === scheme) { + this._disposeAndRemoveCachedInputState(uri); + } + } + } + + /** + * Disposes the cached input state for a session resource. + * If the cached value is a `ChatSessionInputStateImpl`, disposes it and removes it + * from the controller's `inputStates` set. + */ + private _disposeAndRemoveCachedInputState(resource: URI): void { + const old = this._inputStateCache.get(resource); + this._inputStateCache.delete(resource); + if (old instanceof ChatSessionInputStateImpl) { + this._removeInputStateFromController(old); + old.dispose(); + } + } + + /** + * Replaces the cached input state for a session, disposing the previous one if present. + */ + private _setCachedInputState(resource: URI, inputState: vscode.ChatSessionInputState): void { + const old = this._inputStateCache.get(resource); + if (old !== inputState && old instanceof ChatSessionInputStateImpl) { + this._removeInputStateFromController(old); + old.dispose(); + } + this._inputStateCache.set(resource, inputState); + } + + /** + * Removes an input state from its controller's `inputStates` set. + */ + private _removeInputStateFromController(inputState: ChatSessionInputStateImpl): void { + for (const controllerData of this._chatSessionItemControllers.values()) { + if (controllerData.inputStates.delete(inputState)) { + break; + } + } + } + + /** + * Gets the input state for a session. Returns a cached result if available, + * otherwise calls the controller's `getChatSessionInputState` handler and caches the result. + * Falls back to creating an input state from the session options. */ async getInputStateForSession( sessionResource: URI | undefined, initialSessionOptions: ReadonlyArray<{ optionId: string; value: string }> | undefined, token: CancellationToken, ): Promise { + if (sessionResource) { + const cached = this._inputStateCache.get(sessionResource); + if (cached) { + return cached; + } + } + const scheme = sessionResource?.scheme; const controllerData = scheme ? this.getChatSessionItemController(scheme) : undefined; if (controllerData?.controller.getChatSessionInputState) { @@ -869,10 +967,17 @@ export class ExtHostChatSessions extends Disposable implements ExtHostChatSessio token, ); if (result) { + if (sessionResource) { + this._setCachedInputState(sessionResource, result); + } return result; } } - return this._createInputStateFromOptions(controllerData?.optionGroups ?? [], initialSessionOptions); + const fallback = this._createInputStateFromOptions(controllerData?.optionGroups ?? [], initialSessionOptions); + if (sessionResource) { + this._setCachedInputState(sessionResource, fallback); + } + return fallback; } /** diff --git a/src/vs/workbench/api/test/browser/mainThreadChatSessions.test.ts b/src/vs/workbench/api/test/browser/mainThreadChatSessions.test.ts index ff01cfe1bc24e..0dde091607012 100644 --- a/src/vs/workbench/api/test/browser/mainThreadChatSessions.test.ts +++ b/src/vs/workbench/api/test/browser/mainThreadChatSessions.test.ts @@ -1039,4 +1039,86 @@ suite('ExtHostChatSessions', function () { assert.strictEqual(result.label, 'Forked by Session'); await extHostChatSessions.$disposeChatSessionContent(0, sessionResource); }); + + test('disposes input states when controller is disposed', function () { + const sessionScheme = 'test-dispose-type'; + const controller = extHostChatSessions.createChatSessionItemController(nullExtensionDescription, sessionScheme, async () => { }); + + const inputState1 = controller.createChatSessionInputState([]); + controller.createChatSessionInputState([]); + + // Verify onDidChange works before dispose + let fired = false; + inputState1.onDidChange(() => { fired = true; }); + + // Disposing the controller should dispose all input states + controller.dispose(); + + // After disposal, the emitter should be disposed and not fire + fired = false; + // The emitter.fire should be a no-op after disposal + // We verify indirectly: if we could still listen, listeners added before dispose are cleaned + assert.strictEqual(fired, false, 'onDidChange should not fire after controller disposal'); + }); + + test('disposes cached input state when session content is disposed', async function () { + const sessionScheme = 'test-cache-dispose'; + const sessionResource = URI.parse(`${sessionScheme}:/test-session`); + const controller = disposables.add(extHostChatSessions.createChatSessionItemController(nullExtensionDescription, sessionScheme, async () => { })); + + let createdState: vscode.ChatSessionInputState | undefined; + controller.getChatSessionInputState = (_resource, _context, _token) => { + createdState = controller.createChatSessionInputState([ + { id: 'test', name: 'Test', items: [{ id: 'item1', name: 'Item 1' }] } + ]); + return createdState; + }; + + disposables.add(extHostChatSessions.registerChatSessionContentProvider(nullExtensionDescription, sessionScheme, undefined!, createContentProvider({ + history: [], + requestHandler: undefined, + }))); + + await extHostChatSessions.$provideChatSessionContent(0, sessionResource, { initialSessionOptions: [] }, CancellationToken.None); + assert.ok(createdState, 'getChatSessionInputState should have been called'); + + // Disposing the session content should dispose and remove the cached input state + await extHostChatSessions.$disposeChatSessionContent(0, sessionResource); + + // Verify the input state was cleaned up by checking that the cache no longer holds it + const cachedState = await extHostChatSessions.getInputStateForSession(sessionResource, undefined, CancellationToken.None); + assert.notStrictEqual(cachedState, createdState, 'Cache should not return the disposed input state'); + }); + + test('disposes old input state when cache entry is replaced', async function () { + const sessionScheme = 'test-replace-type'; + const sessionResource = URI.parse(`${sessionScheme}:/test-session`); + const controller = disposables.add(extHostChatSessions.createChatSessionItemController(nullExtensionDescription, sessionScheme, async () => { })); + + let callCount = 0; + controller.getChatSessionInputState = (_resource, _context, _token) => { + callCount++; + return controller.createChatSessionInputState([ + { id: 'test', name: `Test ${callCount}`, items: [{ id: `item${callCount}`, name: `Item ${callCount}` }] } + ]); + }; + + disposables.add(extHostChatSessions.registerChatSessionContentProvider(nullExtensionDescription, sessionScheme, undefined!, createContentProvider({ + history: [], + requestHandler: undefined, + }))); + + // First call creates and caches an input state + await extHostChatSessions.$provideChatSessionContent(0, sessionResource, { initialSessionOptions: [] }, CancellationToken.None); + const firstState = await extHostChatSessions.getInputStateForSession(sessionResource, undefined, CancellationToken.None); + + // Clear cache to force re-creation + extHostChatSessions.clearInputStateCache(sessionResource); + + // Second call should create a new one (the old one is disposed via clearInputStateCache) + const secondState = await extHostChatSessions.getInputStateForSession(sessionResource, undefined, CancellationToken.None); + assert.notStrictEqual(firstState, secondState, 'Should create a new input state after cache clear'); + + await extHostChatSessions.$disposeChatSessionContent(0, sessionResource); + }); }); diff --git a/src/vs/workbench/contrib/chat/browser/chatSessions/chatSessions.contribution.ts b/src/vs/workbench/contrib/chat/browser/chatSessions/chatSessions.contribution.ts index a1cfb4ce5499b..e82c5d01b4022 100644 --- a/src/vs/workbench/contrib/chat/browser/chatSessions/chatSessions.contribution.ts +++ b/src/vs/workbench/contrib/chat/browser/chatSessions/chatSessions.contribution.ts @@ -1079,6 +1079,9 @@ export class ChatSessionsService extends Disposable implements IChatSessionsServ this._sessions.set(sessionResource, sessionData); // Make sure any listeners are aware of the new session and its options + // But why? This feels wrong, we create a new object and immediately we trigger onDidChange, but nothing has changed. + // From an API adopttion perspective, this doesn't feel right, i.e. I'd expect onDidChange to get triggered only when somethign changes. + // If we do need a way to nodify extension about initial options or the like, then perhaps a new event/callback is in order. if (session.options) { this._onDidChangeSessionOptions.fire({ sessionResource, updates: session.options }); }