agentHost: preserve concurrent BYOK continuations

Author: sgent-epic

src/vs/platform/agentHost/node/copilot/byokLmProxyService.ts

@@ -77,12 +77,13 @@ const MAX_PENDING_TOOL_CONTINUATIONS = 256;
 type PendingToolCallKind = 'function_call' | 'custom_tool_call';
 
 interface IPendingToolContinuation {
+	readonly scope: string;
 	readonly responseId: string;
 	readonly calls: ReadonlyMap<string, PendingToolCallKind>;
 }
 
 /** Provider state awaiting the SDK's immediate tool-result request. */
-type ByokLmProxyState = Map<string, IPendingToolContinuation>;
+type ByokLmProxyState = Set<IPendingToolContinuation>;
 
 /**
  * Local OpenAI-compatible HTTP proxy that lets the Copilot SDK runtime run
@@ -109,7 +110,7 @@ export class ByokLmProxyService extends LoopbackProxyServer<ByokLmProxyState> im
 	}
 
 	protected createState(): ByokLmProxyState {
-		return new Map();
+		return new Set();
 	}
 
 	async start(): Promise<IByokLmProxyHandle> {
@@ -200,15 +201,13 @@ export class ByokLmProxyService extends LoopbackProxyServer<ByokLmProxyState> im
 			return;
 		}
 
-		const continuationKey = typeof body?.model === 'string' ? this._continuationKey(sessionId, vendor, body.model) : undefined;
-		let bridgeBody = body;
-		if (continuationKey && body.previous_response_id === undefined) {
-			const pending = runtime.state.get(continuationKey);
-			const input = pending && this._recoverToolContinuation(body.input, pending);
-			if (input) {
-				bridgeBody = { ...body, input, previous_response_id: pending.responseId };
-			}
-		}
+		const continuationScope = typeof body?.model === 'string' ? this._continuationScope(sessionId, vendor, body.model) : undefined;
+		const recovered = continuationScope && body.previous_response_id === undefined
+			? this._findToolContinuation(runtime.state, continuationScope, body.input)
+			: undefined;
+		const bridgeBody = recovered
+			? { ...body, input: recovered.input, previous_response_id: recovered.pending.responseId }
+			: body;
 
 		let bridgeRequest;
 		try {
@@ -246,8 +245,11 @@ export class ByokLmProxyService extends LoopbackProxyServer<ByokLmProxyState> im
 				this._writeJsonError(res, 502, result.error, 'api_error');
 				return;
 			}
-			if (continuationKey) {
-				this._updateToolContinuation(runtime.state, continuationKey, result);
+			if (recovered) {
+				runtime.state.delete(recovered.pending);
+			}
+			if (continuationScope) {
+				this._addToolContinuation(runtime.state, continuationScope, result);
 			}
 			if (body.stream === true) {
 				res.writeHead(200, {
@@ -279,10 +281,28 @@ export class ByokLmProxyService extends LoopbackProxyServer<ByokLmProxyState> im
 		}
 	}
 
-	private _continuationKey(sessionId: string, vendor: string, modelId: string): string {
+	private _continuationScope(sessionId: string, vendor: string, modelId: string): string {
 		return JSON.stringify([sessionId, vendor, modelId]);
 	}
 
+	private _findToolContinuation(state: ByokLmProxyState, scope: string, input: IResponsesRequest['input']): { readonly pending: IPendingToolContinuation; readonly input: IResponsesRequest['input'] } | undefined {
+		let match: { readonly pending: IPendingToolContinuation; readonly input: IResponsesRequest['input'] } | undefined;
+		for (const pending of state) {
+			if (pending.scope !== scope) {
+				continue;
+			}
+			const recoveredInput = this._recoverToolContinuation(input, pending);
+			if (!recoveredInput) {
+				continue;
+			}
+			if (match) {
+				return undefined;
+			}
+			match = { pending, input: recoveredInput };
+		}
+		return match;
+	}
+
 	private _recoverToolContinuation(input: IResponsesRequest['input'], pending: IPendingToolContinuation): IResponsesRequest['input'] | undefined {
 		if (!Array.isArray(input)) {
 			return undefined;
@@ -332,34 +352,28 @@ export class ByokLmProxyService extends LoopbackProxyServer<ByokLmProxyState> im
 		}
 	}
 
-	private _updateToolContinuation(state: ByokLmProxyState, key: string, result: IByokLmChatResult): void {
+	private _addToolContinuation(state: ByokLmProxyState, scope: string, result: IByokLmChatResult): void {
 		if (!result.responseId) {
-			state.delete(key);
 			return;
 		}
 		const calls = new Map<string, PendingToolCallKind>();
 		for (const item of result.output) {
 			if (item.type === 'function_call' || item.type === 'custom_tool_call') {
 				if (!item.callId || calls.has(item.callId)) {
-					state.delete(key);
 					return;
 				}
 				calls.set(item.callId, item.type);
 			}
 		}
-		if (calls.size) {
-			// A session can disappear after receiving a tool call, so keep abandoned
-			// continuations from growing for the lifetime of the shared proxy.
-			state.delete(key);
-			state.set(key, { responseId: result.responseId, calls });
-			if (state.size > MAX_PENDING_TOOL_CONTINUATIONS) {
-				const oldestKey = state.keys().next().value;
-				if (oldestKey !== undefined) {
-					state.delete(oldestKey);
-				}
+		if (!calls.size) {
+			return;
+		}
+		state.add({ scope, responseId: result.responseId, calls });
+		if (state.size > MAX_PENDING_TOOL_CONTINUATIONS) {
+			const oldest = state.values().next().value;
+			if (oldest) {
+				state.delete(oldest);
 			}
-		} else {
-			state.delete(key);
 		}
 	}
 

src/vs/platform/agentHost/test/node/byokLmProxyService.test.ts

@@ -422,6 +422,96 @@ suite('ByokLmProxyService', () => {
 		});
 	});
 
+	test('keeps interleaved parent and subagent continuations in the same scope', async () => {
+		const captured: IByokLmChatRequest[] = [];
+		const parentInitial = [{ type: 'message', role: 'user', content: [{ type: 'input_text', text: 'Parent' }] }];
+		const subagentInitial = [{ type: 'message', role: 'user', content: [{ type: 'input_text', text: 'Subagent' }] }];
+		const replay = (input: object[], callId: string) => [
+			...input,
+			{ type: 'function_call', call_id: callId, name: 'tool', arguments: '{}' },
+			{ type: 'function_call_output', call_id: callId, output: 'done' },
+		];
+
+		await withProxy(
+			async request => {
+				captured.push(request);
+				if (captured.length <= 2) {
+					const trajectory = captured.length === 1 ? 'parent' : 'subagent';
+					return {
+						responseId: `resp_${trajectory}`,
+						output: [{ type: 'function_call', callId: `call_${trajectory}`, name: 'tool', argumentsJson: '{}' }],
+					};
+				}
+				return { output: [{ type: 'message', content: [{ type: 'text', text: 'done' }] }] };
+			},
+			async handle => {
+				for (const input of [
+					parentInitial,
+					subagentInitial,
+					replay(parentInitial, 'call_parent'),
+					replay(subagentInitial, 'call_subagent'),
+				]) {
+					const response = await fetch(responsesUrl(handle, 'acme'), {
+						method: 'POST',
+						headers: authHeaders(handle),
+						body: JSON.stringify({ model: 'm', input }),
+					});
+					assert.strictEqual(response.status, 200);
+					await response.text();
+				}
+			},
+		);
+
+		assert.deepStrictEqual(captured.map(request => ({
+			previousResponseId: request.previousResponseId,
+			input: request.input,
+		})), [
+			{ previousResponseId: undefined, input: [{ type: 'message', role: 'user', content: [{ type: 'text', text: 'Parent' }] }] },
+			{ previousResponseId: undefined, input: [{ type: 'message', role: 'user', content: [{ type: 'text', text: 'Subagent' }] }] },
+			{ previousResponseId: 'resp_parent', input: [{ type: 'function_call_output', callId: 'call_parent', output: 'done' }] },
+			{ previousResponseId: 'resp_subagent', input: [{ type: 'function_call_output', callId: 'call_subagent', output: 'done' }] },
+		]);
+	});
+
+	test('preserves full replay when multiple pending responses match', async () => {
+		const captured: IByokLmChatRequest[] = [];
+		const replayedInput = [
+			{ type: 'function_call', call_id: 'call_shared', name: 'tool', arguments: '{}' },
+			{ type: 'function_call_output', call_id: 'call_shared', output: 'done' },
+		];
+
+		await withProxy(
+			async request => {
+				captured.push(request);
+				return captured.length <= 2
+					? { responseId: `resp_${captured.length}`, output: [{ type: 'function_call', callId: 'call_shared', name: 'tool', argumentsJson: '{}' }] }
+					: { output: [{ type: 'message', content: [{ type: 'text', text: 'done' }] }] };
+			},
+			async handle => {
+				for (const input of [[], [], replayedInput]) {
+					const response = await fetch(responsesUrl(handle, 'acme'), {
+						method: 'POST',
+						headers: authHeaders(handle),
+						body: JSON.stringify({ model: 'm', input }),
+					});
+					assert.strictEqual(response.status, 200);
+					await response.text();
+				}
+			},
+		);
+
+		assert.deepStrictEqual({
+			previousResponseId: captured[2]?.previousResponseId,
+			input: captured[2]?.input,
+		}, {
+			previousResponseId: undefined,
+			input: [
+				{ type: 'function_call', callId: 'call_shared', name: 'tool', argumentsJson: '{}' },
+				{ type: 'function_call_output', callId: 'call_shared', output: 'done' },
+			],
+		});
+	});
+
 	test('bounds abandoned tool continuations while preserving recent state', async () => {
 		const captured: IByokLmChatRequest[] = [];
 		const maximumPendingContinuations = 256;
@@ -450,7 +540,7 @@ suite('ByokLmProxyService', () => {
 				};
 
 				// Sessions can disappear after receiving a tool call. Fill the proxy,
-				// refresh its oldest entry, then overflow it with one abandoned session.
+				// add another candidate in its oldest scope, then overflow it.
 				for (let index = 0; index < maximumPendingContinuations; index++) {
 					await post(index, []);
 				}