Fix: CancellationToken not propagating to provideLanguageModelChatResponse on chat stop

[tamuratak]

Close #291713 #277872

The stop button does not propagate cancellation to the provider's CancellationToken in provideLanguageModelChatRequest. This is caused by two issues: (1) the RPC cancel handler for $startChatRequest is deleted immediately because the promise resolves before the provider returns, and (2) request IDs differ across MainThread and ExtensionHost processes, so a single CTS cannot bridge the gap. Fix this by introducing local CancellationTokenSources at both the MainThread and ExtensionHost levels, connected via a new $cancelLanguageModelChatRequest RPC method, and ensure proper cleanup of listeners on disposal.

How rpcProtocol.ts cancellation works and why it breaks:

In rpcProtocol.ts, the _cancelInvokedHandlers map stores a cancel callback keyed by callId (the RPC request number). When a caller sends a cancel message, _receiveCancel looks up the callback and fires it. However, the callback is deleted in promise.then() as soon as the RPC call resolves:

this._cancelInvokedHandlers[callId] = cancel;
// ...
promise.then((r) => {
    delete this._cancelInvokedHandlers[callId]; // deleted on resolve
    // ...
});

$startChatRequest is intentionally designed to return immediately (before streaming completes) so the RPC handler stays alive long enough to forward streaming data. But the promise still resolves right away, which deletes _cancelInvokedHandlers[callId] before the user can press the stop button. After that, any incoming cancel message finds no handler and is silently ignored — this is the root cause of the bug.

Created with GitHub Copilot + MiMO-V2.5-Pro

CC: @connor4312 @lramos15

[Copilot]

Pull request overview

Note

Copilot was unable to run its full agentic suite in this review.

Adds an explicit cross-process cancellation channel for language model chat requests so cancellation can still be signalled after the RPC's built-in token cancel handler has been removed (e.g., once the request promise resolves but streaming continues).

Changes:

• Introduces a new $cancelLanguageModelChatRequest RPC on both MainThreadLanguageModelsShape and ExtHostLanguageModelsShape.
• In ext host and main thread, wraps the incoming CancellationToken in a local CancellationTokenSource keyed by requestId, and registers a cancel listener on the caller-side token to forward cancellation via the new RPC.
• Ensures the local CTS and cancel listeners are disposed on completion, error, and global dispose; reports a CancellationError when the provider finished but the token was cancelled.

Reviewed changes

Copilot reviewed 3 out of 3 changed files in this pull request and generated 3 comments.

File	Description
src/vs/workbench/api/common/extHost.protocol.ts	Adds $cancelLanguageModelChatRequest to both main-thread and ext-host language model shapes.
src/vs/workbench/api/browser/mainThreadLanguageModels.ts	Tracks a per-request CTS, forwards cancellation via the new RPC, and disposes CTS/listeners on completion/error/dispose.
src/vs/workbench/api/common/extHostLanguageModels.ts	Mirrors the same per-request CTS + cancel-listener pattern in the ext host side, and surfaces cancellation as a CancellationError when applicable.

[tamuratak]

The reviewer's concern is based on a misunderstanding. There are two separate $cancelLanguageModelChatRequest methods — one on MainThreadLanguageModelsShape (called from ExtHost) and one on ExtHostLanguageModelsShape (called from MainThread) — each operating on its own request ID namespace.

The call chain is:

1. ExtHost.sendRequest generates requestIdA, passes it to $tryStartChatRequest (MainThread)
2. MainThread.$tryStartChatRequest stores a CTS in _pendingCancelCTS keyed by requestIdA — the same ID it received as a parameter, not a different one
3. When cancelled, the ExtHost cancel listener calls this._proxy.$cancelLanguageModelChatRequest(requestIdA) — which looks up the correct CTS in _pendingCancelCTS by that same requestIdA ✓
4. The CTS token fires, which triggers the sendChatRequest callback's cancel listener, which calls this._proxy.$cancelLanguageModelChatRequest(requestIdM) on the ExtHost side, matching the ExtHost's _pendingCancelTokens by requestIdM ✓

The two ID namespaces are intentionally separate and correctly bridged. Each $cancelLanguageModelChatRequest looks up its own map using the same ID that was used to populate it.

Created with GitHub Copilot + MiMO-V2.5-Pro

[klierronny6-max]

Suggested change

	$cancelLanguageModelChatRequest(requestId: number): void {
	this._pendingCancelCTS.get(requestId)?.cancel();
	}
	$cancelLanguageModelChatRequest(requestId: number): void {
	this._pendingCancelCTS.get(requestId)?.cancel();
	}

[lramos15]

@tamuratak Thanks for the PR! Would you mind taking a look at the Copilot PR feedback and then from there we can do a human review as well