fix(tool-search): preserve context tools with tests

Author: PenguinDOOM

extensions/copilot/src/extension/tools/node/test/toolSearchTool.spec.ts

@@ -10,6 +10,8 @@ import { IToolDeferralService } from '../../../../platform/networking/common/too
 import { IInstantiationService } from '../../../../util/vs/platform/instantiation/common/instantiation';
 import { createExtensionUnitTestingServices } from '../../../test/node/services';
 import { TestChatRequest } from '../../../test/node/testHelpers';
+import { ChatVariablesCollection } from '../../../prompt/common/chatVariablesCollection';
+import { IBuildPromptContext } from '../../../prompt/common/intents';
 import { IToolsService } from '../../common/toolsService';
 import { ToolSearchTool } from '../toolSearchTool';
 import { TestToolsService } from './testToolsService';
@@ -24,6 +26,19 @@ function makeToolInfo(name: string): vscode.LanguageModelToolInformation {
 	} as vscode.LanguageModelToolInformation;
 }
 
+function makePromptContext(availableTools?: readonly vscode.LanguageModelToolInformation[]): IBuildPromptContext {
+	return {
+		query: 'test query',
+		history: [],
+		chatVariables: new ChatVariablesCollection([]),
+		tools: availableTools ? {
+			toolReferences: [],
+			toolInvocationToken: undefined as never,
+			availableTools,
+		} : undefined,
+	};
+}
+
 suite('ToolSearchTool', () => {
 	test('searches only deferred tools enabled for the active request', async () => {
 		const searchedToolNames: string[][] = [];
@@ -157,6 +172,77 @@ suite('ToolSearchTool', () => {
 		expect(searchedToolNames).toEqual([['activate_vs_code_interaction']]);
 	});
 
+	test('preserves request-scoped deferred tools after resolved input is shallow-cloned', async () => {
+		const searchedToolNames: string[][] = [];
+		const nonDeferred = new Set(['read_file']);
+		const tool = new ToolSearchTool(
+			{
+				searchToolsByQuery: async (_query: string, tools: readonly vscode.LanguageModelToolInformation[]) => {
+					searchedToolNames.push(tools.map(candidate => candidate.name));
+					return tools.map(candidate => candidate.name);
+				},
+			} as any,
+			{
+				isNonDeferredTool: (name: string) => nonDeferred.has(name),
+			} as any,
+			{ trace() { } } as any,
+		);
+
+		const resolvedInput = await tool.resolveInput(
+			{ query: 'vscode interaction tools' },
+			makePromptContext([
+				makeToolInfo('read_file'),
+				makeToolInfo('activate_vs_code_interaction'),
+			]),
+			0,
+		);
+
+		const clonedResolvedInput = { ...resolvedInput };
+
+		await tool.invoke(
+			{ input: clonedResolvedInput } as vscode.LanguageModelToolInvocationOptions<{ query: string }>,
+			{ isCancellationRequested: false } as vscode.CancellationToken,
+		);
+
+		expect(searchedToolNames).toEqual([['activate_vs_code_interaction']]);
+		expect((tool as any)._requestScopedDeferredToolsContexts.size).toBe(0);
+	});
+
+	test('preserves request-scoped deferred tools after resolved input is JSON-round-tripped', async () => {
+		const searchedToolNames: string[][] = [];
+		const nonDeferred = new Set(['read_file']);
+		const tool = new ToolSearchTool(
+			{
+				searchToolsByQuery: async (_query: string, tools: readonly vscode.LanguageModelToolInformation[]) => {
+					searchedToolNames.push(tools.map(candidate => candidate.name));
+					return tools.map(candidate => candidate.name);
+				},
+			} as any,
+			{
+				isNonDeferredTool: (name: string) => nonDeferred.has(name),
+			} as any,
+			{ trace() { } } as any,
+		);
+
+		const resolvedInput = await tool.resolveInput(
+			{ query: 'vscode interaction tools' },
+			makePromptContext([
+				makeToolInfo('read_file'),
+				makeToolInfo('activate_vs_code_interaction'),
+			]),
+			0,
+		);
+
+		const jsonRoundTrippedInput = JSON.parse(JSON.stringify(resolvedInput));
+
+		await tool.invoke(
+			{ input: jsonRoundTrippedInput } as vscode.LanguageModelToolInvocationOptions<{ query: string }>,
+			{ isCancellationRequested: false } as vscode.CancellationToken,
+		);
+
+		expect(searchedToolNames).toEqual([['activate_vs_code_interaction']]);
+	});
+
 	test('fails explicitly when invoke runs without request-scoped resolveInput context', async () => {
 		const nonDeferred = new Set(['read_file']);
 		const tool = new ToolSearchTool(

extensions/copilot/src/extension/tools/node/toolSearchTool.ts

@@ -20,13 +20,17 @@ export interface IToolSearchParams {
 }
 
 const DEFAULT_SEARCH_LIMIT = 5;
-const requestScopedDeferredToolsKey = Symbol('requestScopedDeferredTools');
+const requestScopedDeferredToolsContextIdKey = '__toolSearchRequestContextId';
+
+let nextRequestScopedDeferredToolsContextId = 0;
 
 type ResolvedToolSearchParams = IToolSearchParams & {
-	[requestScopedDeferredToolsKey]?: readonly vscode.LanguageModelToolInformation[];
+	[requestScopedDeferredToolsContextIdKey]?: string;
 };
 
 export class ToolSearchTool implements ICopilotModelSpecificTool<IToolSearchParams> {
+	private readonly _requestScopedDeferredToolsContexts = new Map<string, readonly vscode.LanguageModelToolInformation[]>();
+
 	constructor(
 		@IToolEmbeddingsComputer private readonly _toolEmbeddingsComputer: IToolEmbeddingsComputer,
 		@IToolDeferralService private readonly _toolDeferralService: IToolDeferralService,
@@ -42,7 +46,13 @@ export class ToolSearchTool implements ICopilotModelSpecificTool<IToolSearchPara
 			]);
 		}
 
-		const requestScopedDeferredTools = (options.input as ResolvedToolSearchParams)[requestScopedDeferredToolsKey];
+		const requestScopedDeferredToolsContextId = (options.input as ResolvedToolSearchParams)[requestScopedDeferredToolsContextIdKey];
+		const requestScopedDeferredTools = requestScopedDeferredToolsContextId
+			? this._requestScopedDeferredToolsContexts.get(requestScopedDeferredToolsContextId)
+			: undefined;
+		if (requestScopedDeferredToolsContextId && requestScopedDeferredTools) {
+			this._requestScopedDeferredToolsContexts.delete(requestScopedDeferredToolsContextId);
+		}
 
 		if (!requestScopedDeferredTools) {
 			throw new Error('ToolSearchTool: request-scoped deferred tools are unavailable. Ensure resolveInput is called before invoke.');
@@ -67,11 +77,12 @@ export class ToolSearchTool implements ICopilotModelSpecificTool<IToolSearchPara
 
 	async resolveInput(input: IToolSearchParams, promptContext: IBuildPromptContext, _mode: CopilotToolMode): Promise<IToolSearchParams> {
 		const manifest = createRequestToolManifest(promptContext.tools?.availableTools ?? [], this._toolDeferralService);
-		const resolvedInput: ResolvedToolSearchParams = { ...input };
-		Object.defineProperty(resolvedInput, requestScopedDeferredToolsKey, {
-			value: Object.freeze([...manifest.deferredTools]),
-			enumerable: false,
-		});
+		const requestScopedDeferredToolsContextId = `tool-search-${nextRequestScopedDeferredToolsContextId++}`;
+		this._requestScopedDeferredToolsContexts.set(requestScopedDeferredToolsContextId, Object.freeze([...manifest.deferredTools]));
+		const resolvedInput: ResolvedToolSearchParams = {
+			...input,
+			[requestScopedDeferredToolsContextIdKey]: requestScopedDeferredToolsContextId,
+		};
 		return resolvedInput;
 	}
 }