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

[DonJayamanne]

For #288457 (comment)

 getChatSessionInputState called everytime when sending a prompt to a session @mjbvz
Not sure why this happens
I have a session open, send the prompt 2+3, and getChatSessionInputState is invoked.
Send another prompt, and getChatSessionInputState is invoked again.
Not sure why, however from an API adoption point of view, this doesn't seem right.
We're not opening a new session, we're not changing anything.

Root Cause Analysis

Bug: getChatSessionInputState Called on Every Prompt Send

Summary

The extension-side getChatSessionInputState handler is invoked on every prompt send to a contributed chat session, even when the session is already open and no options have changed. This is redundant and potentially expensive — the extension handler may rebuild option groups, query remote state, etc.

Symptoms

1. Open a contributed chat session (e.g., Copilot)
2. Send a prompt (e.g., 2+3) → getChatSessionInputState is called ✓ (expected on first use)
3. Send another prompt (e.g., 4+5) → getChatSessionInputState is called again ✗ (unexpected)
4. Every subsequent prompt triggers the handler, even with no option changes

Root Cause

Call Chain

User sends prompt
    → mainThreadChatAgents2.ts: $invokeAgent()
        → includes chatSessionContext with initialSessionOptions
    → extHostChatAgents2.ts: $invokeAgent() [line 979]
        → if (context.chatSessionContext)  // TRUE for every contributed session request
            → this._chatSessions.getInputStateForSession(sessionResource, ...) [line 981]
                → extHostChatSessions.ts: getInputStateForSession() [line 858]
                    → controllerData.controller.getChatSessionInputState(...) [line 866]
                        → EXTENSION HANDLER CALLED (every time!)

Why It Happens

In extHostChatAgents2.ts, the $invokeAgent method builds a chatSessionContext object (lines 977-995) to pass to the agent's invoke() handler as part of ChatContext. To populate chatSessionContext.inputState, it calls getInputStateForSession() — which unconditionally calls the extension's getChatSessionInputState handler.

There is no caching of the resolved input state. Every prompt re-fetches it from the extension, even though:

• The session already exists and is open
• No options have been changed by the user
• The previous result would still be valid

Key Files

File	Line(s)	Role
src/vs/workbench/api/common/extHostChatAgents2.ts	977-995	Calls getInputStateForSession on every $invokeAgent
src/vs/workbench/api/common/extHostChatSessions.ts	858-876	getInputStateForSession() — unconditionally calls extension handler
src/vs/workbench/api/common/extHostChatSessions.ts	586-606	$provideChatSessionContent() — also calls handler on session open

Other Call Sites of getChatSessionInputState

The handler is also called in these contexts (which are appropriate):

1. Session open ($provideChatSessionContent, line 596) — needed to set initial state
2. New session item ($newChatSessionItem, line 1043) — needed for new sessions
3. Input state refresh ($provideChatSessionInputState, line 1091) — explicit refresh request

The problematic call is specifically the one in getInputStateForSession() triggered from $invokeAgent.

Solution: Per-Session Input State Cache

Approach

Add a ResourceMap<ChatSessionInputState> cache in ExtHostChatSessions keyed by session URI. Return cached results on subsequent prompt sends, and invalidate when the underlying state actually changes.

Cache Lifecycle

Session opens ($provideChatSessionContent)
    → resolve input state → CACHE IT

Prompt 1 ($invokeAgent → getInputStateForSession)
    → CACHE HIT → return immediately (no extension call)

Prompt 2 ($invokeAgent → getInputStateForSession)
    → CACHE HIT → return immediately (no extension call)

User changes option ($provideHandleOptionsChange)
    → INVALIDATE cache for this session

Prompt 3 ($invokeAgent → getInputStateForSession)
    → CACHE MISS → call extension handler → cache result

Session disposed ($disposeChatSessionContent / $releaseSession)
    → DELETE cache entry

Invalidation Points (6 total)

