Retry the agent host IPC connect on a slow-boot channel-registration timeout

src/vs/platform/agentHost/browser/agentHostIpcChannelTransport.ts

@@ -12,6 +12,8 @@
 // upstream to the local agent host process and pipes raw JSON frames over
 // the IPC channel.
 
+import { timeout } from '../../../base/common/async.js';
+import { CancellationTokenSource } from '../../../base/common/cancellation.js';
 import { Emitter } from '../../../base/common/event.js';
 import { Disposable } from '../../../base/common/lifecycle.js';
 import type { IChannel } from '../../../base/parts/ipc/common/ipc.js';
@@ -22,6 +24,33 @@ import { MALFORMED_FRAMES_FORCE_CLOSE_THRESHOLD, MALFORMED_FRAMES_LOG_CAP } from
 
 const REDACTED_TOKEN = '<redacted>';
 
+/** Total wall-clock budget for retrying the upstream `connect` while the agent host registers its IPC channel. */
+const DEFAULT_CONNECT_RETRY_BUDGET_MS = 20_000;
+/** Initial backoff between `connect` retries; doubles up to {@link DEFAULT_CONNECT_RETRY_MAX_DELAY_MS}. */
+const DEFAULT_CONNECT_RETRY_INITIAL_DELAY_MS = 250;
+/** Upper bound on the backoff between `connect` retries. */
+const DEFAULT_CONNECT_RETRY_MAX_DELAY_MS = 2_000;
+
+/** Connect-retry tunables for {@link AgentHostIpcChannelTransport}; overridable in tests. */
+export interface IAgentHostIpcChannelTransportOptions {
+	readonly connectRetryBudgetMs?: number;
+	readonly connectRetryInitialDelayMs?: number;
+	readonly connectRetryMaxDelayMs?: number;
+	/** Delay primitive between retries; defaults to a real timer. */
+	readonly sleep?: (ms: number) => Promise<void>;
+}
+
+/**
+ * The IPC `ChannelServer` rejects calls to a channel it hasn't registered yet
+ * with an error named `'Unknown channel'` once its timeout elapses. For the
+ * agent host that means the host process is still booting and hasn't registered
+ * `agentHostProtocol` — a transient, retryable condition rather than a hard
+ * failure.
+ */
+function isUnknownChannelError(error: unknown): boolean {
+	return error instanceof Error && error.name === 'Unknown channel';
+}
+
 /**
  * Wraps an {@link IChannel} as an {@link IClientTransport} for the agent
  * host protocol. Frames are passed as JSON strings to avoid the IPC layer's
@@ -45,11 +74,24 @@ export class AgentHostIpcChannelTransport extends Disposable implements IClientT
 	private _closeFired = false;
 	private _malformedFrames = 0;
 
+	private readonly _connectRetryBudgetMs: number;
+	private readonly _connectRetryInitialDelayMs: number;
+	private readonly _connectRetryMaxDelayMs: number;
+	private readonly _sleep: (ms: number) => Promise<void>;
+
+	/** Cancels an in-flight connect-retry backoff when the transport is disposed. */
+	private readonly _connectCts = new CancellationTokenSource();
+
 	constructor(
 		private readonly _channel: IChannel,
 		private readonly _ahpLogger?: AhpJsonlLogger,
+		options?: IAgentHostIpcChannelTransportOptions,
 	) {
 		super();
+		this._connectRetryBudgetMs = options?.connectRetryBudgetMs ?? DEFAULT_CONNECT_RETRY_BUDGET_MS;
+		this._connectRetryInitialDelayMs = options?.connectRetryInitialDelayMs ?? DEFAULT_CONNECT_RETRY_INITIAL_DELAY_MS;
+		this._connectRetryMaxDelayMs = options?.connectRetryMaxDelayMs ?? DEFAULT_CONNECT_RETRY_MAX_DELAY_MS;
+		this._sleep = options?.sleep ?? (ms => timeout(ms, this._connectCts.token));
 	}
 
 	get isOpen(): boolean {
@@ -61,13 +103,56 @@ export class AgentHostIpcChannelTransport extends Disposable implements IClientT
 			throw new Error('Transport is disposed');
 		}
 		// Subscribe before connecting so we don't miss any frames the upstream
-		// host emits between open and our listener attaching.
+		// host emits between open and our listener attaching. Event listens to a
+		// not-yet-registered IPC channel are buffered by the ChannelServer and
+		// flushed once it registers, so subscribing once — before any connect
+		// retry — is correct even while the host is still booting.
 		this._register(this._channel.listen<string>('frame')(text => this._handleFrame(text)));
 		this._register(this._channel.listen<void>('close')(() => this._fireClose()));
