agentHost: preserve BYOK state through tool continuations

Author: sgent-epic

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

@@ -6,6 +6,7 @@
 import type * as http from 'http';
 import { createDecorator } from '../../../instantiation/common/instantiation.js';
 import { ILogService } from '../../../log/common/log.js';
+import type { IByokLmChatResult } from '../../common/agentHostByokLm.js';
 import { IByokLmBridgeRegistry } from '../byokLmBridgeRegistry.js';
 import { parseProxyBearer } from '../claude/claudeProxyAuth.js';
 import {
@@ -72,12 +73,15 @@ const PROXY_USER_FACING_NAME = 'ByokLmProxyService';
 const VENDOR_PATH_PREFIX = '/v/';
 const RESPONSES_SUFFIX = '/responses';
 
-/**
- * The BYOK proxy keeps no per-bind mutable state: the active renderer bridge is
- * resolved from {@link IByokLmBridgeRegistry} at request time, and the nonce
- * lives on the runtime owned by {@link LoopbackProxyServer}.
- */
-type ByokLmProxyState = undefined;
+type PendingToolCallKind = 'function_call' | 'custom_tool_call';
+
+interface IPendingToolContinuation {
+	readonly responseId: string;
+	readonly calls: ReadonlyMap<string, PendingToolCallKind>;
+}
+
+/** Provider state awaiting the SDK's immediate tool-result request. */
+type ByokLmProxyState = Map<string, IPendingToolContinuation>;
 
 /**
  * Local OpenAI-compatible HTTP proxy that lets the Copilot SDK runtime run
@@ -104,8 +108,7 @@ export class ByokLmProxyService extends LoopbackProxyServer<ByokLmProxyState> im
 	}
 
 	protected createState(): ByokLmProxyState {
-		// No per-bind state — the bridge is resolved from the registry per request.
-		return undefined;
+		return new Map();
 	}
 
 	async start(): Promise<IByokLmProxyHandle> {
@@ -145,14 +148,15 @@ export class ByokLmProxyService extends LoopbackProxyServer<ByokLmProxyState> im
 		// Inbound requests carry `Bearer <nonce>.<sessionId>`; the runtime is
 		// handed `<nonce>.<sessionId>` at session launch.
 		const auth = parseProxyBearer(req.headers, runtime.nonce);
-		if (!auth.valid || !auth.sessionId) {
+		const sessionId = auth.sessionId;
+		if (!auth.valid || !sessionId) {
 			this._writeJsonError(res, 401, 'Invalid authentication', 'authentication_error');
 			return;
 		}
 
 		const vendor = this._parseVendorFromResponsesPath(pathname);
 		if (method === 'POST' && vendor !== undefined) {
-			await this._handleResponses(req, res, runtime, vendor);
+			await this._handleResponses(req, res, runtime, vendor, sessionId);
 			return;
 		}
 
@@ -185,7 +189,7 @@ export class ByokLmProxyService extends LoopbackProxyServer<ByokLmProxyState> im
 		return vendor;
 	}
 
-	private async _handleResponses(req: http.IncomingMessage, res: http.ServerResponse, runtime: ILoopbackProxyRuntime<ByokLmProxyState>, vendor: string): Promise<void> {
+	private async _handleResponses(req: http.IncomingMessage, res: http.ServerResponse, runtime: ILoopbackProxyRuntime<ByokLmProxyState>, vendor: string, sessionId: string): Promise<void> {
 		let body: IResponsesRequest;
 		try {
 			const raw = await readProxyRequestBody(req);
@@ -195,9 +199,19 @@ 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 };
+			}
+		}
+
 		let bridgeRequest;
 		try {
-			bridgeRequest = responsesRequestToBridge(vendor, body);
+			bridgeRequest = responsesRequestToBridge(vendor, bridgeBody);
 		} catch (err) {
 			const message = err instanceof ResponsesTranslationError ? err.message : String(err);
 			this._writeJsonError(res, 400, message, 'invalid_request_error');
@@ -231,6 +245,9 @@ 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 (body.stream === true) {
 				res.writeHead(200, {
 					'Content-Type': 'text/event-stream',
@@ -261,6 +278,81 @@ export class ByokLmProxyService extends LoopbackProxyServer<ByokLmProxyState> im
 		}
 	}
 
+	private _continuationKey(sessionId: string, vendor: string, modelId: string): string {
+		return JSON.stringify([sessionId, vendor, modelId]);
+	}
+
+	private _recoverToolContinuation(input: IResponsesRequest['input'], pending: IPendingToolContinuation): IResponsesRequest['input'] | undefined {
+		if (!Array.isArray(input)) {
+			return undefined;
+		}
+
+		// A previous_response_id request must contain only the new tool outputs.
+		let start = input.length;
+		while (start > 0 && this._toolOutputKind((input[start - 1] as { readonly type?: unknown } | null)?.type)) {
+			start--;
+		}
+		if (input.length - start !== pending.calls.size) {
+			return undefined;
+		}
+
+		const outputs = input.slice(start);
+		const seen = new Set<string>();
+		for (const value of outputs) {
+			if (typeof value !== 'object' || value === null || Array.isArray(value)) {
+				return undefined;
+			}
+			const item = value as { readonly type?: unknown; readonly call_id?: unknown; readonly output?: unknown };
+			const kind = this._toolOutputKind(item.type);
+			const callId = item.call_id;
+			if (
+				!kind
+				|| typeof callId !== 'string'
+				|| !callId
+				|| seen.has(callId)
+				|| pending.calls.get(callId) !== kind
+				|| (item.output !== undefined && typeof item.output !== 'string')
+			) {
+				return undefined;
+			}
+			seen.add(callId);
+		}
+		return outputs;
+	}
+
+	private _toolOutputKind(type: unknown): PendingToolCallKind | undefined {
+		switch (type) {
+			case 'function_call_output':
+				return 'function_call';
+			case 'custom_tool_call_output':
+				return 'custom_tool_call';
+			default:
+				return undefined;
+		}
+	}
+
+	private _updateToolContinuation(state: ByokLmProxyState, key: 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) {
+			state.set(key, { responseId: result.responseId, calls });
+		} else {
+			state.delete(key);
+		}
+	}
+
 	private _writeJsonError(res: http.ServerResponse, status: number, message: string, type = 'api_error'): void {
 		if (res.headersSent || res.writableEnded) {
 			return;

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

@@ -57,8 +57,8 @@ suite('ByokLmProxyService', () => {
 		return `${handle.providerBaseUrl(vendor)}/responses`;
 	}
 
-	function authHeaders(handle: IByokLmProxyHandle): Record<string, string> {
-		return { 'Content-Type': 'application/json', 'Authorization': `Bearer ${handle.nonce}.${sessionId}` };
+	function authHeaders(handle: IByokLmProxyHandle, selectedSessionId = sessionId): Record<string, string> {
+		return { 'Content-Type': 'application/json', 'Authorization': `Bearer ${handle.nonce}.${selectedSessionId}` };
 	}
 
 	test('serves the unauthenticated health check', async () => {
@@ -298,6 +298,88 @@ suite('ByokLmProxyService', () => {
 		]);
 	});
 
+	test('recovers only the exact scoped tool continuation without overriding explicit state', async () => {
+		const captured: IByokLmChatRequest[] = [];
+		const initialInput = [{ type: 'message', role: 'user', content: [{ type: 'input_text', text: 'Use both tools.' }] }];
+		const outputs = [
+			{ type: 'function_call_output', call_id: 'function_1', output: 'Rain' },
+			{ type: 'custom_tool_call_output', call_id: 'custom_1', output: 'Applied patch.' },
+		];
+		const replayedInput = [
+			...initialInput,
+			{ type: 'function_call', call_id: 'function_1', name: 'get_weather', arguments: '{}' },
+			{ type: 'custom_tool_call', call_id: 'custom_1', name: 'apply_patch', input: '*** Begin Patch\n*** End Patch' },
+			...outputs,
+		];
+
+		await withProxy(
+			async request => {
+				captured.push(request);
+				if (captured.length === 1) {
+					return {
+						responseId: 'resp_provider_1',
+						output: [
+							{ type: 'function_call', callId: 'function_1', name: 'get_weather', argumentsJson: '{}' },
+							{ type: 'custom_tool_call', callId: 'custom_1', name: 'apply_patch', input: '*** Begin Patch\n*** End Patch' },
+						],
+					};
+				}
+				if (captured.length === 6) {
+					return { responseId: 'resp_provider_2', output: [{ type: 'message', content: [{ type: 'text', text: 'done' }] }] };
+				}
+				return { output: [], error: 'not a valid continuation' };
+			},
+			async handle => {
+				const post = (vendor: string, model: string, requestSessionId: string, input: unknown, previousResponseId?: string) => fetch(responsesUrl(handle, vendor), {
+					method: 'POST',
+					headers: authHeaders(handle, requestSessionId),
+					body: JSON.stringify({ model, input, ...(previousResponseId ? { previous_response_id: previousResponseId } : {}) }),
+				});
+
+				let response = await post('acme', 'm', sessionId, initialInput);
+				assert.strictEqual(response.status, 200);
+				await response.text();
+
+				response = await post('acme', 'm', sessionId, replayedInput, 'resp_explicit');
+				assert.strictEqual(response.status, 502);
+				await response.text();
+
+				for (const [vendor, model, requestSessionId] of [
+					['acme', 'm', 'sess-2'],
+					['other', 'm', sessionId],
+					['acme', 'other-model', sessionId],
+				]) {
+					response = await post(vendor, model, requestSessionId, replayedInput);
+					assert.strictEqual(response.status, 502);
+					await response.text();
+				}
+
+				response = await post('acme', 'm', sessionId, replayedInput);
+				assert.strictEqual(response.status, 200);
+				await response.text();
+
+				response = await post('acme', 'm', sessionId, replayedInput);
+				assert.strictEqual(response.status, 502);
+				await response.text();
+			},
+		);
+
+		assert.strictEqual(captured[1]?.previousResponseId, 'resp_explicit');
+		assert.deepStrictEqual(captured.slice(2, 5).map(request => request.previousResponseId), [undefined, undefined, undefined]);
+		assert.deepStrictEqual({
+			previousResponseId: captured[5]?.previousResponseId,
+			input: captured[5]?.input,
+			clearedPreviousResponseId: captured[6]?.previousResponseId,
+		}, {
+			previousResponseId: 'resp_provider_1',
+			input: [
+				{ type: 'function_call_output', callId: 'function_1', output: 'Rain' },
+				{ type: 'custom_tool_call_output', callId: 'custom_1', output: 'Applied patch.' },
+			],
+			clearedPreviousResponseId: undefined,
+		});
+	});
+
 	test('decodes a url-encoded vendor path segment', async () => {
 		let captured: IByokLmChatRequest | undefined;
 		await withProxy(

src/vs/platform/agentHost/test/node/providerIntegration/copilotByokResponses.integrationTest.ts

@@ -6,7 +6,7 @@
 import assert from 'assert';
 import { mkdtemp, rm } from 'fs/promises';
 import { tmpdir } from 'os';
-import { CopilotClient } from '@github/copilot-sdk';
+import { CopilotClient, defineTool } from '@github/copilot-sdk';
 import { Emitter } from '../../../../../base/common/event.js';
 import { ensureNoDisposablesAreLeakedInTestSuite } from '../../../../../base/test/common/utils.js';
 import { NullLogService } from '../../../../log/common/log.js';
@@ -20,7 +20,7 @@ suite('Agent Host Provider Integration - Copilot BYOK Responses', function () {
 
 	const store = ensureNoDisposablesAreLeakedInTestSuite();
 
-	test('bundled SDK consumes structured reasoning and text from the proxy', async function () {
+	test('bundled SDK preserves provider state through a tool continuation', async function () {
 		this.timeout(120_000);
 
 		const sessionId = 'byok-responses-integration';
@@ -31,19 +31,18 @@ suite('Agent Host Provider Integration - Copilot BYOK Responses', function () {
 		const registration = registry.register('client', {
 			chat: async request => {
 				captured.push(request);
-				if (captured.length > 1) {
+				if (captured.length === 1) {
 					return {
-						responseId: 'resp_provider_2',
-						output: [{ type: 'message', content: [{ type: 'text', text: 'second' }] }],
+						responseId: 'resp_provider_1',
+						output: [
+							{ type: 'reasoning', id: 'rs_provider', summary: ['Calling echo'], encryptedContent: 'opaque' },
+							{ type: 'function_call', callId: 'call_1', name: 'echo', argumentsJson: '{}' },
+						],
 					};
 				}
 				return {
-					responseId: 'resp_provider',
-					output: [
-						{ type: 'reasoning', id: 'rs_provider', summary: ['considered options'], encryptedContent: 'opaque' },
-						{ type: 'message', content: [{ type: 'text', text: 'hello' }] },
-					],
-					usage: { inputTokens: 1, outputTokens: 2, reasoningTokens: 1 },
+					responseId: 'resp_provider_2',
+					output: [{ type: 'message', content: [{ type: 'text', text: 'final response' }] }],
 				};
 			},
 			onDidChangeModels: models.event,
@@ -69,47 +68,52 @@ suite('Agent Host Provider Integration - Copilot BYOK Responses', function () {
 				sessionId,
 				model: 'test-model',
 				reasoningEffort: 'medium',
-				availableTools: [],
+				tools: [
+					defineTool('echo', {
+						description: 'Returns a fixed echo result.',
+						parameters: { type: 'object', properties: {}, additionalProperties: false },
+						handler: async () => 'echo result',
+						skipPermission: true,
+						defer: 'never',
+					}),
+				],
+				availableTools: ['custom:echo'],
 				provider: {
 					type: 'openai',
 					wireApi: 'responses',
 					baseUrl: handle.providerBaseUrl('acme'),
 					bearerToken: `${handle.nonce}.${sessionId}`,
 				},
 			});
-			const reasoning: string[] = [];
-			session.on('assistant.reasoning', event => reasoning.push(event.data.content));
-
-			const result = await session.sendAndWait({ prompt: 'Reply exactly hello.' }, 30_000);
-			const secondResult = await session.sendAndWait({ prompt: 'Reply exactly second.' }, 30_000);
-			const replayedReasoning = captured[1]?.input.find(item => item.type === 'reasoning');
+			const result = await session.sendAndWait({ prompt: 'Call echo once, then reply exactly final response.' }, 30_000);
 
 			assert.deepStrictEqual({
 				result: result?.type === 'assistant.message' ? result.data.content : undefined,
-				secondResult: secondResult?.type === 'assistant.message' ? secondResult.data.content : undefined,
-				reasoning,
+				requestCount: captured.length,
 				firstRequest: {
 					vendor: captured[0]?.vendor,
 					modelId: captured[0]?.modelId,
 					inputTypes: captured[0]?.input.map(item => item.type),
 					reasoningEffort: captured[0]?.reasoningEffort,
 				},
-				replayedReasoning,
+				secondRequest: {
+					previousResponseId: captured[1]?.previousResponseId,
+					input: captured[1]?.input.map(item => item.type === 'function_call_output'
+						? { type: item.type, callId: item.callId, output: item.output }
+						: { type: item.type }),
+				},
 			}, {
-				result: 'hello',
-				secondResult: 'second',
-				reasoning: ['considered options'],
+				result: 'final response',
+				requestCount: 2,
 				firstRequest: {
 					vendor: 'acme',
 					modelId: 'test-model',
 					inputTypes: ['message'],
 					reasoningEffort: 'medium',
 				},
-				replayedReasoning: {
-					type: 'reasoning',
-					id: 'rs_provider',
-					summary: ['considered options'],
-					encryptedContent: 'opaque',
+				secondRequest: {
+					previousResponseId: 'resp_provider_1',
+					input: [{ type: 'function_call_output', callId: 'call_1', output: 'echo result' }],
 				},
 			});