#	Event	Location	Action
1	Session content disposed	$disposeChatSessionContent	Delete entry for session URI
2	Session released (agent side)	$releaseSession in extHostChatAgents2.ts	Delete entry for session URI
3	Options changed from main thread	$provideHandleOptionsChange	Delete entry for session URI
4	Extension mutates inputState.groups	onChangedDelegate in createChatSessionInputState	Clear all entries for controller's scheme
5	Controller unregistered	Controller dispose callback	Clear all entries for controller's scheme
6	Provider options refreshed	$provideChatSessionProviderOptions	Clear all entries for provider's scheme

Why URI-Only Cache Key Is Safe

initialSessionOptions (passed alongside the session URI) is derived from chatSessionsService.getSessionOptions(sessionResource) — it's the current option state for that specific session. When a user changes an option, $provideHandleOptionsChange fires, which invalidates the cache. So the options are implicitly part of the cache validity — they don't need to be part of the key.

Changes Made

1. src/vs/workbench/api/common/extHostChatSessions.ts:

   • Added _inputStateCache: ResourceMap<ChatSessionInputState>
   • Added clearInputStateCache(resource) public method
   • Added _clearInputStateCacheForScheme(scheme) private helper
   • Modified getInputStateForSession() to check/populate cache
   • Added cache seeding in $provideChatSessionContent()
   • Added cache invalidation at 5 points (dispose, options change, group mutation, controller unregister, provider options refresh)

2. src/vs/workbench/api/common/extHostChatAgents2.ts:

   • Added clearInputStateCache call in $releaseSession()

Leaks - objects created by createChatSessionInputState never disposed

Fix: ChatSessionInputState disposal leak

Problem Statement

