agentHost: preserve BYOK state through tool continuations

extensions/copilot/src/extension/byok/node/test/openAIEndpoint.spec.ts

@@ -10,7 +10,12 @@ import { ConfigKey, IConfigurationService } from '../../../../platform/configura
 import { IChatModelInformation, ModelSupportedEndpoint } from '../../../../platform/endpoint/common/endpointProvider';
 import { CustomDataPartMimeTypes } from '../../../../platform/endpoint/common/endpointTypes';
 import { ChatEndpoint } from '../../../../platform/endpoint/node/chatEndpoint';
+import { ILogService } from '../../../../platform/log/common/logService';
+import { IResponseDelta } from '../../../../platform/networking/common/fetch';
 import { ICreateEndpointBodyOptions, IEndpointBody, IMakeChatRequestOptions } from '../../../../platform/networking/common/networking';
+import { ITelemetryService } from '../../../../platform/telemetry/common/telemetry';
+import { TelemetryData } from '../../../../platform/telemetry/common/telemetryData';
+import { createFakeStreamResponse } from '../../../../platform/test/node/fetcher';
 import { ITestingServicesAccessor } from '../../../../platform/test/node/services';
 import { CancellationToken } from '../../../../util/vs/base/common/cancellation';
 import { DisposableStore } from '../../../../util/vs/base/common/lifecycle';
@@ -577,6 +582,46 @@ describe('OpenAIEndpoint - Reasoning Properties', () => {
 			expect(body.previous_response_id).toBeUndefined();
 			expect(body.store).toBe(false);
 		});
