Use better type for ChatParticipantToolToken

src/vs/workbench/api/common/extHostChatAgents2.ts

@@ -450,13 +450,15 @@ export class ExtHostChatAgents2 extends Disposable implements ExtHostChatAgentsS
 		const { request, location, history } = await this._createRequest(requestDto, context, detector.extension);
 
 		const model = await this.getModelForRequest(request, detector.extension);
+		const toolInvocationToken = this._tools.getToolInvocationToken(request.sessionId);
 		const extRequest = typeConvert.ChatAgentRequest.to(
 			request,
 			location,
 			model,
 			this.getDiagnosticsWhenEnabled(detector.extension),
 			this.getToolsForRequest(detector.extension, request),
 			detector.extension,
+			toolInvocationToken,
 			this._logService);
 
 		return detector.provider.provideParticipantDetection(
@@ -541,13 +543,15 @@ export class ExtHostChatAgents2 extends Disposable implements ExtHostChatAgentsS
 			stream = new ChatAgentResponseStream(agent.extension, request, this._proxy, this._commands.converter, sessionDisposables);
 
 			const model = await this.getModelForRequest(request, agent.extension);
+			const toolInvocationToken = this._tools.getToolInvocationToken(request.sessionId);
 			const extRequest = typeConvert.ChatAgentRequest.to(
 				request,
 				location,
 				model,
 				this.getDiagnosticsWhenEnabled(agent.extension),
 				this.getToolsForRequest(agent.extension, request),
 				agent.extension,
+				toolInvocationToken,
 				this._logService
 			);
 			inFlightRequest = { requestId: requestDto.requestId, extRequest };
@@ -653,6 +657,7 @@ export class ExtHostChatAgents2 extends Disposable implements ExtHostChatAgentsS
 	$releaseSession(sessionId: string): void {
 		this._sessionDisposables.deleteAndDispose(sessionId);
 		this._onDidDisposeChatSession.fire(sessionId);
+		this._tools.releaseSession(sessionId);
 	}
 
 	async $provideFollowups(requestDto: Dto<IChatAgentRequest>, handle: number, result: IChatAgentResult, context: { history: IChatAgentHistoryEntryDto[] }, token: CancellationToken): Promise<IChatFollowup[]> {

src/vs/workbench/api/common/extHostLanguageModelTools.ts

@@ -11,15 +11,15 @@ import { IDisposable, toDisposable } from '../../../base/common/lifecycle.js';
 import { revive } from '../../../base/common/marshalling.js';
 import { generateUuid } from '../../../base/common/uuid.js';
 import { IExtensionDescription } from '../../../platform/extensions/common/extensions.js';
-import { IPreparedToolInvocation, isToolInvocationContext, IToolInvocation, IToolInvocationContext, IToolResult } from '../../contrib/chat/common/languageModelToolsService.js';
+import { IPreparedToolInvocation, IToolInvocation, IToolInvocationContext, IToolResult } from '../../contrib/chat/common/languageModelToolsService.js';
 import { ExtensionEditToolId, InternalEditToolId } from '../../contrib/chat/common/tools/editFileTool.js';
 import { InternalFetchWebPageToolId } from '../../contrib/chat/common/tools/tools.js';
+import { SearchExtensionsToolId } from '../../contrib/extensions/common/searchExtensionsTool.js';
 import { checkProposedApiEnabled, isProposedApiEnabled } from '../../services/extensions/common/extensions.js';
 import { Dto, SerializableObjectWithBuffers } from '../../services/extensions/common/proxyIdentifier.js';
 import { ExtHostLanguageModelToolsShape, IMainContext, IToolDataDto, MainContext, MainThreadLanguageModelToolsShape } from './extHost.protocol.js';
 import { ExtHostLanguageModels } from './extHostLanguageModels.js';
 import * as typeConvert from './extHostTypeConverters.js';
-import { SearchExtensionsToolId } from '../../contrib/extensions/common/searchExtensionsTool.js';
 
 class Tool {
 
@@ -61,6 +61,8 @@ export class ExtHostLanguageModelTools implements ExtHostLanguageModelToolsShape
 	/** A map of all known tools, from other EHs or registered in vscode core */
 	private readonly _allTools = new Map<string, Tool>();
 
+	private readonly _sessionToolInvocationTokens = new Map<string, vscode.ChatParticipantToolToken>();
+
 	constructor(
 		mainContext: IMainContext,
 		private readonly _languageModels: ExtHostLanguageModels,
@@ -74,6 +76,30 @@ export class ExtHostLanguageModelTools implements ExtHostLanguageModelToolsShape
 		});
 	}
 
+	releaseSession(sessionId: string): void {
+		this._sessionToolInvocationTokens.delete(sessionId);
+	}
+
+	getToolInvocationToken(sessionId: string): vscode.ChatParticipantToolToken {
+		if (!this._sessionToolInvocationTokens.has(sessionId)) {
+			const token: vscode.ChatParticipantToolToken = Object.freeze({
+				value: Symbol('ChatParticipantToolToken') as any
+			});
+			this._sessionToolInvocationTokens.set(sessionId, token);
+		}
+
+		return this._sessionToolInvocationTokens.get(sessionId)!;
+	}
+
+	private getContextForToken(token: vscode.ChatParticipantToolToken): IToolInvocationContext {
+		for (const [sessionId, sessionToken] of this._sessionToolInvocationTokens.entries()) {
+			if (token === sessionToken) {
+				return { sessionId };
+			}
+		}
+		throw new Error(`Invalid tool invocation token`);
+	}
+
 	async $countTokensForInvocation(callId: string, input: string, token: CancellationToken): Promise<number> {
 		const fn = this._tokenCountFuncs.get(callId);
 		if (!fn) {
@@ -90,9 +116,7 @@ export class ExtHostLanguageModelTools implements ExtHostLanguageModelToolsShape
 		}
 
 		try {
-			if (options.toolInvocationToken && !isToolInvocationContext(options.toolInvocationToken)) {
-				throw new Error(`Invalid tool invocation token`);
-			}
+			const invocationContext = options.toolInvocationToken && this.getContextForToken(options.toolInvocationToken);
 
 			if ((toolId === InternalEditToolId || toolId === ExtensionEditToolId) && !isProposedApiEnabled(extension, 'chatParticipantPrivate')) {
 				throw new Error(`Invalid tool: ${toolId}`);
@@ -104,7 +128,7 @@ export class ExtHostLanguageModelTools implements ExtHostLanguageModelToolsShape
 				callId,
 				parameters: options.input,
 				tokenBudget: options.tokenizationOptions?.tokenBudget,
-				context: options.toolInvocationToken as IToolInvocationContext | undefined,
+				context: invocationContext,
 				chatRequestId: isProposedApiEnabled(extension, 'chatParticipantPrivate') ? options.chatRequestId : undefined,
 				chatInteractionId: isProposedApiEnabled(extension, 'chatParticipantPrivate') ? options.chatInteractionId : undefined,
 			}, token);
@@ -159,7 +183,7 @@ export class ExtHostLanguageModelTools implements ExtHostLanguageModelToolsShape
 
 		const options: vscode.LanguageModelToolInvocationOptions<Object> = {
 			input: dto.parameters,
-			toolInvocationToken: dto.context as vscode.ChatParticipantToolToken | undefined,
+			toolInvocationToken: dto.context && this.getToolInvocationToken(dto.context.sessionId),
 		};
 		if (isProposedApiEnabled(item.extension, 'chatParticipantPrivate')) {
 			options.chatRequestId = dto.chatRequestId;

src/vs/workbench/api/common/extHostTypeConverters.ts

@@ -2921,7 +2921,16 @@ export namespace ChatResponsePart {
 }
 
 export namespace ChatAgentRequest {
-	export function to(request: IChatAgentRequest, location2: vscode.ChatRequestEditorData | vscode.ChatRequestNotebookData | undefined, model: vscode.LanguageModelChat, diagnostics: readonly [vscode.Uri, readonly vscode.Diagnostic[]][], tools: Map<string, boolean>, extension: IRelaxedExtensionDescription, logService: ILogService): vscode.ChatRequest {
+	export function to(
+		request: IChatAgentRequest,
+		location2: vscode.ChatRequestEditorData | vscode.ChatRequestNotebookData | undefined,
+		model: vscode.LanguageModelChat,
+		diagnostics: readonly [vscode.Uri, readonly vscode.Diagnostic[]][],
+		tools: Map<string, boolean>,
+		extension: IRelaxedExtensionDescription,
+		toolInvocationToken: vscode.ChatParticipantToolToken,
+		logService: ILogService
+	): vscode.ChatRequest {
 		const toolReferences = request.variables.variables.filter(v => v.kind === 'tool');
 		const variableReferences = request.variables.variables.filter(v => v.kind !== 'tool');
 		const requestWithAllProps: vscode.ChatRequest = {
@@ -2939,7 +2948,7 @@ export namespace ChatAgentRequest {
 			acceptedConfirmationData: request.acceptedConfirmationData,
 			rejectedConfirmationData: request.rejectedConfirmationData,
 			location2,
-			toolInvocationToken: Object.freeze({ sessionId: request.sessionId }) as never,
+			toolInvocationToken,
 			tools,
 			model,
 			editedFileEvents: request.editedFileEvents,
@@ -2961,7 +2970,6 @@ export namespace ChatAgentRequest {
 			delete (requestWithAllProps as any).tools;
 		}
 
-
 		return requestWithAllProps;
 	}
 }

src/vs/workbench/contrib/chat/common/languageModelToolsService.ts

@@ -119,10 +119,6 @@ export interface IToolInvocationContext {
 	sessionId: string;
 }
 
-export function isToolInvocationContext(obj: any): obj is IToolInvocationContext {
-	return typeof obj === 'object' && typeof obj.sessionId === 'string';
-}
-
 export interface IToolResultInputOutputDetails {
 	readonly input: string;
 	readonly output: ({ type: 'text'; value: string } | { type: 'data'; mimeType: string; value64: string })[];

src/vscode-dts/vscode.d.ts

@@ -20621,7 +20621,12 @@ declare module 'vscode' {
 	/**
 	 * A token that can be passed to {@link lm.invokeTool} when invoking a tool inside the context of handling a chat request.
 	 */
-	export type ChatParticipantToolToken = never;
+	export interface ChatParticipantToolToken {
+		/**
+		 * A unique symbol to prevent accidental construction of this type (see note on {@link LanguageModelToolInvocationOptions.toolInvocationToken}).
+		 */
+		readonly value: unique symbol;
+	}
 
 	/**
 	 * Options provided for tool invocation.