agentHost: cancel late callbacks during abort

src/vs/platform/agentHost/common/pendingRequestRegistry.ts

@@ -6,6 +6,8 @@
 import { DeferredPromise } from '../../../base/common/async.js';
 import { CancellationError } from '../../../base/common/errors.js';
 
+type MetadataArgument<TMeta> = [TMeta] extends [void] ? [metadata?: TMeta] : [metadata: TMeta];
+
 /**
  * Registry of parked deferred promises keyed by string id. Used to
  * track request/response round-trips where a callback fires a signal
@@ -18,30 +20,31 @@ import { CancellationError } from '../../../base/common/errors.js';
  * miss its response and the awaited promise would deadlock — a
  * regression caught in Claude phase 7.
  */
-export class PendingRequestRegistry<T> {
-	private readonly _entries = new Map<string, DeferredPromise<T>>();
+export class PendingRequestRegistry<TResult, TMeta = void> {
+	private readonly _entries = new Map<string, { deferred: DeferredPromise<TResult>; metadata: TMeta }>();
 	/**
 	 * Results delivered via {@link respondOrBuffer} before any deferred was
 	 * parked under the same key. A subsequent {@link register} consumes the
 	 * buffered value and resolves immediately, tolerating a completion that
 	 * races ahead of the handler that awaits it.
 	 */
-	private readonly _earlyResults = new Map<string, T>();
+	private readonly _earlyResults = new Map<string, TResult>();
 
-	registerAndFire(key: string, fire: () => void): Promise<T> {
+	/** Atomically park a deferred and optional metadata, then invoke `fire` to prevent synchronous responses racing registration. */
+	registerAndFire(key: string, fire: () => void, ...metadata: MetadataArgument<TMeta>): Promise<TResult> {
 		if (this._earlyResults.has(key)) {
-			const buffered = this._earlyResults.get(key) as T;
+			const buffered = this._earlyResults.get(key) as TResult;
 			this._earlyResults.delete(key);
 			return Promise.resolve(buffered);
 		}
-		const deferred = new DeferredPromise<T>();
-		this._entries.set(key, deferred);
+		const deferred = new DeferredPromise<TResult>();
+		this._entries.set(key, { deferred, metadata: metadata[0] as TMeta });
 		fire();
 		return deferred.p;
 	}
 
 	/**
-	 * Park a deferred under `key` and return its promise. Use when there
+	 * Park a deferred and optional metadata under `key`, then return its promise. Use when there
 	 * is no synchronous responder to guard against — the request that
 	 * eventually feeds {@link respond} originates from a different code
 	 * path (e.g. an MCP handler invoked by the SDK whose completion
@@ -52,31 +55,60 @@ export class PendingRequestRegistry<T> {
 	 * a {@link CancellationError} so its awaiter unwinds instead of
 	 * leaking forever.
 	 */
-	register(key: string): Promise<T> {
+	register(key: string, ...metadata: MetadataArgument<TMeta>): Promise<TResult> {
 		if (this._earlyResults.has(key)) {
-			const buffered = this._earlyResults.get(key) as T;
+			const buffered = this._earlyResults.get(key) as TResult;
 			this._earlyResults.delete(key);
 			return Promise.resolve(buffered);
 		}
 		const existing = this._entries.get(key);
-		if (existing && !existing.isSettled) {
-			existing.error(new CancellationError());
+		if (existing && !existing.deferred.isSettled) {
+			existing.deferred.error(new CancellationError());
 		}
-		const deferred = new DeferredPromise<T>();
-		this._entries.set(key, deferred);
+		const deferred = new DeferredPromise<TResult>();
+		this._entries.set(key, { deferred, metadata: metadata[0] as TMeta });
 		return deferred.p;
 	}
 
-	respond(key: string, value: T): boolean {
-		const deferred = this._entries.get(key);
-		if (!deferred) {
+	/** Resolve the deferred parked under `key`, returning whether an entry was found. */
+	respond(key: string, value: TResult): boolean {
+		const entry = this._entries.get(key);
+		if (!entry) {
 			return false;
 		}
 		this._entries.delete(key);
-		deferred.complete(value);
+		entry.deferred.complete(value);
 		return true;
 	}
 
+	/** Return the metadata associated with the request parked under `key`. */
+	getMetadata(key: string): TMeta | undefined {
+		return this._entries.get(key)?.metadata;
+	}
+
+	/** Whether a request is currently parked under `key`. */
+	has(key: string): boolean {
+		return this._entries.has(key);
+	}
+
+	/** Iterate the keys and metadata of all currently parked requests. */
+	*entries(): IterableIterator<readonly [key: string, metadata: TMeta]> {
+		for (const [key, entry] of this._entries) {
+			yield [key, entry.metadata];
+		}
+	}
+
+	/** Resolve every parked request whose metadata matches `predicate`, returning the number resolved. */
+	respondWhere(predicate: (metadata: TMeta, key: string) => boolean, value: TResult): number {
+		let count = 0;
+		for (const [key, entry] of this._entries) {
+			if (predicate(entry.metadata, key) && this.respond(key, value)) {
+				count++;
+			}
+		}
+		return count;
+	}
+
 	/**
 	 * Like {@link respond}, but if no deferred is parked under `key`, buffer
 	 * the value so a subsequent {@link register} / {@link registerAndFire}
@@ -85,7 +117,7 @@ export class PendingRequestRegistry<T> {
 	 * Copilot client-tool round-trip, whose SDK handler and the workbench
 	 * completion race).
 	 */
-	respondOrBuffer(key: string, value: T): void {
+	respondOrBuffer(key: string, value: TResult): void {
 		if (!this.respond(key, value)) {
 			this._earlyResults.set(key, value);
 		}
@@ -104,10 +136,10 @@ export class PendingRequestRegistry<T> {
 	 * rather than an error. Use {@link rejectAll} when callers must observe
 	 * a thrown error instead (cancellation, dispose).
 	 */
-	denyAll(denyValue: T): void {
-		for (const [, deferred] of this._entries) {
-			if (!deferred.isSettled) {
-				deferred.complete(denyValue);
+	denyAll(denyValue: TResult): void {
+		for (const [, entry] of this._entries) {
+			if (!entry.deferred.isSettled) {
+				entry.deferred.complete(denyValue);
 			}
 		}
 		this._entries.clear();
@@ -125,9 +157,9 @@ export class PendingRequestRegistry<T> {
 	 * where the awaited consumer must observe an error to unwind.
 	 */
 	rejectAll(error: Error): void {
-		for (const [, deferred] of this._entries) {
-			if (!deferred.isSettled) {
-				deferred.error(error);
+		for (const [, entry] of this._entries) {
+			if (!entry.deferred.isSettled) {
+				entry.deferred.error(error);
 			}
 		}
 		this._entries.clear();