Ensure we do not support the tools in an untrusted workspace (#25142)

Author: DonJayamanne

src/client/chat/configurePythonEnvTool.ts

@@ -19,6 +19,7 @@ import { TerminalCodeExecutionProvider } from '../terminals/codeExecution/termin
 import {
     getEnvDetailsForResponse,
     getToolResponseIfNotebook,
+    getUntrustedWorkspaceResponse,
     IResourceReference,
     isCancellationError,
     raceCancellationError,
@@ -53,6 +54,9 @@ export class ConfigurePythonEnvTool implements LanguageModelTool<IResourceRefere
         options: LanguageModelToolInvocationOptions<IResourceReference>,
         token: CancellationToken,
     ): Promise<LanguageModelToolResult> {
+        if (!workspace.isTrusted) {
+            return getUntrustedWorkspaceResponse();
+        }
         const resource = resolveFilePath(options.input.resourcePath);
         const notebookResponse = getToolResponseIfNotebook(resource);
         if (notebookResponse) {

src/client/chat/createVirtualEnvTool.ts

@@ -22,6 +22,7 @@ import {
     doesWorkspaceHaveVenvOrCondaEnv,
     getDisplayVersion,
     getEnvDetailsForResponse,
+    getUntrustedWorkspaceResponse,
     IResourceReference,
     isCancellationError,
     raceCancellationError,
@@ -72,6 +73,9 @@ export class CreateVirtualEnvTool implements LanguageModelTool<ICreateVirtualEnv
         options: LanguageModelToolInvocationOptions<ICreateVirtualEnvToolParams>,
         token: CancellationToken,
     ): Promise<LanguageModelToolResult> {
+        if (!workspace.isTrusted) {
+            return getUntrustedWorkspaceResponse();
+        }
         const resource = resolveFilePath(options.input.resourcePath);
         let info = await this.getPreferredEnvForCreation(resource);
         if (!info) {

src/client/chat/getExecutableTool.ts

@@ -10,6 +10,7 @@ import {
     LanguageModelToolInvocationPrepareOptions,
     LanguageModelToolResult,
     PreparedToolInvocation,
+    workspace,
 } from 'vscode';
 import { PythonExtension } from '../api/types';
 import { IServiceContainer } from '../ioc/types';
@@ -19,6 +20,7 @@ import {
     getEnvDisplayName,
     getEnvironmentDetails,
     getToolResponseIfNotebook,
+    getUntrustedWorkspaceResponse,
     IResourceReference,
     raceCancellationError,
 } from './utils';
@@ -45,6 +47,10 @@ export class GetExecutableTool implements LanguageModelTool<IResourceReference>
         options: LanguageModelToolInvocationOptions<IResourceReference>,
         token: CancellationToken,
     ): Promise<LanguageModelToolResult> {
+        if (!workspace.isTrusted) {
+            return getUntrustedWorkspaceResponse();
+        }
+
         const resourcePath = resolveFilePath(options.input.resourcePath);
         const notebookResponse = getToolResponseIfNotebook(resourcePath);
         if (notebookResponse) {

src/client/chat/getPythonEnvTool.ts

@@ -10,13 +10,20 @@ import {
     LanguageModelToolInvocationPrepareOptions,
     LanguageModelToolResult,
     PreparedToolInvocation,
+    workspace,
 } from 'vscode';
 import { PythonExtension } from '../api/types';
 import { IServiceContainer } from '../ioc/types';
 import { ICodeExecutionService } from '../terminals/types';
 import { TerminalCodeExecutionProvider } from '../terminals/codeExecution/terminalCodeExecution';
 import { IProcessServiceFactory, IPythonExecutionFactory } from '../common/process/types';
-import { getEnvironmentDetails, getToolResponseIfNotebook, IResourceReference, raceCancellationError } from './utils';
+import {
+    getEnvironmentDetails,
+    getToolResponseIfNotebook,
+    getUntrustedWorkspaceResponse,
+    IResourceReference,
+    raceCancellationError,
+} from './utils';
 import { resolveFilePath } from './utils';
 import { getPythonPackagesResponse } from './listPackagesTool';
 import { ITerminalHelper } from '../common/terminal/types';
@@ -45,6 +52,10 @@ export class GetEnvironmentInfoTool implements LanguageModelTool<IResourceRefere
         options: LanguageModelToolInvocationOptions<IResourceReference>,
         token: CancellationToken,
     ): Promise<LanguageModelToolResult> {
+        if (!workspace.isTrusted) {
+            return getUntrustedWorkspaceResponse();
+        }
+
         const resourcePath = resolveFilePath(options.input.resourcePath);
         const notebookResponse = getToolResponseIfNotebook(resourcePath);
         if (notebookResponse) {

src/client/chat/installPackagesTool.ts

@@ -10,12 +10,14 @@ import {
     LanguageModelToolInvocationPrepareOptions,
     LanguageModelToolResult,
     PreparedToolInvocation,
+    workspace,
 } from 'vscode';
 import { PythonExtension } from '../api/types';
 import { IServiceContainer } from '../ioc/types';
 import {
     getEnvDisplayName,
     getToolResponseIfNotebook,
+    getUntrustedWorkspaceResponse,
     IResourceReference,
     isCancellationError,
     isCondaEnv,
@@ -43,6 +45,10 @@ export class InstallPackagesTool implements LanguageModelTool<IInstallPackageArg
         options: LanguageModelToolInvocationOptions<IInstallPackageArgs>,
         token: CancellationToken,
     ): Promise<LanguageModelToolResult> {
+        if (!workspace.isTrusted) {
+            return getUntrustedWorkspaceResponse();
+        }
+
         const resourcePath = resolveFilePath(options.input.resourcePath);
         const packageCount = options.input.packageList.length;
         const packagePlurality = packageCount === 1 ? 'package' : 'packages';

src/client/chat/selectEnvTool.ts

@@ -24,6 +24,7 @@ import {
     doesWorkspaceHaveVenvOrCondaEnv,
     getEnvDetailsForResponse,
     getToolResponseIfNotebook,
+    getUntrustedWorkspaceResponse,
     IResourceReference,
 } from './utils';
 import { resolveFilePath } from './utils';
@@ -61,6 +62,10 @@ export class SelectPythonEnvTool implements LanguageModelTool<ISelectPythonEnvTo
         options: LanguageModelToolInvocationOptions<ISelectPythonEnvToolArguments>,
         token: CancellationToken,
     ): Promise<LanguageModelToolResult> {
+        if (!workspace.isTrusted) {
+            return getUntrustedWorkspaceResponse();
+        }
+
         const resource = resolveFilePath(options.input.resourcePath);
         let selected: boolean | undefined = false;
         const hasVenvOrCondaEnvInWorkspaceFolder = doesWorkspaceHaveVenvOrCondaEnv(resource, this.api);

src/client/chat/utils.ts

@@ -126,6 +126,10 @@ export async function getEnvironmentDetails(
     return message.join('\n');
 }
 
+export function getUntrustedWorkspaceResponse() {
+    return new LanguageModelToolResult([new LanguageModelTextPart('Cannot use this tool in an untrusted workspace.')]);
+}
+
 export async function getTerminalCommand(
     environment: ResolvedEnvironment,
     resource: Uri | undefined,
@@ -239,6 +243,9 @@ export async function getEnvDetailsForResponse(
     resource: Uri | undefined,
     token: CancellationToken,
 ): Promise<LanguageModelToolResult> {
+    if (!workspace.isTrusted) {
+        throw new Error('Cannot use this tool in an untrusted workspace.');
+    }
     const envPath = api.getActiveEnvironmentPath(resource);
     environment = environment || (await raceCancellationError(api.resolveEnvironment(envPath), token));
     if (!environment || !environment.version) {