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

[easyatm]

I'm also hitting this bug in my own Language Model Chat Provider extension — the CancellationToken passed to provideLanguageModelChatResponse never fires when the user clicks the stop button, making it impossible to cleanly cancel ongoing LLM requests.

This PR's analysis (the RPC cancel handler being deleted as soon as $startChatRequest resolves) matches exactly what I've observed. The fix looks correct and well-thought-out.

@lramos15 @roblourens
Could you prioritize the human review and get this merged? It's a real pain point for anyone building third-party LM providers.

[vritant24]

Thanks for the change, it makes sense but I think this can be changed a little to be more maintainable and robust. I also request that you add Unit tests.

If you need any more assistance, please let me know and I'm happy to help

[Copilot]

Pull request overview

Copilot reviewed 7 out of 7 changed files in this pull request and generated 1 comment.

[vritant24]

@tamuratak can you address the last copilot review? Everything else looks great at the moment

[Copilot]

Pull request overview

Copilot reviewed 7 out of 7 changed files in this pull request and generated 1 comment.

[vritant24]

@tamuratak Thanks for addressing the review. I ran the PR Checks but there seem to be some failing

[tamuratak]

Removed console.warn

[lramos15]

I find the abstraction created by CancellableRequestMap to be unnecessary. We've done this pattern of firing off an RPC call that resolves and then continues async work before so we should follow those patterns

Please see speech ($cancelSpeechToTextSession), notebooks ($cancelCells), and testing ($cancelExtensionTestRun) for examples.

I'd suggest:

1. Keep _pendingProgress as a plain Map (as it was before)
2. Use DisposableMap<number, CancellationTokenSource> for both directions of cancellation — this is exactly what speech does with this.sessions = new Map<number, CTS>()
3. Register the cancel listener inline when setting up the request, dispose it alongside the CTS on cleanup

[Copilot]

Pull request overview

Copilot reviewed 5 out of 5 changed files in this pull request and generated 1 comment.

[Copilot]

Pull request overview

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

[Copilot]

Pull request overview

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

[tamuratak]

done