-		await this._channel.call('connect');
+		await this._connectWithRetry();
 		this._isOpen = true;
 	}
 
+	/**
+	 * Opens the upstream connection, retrying while the agent host is still
+	 * registering its `agentHostProtocol` IPC channel. On a slow host boot the
+	 * channel can be registered only after the IPC ChannelServer's unknown-channel
+	 * timeout, which rejects `call('connect')` with a transient "Unknown channel"
+	 * error even though the channel appears moments later. The local transport
+	 * cannot reconnect once the protocol client gives up, so treating that
+	 * transient timeout as fatal is what leaves the agent host missing from the
+	 * picker until a window reload — retry with backoff up to a bounded budget
+	 * instead, and surface any other error (or budget exhaustion) unchanged.
+	 */
+	private async _connectWithRetry(): Promise<void> {
+		const deadline = Date.now() + this._connectRetryBudgetMs;
+		let delay = this._connectRetryInitialDelayMs;
+		for (; ;) {
+			try {
+				await this._channel.call('connect');
+				return;
+			} catch (error) {
+				// Retry only the transient "channel not registered yet" timeout, and
+				// only while the transport is live and within the wall-clock budget.
+				const remaining = deadline - Date.now();
+				if (this._store.isDisposed || !isUnknownChannelError(error) || remaining <= 0) {
+					throw error;
+				}
+				// Clamp the backoff to the remaining budget so a throttled timer can't
+				// push the next attempt (and its own IPC timeout) past the budget, then
+				// recheck the deadline before issuing another call.
+				await this._sleep(Math.min(delay, remaining));
+				if (this._store.isDisposed) {
+					throw new Error('Transport is disposed');
+				}
+				if (Date.now() >= deadline) {
+					throw error;
+				}
+				delay = Math.min(delay * 2, this._connectRetryMaxDelayMs);
+			}
+		}
+	}
+
 	send(message: ProtocolMessage | AhpServerNotification | JsonRpcResponse): void {
 		if (!this._isOpen || this._closeFired) {
 			// Surface the failure via the close event; callers observe that.
@@ -82,6 +167,8 @@ export class AgentHostIpcChannelTransport extends Disposable implements IClientT
 	}
 
 	override dispose(): void {
+		// Cancel any in-flight connect-retry backoff so teardown doesn't wait it out.
+		this._connectCts.dispose(true);
 		if (this._isOpen && !this._closeFired) {
 			// Best-effort close — ignore any rejection since we're tearing down.
 			this._channel.call('close').catch(() => { });

src/vs/platform/agentHost/test/browser/agentHostIpcChannelTransport.test.ts

@@ -19,13 +19,16 @@ class FakeChannel extends Disposable implements IChannel {
 	readonly frameEmitter = this._register(new Emitter<string>());
 	readonly closeEmitter = this._register(new Emitter<void>());
 	readonly calls: { command: string; arg: unknown }[] = [];
+	readonly listenCounts = { frame: 0, close: 0 };
 	connectResult: Promise<void> = Promise.resolve();
+	/** When set, drives each `call('connect')` (e.g. to reject transiently before resolving). */
+	connectHandler: (() => Promise<void>) | undefined;
 	sendResult: Promise<void> = Promise.resolve();
 
 	call<T>(command: string, arg?: unknown): Promise<T> {
 		this.calls.push({ command, arg });
 		if (command === 'connect') {
-			return this.connectResult as Promise<T>;
+			return (this.connectHandler ? this.connectHandler() : this.connectResult) as Promise<T>;
 		}
 		if (command === 'send') {
 			return this.sendResult as Promise<T>;
@@ -35,15 +38,24 @@ class FakeChannel extends Disposable implements IChannel {
 
 	listen<T>(event: string): Event<T> {
 		if (event === 'frame') {
+			this.listenCounts.frame++;
 			return this.frameEmitter.event as Event<unknown> as Event<T>;
 		}
 		if (event === 'close') {
+			this.listenCounts.close++;
 			return this.closeEmitter.event as Event<unknown> as Event<T>;
 		}
 		throw new Error(`Unknown event: ${event}`);
 	}
 }
 
+/** Builds the error the IPC ChannelServer sends for a call to a not-yet-registered channel. */
+function unknownChannelError(): Error {
+	const error = new Error(`Channel name 'agentHostProtocol' timed out after 1000ms`);
+	error.name = 'Unknown channel';
+	return error;
+}
+
 suite('AgentHostIpcChannelTransport', () => {
 	const ds = ensureNoDisposablesAreLeakedInTestSuite();
 
@@ -78,6 +90,84 @@ suite('AgentHostIpcChannelTransport', () => {
 		assert.strictEqual(transport.isOpen, false);
 	});
 
+	test('retries a transient "Unknown channel" timeout until the host registers its channel', async () => {
+		const channel = ds.add(new FakeChannel());
+		let attempts = 0;
+		channel.connectHandler = () => {
+			attempts++;
+			return attempts < 3 ? Promise.reject(unknownChannelError()) : Promise.resolve();
+		};
+		const transport = ds.add(new AgentHostIpcChannelTransport(channel, undefined, { sleep: () => Promise.resolve() }));
+
+		await transport.connect();
+
+		// Retried until the third attempt succeeded, and frame/close were
+		// subscribed exactly once despite the retries.
+		assert.deepStrictEqual(
+			{ attempts, isOpen: transport.isOpen, listenCounts: channel.listenCounts },
+			{ attempts: 3, isOpen: true, listenCounts: { frame: 1, close: 1 } },
+		);
+	});
+
+	test('does not retry a non-transient connect error', async () => {
+		const channel = ds.add(new FakeChannel());
+		let attempts = 0;
+		channel.connectHandler = () => { attempts++; return Promise.reject(new Error('boom')); };
+		const transport = ds.add(new AgentHostIpcChannelTransport(channel, undefined, { sleep: () => Promise.resolve() }));
+
+		await assert.rejects(() => transport.connect(), /boom/);
+		assert.strictEqual(attempts, 1);
+		assert.strictEqual(transport.isOpen, false);
+	});
+
+	test('gives up once the connect-retry budget is exhausted', async () => {
+		const channel = ds.add(new FakeChannel());
+		let attempts = 0;
+		channel.connectHandler = () => { attempts++; return Promise.reject(unknownChannelError()); };
+		const transport = ds.add(new AgentHostIpcChannelTransport(channel, undefined, { connectRetryBudgetMs: 0, sleep: () => Promise.resolve() }));
+
+		await assert.rejects(() => transport.connect(), /Unknown channel/);
+		assert.strictEqual(attempts, 1);
+		assert.strictEqual(transport.isOpen, false);
+	});
+
+	test('bounds the retry budget even when a backoff overshoots the deadline', async () => {
+		const channel = ds.add(new FakeChannel());
+		let attempts = 0;
+		channel.connectHandler = () => { attempts++; return Promise.reject(unknownChannelError()); };
+		const sleeps: number[] = [];
+		const transport = ds.add(new AgentHostIpcChannelTransport(channel, undefined, {
+			connectRetryBudgetMs: 20,
+			connectRetryInitialDelayMs: 10_000, // unclamped this would ignore the budget entirely
+			sleep: ms => { sleeps.push(ms); return new Promise<void>(resolve => setTimeout(resolve, ms)); },
+		}));
+
+		await assert.rejects(() => transport.connect(), /Unknown channel/);
+
+		// No attempt is issued past the deadline, and any backoff is clamped to the
+		// remaining budget rather than the 10s initial delay.
+		assert.strictEqual(attempts, 1);
+		assert.ok(sleeps.every(ms => ms <= 20), `backoff should be clamped to the remaining budget, got ${JSON.stringify(sleeps)}`);
+	});
+
+	test('stops retrying once the transport is disposed', async () => {
+		const channel = ds.add(new FakeChannel());
+		let attempts = 0;
+		channel.connectHandler = () => { attempts++; return Promise.reject(unknownChannelError()); };
+		// A backoff the test controls, so we can dispose while it is pending.
+		let releaseSleep!: () => void;
+		const sleepGate = new Promise<void>(resolve => { releaseSleep = resolve; });
+		const transport = new AgentHostIpcChannelTransport(channel, undefined, { sleep: () => sleepGate });
+
+		const connectPromise = transport.connect();
+		await Promise.resolve(); // let the first (rejected) connect attempt run
+		transport.dispose();
+		releaseSleep();
+
+		await assert.rejects(() => connectPromise);
+		assert.strictEqual(attempts, 1);
+	});
+
 	test('drops send when transport is not open', async () => {
 		const channel = ds.add(new FakeChannel());
 		const transport = ds.add(new AgentHostIpcChannelTransport(channel));