ChatSessionInputStateImpl objects are created but never disposed. Each instance holds an Emitter<void> (#onDidChangeEmitter) which is never cleaned up. Additionally, these objects accumulate in the inputStates: Set<ChatSessionInputStateImpl> on the controller data — they are added via inputStates.add(inputState) but never removed.

Root Cause Analysis

Where ChatSessionInputStateImpl is created

There are two categories of creation:

1. "Managed" instances — created via controller.createChatSessionInputState() (line 516-548 in extHostChatSessions.ts). These are:

   • Registered in inputStates set (line 546)
   • Given an onChangedDelegate callback that fires $updateChatSessionInputState over the proxy
   • Created by the extension's getChatSessionInputState handler

2. "Unmanaged" instances — created via new ChatSessionInputStateImpl(groups) directly (lines 425, 857, 871, 1101). These are:

   • Created as fallback/simple input states (no onChangedDelegate)
   • Not tracked in inputStates
   • Used for legacy provider paths and _createInputStateFromOptions

Lifecycle of managed instances

getChatSessionInputState called (by extension handler)
  → extension calls controller.createChatSessionInputState(groups)
    → new ChatSessionInputStateImpl(groups, onChangedDelegate)
    → inputStates.add(inputState)                        ← ADD
    → returned to extension, then to VS Code
      
VS Code holds it in _inputStateCache (ResourceMap)
Extension holds it in newInputStates (WeakRef[])
Extension subscribes to onDidChange (event listener leak)

Session disposed / new session created:
  → _inputStateCache entry cleared (cache invalidation)
  → BUT inputStates.delete(inputState) NEVER happens      ← LEAK
  → #onDidChangeEmitter never disposed                     ← LEAK
  → onDidChange listener from extension never removed      ← LEAK

Multiple leak vectors

1. inputStates set grows unbounded: Every call to createChatSessionInputState adds to the set; nothing removes from it. The set is per-controller and lives forever. When $provideHandleOptionsChange fires, it iterates ALL inputStates (line 709) and updates them — stale ones too.

2. #onDidChangeEmitter never disposed: ChatSessionInputStateImpl has a new Emitter<void>() but no dispose() method. The emitter leaks.

3. Extension-side onDidChange listener leaks: The Copilot extension subscribes state.onDidChange(...) (line 282) but the returned IDisposable is never stored or disposed.

4. _inputStateCache retains references: The cache is cleaned up in some paths but since the objects themselves are never disposed, any remaining reference keeps the emitter alive.

When are new input states created?

• Every time a new (untitled) chat session is opened (via $provideChatSessionContent)
• Every time a new session item is created (via $newChatSessionItem)
• Every time $provideChatSessionInputState is called from the main thread
• On cache miss in getInputStateForSession (every prompt send, if not cached)

So over the lifetime of a VS Code window, if a user opens multiple chat sessions, each creates a new ChatSessionInputStateImpl that is never cleaned up.

Proposed Solution

1. Make ChatSessionInputStateImpl disposable

Add a dispose() method that:

• Disposes #onDidChangeEmitter
• Clears internal state

2. Remove from inputStates set when no longer needed

When a session's input state is replaced (e.g., new getChatSessionInputState call for the same session) or the session is disposed, remove the old input state from inputStates and dispose it.

3. Specific changes

A. ChatSessionInputStateImpl — add dispose

class ChatSessionInputStateImpl implements vscode.ChatSessionInputState {
    // ... existing ...
    dispose(): void {
        this.#onDidChangeEmitter.dispose();
    }
}

B. _inputStateCache — dispose evicted entries

When an entry is removed from _inputStateCache, if the evicted value is a ChatSessionInputStateImpl that is "unmanaged" (not in the inputStates set), dispose it. For managed ones, the lifecycle is tied to the inputStates set.

Actually, looking more carefully: the cache stores whatever comes back from getChatSessionInputState, which could be a managed instance. The managed instances are also in inputStates. So:

• The _inputStateCache should NOT own disposal — it's just a lookup cache.
• The inputStates set is the owner of managed instances.
• Unmanaged instances (created via _createInputStateFromOptions) are not tracked anywhere and should either be tracked or made lightweight (no emitter).

C. inputStates set — remove and dispose when session completes

When $disposeChatSessionContent is called or when the controller is disposed, clear and dispose all inputStates:

• In controller dispose: iterate inputStates, call dispose() on each, then clear().
• When a new input state replaces an old one for the same logical session: remove the old from the set and dispose it.

D. Track association between session URI and its input state

Currently there's no mapping from session URI → its inputState in the inputStates set. The _inputStateCache serves this purpose partially. We need to ensure that when a session's input state is replaced or the session is disposed, the corresponding entry in inputStates is removed and disposed.

Approach: Add a ResourceMap<ChatSessionInputStateImpl> to track the active input state per session. When a new one is created for the same session, dispose the old one. On session dispose, look up and dispose.

Alternatively, since the _inputStateCache already maps URI → inputState, we can dispose the old value when setting a new one, and dispose on delete. But we need to be careful: the _inputStateCache stores the interface type, not the impl.

Simplest correct approach

1. Make ChatSessionInputStateImpl disposable (dispose the emitter)
2. When removing entries from inputStates set, dispose them:
   • Controller dispose: dispose all inputStates
   • When creating a new input state via createChatSessionInputState: no change (we don't know which session it's for yet)
3. When removing entries from _inputStateCache: dispose the removed value if it's a ChatSessionInputStateImpl
4. In $disposeChatSessionContent: already clears cache → will now dispose
5. In clearInputStateCache: already clears cache → will now dispose

The key insight is that the _inputStateCache is the right place to manage the lifecycle of per-session input states. When a cache entry is evicted (replaced or deleted), the old value should be disposed AND removed from inputStates.

Files to Change

1. src/vs/workbench/api/common/extHostChatSessions.ts

   • Add dispose() to ChatSessionInputStateImpl
   • When deleting from _inputStateCache, also dispose the old value and remove from inputStates
   • When controller is disposed, dispose all its inputStates
   • When setting a new value in _inputStateCache, dispose the old value

2. src/vs/workbench/api/test/browser/mainThreadChatSessions.test.ts (or equivalent)

   • Add test verifying that input states are disposed when sessions are disposed

Todos

• make-disposable: Make ChatSessionInputStateImpl implement dispose pattern
• dispose-on-cache-evict: Dispose old input states when evicted from cache
• dispose-on-controller-teardown: Dispose all inputStates when controller is disposed
• remove-from-inputstates: Remove disposed input states from the inputStates set
• add-tests: Add tests for proper disposal
• compile-check: Verify TypeScript compilation

[Copilot]

Pull request overview

This PR adds extension-host-side caching for ChatSessionInputState keyed by session resource URIs, aiming to keep chat session input state stable across invocations and reduce repeated controller calls.

Changes:

• Added a per-session ResourceMap cache for resolved ChatSessionInputState objects in ExtHostChatSessions.
• Added cache seeding and invalidation/cleanup paths (session dispose, controller unregister, provider option-group changes).
• Cleared cached input state when a chat session is released from ExtHostChatAgents2.

Show a summary per file

File	Description
src/vs/workbench/api/common/extHostChatSessions.ts	Introduces _inputStateCache, cache seeding, and cache invalidation APIs for chat session input state resolution.
src/vs/workbench/api/common/extHostChatAgents2.ts	Clears input-state cache on session release to avoid retaining stale state.

Copilot's findings

Comments suppressed due to low confidence (1)

src/vs/workbench/api/common/extHostChatSessions.ts:909

• The input-state cache can return stale option selections after session option updates. $provideHandleOptionsChange updates existing ChatSessionInputStateImpl instances via _setGroups(...) + _fireDidChange(), which bypasses the groups setter and therefore bypasses the new cache invalidation logic. With caching enabled, subsequent getInputStateForSession(...) calls can keep returning an outdated cached state instead of recomputing/refreshing based on the updated options. Consider invalidating the cache for the affected sessionResource (or at least the scheme) when handling option changes, or ensure option-change propagation updates the cached state as well.

		if (sessionResource) {
			const cached = this._inputStateCache.get(sessionResource);
			if (cached) {
				return cached;
			}
		}

• Files reviewed: 2/2 changed files
• Comments generated: 3

[Copilot]

$provideChatSessionContent seeds _inputStateCache before awaiting provider.provider.provideChatSessionContent(...). If that provider call throws or the token is cancelled before the session is fully created, the cache entry for sessionResource will remain and can retain references unnecessarily. Consider seeding the cache only after provideChatSessionContent succeeds, or deleting the cache entry in an error/cancellation path.

[Copilot]

When inputState.groups is set, this callback updates entry.optionGroups (controller-wide) but only invalidates cache entries whose cached value is the same inputState object. If multiple session resources for the same scheme have cached input states, they can become inconsistent with the new optionGroups/defaults. Consider clearing cached input states for the whole scheme (or otherwise ensuring all affected cached states are refreshed) when controller-level option groups change.

This issue also appears on line 904 of the same file.

Suggested change

	// 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);
	}
	}
	// Option groups are controller-wide, so clear cached input states to force refresh.
	this._inputStateCache.clear();

[Copilot]

New caching behavior in getInputStateForSession (cache hits, scheme-wide invalidation, and seeding on session open) is not covered by existing tests. There is already an ExtHostChatSessions test suite (src/vs/workbench/api/test/browser/mainThreadChatSessions.test.ts), so it would be good to add regression tests that (1) verify cached values are reused to avoid repeated controller calls and (2) verify cache invalidation when session options/input state change.

[github-actions]

Screenshot Changes

Base: 1dfb3b9d Current: 8b4331f4

Changed (2)

agentSessionsViewer/WithBadge/Dark

Before	After

agentSessionsViewer/WithBadge/Light

Before	After