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

@@ -372,6 +372,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,
@@ -453,6 +458,7 @@ export class ExtHostChatSessions extends Disposable implements ExtHostChatSessio
 		const disposable: vscode.Disposable = {
 			dispose: () => {
 				this._chatSessionItemControllers.delete(controllerHandle);
+				this._clearInputStateCacheForScheme(chatSessionType);
 				disposables.dispose();
 				this._proxy.$unregisterChatSessionItemController(controllerHandle);
 			}
@@ -518,6 +524,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 +559,7 @@ export class ExtHostChatSessions extends Disposable implements ExtHostChatSessio
 
 		disposables.add(toDisposable(() => {
 			this._chatSessionItemControllers.delete(controllerHandle);
+			this._clearInputStateCacheForScheme(id);
 			this._proxy.$unregisterChatSessionItemController(controllerHandle);
 		}));
 
@@ -605,6 +618,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._inputStateCache.set(sessionResource, inputState);
+
 		const session = await provider.provider.provideChatSessionContent(sessionResource, token, {
 			sessionOptions: context?.initialSessionOptions ?? [],
 			inputState,
@@ -735,6 +751,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 +770,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._inputStateCache.delete(uri);
 	}
 
 	async $invokeChatSessionRequestHandler(handle: number, sessionResource: UriComponents, request: IChatAgentRequest, history: any[], token: CancellationToken): Promise<IChatAgentResult> {
@@ -852,14 +872,42 @@ 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._inputStateCache.delete(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._inputStateCache.delete(uri);
+			}
+		}
+	}
+
+	/**
+	 * 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 +917,17 @@ export class ExtHostChatSessions extends Disposable implements ExtHostChatSessio
 				token,
 			);
 			if (result) {
+				if (sessionResource) {
+					this._inputStateCache.set(sessionResource, result);
+				}
 				return result;
 			}
 		}
-		return this._createInputStateFromOptions(controllerData?.optionGroups ?? [], initialSessionOptions);
+		const fallback = this._createInputStateFromOptions(controllerData?.optionGroups ?? [], initialSessionOptions);
+		if (sessionResource) {
+			this._inputStateCache.set(sessionResource, fallback);
+		}
+		return fallback;
 	}
 
 	/**