Skip to content

Commit 70e17fc

Browse files
authored
Ensure we do not support the tools in an untrusted workspace (#25142)
1 parent 9b17438 commit 70e17fc

File tree

7 files changed

+44
-1
lines changed

7 files changed

+44
-1
lines changed

src/client/chat/configurePythonEnvTool.ts

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,7 @@ import { TerminalCodeExecutionProvider } from '../terminals/codeExecution/termin
1919
import {
2020
getEnvDetailsForResponse,
2121
getToolResponseIfNotebook,
22+
getUntrustedWorkspaceResponse,
2223
IResourceReference,
2324
isCancellationError,
2425
raceCancellationError,
@@ -53,6 +54,9 @@ export class ConfigurePythonEnvTool implements LanguageModelTool
5354
options: LanguageModelToolInvocationOptions<IResourceReference>,
5455
token: CancellationToken,
5556
): Promise<LanguageModelToolResult> {
57+
if (!workspace.isTrusted) {
58+
return getUntrustedWorkspaceResponse();
59+
}
5660
const resource = resolveFilePath(options.input.resourcePath);
5761
const notebookResponse = getToolResponseIfNotebook(resource);
5862
if (notebookResponse) {

src/client/chat/createVirtualEnvTool.ts

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,7 @@ import {
2222
doesWorkspaceHaveVenvOrCondaEnv,
2323
getDisplayVersion,
2424
getEnvDetailsForResponse,
25+
getUntrustedWorkspaceResponse,
2526
IResourceReference,
2627
isCancellationError,
2728
raceCancellationError,
@@ -72,6 +73,9 @@ export class CreateVirtualEnvTool implements LanguageModelTool
7273
options: LanguageModelToolInvocationOptions<ICreateVirtualEnvToolParams>,
7374
token: CancellationToken,
7475
): Promise<LanguageModelToolResult> {
76+
if (!workspace.isTrusted) {
77+
return getUntrustedWorkspaceResponse();
78+
}
7579
const resource = resolveFilePath(options.input.resourcePath);
7680
let info = await this.getPreferredEnvForCreation(resource);
7781
if (!info) {

src/client/chat/getExecutableTool.ts

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,7 @@ import {
1010
LanguageModelToolInvocationPrepareOptions,
1111
LanguageModelToolResult,
1212
PreparedToolInvocation,
13+
workspace,
1314
} from 'vscode';
1415
import { PythonExtension } from '../api/types';
1516
import { IServiceContainer } from '../ioc/types';
@@ -19,6 +20,7 @@ import {
1920
getEnvDisplayName,
2021
getEnvironmentDetails,
2122
getToolResponseIfNotebook,
23+
getUntrustedWorkspaceResponse,
2224
IResourceReference,
2325
raceCancellationError,
2426
} from './utils';
@@ -45,6 +47,10 @@ export class GetExecutableTool implements LanguageModelTool
4547
options: LanguageModelToolInvocationOptions<IResourceReference>,
4648
token: CancellationToken,
4749
): Promise<LanguageModelToolResult> {
50+
if (!workspace.isTrusted) {
51+
return getUntrustedWorkspaceResponse();
52+
}
53+
4854
const resourcePath = resolveFilePath(options.input.resourcePath);
4955
const notebookResponse = getToolResponseIfNotebook(resourcePath);
5056
if (notebookResponse) {

src/client/chat/getPythonEnvTool.ts

Lines changed: 12 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -10,13 +10,20 @@ import {
1010
LanguageModelToolInvocationPrepareOptions,
1111
LanguageModelToolResult,
1212
PreparedToolInvocation,
13+
workspace,
1314
} from 'vscode';
1415
import { PythonExtension } from '../api/types';
1516
import { IServiceContainer } from '../ioc/types';
1617
import { ICodeExecutionService } from '../terminals/types';
1718
import { TerminalCodeExecutionProvider } from '../terminals/codeExecution/terminalCodeExecution';
1819
import { IProcessServiceFactory, IPythonExecutionFactory } from '../common/process/types';
19-
import { getEnvironmentDetails, getToolResponseIfNotebook, IResourceReference, raceCancellationError } from './utils';
20+
import {
21+
getEnvironmentDetails,
22+
getToolResponseIfNotebook,
23+
getUntrustedWorkspaceResponse,
24+
IResourceReference,
25+
raceCancellationError,
26+
} from './utils';
2027
import { resolveFilePath } from './utils';
2128
import { getPythonPackagesResponse } from './listPackagesTool';
2229
import { ITerminalHelper } from '../common/terminal/types';
@@ -45,6 +52,10 @@ export class GetEnvironmentInfoTool implements LanguageModelTool
4552
options: LanguageModelToolInvocationOptions<IResourceReference>,
4653
token: CancellationToken,
4754
): Promise<LanguageModelToolResult> {
55+
if (!workspace.isTrusted) {
56+
return getUntrustedWorkspaceResponse();
57+
}
58+
4859
const resourcePath = resolveFilePath(options.input.resourcePath);
4960
const notebookResponse = getToolResponseIfNotebook(resourcePath);
5061
if (notebookResponse) {

src/client/chat/installPackagesTool.ts

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -10,12 +10,14 @@ import {
1010
LanguageModelToolInvocationPrepareOptions,
1111
LanguageModelToolResult,
1212
PreparedToolInvocation,
13+
workspace,
1314
} from 'vscode';
1415
import { PythonExtension } from '../api/types';
1516
import { IServiceContainer } from '../ioc/types';
1617
import {
1718
getEnvDisplayName,
1819
getToolResponseIfNotebook,
20+
getUntrustedWorkspaceResponse,
1921
IResourceReference,
2022
isCancellationError,
2123
isCondaEnv,
@@ -43,6 +45,10 @@ export class InstallPackagesTool implements LanguageModelTool
4345
options: LanguageModelToolInvocationOptions<IInstallPackageArgs>,
4446
token: CancellationToken,
4547
): Promise<LanguageModelToolResult> {
48+
if (!workspace.isTrusted) {
49+
return getUntrustedWorkspaceResponse();
50+
}
51+
4652
const resourcePath = resolveFilePath(options.input.resourcePath);
4753
const packageCount = options.input.packageList.length;
4854
const packagePlurality = packageCount === 1 ? 'package' : 'packages';

src/client/chat/selectEnvTool.ts

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -24,6 +24,7 @@ import {
2424
doesWorkspaceHaveVenvOrCondaEnv,
2525
getEnvDetailsForResponse,
2626
getToolResponseIfNotebook,
27+
getUntrustedWorkspaceResponse,
2728
IResourceReference,
2829
} from './utils';
2930
import { resolveFilePath } from './utils';
@@ -61,6 +62,10 @@ export class SelectPythonEnvTool implements LanguageModelTool
6162
options: LanguageModelToolInvocationOptions<ISelectPythonEnvToolArguments>,
6263
token: CancellationToken,
6364
): Promise<LanguageModelToolResult> {
65+
if (!workspace.isTrusted) {
66+
return getUntrustedWorkspaceResponse();
67+
}
68+
6469
const resource = resolveFilePath(options.input.resourcePath);
6570
let selected: boolean | undefined = false;
6671
const hasVenvOrCondaEnvInWorkspaceFolder = doesWorkspaceHaveVenvOrCondaEnv(resource, this.api);

src/client/chat/utils.ts

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -126,6 +126,10 @@ export async function getEnvironmentDetails(
126126
return message.join('\n');
127127
}
128128

129+
export function getUntrustedWorkspaceResponse() {
130+
return new LanguageModelToolResult([new LanguageModelTextPart('Cannot use this tool in an untrusted workspace.')]);
131+
}
132+
129133
export async function getTerminalCommand(
130134
environment: ResolvedEnvironment,
131135
resource: Uri | undefined,
@@ -239,6 +243,9 @@ export async function getEnvDetailsForResponse(
239243
resource: Uri | undefined,
240244
token: CancellationToken,
241245
): Promise<LanguageModelToolResult> {
246+
if (!workspace.isTrusted) {
247+
throw new Error('Cannot use this tool in an untrusted workspace.');
248+
}
242249
const envPath = api.getActiveEnvironmentPath(resource);
243250
environment = environment || (await raceCancellationError(api.resolveEnvironment(envPath), token));
244251
if (!environment || !environment.version) {

0 commit comments

Comments
 (0)