feat(chat): implement input state caching and management for chat sessions

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);

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</* proxyId */ string, { readonly originalCommandId: string; readonly controllerHandle: number }>();
 
+	/**
+	 * Cache of resolved input states per session resource
+	 */
+	private readonly _inputStateCache = new ResourceMap<vscode.ChatSessionInputState>();
+
 	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<ChatSessionInputStateImpl>();
+		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,15 +783,17 @@ export class ExtHostChatSessions extends Disposable implements ExtHostChatSessio
 	}
 
 	async $disposeChatSessionContent(providerHandle: number, sessionResource: UriComponents): Promise<void> {
-		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;
 		}
 
 		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<IChatAgentResult> {
@@ -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<vscode.ChatSessionInputState> {
+		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;
 	}
 
 	/**

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);
+	});
 });

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 });
 		}