+
+		it.each([
+			{ zeroDataRetentionEnabled: true, expectedMarker: undefined },
+			{ zeroDataRetentionEnabled: false, expectedMarker: 'resp_123' },
+		])('publishes only resumable Responses state when ZDR is $zeroDataRetentionEnabled', async ({ zeroDataRetentionEnabled, expectedMarker }) => {
+			const endpoint = instaService.createInstance(OpenAIEndpoint,
+				{
+					...modelMetadata,
+					vendor: 'OpenAI',
+					zeroDataRetentionEnabled,
+				},
+				'test-api-key',
+				'https://api.openai.com/v1/responses');
+			const response = createFakeStreamResponse(`data: ${JSON.stringify({
+				type: 'response.completed',
+				response: {
+					id: 'resp_123',
+					model: modelMetadata.id,
+					created_at: 123,
+					usage: { input_tokens: 0, output_tokens: 0, total_tokens: 0 },
+					output: [],
+				},
+			})}\n\n`);
+			const deltas: IResponseDelta[] = [];
+			const stream = await endpoint.processResponseFromChatEndpoint(
+				accessor.get(ITelemetryService),
+				accessor.get(ILogService),
+				response,
+				1,
+				async (_text, _index, delta) => {
+					deltas.push(delta);
+					return undefined;
+				},
+				TelemetryData.createAndMarkAsIssued(),
+			);
+
+			for await (const _ of stream) { }
+
+			expect(deltas.at(-1)?.statefulMarker).toBe(expectedMarker);
+		});
 	});
 
 	describe('reasoning effort forwarding', () => {

extensions/copilot/src/platform/endpoint/node/chatEndpoint.ts

@@ -521,7 +521,7 @@ export class ChatEndpoint implements IChatEndpoint {
 	): Promise<AsyncIterableObject<ChatCompletion>> {
 		if (this.useResponsesApi) {
 			const compactionThreshold = getResponsesApiCompactionThreshold(this._configurationService, this._expService, this);
-			return processResponseFromChatEndpoint(this._instantiationService, telemetryService, logService, response, expectedNumChoices, finishCallback, telemetryData, compactionThreshold);
+			return processResponseFromChatEndpoint(this._instantiationService, telemetryService, logService, response, expectedNumChoices, finishCallback, telemetryData, compactionThreshold, this.modelMetadata.zeroDataRetentionEnabled !== true);
 		} else if (this.useMessagesApi) {
 			return processResponseFromMessagesEndpoint(this._instantiationService, telemetryService, logService, response, finishCallback, telemetryData);
 		} else if (!this._supportsStreaming) {

extensions/copilot/src/platform/endpoint/node/responsesApi.ts

@@ -959,7 +959,7 @@ function keepLatestCompactionOutput(output: OpenAI.Responses.ResponseOutputItem[
 	return output.filter((item, idx) => !isCompactionOutputItem(item) || idx === latestCompactionOutput.outputIndex);
 }
 
-export async function processResponseFromChatEndpoint(instantiationService: IInstantiationService, telemetryService: ITelemetryService, logService: ILogService, response: Response, expectedNumChoices: number, finishCallback: FinishedCallback, telemetryData: TelemetryData, compactionThreshold?: number): Promise<AsyncIterableObject<ChatCompletion>> {
+export async function processResponseFromChatEndpoint(instantiationService: IInstantiationService, telemetryService: ITelemetryService, logService: ILogService, response: Response, expectedNumChoices: number, finishCallback: FinishedCallback, telemetryData: TelemetryData, compactionThreshold?: number, emitStatefulMarker = true): Promise<AsyncIterableObject<ChatCompletion>> {
 	return new AsyncIterableObject<ChatCompletion>(async feed => {
 		const requestId = response.headers.get('X-Request-ID') ?? generateUuid();
 		const ghRequestId = response.headers.get('x-github-request-id') ?? '';
@@ -977,7 +977,7 @@ export async function processResponseFromChatEndpoint(instantiationService: IIns
 				const parsedData = JSON.parse(ev.data);
 				const responseStreamEvent: OpenAI.Responses.ResponseStreamEvent = { type: ev.type, ...parsedData };
 				dumper.logEvent(responseStreamEvent);
-				const completion = processor.push(responseStreamEvent, finishCallback);
+				const completion = processor.push(responseStreamEvent, finishCallback, emitStatefulMarker);
 				if (completion) {
 					sendCompletionOutputTelemetry(telemetryService, logService, completion, telemetryData);
 					feed.emitOne(completion);
@@ -1233,7 +1233,7 @@ export class OpenAIResponsesProcessor {
 		});
 	}
 
-	public push(chunk: OpenAI.Responses.ResponseStreamEvent, _onProgress: FinishedCallback): ChatCompletion | undefined {
+	public push(chunk: OpenAI.Responses.ResponseStreamEvent, _onProgress: FinishedCallback, emitStatefulMarker = true): ChatCompletion | undefined {
 		const onProgress = (delta: IResponseDelta): undefined => {
 			this.textAccumulator += delta.text;
 			_onProgress(this.textAccumulator, 0, delta);
@@ -1423,7 +1423,7 @@ export class OpenAIResponsesProcessor {
 				}
 				onProgress({
 					text: '',
-					statefulMarker: chunk.response.id,
+					statefulMarker: emitStatefulMarker ? chunk.response.id : undefined,
 					contextManagement: shouldEmitResolvedCompaction ? latestCompactionItem : undefined,
 				});
 				return {

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

@@ -122,6 +122,7 @@ export type IByokLmOutputItem =
 
 export interface IByokLmChatResult {
 	readonly output: IByokLmOutputItem[];
+	/** Opaque provider state identifier. Present only when a subsequent request can resume it. */
 	readonly responseId?: string;
 	readonly usage?: {
 		readonly inputTokens?: number;

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 {
@@ -71,13 +72,18 @@ export interface IByokLmProxyService {
 const PROXY_USER_FACING_NAME = 'ByokLmProxyService';
 const VENDOR_PATH_PREFIX = '/v/';
 const RESPONSES_SUFFIX = '/responses';
+const MAX_PENDING_TOOL_CONTINUATIONS = 256;
 
-/**
- * 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 scope: string;
+	readonly responseId: string;
+	readonly calls: ReadonlyMap<string, PendingToolCallKind>;
+}
+
+/** Provider state awaiting the SDK's immediate tool-result request. */
+type ByokLmProxyState = Set<IPendingToolContinuation>;
 
 /**
  * Local OpenAI-compatible HTTP proxy that lets the Copilot SDK runtime run
@@ -104,8 +110,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 Set();
 	}
 
 	async start(): Promise<IByokLmProxyHandle> {
@@ -145,14 +150,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 +191,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 +201,22 @@ export class ByokLmProxyService extends LoopbackProxyServer<ByokLmProxyState> im
 			return;
 		}
 
+		const continuationScope = typeof body?.model === 'string' ? this._continuationScope(sessionId, vendor, body.model) : undefined;
+		const explicitResponseId = body?.previous_response_id;
+		const explicit = continuationScope && explicitResponseId !== undefined
+			? Array.from(runtime.state).find(pending => pending.scope === continuationScope && pending.responseId === explicitResponseId)
+			: undefined;
+		const recovered = continuationScope && explicitResponseId === undefined
+			? this._findToolContinuation(runtime.state, continuationScope, body.input)
+			: undefined;
+		const consumed = recovered?.pending ?? explicit;
+		const bridgeBody = recovered
+			? { ...body, input: recovered.input, previous_response_id: recovered.pending.responseId }
+			: body;
+
 		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 +250,12 @@ export class ByokLmProxyService extends LoopbackProxyServer<ByokLmProxyState> im
 				this._writeJsonError(res, 502, result.error, 'api_error');
 				return;
 			}
+			if (consumed) {
+				runtime.state.delete(consumed);
+			}
+			if (continuationScope) {
+				this._addToolContinuation(runtime.state, continuationScope, result);
+			}
 			if (body.stream === true) {
 				res.writeHead(200, {
 					'Content-Type': 'text/event-stream',
@@ -261,6 +286,102 @@ export class ByokLmProxyService extends LoopbackProxyServer<ByokLmProxyState> im
 		}
 	}
 
+	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;
+		}
+
+		// 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 _addToolContinuation(state: ByokLmProxyState, scope: string, result: IByokLmChatResult): void {
+		if (!result.responseId) {
+			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)) {
+					return;
+				}
+				calls.set(item.callId, item.type);
+			}
+		}
+		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);
+			}
+		}
+	}
+
 	private _writeJsonError(res: http.ServerResponse, status: number, message: string, type = 'api_error'): void {
 		if (res.headersSent || res.writableEnded) {
 			return;