diff --git a/src/vs/workbench/contrib/chat/browser/chat.shared.contribution.ts b/src/vs/workbench/contrib/chat/browser/chat.shared.contribution.ts index c246d3438768b..9221a11be9287 100644 --- a/src/vs/workbench/contrib/chat/browser/chat.shared.contribution.ts +++ b/src/vs/workbench/contrib/chat/browser/chat.shared.contribution.ts @@ -1147,7 +1147,7 @@ configurationRegistry.registerConfiguration({ }, [ChatConfiguration.ToolRiskAssessmentEnabled]: { type: 'boolean', - description: nls.localize('chat.tools.riskAssessment.enabled', "When enabled, terminal tool confirmations show an LLM-generated risk level (Safe / Caution / Review carefully) and a short explanation."), + description: nls.localize('chat.tools.riskAssessment.enabled', "When enabled, tool confirmations show an LLM-generated risk level (Safe / Caution / Review carefully) and a short explanation."), default: true, experiment: { mode: 'auto' diff --git a/src/vs/workbench/contrib/chat/browser/tools/chatToolRiskAssessmentService.ts b/src/vs/workbench/contrib/chat/browser/tools/chatToolRiskAssessmentService.ts index a8183c09e88d4..71d99b6cdac2c 100644 --- a/src/vs/workbench/contrib/chat/browser/tools/chatToolRiskAssessmentService.ts +++ b/src/vs/workbench/contrib/chat/browser/tools/chatToolRiskAssessmentService.ts @@ -28,18 +28,24 @@ export interface IToolRiskAssessment { export const IChatToolRiskAssessmentService = createDecorator('chatToolRiskAssessmentService'); +/** + * Which rubric the model uses to assess a tool call: `terminal` for a shell command, `generic` + * for file edits, reads, fetches, and everything else. When omitted, auto-detected from the tool id. + */ +export type ToolRiskPromptKind = 'terminal' | 'generic'; + export interface IChatToolRiskAssessmentService { readonly _serviceBrand: undefined; /** Returns whether the feature is enabled by configuration. */ isEnabled(): boolean; /** Synchronously read a previously cached assessment, or undefined if none. */ - getCached(tool: IToolData, parameters: unknown): IToolRiskAssessment | undefined; + getCached(tool: IToolData, parameters: unknown, kind?: ToolRiskPromptKind): IToolRiskAssessment | undefined; /** - * Get a cached or freshly-computed risk assessment for a tool call. - * Returns `undefined` when the feature is disabled, no model is available, - * or the assessment cannot be parsed. + * Get a cached or freshly-computed risk assessment for a tool call. Returns `undefined` when no + * model is available or the assessment cannot be parsed, or when the feature is disabled unless + * `options.ignoreEnablement` is set (used by the Autopilot risk gate). */ - assess(tool: IToolData, parameters: unknown, token: CancellationToken): Promise; + assess(tool: IToolData, parameters: unknown, token: CancellationToken, kind?: ToolRiskPromptKind, options?: { ignoreEnablement?: boolean }): Promise; } const MAX_PARAM_BYTES = 2000; @@ -64,16 +70,17 @@ export class ChatToolRiskAssessmentService implements IChatToolRiskAssessmentSer return this._configurationService.getValue(ChatConfiguration.ToolRiskAssessmentEnabled) !== false; } - getCached(tool: IToolData, parameters: unknown): IToolRiskAssessment | undefined { - return this._cache.get(this._cacheKey(tool, parameters))?.assessment; + getCached(tool: IToolData, parameters: unknown, kind?: ToolRiskPromptKind): IToolRiskAssessment | undefined { + return this._cache.get(this._cacheKey(tool, parameters, resolveRiskPromptKind(tool, kind)))?.assessment; } - async assess(tool: IToolData, parameters: unknown, token: CancellationToken): Promise { - if (!this.isEnabled()) { + async assess(tool: IToolData, parameters: unknown, token: CancellationToken, kind?: ToolRiskPromptKind, options?: { ignoreEnablement?: boolean }): Promise { + if (!options?.ignoreEnablement && !this.isEnabled()) { return undefined; } - const key = this._cacheKey(tool, parameters); + const resolvedKind = resolveRiskPromptKind(tool, kind); + const key = this._cacheKey(tool, parameters, resolvedKind); const cached = this._cache.get(key); if (cached) { @@ -87,7 +94,7 @@ export class ChatToolRiskAssessmentService implements IChatToolRiskAssessmentSer const promise = (async () => { try { - const assessment = await this._invokeModel(tool, parameters, token); + const assessment = await this._invokeModel(tool, parameters, resolvedKind, token); if (token.isCancellationRequested) { return undefined; } @@ -104,11 +111,11 @@ export class ChatToolRiskAssessmentService implements IChatToolRiskAssessmentSer return promise; } - private _cacheKey(tool: IToolData, parameters: unknown): string { - return tool.id + '::' + stableStringify(normalizeRiskCacheParameters(tool, parameters)); + private _cacheKey(tool: IToolData, parameters: unknown, kind: ToolRiskPromptKind): string { + return kind + '::' + tool.id + '::' + stableStringify(normalizeRiskCacheParameters(parameters, kind)); } - private async _invokeModel(tool: IToolData, parameters: unknown, token: CancellationToken): Promise { + private async _invokeModel(tool: IToolData, parameters: unknown, kind: ToolRiskPromptKind, token: CancellationToken): Promise { const modelId = this._configurationService.getValue(ChatConfiguration.ToolRiskAssessmentModel) || 'copilot-utility-small'; const models = await this._languageModelsService.selectLanguageModels({ vendor: 'copilot', id: modelId }); @@ -116,7 +123,7 @@ export class ChatToolRiskAssessmentService implements IChatToolRiskAssessmentSer return undefined; } - const prompt = buildPrompt(tool, parameters); + const prompt = buildPrompt(tool, parameters, kind); const response = await this._languageModelsService.sendChatRequest( models[0], undefined, @@ -149,20 +156,35 @@ export class ChatToolRiskAssessmentService implements IChatToolRiskAssessmentSer } } +/** + * Resolve which rubric to assess a tool call under. An explicit kind wins; otherwise it is + * auto-detected from the tool id so `run_in_terminal` keeps the terminal rubric. + */ +function resolveRiskPromptKind(tool: IToolData, kind: ToolRiskPromptKind | undefined): ToolRiskPromptKind { + return kind ?? (tool.id === TerminalToolId.RunInTerminal ? 'terminal' : 'generic'); +} + /** * Compute the subset of tool parameters that are relevant to the risk * assessment, used as the cache key so re-invocations of the same tool call * hit the cache even when model-generated descriptive fields differ. */ -function normalizeRiskCacheParameters(tool: IToolData, parameters: unknown): unknown { - if (tool.id === TerminalToolId.RunInTerminal && parameters && typeof parameters === 'object') { +function normalizeRiskCacheParameters(parameters: unknown, kind: ToolRiskPromptKind): unknown { + if (kind === 'terminal' && parameters && typeof parameters === 'object') { const p = parameters as Record; return { command: p.command }; } return parameters; } -function buildPrompt(tool: IToolData, parameters: unknown): string { +function buildPrompt(tool: IToolData, parameters: unknown, kind: ToolRiskPromptKind): string { + const argsJson = serializeParameters(parameters); + return kind === 'terminal' + ? buildTerminalPrompt(tool, argsJson) + : buildGenericToolPrompt(tool, argsJson); +} + +function serializeParameters(parameters: unknown): string { let argsJson: string; try { argsJson = JSON.stringify(parameters ?? {}); @@ -172,6 +194,10 @@ function buildPrompt(tool: IToolData, parameters: unknown): string { if (argsJson.length > MAX_PARAM_BYTES) { argsJson = argsJson.slice(0, MAX_PARAM_BYTES) + '...[truncated]'; } + return argsJson; +} + +function buildTerminalPrompt(tool: IToolData, argsJson: string): string { return [ `You assess what one terminal command does for a code-editing AI agent, and how risky it is.`, `Reply with STRICT JSON only (no prose, no markdown fences):`, @@ -232,6 +258,61 @@ function buildPrompt(tool: IToolData, parameters: unknown): string { ].join('\n'); } +function buildGenericToolPrompt(tool: IToolData, argsJson: string): string { + return [ + `You assess what one tool call does for a code-editing AI agent, and how risky it is.`, + `The tool may edit files, read files, fetch data, or perform some other action.`, + `Reply with STRICT JSON only (no prose, no markdown fences):`, + `{`, + ` "risk": "green" | "orange" | "red",`, + ` "explanation": ""`, + `}`, + ``, + `Rules for "risk" — apply in order; take the FIRST match:`, + ` 1. permanently destroys source code or user data with no recovery`, + ` (irrecoverable deletion, wiping a database, unrecoverable overwrite) -> red`, + ` 2. executes code downloaded on the fly from an arbitrary or untrusted URL -> red`, + ` 3. sends data to a remote server or changes remote state (POST/PUT, upload, deploy) -> orange`, + ` 4. modifies local files or workspace state (edits, creates, reversible deletes)`, + ` or installs packages from a standard registry -> orange`, + ` 5. otherwise (reads files, lists, searches, fetches public read-only data) -> green`, + ``, + `Read-only operations are always GREEN. Editing or creating a workspace file is`, + `ORANGE (reversible via undo or version control), never red. RED is reserved for`, + `actions whose effects cannot be undone. Installing a package from a normal`, + `registry is ORANGE; only running code piped straight from an arbitrary URL is RED.`, + ``, + `Examples:`, + ` read a file's contents -> green`, + ` list files in a directory -> green`, + ` search the workspace for a symbol -> green`, + ` fetch a public web page (GET) -> green`, + ` edit an existing source file -> orange`, + ` create a new file in the workspace -> orange`, + ` install a package -> orange`, + ` POST data to an external API -> orange`, + ` wipe a database table -> red`, + ` run code from an untrusted URL -> red`, + ``, + `Write "explanation" in this exact shape:`, + // allow-any-unicode-next-line + ` - green : " ." e.g. "Reads the contents of package.json."`, + // allow-any-unicode-next-line + ` - orange: " — ." e.g. "Edits src/app.ts — changes workspace source."`, + // allow-any-unicode-next-line + ` - red : " — ." e.g. "Deletes src/app.ts — permanently removes source."`, + ``, + `Strict explanation rules:`, + ` - Cite the ACTUAL files, paths, URLs, or values from the arguments below.`, + ` - Never use generic phrases like "may have side effects". Always name WHAT is read or changed.`, + ` - Plain prose. No quotes around the sentence. No markdown fences.`, + ``, + `Tool: ${tool.displayName} (id: ${tool.id})`, + `Description: ${tool.modelDescription || tool.userDescription || ''}`, + `Arguments (JSON): ${argsJson}`, + ].join('\n'); +} + function parseAssessment(rawText: string, tool: IToolData): IToolRiskAssessment | undefined { let text = rawText.trim(); if (text.startsWith('```')) { diff --git a/src/vs/workbench/contrib/chat/browser/tools/languageModelToolsService.ts b/src/vs/workbench/contrib/chat/browser/tools/languageModelToolsService.ts index b0f882aba191f..90485807192f1 100644 --- a/src/vs/workbench/contrib/chat/browser/tools/languageModelToolsService.ts +++ b/src/vs/workbench/contrib/chat/browser/tools/languageModelToolsService.ts @@ -42,7 +42,7 @@ import { ChatContextKeys } from '../../common/actions/chatContextKeys.js'; import { ChatRequestToolReferenceEntry, toToolSetVariableEntry, toToolVariableEntry } from '../../common/attachments/chatVariableEntries.js'; import { IVariableReference } from '../../common/chatModes.js'; import { ConfirmedReason, IChatService, IChatToolInvocation, ToolConfirmKind } from '../../common/chatService/chatService.js'; -import { ChatConfiguration, isAutoApproveLevel } from '../../common/constants.js'; +import { ChatConfiguration, isAutoApproveLevel, isAutopilotLevel } from '../../common/constants.js'; import { localChatSessionType } from '../../common/chatSessionsService.js'; import { ILanguageModelChatMetadata } from '../../common/languageModels.js'; import { IChatModel, IChatRequestModel } from '../../common/model/chatModel.js'; @@ -51,10 +51,12 @@ import { chatSessionResourceToId, getChatSessionType } from '../../common/model/ import { HookType } from '../../common/promptSyntax/hookTypes.js'; import { CopilotChatSettingId, CopilotToolId } from '../../common/tools/copilotToolIds.js'; import { ILanguageModelToolsConfirmationService } from '../../common/tools/languageModelToolsConfirmationService.js'; +import { TerminalToolId } from '../../common/tools/terminalToolIds.js'; import { CountTokensCallback, createToolSchemaUri, IBeginToolCallOptions, IExternalPreToolUseHookResult, ILanguageModelToolsService, IPreparedToolInvocation, isToolSet, IToolAndToolSetEnablementMap, IToolData, IToolImpl, IToolInvocation, IToolInvokedEvent, IToolResult, IToolResultInputOutputDetails, IToolSet, SpecedToolAliases, stringifyPromptTsxPart, ToolDataSource, ToolInvocationPresentation, toolMatchesModel, ToolSet, ToolSetForModel, VSCodeToolReference } from '../../common/tools/languageModelToolsService.js'; import { IToolResultCompressor } from '../../common/tools/toolResultCompressor.js'; import { getToolConfirmationAlert } from '../accessibility/chatAccessibilityProvider.js'; import { IChatWidgetService } from '../chat.js'; +import { IChatToolRiskAssessmentService, ToolRiskLevel } from './chatToolRiskAssessmentService.js'; const jsonSchemaRegistry = Registry.as(JSONContributionRegistry.Extensions.JSONContribution); @@ -73,6 +75,13 @@ export const enum AutoApproveStorageKeys { const SkipAutoApproveConfirmationKey = 'vscode.chat.tools.global.autoApprove.testMode'; +/** + * Marks a {@link ToolConfirmKind.ConfirmationNotNeeded} decision that came from the session + * auto-approving everything, rather than a per-tool setting or an explicit user action. Shared so + * `shouldAutoConfirm`, the Autopilot risk gate, and approval telemetry use the same string. + */ +const autoApproveAllReason = 'auto-approve-all'; + // This tool will always require user confirmation even in auto approval mode. // Users cannot auto approve this tool via settings either, as this is a tool used before the agentic loop. const toolIdsThatCannotBeAutoApproved = new Set([ @@ -80,6 +89,15 @@ const toolIdsThatCannotBeAutoApproved = new Set([ 'vscode_get_modified_files_confirmation', ]); +// Fetch uses two tools: the model-facing 'copilot_fetchWebPage' and the internal +// 'vscode_fetchWebPage_internal' it delegates to. Both auto-approve themselves, so the Autopilot +// risk gate classifies them to catch dangerous fetches (leaking secrets to an attacker URL, +// hitting internal hosts). +const fetchWebPageToolIds = new Set([ + 'copilot_fetchWebPage', + 'vscode_fetchWebPage_internal', +]); + export const globalAutoApproveDescription = localize2( { key: 'autoApprove3.markdown', @@ -140,6 +158,7 @@ export class LanguageModelToolsService extends Disposable implements ILanguageMo @ICommandService private readonly _commandService: ICommandService, @IChatWidgetService private readonly _chatWidgetService: IChatWidgetService, @IToolResultCompressor private readonly _toolResultCompressor: IToolResultCompressor, + @IChatToolRiskAssessmentService private readonly _riskAssessmentService: IChatToolRiskAssessmentService, ) { super(); @@ -591,9 +610,13 @@ export class LanguageModelToolsService extends Disposable implements ILanguageMo preparedInvocation = await this.prepareToolInvocationWithHookResult(tool, dto, preToolUseHookResult, token); prepareTimeWatch.stop(); - const { autoConfirmed, preparedInvocation: updatedPreparedInvocation } = await this.resolveAutoConfirmFromHook(preToolUseHookResult, tool, dto, preparedInvocation, dto.context?.sessionResource); + const { autoConfirmed: resolvedAutoConfirmed, preparedInvocation: updatedPreparedInvocation } = await this.resolveAutoConfirmFromHook(preToolUseHookResult, tool, dto, preparedInvocation, dto.context?.sessionResource); preparedInvocation = updatedPreparedInvocation; + // In Autopilot, run the risk classifier on an auto-approved call that would + // otherwise show a confirmation. A "red" rating skips the call; anything else + // (including a classifier failure) keeps the original auto-confirmation. + const { autoConfirmed, skipExplanation: riskSkipExplanation } = await this._maybeApplyAutopilotRiskGate(tool, dto, preparedInvocation, resolvedAutoConfirmed, token); // Important: a tool invocation that will be autoconfirmed should never // be in the chat response in the `NeedsConfirmation` state, even briefly, @@ -612,6 +635,27 @@ export class LanguageModelToolsService extends Disposable implements ILanguageMo } dto.toolSpecificData = toolInvocation?.toolSpecificData; + + // Enforce a risk skip here, before the confirmation flow below: run_in_terminal + // suppresses its own confirmation under Autopilot and never reaches it. The tool + // is not run, and an info note explains why. + if (riskSkipExplanation) { + this._logToolApprovalTelemetry(tool, dto, { type: ToolConfirmKind.Skipped }); + // Terminal and edit tools hide their invocation part once complete, so show the + // reason as a separate info note. + this._chatService.appendProgress(request, { + kind: 'info', + content: new MarkdownString(localize('autopilotRiskSkipped', "Autopilot skipped \"{0}\" because it was assessed as high-risk: {1}", tool.data.displayName, riskSkipExplanation)), + }); + toolResult = { + content: [{ + kind: 'text', + value: `Autopilot skipped this tool call because it was automatically assessed as high-risk: ${riskSkipExplanation} The action was not performed. Do not retry it as-is — choose a safer approach or leave it for the user to run manually.` + }] + }; + return toolResult; + } + if (preparedInvocation?.confirmationMessages?.title) { if (!IChatToolInvocation.executionConfirmedOrDenied(toolInvocation) && !autoConfirmed) { this.playAccessibilitySignal([toolInvocation], dto.context?.sessionResource); @@ -760,7 +804,7 @@ export class LanguageModelToolsService extends Disposable implements ILanguageMo [ToolConfirmKind.UserAction]: 'userAction', [ToolConfirmKind.Skipped]: 'skipped', }; - const allowedConfirmationNotNeededReasons = new Set(['auto-approve-all', 'inlineChat']); + const allowedConfirmationNotNeededReasons = new Set([autoApproveAllReason, 'inlineChat']); let confirmationNotNeededReason: string | undefined; if (reason.type === ToolConfirmKind.ConfirmationNotNeeded && reason.reason) { const raw = typeof reason.reason === 'string' ? reason.reason : reason.reason.value; @@ -875,6 +919,95 @@ export class LanguageModelToolsService extends Disposable implements ILanguageMo return { autoConfirmed, preparedInvocation }; } + /** + * In Autopilot, runs the risk classifier on an auto-approved call and skips it when the rating + * is {@link ToolRiskLevel.Red}. Any other result returns the original auto-confirmation + * unchanged. + * + * To keep the classifier off the hot path, it only runs when all of these hold: + * - the call was auto-approved by the session approving everything, or is a `run_in_terminal` / + * fetch call that self-approved (these can run risky commands or prompt-injected URLs without + * ever showing a confirmation); + * - it would otherwise show a confirmation (the self-approving tools above are the exception); + * - the session is a local panel session at the Autopilot level with Advanced Autopilot on. + * + * This is independent of `chat.tools.riskAssessment.enabled`, which only controls the + * confirmation risk badge. CLI and agent-host sessions handle their own confirmations and are + * excluded. + * + * Fails open: a cancelled, unavailable, or failed assessment keeps the original + * auto-confirmation so Autopilot keeps moving. + */ + private async _maybeApplyAutopilotRiskGate( + tool: IToolEntry, + dto: IToolInvocation, + preparedInvocation: IPreparedToolInvocation | undefined, + autoConfirmed: ConfirmedReason | undefined, + token: CancellationToken, + ): Promise<{ autoConfirmed: ConfirmedReason | undefined; skipExplanation?: string }> { + const isTerminalTool = tool.data.id === TerminalToolId.RunInTerminal; + const isFetchTool = fetchWebPageToolIds.has(tool.data.id); + const isAlwaysClassifyTool = isTerminalTool || isFetchTool; + + // Normally only gate calls the session auto-approved wholesale (the `autoApproveAllReason` + // sentinel). A per-tool setting, user action, or hook carries a concrete reason and is + // respected as-is. + // + // Exception: run_in_terminal and fetch self-approve without a confirmation, so a risky command + // or a prompt-injected URL would run unclassified. Gate them when they arrive self-approved + // (no reason and no confirmation of their own); an explicit allow carries a concrete reason + // instead of `undefined`, so it stays respected. + const isBlanketSessionApprove = autoConfirmed?.type === ToolConfirmKind.ConfirmationNotNeeded + && autoConfirmed.reason === autoApproveAllReason; + const isSelfApprovedAlwaysClassify = isAlwaysClassifyTool + && autoConfirmed === undefined + && !preparedInvocation?.confirmationMessages?.title; + if (!isBlanketSessionApprove && !isSelfApprovedAlwaysClassify) { + return { autoConfirmed }; + } + + // Only gate calls that would otherwise show a confirmation, plus the self-approving tools above. + if (!isAlwaysClassifyTool && !preparedInvocation?.confirmationMessages?.title) { + return { autoConfirmed }; + } + + // Check the Advanced Autopilot flag first: it is default-off, so the common case bails before + // the session lookups below. This does not consult `chat.tools.riskAssessment.enabled`, which + // only controls the confirmation risk badge. + if (this._configurationService.getValue(ChatConfiguration.AutopilotAdvancedEnabled) !== true) { + return { autoConfirmed }; + } + + // Scope to local panel sessions at the Autopilot level. CLI and agent-host sessions handle + // their own confirmations. + const sessionResource = dto.context?.sessionResource; + if (!sessionResource || getChatSessionType(sessionResource) !== localChatSessionType) { + return { autoConfirmed }; + } + if (!this._isSessionInAutopilotLevel(sessionResource)) { + return { autoConfirmed }; + } + + try { + // ignoreEnablement: assess even when the risk-badge setting is off. + const assessment = await this._riskAssessmentService.assess(tool.data, dto.parameters, token, undefined, { ignoreEnablement: true }); + if (token.isCancellationRequested) { + return { autoConfirmed }; + } + if (assessment?.risk === ToolRiskLevel.Red) { + const fallbackExplanation = localize('autopilotRiskSkipFallback', "The action was assessed as potentially destructive or irreversible."); + const explanation = assessment.explanation.trim() || fallbackExplanation; + this._logService.info(`[LanguageModelToolsService#invokeTool] Autopilot skipping high-risk tool ${tool.data.id}: ${explanation}`); + return { autoConfirmed: { type: ToolConfirmKind.Skipped }, skipExplanation: explanation }; + } + } catch (err) { + this._logService.warn(`[LanguageModelToolsService#invokeTool] Autopilot risk assessment failed for tool ${tool.data.id}, allowing: ${toErrorMessage(err)}`); + } + + // Green/orange, no assessment, or a failure: keep the original auto-confirmation (fail open). + return { autoConfirmed }; + } + private async prepareToolInvocation(tool: IToolEntry, dto: IToolInvocation, forceConfirmationReason: string | undefined, token: CancellationToken): Promise { let prepared: IPreparedToolInvocation | undefined; if (tool.impl!.prepareToolInvocation) { @@ -1155,6 +1288,29 @@ export class LanguageModelToolsService extends Disposable implements ILanguageMo return isAutoApproveLevel(request?.modeInfo?.permissionLevel) || this._isSessionLiveAutoApproveLevel(chatSessionResource); } + /** + * True if the session's live permission picker level is Autopilot. Like + * {@link _isSessionLiveAutoApproveLevel}, but excludes plain Auto-Approve. + */ + private _isSessionLiveAutopilotLevel(chatSessionResource: URI): boolean { + const widget = this._chatWidgetService.getWidgetBySessionResource(chatSessionResource) + ?? this._chatWidgetService.lastFocusedWidget; + return !!widget && isAutopilotLevel(widget.input.currentModeInfo.permissionLevel); + } + + /** + * True if the session is at the Autopilot level (not plain Auto-Approve), via either the last + * request's stamped level or the live picker level. + */ + private _isSessionInAutopilotLevel(chatSessionResource: URI | undefined): boolean { + if (!chatSessionResource) { + return false; + } + const model = this._chatService.getSession(chatSessionResource); + const request = model?.getRequests().at(-1); + return isAutopilotLevel(request?.modeInfo?.permissionLevel) || this._isSessionLiveAutopilotLevel(chatSessionResource); + } + private getEligibleForAutoApprovalSpecialCase(toolData: IToolData): string | undefined { if (toolData.id === 'vscode_fetchWebPage_internal') { return 'fetch'; @@ -1210,7 +1366,7 @@ export class LanguageModelToolsService extends Disposable implements ILanguageMo if (chatSessionResource && !this._isAutoApprovePolicyRestricted() && this._isSessionInAutoApproveLevel(chatSessionResource)) { // CLI sessions still need their multi-option dialogs (e.g. uncommitted changes). if (!(toolIdsThatCannotBeAutoApproved.has(tool.data.id) && getChatSessionType(chatSessionResource) !== localChatSessionType)) { - return { type: ToolConfirmKind.ConfirmationNotNeeded, reason: 'auto-approve-all' }; + return { type: ToolConfirmKind.ConfirmationNotNeeded, reason: autoApproveAllReason }; } } @@ -1251,7 +1407,7 @@ export class LanguageModelToolsService extends Disposable implements ILanguageMo const sessionAutoApprove = chatSessionResource && !this._isAutoApprovePolicyRestricted() && this._isSessionInAutoApproveLevel(chatSessionResource); if (sessionAutoApprove) { if (!(toolIdsThatCannotBeAutoApproved.has(toolId) && getChatSessionType(chatSessionResource!) !== localChatSessionType)) { - return { type: ToolConfirmKind.ConfirmationNotNeeded, reason: 'auto-approve-all' }; + return { type: ToolConfirmKind.ConfirmationNotNeeded, reason: autoApproveAllReason }; } } diff --git a/src/vs/workbench/contrib/chat/browser/widget/chatContentParts/chatElicitationContentPart.ts b/src/vs/workbench/contrib/chat/browser/widget/chatContentParts/chatElicitationContentPart.ts index bb3ebb096eafd..e964961986115 100644 --- a/src/vs/workbench/contrib/chat/browser/widget/chatContentParts/chatElicitationContentPart.ts +++ b/src/vs/workbench/contrib/chat/browser/widget/chatContentParts/chatElicitationContentPart.ts @@ -4,7 +4,6 @@ *--------------------------------------------------------------------------------------------*/ import { IMarkdownString, isMarkdownString, MarkdownString } from '../../../../../../base/common/htmlContent.js'; -import { CancellationTokenSource } from '../../../../../../base/common/cancellation.js'; import { Disposable, IDisposable, toDisposable } from '../../../../../../base/common/lifecycle.js'; import { autorun } from '../../../../../../base/common/observable.js'; import { IInstantiationService } from '../../../../../../platform/instantiation/common/instantiation.js'; @@ -19,7 +18,7 @@ import { AcceptElicitationRequestActionId } from '../../actions/chatElicitationA import { IChatToolRiskAssessmentService } from '../../tools/chatToolRiskAssessmentService.js'; import { ChatConfirmationWidget, IChatConfirmationButton } from './chatConfirmationWidget.js'; import { IChatContentPart, IChatContentPartRenderContext } from './chatContentParts.js'; -import { ToolRiskBadgeWidget } from './toolInvocationParts/toolRiskBadgeWidget.js'; +import { createToolRiskBadge } from './toolInvocationParts/toolRiskBadgeHelper.js'; import { IAction } from '../../../../../../base/common/actions.js'; export class ChatElicitationContentPart extends Disposable implements IChatContentPart { @@ -134,38 +133,8 @@ export class ChatElicitationContentPart extends Disposable implements IChatConte if (elicitation.kind !== 'elicitation2' || !elicitation.riskAssessment) { return undefined; } - if (!this.riskAssessmentService.isEnabled()) { - return undefined; - } const { toolId, parameters } = elicitation.riskAssessment; - const tool = this.languageModelToolsService.getTool(toolId); - if (!tool) { - return undefined; - } - const widget = this._register(this.instantiationService.createInstance(ToolRiskBadgeWidget)); - const cached = this.riskAssessmentService.getCached(tool, parameters); - if (cached) { - widget.setAssessment(cached); - } else { - widget.setLoading(); - const cts = this._register(new CancellationTokenSource()); - (async () => { - try { - const result = await this.riskAssessmentService.assess(tool, parameters, cts.token); - if (cts.token.isCancellationRequested) { - return; - } - if (!result) { - widget.setHidden(); - return; - } - widget.setAssessment(result); - } catch { - widget.setHidden(); - } - })(); - } - return widget.domNode; + return createToolRiskBadge(this._store, this.instantiationService, this.riskAssessmentService, this.languageModelToolsService, toolId, parameters)?.domNode; } hasSameContent(other: IChatProgressRenderableResponseContent): boolean { diff --git a/src/vs/workbench/contrib/chat/browser/widget/chatContentParts/toolInvocationParts/abstractToolConfirmationSubPart.ts b/src/vs/workbench/contrib/chat/browser/widget/chatContentParts/toolInvocationParts/abstractToolConfirmationSubPart.ts index 94c5130b60cc0..4555ca617cd62 100644 --- a/src/vs/workbench/contrib/chat/browser/widget/chatContentParts/toolInvocationParts/abstractToolConfirmationSubPart.ts +++ b/src/vs/workbench/contrib/chat/browser/widget/chatContentParts/toolInvocationParts/abstractToolConfirmationSubPart.ts @@ -15,9 +15,11 @@ import { ChatContextKeys } from '../../../../common/actions/chatContextKeys.js'; import { ConfirmedReason, IChatToolInvocation, ToolConfirmKind } from '../../../../common/chatService/chatService.js'; import { ILanguageModelToolsService } from '../../../../common/tools/languageModelToolsService.js'; import { IChatWidgetService } from '../../../chat.js'; +import { IChatToolRiskAssessmentService } from '../../../tools/chatToolRiskAssessmentService.js'; import { ChatCustomConfirmationWidget, IChatConfirmationButton } from '../chatConfirmationWidget.js'; import { IChatContentPartRenderContext } from '../chatContentParts.js'; import { BaseChatToolInvocationSubPart } from './chatToolInvocationSubPart.js'; +import { createToolRiskBadge } from './toolRiskBadgeHelper.js'; export interface IToolConfirmationConfig { allowActionId: string; @@ -50,6 +52,7 @@ export abstract class AbstractToolConfirmationSubPart extends BaseChatToolInvoca @IContextKeyService protected readonly contextKeyService: IContextKeyService, @IChatWidgetService protected readonly chatWidgetService: IChatWidgetService, @ILanguageModelToolsService protected readonly languageModelToolsService: ILanguageModelToolsService, + @IChatToolRiskAssessmentService protected readonly riskAssessmentService: IChatToolRiskAssessmentService, ) { super(toolInvocation); @@ -114,6 +117,11 @@ export abstract class AbstractToolConfirmationSubPart extends BaseChatToolInvoca const contentElement = this.createContentElement(); const tool = languageModelToolsService.getTool(toolInvocation.toolId); + // Risk badges describe a pending action, so they only show on the pre-execution + // confirmation, not after the tool has run. + const riskBadge = state.type === IChatToolInvocation.StateKind.WaitingForConfirmation + ? this.createRiskBadgeDomNode(state.parameters) + : undefined; const confirmWidget = this._register(this.instantiationService.createInstance( ChatCustomConfirmationWidget<(() => void)>, this.context, @@ -123,6 +131,7 @@ export abstract class AbstractToolConfirmationSubPart extends BaseChatToolInvoca subtitle: config.subtitle, buttons, message: contentElement, + footerBanner: riskBadge, toolbarData: { arg: toolInvocation, partType: config.partType, @@ -200,6 +209,15 @@ export abstract class AbstractToolConfirmationSubPart extends BaseChatToolInvoca return []; } + /** + * Create the risk-assessment badge DOM node for this confirmation, or + * `undefined` when the feature is disabled or the tool is unknown. Returned + * as a `footerBanner` for the confirmation widget. + */ + protected createRiskBadgeDomNode(parameters: unknown): HTMLElement | undefined { + return createToolRiskBadge(this._store, this.instantiationService, this.riskAssessmentService, this.languageModelToolsService, this.toolInvocation.toolId, parameters)?.domNode; + } + /** * When true, "Allow Once" stays the primary button even when a * session-scoped action is available. Subclasses override this diff --git a/src/vs/workbench/contrib/chat/browser/widget/chatContentParts/toolInvocationParts/chatMissingSandboxDepsConfirmationSubPart.ts b/src/vs/workbench/contrib/chat/browser/widget/chatContentParts/toolInvocationParts/chatMissingSandboxDepsConfirmationSubPart.ts index e31d3daf423d5..70ac5dcf41ed3 100644 --- a/src/vs/workbench/contrib/chat/browser/widget/chatContentParts/toolInvocationParts/chatMissingSandboxDepsConfirmationSubPart.ts +++ b/src/vs/workbench/contrib/chat/browser/widget/chatContentParts/toolInvocationParts/chatMissingSandboxDepsConfirmationSubPart.ts @@ -14,6 +14,7 @@ import { IChatToolInvocation, type IChatTerminalToolInvocationData } from '../.. import { ILanguageModelToolsService } from '../../../../common/tools/languageModelToolsService.js'; import { AcceptToolConfirmationActionId, SkipToolConfirmationActionId } from '../../../actions/chatToolActions.js'; import { IChatCodeBlockInfo, IChatWidgetService } from '../../../chat.js'; +import { IChatToolRiskAssessmentService } from '../../../tools/chatToolRiskAssessmentService.js'; import { IChatContentPartRenderContext } from '../chatContentParts.js'; import { AbstractToolConfirmationSubPart } from './abstractToolConfirmationSubPart.js'; @@ -30,8 +31,9 @@ export class ChatMissingSandboxDepsConfirmationSubPart extends AbstractToolConfi @IContextKeyService contextKeyService: IContextKeyService, @IChatWidgetService chatWidgetService: IChatWidgetService, @ILanguageModelToolsService languageModelToolsService: ILanguageModelToolsService, + @IChatToolRiskAssessmentService riskAssessmentService: IChatToolRiskAssessmentService, ) { - super(toolInvocation, context, instantiationService, keybindingService, contextKeyService, chatWidgetService, languageModelToolsService); + super(toolInvocation, context, instantiationService, keybindingService, contextKeyService, chatWidgetService, languageModelToolsService, riskAssessmentService); this.render({ allowActionId: AcceptToolConfirmationActionId, diff --git a/src/vs/workbench/contrib/chat/browser/widget/chatContentParts/toolInvocationParts/chatModifiedFilesConfirmationSubPart.ts b/src/vs/workbench/contrib/chat/browser/widget/chatContentParts/toolInvocationParts/chatModifiedFilesConfirmationSubPart.ts index d71f6508afe7a..4c3e91c09eb30 100644 --- a/src/vs/workbench/contrib/chat/browser/widget/chatContentParts/toolInvocationParts/chatModifiedFilesConfirmationSubPart.ts +++ b/src/vs/workbench/contrib/chat/browser/widget/chatContentParts/toolInvocationParts/chatModifiedFilesConfirmationSubPart.ts @@ -22,6 +22,7 @@ import { ILanguageModelToolsService } from '../../../../common/tools/languageMod import { ModifiedFileEntryState } from '../../../../common/editing/chatEditingService.js'; import { ChatContextKeys } from '../../../../common/actions/chatContextKeys.js'; import { IChatCodeBlockInfo, IChatWidgetService } from '../../../chat.js'; +import { IChatToolRiskAssessmentService } from '../../../tools/chatToolRiskAssessmentService.js'; import { IChatContentPartRenderContext } from '../chatContentParts.js'; import { ChatCustomConfirmationWidget, IChatConfirmationButton } from '../chatConfirmationWidget.js'; import { CollapsibleListPool, IChatCollapsibleListItem } from '../chatReferencesContentPart.js'; @@ -44,8 +45,9 @@ export class ChatModifiedFilesConfirmationSubPart extends AbstractToolConfirmati @IMarkdownRendererService private readonly markdownRendererService: IMarkdownRendererService, @IEditorService private readonly editorService: IEditorService, @ICommandService private readonly commandService: ICommandService, + @IChatToolRiskAssessmentService riskAssessmentService: IChatToolRiskAssessmentService, ) { - super(toolInvocation, context, instantiationService, keybindingService, contextKeyService, chatWidgetService, languageModelToolsService); + super(toolInvocation, context, instantiationService, keybindingService, contextKeyService, chatWidgetService, languageModelToolsService, riskAssessmentService); const state = toolInvocation.state.get(); if (state.type !== IChatToolInvocation.StateKind.WaitingForConfirmation || !state.confirmationMessages?.title) { @@ -67,6 +69,7 @@ export class ChatModifiedFilesConfirmationSubPart extends AbstractToolConfirmati subtitle: typeof toolInvocation.originMessage === 'string' ? toolInvocation.originMessage : toolInvocation.originMessage?.value, buttons: this.createButtons(data.options), message: this.createWidgetContentElement(state.confirmationMessages.message, data), + footerBanner: this.createRiskBadgeDomNode(state.parameters), } )); diff --git a/src/vs/workbench/contrib/chat/browser/widget/chatContentParts/toolInvocationParts/chatTerminalToolConfirmationSubPart.ts b/src/vs/workbench/contrib/chat/browser/widget/chatContentParts/toolInvocationParts/chatTerminalToolConfirmationSubPart.ts index 518aa7c2bd35a..9b4b9c61927fa 100644 --- a/src/vs/workbench/contrib/chat/browser/widget/chatContentParts/toolInvocationParts/chatTerminalToolConfirmationSubPart.ts +++ b/src/vs/workbench/contrib/chat/browser/widget/chatContentParts/toolInvocationParts/chatTerminalToolConfirmationSubPart.ts @@ -8,7 +8,6 @@ import { HoverStyle } from '../../../../../../../base/browser/ui/hover/hover.js' import { HoverPosition } from '../../../../../../../base/browser/ui/hover/hoverWidget.js'; import { Separator } from '../../../../../../../base/common/actions.js'; import { asArray } from '../../../../../../../base/common/arrays.js'; -import { CancellationTokenSource } from '../../../../../../../base/common/cancellation.js'; import { Codicon } from '../../../../../../../base/common/codicons.js'; import { ErrorNoTelemetry } from '../../../../../../../base/common/errors.js'; import { createCommandUri, escapeMarkdownSyntaxTokens, MarkdownString, type IMarkdownString } from '../../../../../../../base/common/htmlContent.js'; @@ -41,7 +40,7 @@ import { IChatContentPartRenderContext } from '../chatContentParts.js'; import { ChatMarkdownContentPart } from '../chatMarkdownContentPart.js'; import { CodeBlockPart, ICodeBlockRenderOptions } from '../codeBlockPart.js'; import { BaseChatToolInvocationSubPart } from './chatToolInvocationSubPart.js'; -import { ToolRiskBadgeWidget } from './toolRiskBadgeWidget.js'; +import { createToolRiskBadge } from './toolRiskBadgeHelper.js'; export const enum TerminalToolConfirmationStorageKeys { TerminalAutoApproveWarningAccepted = 'chat.tools.terminal.autoApprove.warningAccepted' @@ -191,7 +190,7 @@ export class ChatTerminalToolConfirmationSubPart extends BaseChatToolInvocationS position: { hoverPosition: HoverPosition.LEFT }, })); - const riskBadge = this._createRiskBadge(state.parameters); + const riskBadge = createToolRiskBadge(this._store, this.instantiationService, this.riskAssessmentService, this.languageModelToolsService, this.toolInvocation.toolId, state.parameters, 'terminal'); const confirmWidget = this._register(this.instantiationService.createInstance( ChatCustomConfirmationWidget, @@ -507,43 +506,6 @@ export class ChatTerminalToolConfirmationSubPart extends BaseChatToolInvocationS return promptResult.result === true; } - private _createRiskBadge(parameters: unknown): ToolRiskBadgeWidget | undefined { - if (!this.riskAssessmentService.isEnabled()) { - return undefined; - } - const tool = this.languageModelToolsService.getTool(this.toolInvocation.toolId); - if (!tool) { - return undefined; - } - const widget = this._register(this.instantiationService.createInstance(ToolRiskBadgeWidget)); - const cached = this.riskAssessmentService.getCached(tool, parameters); - if (cached) { - widget.setAssessment(cached); - } else { - widget.setLoading(); - const cts = new CancellationTokenSource(); - this._register(toDisposable(() => cts.dispose(true))); - (async () => { - try { - const result = await this.riskAssessmentService.assess(tool, parameters, cts.token); - if (cts.token.isCancellationRequested || widget.isDisposed) { - return; - } - if (!result) { - widget.setHidden(); - return; - } - widget.setAssessment(result); - } catch { - if (!widget.isDisposed) { - widget.setHidden(); - } - } - })(); - } - return widget; - } - private _appendMarkdownPart(container: HTMLElement, message: string | IMarkdownString, codeBlockRenderOptions: ICodeBlockRenderOptions) { const part = this._register(this.instantiationService.createInstance(ChatMarkdownContentPart, { diff --git a/src/vs/workbench/contrib/chat/browser/widget/chatContentParts/toolInvocationParts/chatToolConfirmationSubPart.ts b/src/vs/workbench/contrib/chat/browser/widget/chatContentParts/toolInvocationParts/chatToolConfirmationSubPart.ts index 2ada03f701601..1cf66bb3a5718 100644 --- a/src/vs/workbench/contrib/chat/browser/widget/chatContentParts/toolInvocationParts/chatToolConfirmationSubPart.ts +++ b/src/vs/workbench/contrib/chat/browser/widget/chatContentParts/toolInvocationParts/chatToolConfirmationSubPart.ts @@ -25,6 +25,7 @@ import { createToolSchemaUri, ILanguageModelToolsService, IToolConfirmationMessa import { ILanguageModelToolsConfirmationService } from '../../../../common/tools/languageModelToolsConfirmationService.js'; import { AcceptToolConfirmationActionId, SkipToolConfirmationActionId } from '../../../actions/chatToolActions.js'; import { IChatCodeBlockInfo, IChatWidgetService } from '../../../chat.js'; +import { IChatToolRiskAssessmentService } from '../../../tools/chatToolRiskAssessmentService.js'; import { renderFileWidgets } from '../chatInlineAnchorWidget.js'; import { CodeBlockPart, ICodeBlockRenderOptions } from '../codeBlockPart.js'; import { IChatContentPartRenderContext } from '../chatContentParts.js'; @@ -58,13 +59,14 @@ export class ToolConfirmationSubPart extends AbstractToolConfirmationSubPart { @ILanguageModelToolsService languageModelToolsService: ILanguageModelToolsService, @IChatMarkdownAnchorService private readonly chatMarkdownAnchorService: IChatMarkdownAnchorService, @ILanguageModelToolsConfirmationService private readonly confirmationService: ILanguageModelToolsConfirmationService, + @IChatToolRiskAssessmentService riskAssessmentService: IChatToolRiskAssessmentService, ) { const state = toolInvocation.state.get(); if (state.type !== IChatToolInvocation.StateKind.WaitingForConfirmation || !state.confirmationMessages?.title) { throw new Error('Confirmation messages are missing'); } - super(toolInvocation, context, instantiationService, keybindingService, contextKeyService, chatWidgetService, languageModelToolsService); + super(toolInvocation, context, instantiationService, keybindingService, contextKeyService, chatWidgetService, languageModelToolsService, riskAssessmentService); this.render({ allowActionId: AcceptToolConfirmationActionId, diff --git a/src/vs/workbench/contrib/chat/browser/widget/chatContentParts/toolInvocationParts/chatToolPostExecuteConfirmationPart.ts b/src/vs/workbench/contrib/chat/browser/widget/chatContentParts/toolInvocationParts/chatToolPostExecuteConfirmationPart.ts index da2b868f353e9..654129d02122e 100644 --- a/src/vs/workbench/contrib/chat/browser/widget/chatContentParts/toolInvocationParts/chatToolPostExecuteConfirmationPart.ts +++ b/src/vs/workbench/contrib/chat/browser/widget/chatContentParts/toolInvocationParts/chatToolPostExecuteConfirmationPart.ts @@ -16,6 +16,7 @@ import { ILanguageModelToolsConfirmationService } from '../../../../common/tools import { ILanguageModelToolsService, IToolResultDataPart, IToolResultPromptTsxPart, IToolResultTextPart, stringifyPromptTsxPart } from '../../../../common/tools/languageModelToolsService.js'; import { AcceptToolPostConfirmationActionId, SkipToolPostConfirmationActionId } from '../../../actions/chatToolActions.js'; import { IChatCodeBlockInfo, IChatWidgetService } from '../../../chat.js'; +import { IChatToolRiskAssessmentService } from '../../../tools/chatToolRiskAssessmentService.js'; import { IChatContentPartRenderContext } from '../chatContentParts.js'; import { ChatCollapsibleIOPart } from '../chatToolInputOutputContentPart.js'; import { ChatToolOutputContentSubPart } from '../chatToolOutputContentSubPart.js'; @@ -36,8 +37,9 @@ export class ChatToolPostExecuteConfirmationPart extends AbstractToolConfirmatio @IChatWidgetService chatWidgetService: IChatWidgetService, @ILanguageModelToolsService languageModelToolsService: ILanguageModelToolsService, @ILanguageModelToolsConfirmationService private readonly confirmationService: ILanguageModelToolsConfirmationService, + @IChatToolRiskAssessmentService riskAssessmentService: IChatToolRiskAssessmentService, ) { - super(toolInvocation, context, instantiationService, keybindingService, contextKeyService, chatWidgetService, languageModelToolsService); + super(toolInvocation, context, instantiationService, keybindingService, contextKeyService, chatWidgetService, languageModelToolsService, riskAssessmentService); const subtitle = toolInvocation.pastTenseMessage || toolInvocation.invocationMessage; this.render({ allowActionId: AcceptToolPostConfirmationActionId, diff --git a/src/vs/workbench/contrib/chat/browser/widget/chatContentParts/toolInvocationParts/toolRiskBadgeHelper.ts b/src/vs/workbench/contrib/chat/browser/widget/chatContentParts/toolInvocationParts/toolRiskBadgeHelper.ts new file mode 100644 index 0000000000000..e1dedd48d2084 --- /dev/null +++ b/src/vs/workbench/contrib/chat/browser/widget/chatContentParts/toolInvocationParts/toolRiskBadgeHelper.ts @@ -0,0 +1,71 @@ +/*--------------------------------------------------------------------------------------------- + * Copyright (c) Microsoft Corporation. All rights reserved. + * Licensed under the MIT License. See License.txt in the project root for license information. + *--------------------------------------------------------------------------------------------*/ + +import { CancellationTokenSource } from '../../../../../../../base/common/cancellation.js'; +import { DisposableStore, toDisposable } from '../../../../../../../base/common/lifecycle.js'; +import { IInstantiationService } from '../../../../../../../platform/instantiation/common/instantiation.js'; +import { IChatToolRiskAssessmentService, ToolRiskPromptKind } from '../../../tools/chatToolRiskAssessmentService.js'; +import { ILanguageModelToolsService } from '../../../../common/tools/languageModelToolsService.js'; +import { ToolRiskBadgeWidget } from './toolRiskBadgeWidget.js'; + +/** + * Creates a {@link ToolRiskBadgeWidget} for a tool confirmation surface, or `undefined` when the + * feature is disabled or the tool is unknown. A cached assessment renders synchronously; otherwise + * the badge shows a loading state and assesses asynchronously, hiding itself on failure. + * + * The widget and its assessment token are registered on `store`, so disposing the store cancels + * any in-flight assessment. The widget is returned so terminal confirmations can attach + * `setDetails` / `onDidHide`; most callers only need its `domNode` as a `footerBanner`. + * + * `kind` selects the rubric (terminal vs. generic); when omitted it is auto-detected from the tool id. + */ +export function createToolRiskBadge( + store: DisposableStore, + instantiationService: IInstantiationService, + riskAssessmentService: IChatToolRiskAssessmentService, + languageModelToolsService: ILanguageModelToolsService, + toolId: string, + parameters: unknown, + kind?: ToolRiskPromptKind, +): ToolRiskBadgeWidget | undefined { + // Check the feature flag before the tool lookup so it is skipped when disabled. + if (!riskAssessmentService.isEnabled()) { + return undefined; + } + + const tool = languageModelToolsService.getTool(toolId); + if (!tool) { + return undefined; + } + + const widget = store.add(instantiationService.createInstance(ToolRiskBadgeWidget)); + const cached = riskAssessmentService.getCached(tool, parameters, kind); + if (cached) { + widget.setAssessment(cached); + return widget; + } + + widget.setLoading(); + const cts = new CancellationTokenSource(); + store.add(toDisposable(() => cts.dispose(true))); + (async () => { + try { + const result = await riskAssessmentService.assess(tool, parameters, cts.token, kind); + if (cts.token.isCancellationRequested || widget.isDisposed) { + return; + } + if (!result) { + widget.setHidden(); + return; + } + widget.setAssessment(result); + } catch { + if (!widget.isDisposed) { + widget.setHidden(); + } + } + })(); + return widget; +} diff --git a/src/vs/workbench/contrib/chat/common/constants.ts b/src/vs/workbench/contrib/chat/common/constants.ts index 2e073a826c561..7a715331fe856 100644 --- a/src/vs/workbench/contrib/chat/common/constants.ts +++ b/src/vs/workbench/contrib/chat/common/constants.ts @@ -130,6 +130,15 @@ export function isAutoApproveLevel(level: ChatPermissionLevel | undefined): bool return level === ChatPermissionLevel.AutoApprove || level === ChatPermissionLevel.Autopilot; } +/** + * True for {@link ChatPermissionLevel.Autopilot} only. Unlike {@link isAutoApproveLevel}, this + * excludes {@link ChatPermissionLevel.AutoApprove}, so it can gate Autopilot-only behavior such as + * risk-based skipping of tool calls. + */ +export function isAutopilotLevel(level: ChatPermissionLevel | undefined): boolean { + return level === ChatPermissionLevel.Autopilot; +} + // Thinking display modes for pinned content export enum ThinkingDisplayMode { Collapsed = 'collapsed', diff --git a/src/vs/workbench/contrib/chat/test/browser/tools/languageModelToolsService.test.ts b/src/vs/workbench/contrib/chat/test/browser/tools/languageModelToolsService.test.ts index 0ec94c6fe35af..2eb04ea38e088 100644 --- a/src/vs/workbench/contrib/chat/test/browser/tools/languageModelToolsService.test.ts +++ b/src/vs/workbench/contrib/chat/test/browser/tools/languageModelToolsService.test.ts @@ -6,7 +6,7 @@ import * as assert from 'assert'; import { Barrier } from '../../../../../../base/common/async.js'; import { VSBuffer } from '../../../../../../base/common/buffer.js'; -import { CancellationToken } from '../../../../../../base/common/cancellation.js'; +import { CancellationToken, CancellationTokenSource } from '../../../../../../base/common/cancellation.js'; import { CancellationError, isCancellationError } from '../../../../../../base/common/errors.js'; import { URI } from '../../../../../../base/common/uri.js'; import { ensureNoDisposablesAreLeakedInTestSuite } from '../../../../../../base/test/common/utils.js'; @@ -23,8 +23,9 @@ import { ConfirmationOptionKind } from '../../../../../../platform/agentHost/com import { ITelemetryService } from '../../../../../../platform/telemetry/common/telemetry.js'; import { workbenchInstantiationService } from '../../../../../test/browser/workbenchTestServices.js'; import { LanguageModelToolsService } from '../../../browser/tools/languageModelToolsService.js'; +import { IChatToolRiskAssessmentService, IToolRiskAssessment, ToolRiskLevel, ToolRiskPromptKind } from '../../../browser/tools/chatToolRiskAssessmentService.js'; import { ChatModel, IChatModel } from '../../../common/model/chatModel.js'; -import { IChatService, IChatToolInputInvocationData, IChatToolInvocation, ToolConfirmKind } from '../../../common/chatService/chatService.js'; +import { IChatService, IChatProgress, IChatInfoMessage, IChatToolInputInvocationData, IChatToolInvocation, ToolConfirmKind } from '../../../common/chatService/chatService.js'; import { ChatConfiguration, ChatPermissionLevel } from '../../../common/constants.js'; import { SpecedToolAliases, isToolResultInputOutputDetails, IToolData, IToolImpl, IToolInvocation, ToolDataSource, ToolSet, IToolResultTextPart } from '../../../common/tools/languageModelToolsService.js'; import { MockChatService } from '../../common/chatService/mockChatService.js'; @@ -70,6 +71,45 @@ class TestTelemetryService implements Partial { } } +/** + * Configurable stub for {@link IChatToolRiskAssessmentService}. `enabled` models the + * `chat.tools.riskAssessment.enabled` confirmation-badge setting; tests that exercise the + * gate set `assessment`, `assessError`, or `onAssess` and inspect `assessCalls`. Note the + * Autopilot gate is independent of `enabled` (it passes `ignoreEnablement`), so the gate's + * opt-in is driven by Advanced Autopilot, not this field. + */ +class TestChatToolRiskAssessmentService implements IChatToolRiskAssessmentService { + declare readonly _serviceBrand: undefined; + + public enabled = false; + public assessment: IToolRiskAssessment | undefined = undefined; + public assessError: Error | undefined = undefined; + /** Invoked synchronously at the start of {@link assess} so tests can cancel mid-flight. */ + public onAssess: (() => void) | undefined = undefined; + public readonly assessCalls: { toolId: string; parameters: unknown; kind?: ToolRiskPromptKind }[] = []; + + isEnabled(): boolean { + return this.enabled; + } + + getCached(): IToolRiskAssessment | undefined { + return undefined; + } + + async assess(tool: IToolData, parameters: unknown, _token: CancellationToken, kind?: ToolRiskPromptKind, options?: { ignoreEnablement?: boolean }): Promise { + this.assessCalls.push({ toolId: tool.id, parameters, kind }); + this.onAssess?.(); + // Mirror the real service: honor the badge setting unless the caller opts out. + if (!options?.ignoreEnablement && !this.enabled) { + return undefined; + } + if (this.assessError) { + throw this.assessError; + } + return this.assessment; + } +} + function registerToolForTest(service: LanguageModelToolsService, store: any, id: string, impl: IToolImpl, data?: Partial) { const toolData: IToolData = { id, @@ -121,6 +161,7 @@ interface TestToolsServiceSetup { chatService: MockChatService; service: LanguageModelToolsService; contextKeyService: IContextKeyService; + riskAssessmentService: TestChatToolRiskAssessmentService; } interface TestToolsServiceOptions { @@ -152,6 +193,8 @@ function createTestToolsService(store: ReturnType Promise<{ content: { value: string }[] }>; wasInvoked: () => boolean } { + const withConfirmation = opts?.withConfirmation ?? true; + const permissionLevel = opts?.permissionLevel ?? ChatPermissionLevel.Autopilot; + const advancedEnabled = opts?.advancedEnabled ?? true; + const toolId = opts?.toolId ?? 'riskGateTool'; + + setup.configurationService.setUserConfiguration(ChatConfiguration.AutopilotAdvancedEnabled, advancedEnabled); + setup.configurationService.setUserConfiguration('chat.tools.global.autoApprove', false); + + let invoked = false; + const tool = registerToolForTest(setup.service, store, toolId, { + prepareToolInvocation: async () => (withConfirmation ? { confirmationMessages: { title: 'Confirm?', message: 'Proceed?' } } : {}), + invoke: async () => { invoked = true; return { content: [{ kind: 'text', value: 'ran' }] }; }, + }); + + const sessionId = 'riskGateSession'; + stubGetSession(setup.chatService, sessionId, { requestId: 'req-risk', modeInfo: { permissionLevel } }); + + return { + invoke: (token: CancellationToken = CancellationToken.None) => setup.service.invokeTool(tool.makeDto({ x: 1 }, { sessionId }), async () => 0, token) as Promise<{ content: { value: string }[] }>, + wasInvoked: () => invoked, + }; } suite('LanguageModelToolsService', () => { @@ -1612,6 +1689,275 @@ suite('LanguageModelToolsService', () => { assert.strictEqual(result.content[0].value, 'terminal executed'); }); + test('autopilot risk gate skips a tool assessed as high-risk (red)', async () => { + const setup = createTestToolsService(store); + setup.riskAssessmentService.enabled = true; + setup.riskAssessmentService.assessment = { risk: ToolRiskLevel.Red, explanation: 'Deletes source files irreversibly.' }; + const t = setupRiskGateTool(setup, store); + + const result = await t.invoke(); + + assert.deepStrictEqual( + { + invoked: t.wasInvoked(), + assessCalls: setup.riskAssessmentService.assessCalls.length, + mentionsRisk: String(result.content[0].value).includes('Deletes source files irreversibly.'), + }, + { invoked: false, assessCalls: 1, mentionsRisk: true }, + ); + }); + + test('autopilot risk gate allows a low-risk (green) tool call', async () => { + const setup = createTestToolsService(store); + setup.riskAssessmentService.enabled = true; + setup.riskAssessmentService.assessment = { risk: ToolRiskLevel.Green, explanation: 'Reads a file.' }; + const t = setupRiskGateTool(setup, store); + + const result = await t.invoke(); + + assert.deepStrictEqual( + { invoked: t.wasInvoked(), assessCalls: setup.riskAssessmentService.assessCalls.length, value: result.content[0].value }, + { invoked: true, assessCalls: 1, value: 'ran' }, + ); + }); + + test('autopilot risk gate allows a medium-risk (orange) tool call (red-only threshold)', async () => { + const setup = createTestToolsService(store); + setup.riskAssessmentService.enabled = true; + setup.riskAssessmentService.assessment = { risk: ToolRiskLevel.Orange, explanation: 'Edits a file.' }; + const t = setupRiskGateTool(setup, store); + + const result = await t.invoke(); + + assert.deepStrictEqual( + { invoked: t.wasInvoked(), assessCalls: setup.riskAssessmentService.assessCalls.length, value: result.content[0].value }, + { invoked: true, assessCalls: 1, value: 'ran' }, + ); + }); + + test('autopilot risk gate fails open when the classifier returns no assessment', async () => { + const setup = createTestToolsService(store); + setup.riskAssessmentService.enabled = true; + setup.riskAssessmentService.assessment = undefined; + const t = setupRiskGateTool(setup, store); + + const result = await t.invoke(); + + assert.deepStrictEqual( + { invoked: t.wasInvoked(), assessCalls: setup.riskAssessmentService.assessCalls.length, value: result.content[0].value }, + { invoked: true, assessCalls: 1, value: 'ran' }, + ); + }); + + test('autopilot risk gate fails open when the classifier throws', async () => { + const setup = createTestToolsService(store); + setup.riskAssessmentService.enabled = true; + setup.riskAssessmentService.assessError = new Error('network down'); + const t = setupRiskGateTool(setup, store); + + const result = await t.invoke(); + + assert.deepStrictEqual( + { invoked: t.wasInvoked(), value: result.content[0].value }, + { invoked: true, value: 'ran' }, + ); + }); + + test('autopilot risk gate does not assess tool calls that have no confirmation', async () => { + const setup = createTestToolsService(store); + setup.riskAssessmentService.enabled = true; + setup.riskAssessmentService.assessment = { risk: ToolRiskLevel.Red, explanation: 'should not matter' }; + const t = setupRiskGateTool(setup, store, { withConfirmation: false }); + + const result = await t.invoke(); + + assert.deepStrictEqual( + { invoked: t.wasInvoked(), assessCalls: setup.riskAssessmentService.assessCalls.length, value: result.content[0].value }, + { invoked: true, assessCalls: 0, value: 'ran' }, + ); + }); + + test('autopilot risk gate classifies a terminal command even when it has no confirmation', async () => { + // run_in_terminal suppresses its own confirmation under auto-approve sessions, so the + // gate must classify it anyway; a red command is skipped despite the missing confirmation. + const setup = createTestToolsService(store); + setup.riskAssessmentService.enabled = true; + setup.riskAssessmentService.assessment = { risk: ToolRiskLevel.Red, explanation: 'Force-pushes main, overwriting history.' }; + const t = setupRiskGateTool(setup, store, { withConfirmation: false, toolId: 'run_in_terminal' }); + + const result = await t.invoke(); + + assert.deepStrictEqual( + { + invoked: t.wasInvoked(), + assessCalls: setup.riskAssessmentService.assessCalls.length, + isRiskMessage: String(result.content[0].value).startsWith('Autopilot skipped this tool call'), + }, + { invoked: false, assessCalls: 1, isRiskMessage: true }, + ); + }); + + test('autopilot risk gate runs a non-red terminal command that has no confirmation', async () => { + // A terminal command is always classified in Autopilot, but a non-red verdict still runs. + const setup = createTestToolsService(store); + setup.riskAssessmentService.enabled = true; + setup.riskAssessmentService.assessment = { risk: ToolRiskLevel.Orange, explanation: 'Installs a package.' }; + const t = setupRiskGateTool(setup, store, { withConfirmation: false, toolId: 'run_in_terminal' }); + + const result = await t.invoke(); + + assert.deepStrictEqual( + { invoked: t.wasInvoked(), assessCalls: setup.riskAssessmentService.assessCalls.length, value: result.content[0].value }, + { invoked: true, assessCalls: 1, value: 'ran' }, + ); + }); + + test('autopilot risk gate classifies a fetch web page call even when it has no confirmation', async () => { + // Fetch web page tools auto-approve themselves (URL in the prompt / trusted domain) and so + // surface no confirmation; the gate must classify them anyway so a dangerous URL (e.g. one + // injected into the prompt to exfiltrate secrets) is still skipped when assessed red. + const setup = createTestToolsService(store); + setup.riskAssessmentService.enabled = true; + setup.riskAssessmentService.assessment = { risk: ToolRiskLevel.Red, explanation: 'Sends workspace secrets to an untrusted host.' }; + const t = setupRiskGateTool(setup, store, { withConfirmation: false, toolId: 'vscode_fetchWebPage_internal' }); + + const result = await t.invoke(); + + assert.deepStrictEqual( + { + invoked: t.wasInvoked(), + assessCalls: setup.riskAssessmentService.assessCalls.length, + isRiskMessage: String(result.content[0].value).startsWith('Autopilot skipped this tool call'), + }, + { invoked: false, assessCalls: 1, isRiskMessage: true }, + ); + }); + + test('autopilot risk gate runs a non-red fetch web page call that has no confirmation', async () => { + // A fetch is always classified in Autopilot, but a non-red verdict still runs. + const setup = createTestToolsService(store); + setup.riskAssessmentService.enabled = true; + setup.riskAssessmentService.assessment = { risk: ToolRiskLevel.Green, explanation: 'Fetches public documentation.' }; + const t = setupRiskGateTool(setup, store, { withConfirmation: false, toolId: 'copilot_fetchWebPage' }); + + const result = await t.invoke(); + + assert.deepStrictEqual( + { invoked: t.wasInvoked(), assessCalls: setup.riskAssessmentService.assessCalls.length, value: result.content[0].value }, + { invoked: true, assessCalls: 1, value: 'ran' }, + ); + }); + + test('autopilot risk gate is inert when Advanced Autopilot is disabled', async () => { + const setup = createTestToolsService(store); + setup.riskAssessmentService.enabled = true; + setup.riskAssessmentService.assessment = { risk: ToolRiskLevel.Red, explanation: 'should not matter' }; + const t = setupRiskGateTool(setup, store, { advancedEnabled: false }); + + const result = await t.invoke(); + + assert.deepStrictEqual( + { invoked: t.wasInvoked(), assessCalls: setup.riskAssessmentService.assessCalls.length, value: result.content[0].value }, + { invoked: true, assessCalls: 0, value: 'ran' }, + ); + }); + + test('autopilot risk gate does not apply at the plain Auto-Approve level', async () => { + const setup = createTestToolsService(store); + setup.riskAssessmentService.enabled = true; + setup.riskAssessmentService.assessment = { risk: ToolRiskLevel.Red, explanation: 'should not matter' }; + const t = setupRiskGateTool(setup, store, { permissionLevel: ChatPermissionLevel.AutoApprove }); + + const result = await t.invoke(); + + assert.deepStrictEqual( + { invoked: t.wasInvoked(), assessCalls: setup.riskAssessmentService.assessCalls.length, value: result.content[0].value }, + { invoked: true, assessCalls: 0, value: 'ran' }, + ); + }); + + test('autopilot risk gate runs even when the risk assessment badge setting is disabled', async () => { + // The gate is independent of chat.tools.riskAssessment.enabled (which only controls the + // confirmation risk badge): a red verdict still skips the call. Also verifies the gate + // passes ignoreEnablement — without it the stub would return undefined and the tool would run. + const setup = createTestToolsService(store); + setup.riskAssessmentService.enabled = false; + setup.riskAssessmentService.assessment = { risk: ToolRiskLevel.Red, explanation: 'Deletes source files irreversibly.' }; + const t = setupRiskGateTool(setup, store); + + const result = await t.invoke(); + + assert.deepStrictEqual( + { + invoked: t.wasInvoked(), + assessCalls: setup.riskAssessmentService.assessCalls.length, + isRiskMessage: String(result.content[0].value).startsWith('Autopilot skipped this tool call'), + }, + { invoked: false, assessCalls: 1, isRiskMessage: true }, + ); + }); + + test('autopilot risk gate skips on red even when the classifier explanation is empty', async () => { + const setup = createTestToolsService(store); + setup.riskAssessmentService.enabled = true; + setup.riskAssessmentService.assessment = { risk: ToolRiskLevel.Red, explanation: '' }; + const t = setupRiskGateTool(setup, store); + + const result = await t.invoke(); + + // The skip must still read as an automated risk-skip, never the user-skip fallback message. + assert.deepStrictEqual( + { + invoked: t.wasInvoked(), + assessCalls: setup.riskAssessmentService.assessCalls.length, + isRiskMessage: String(result.content[0].value).startsWith('Autopilot skipped this tool call'), + isUserSkipMessage: String(result.content[0].value).includes('The user chose to skip'), + }, + { invoked: false, assessCalls: 1, isRiskMessage: true, isUserSkipMessage: false }, + ); + }); + + test('autopilot risk gate does not skip when cancelled during assessment', async () => { + const setup = createTestToolsService(store); + setup.riskAssessmentService.enabled = true; + setup.riskAssessmentService.assessment = { risk: ToolRiskLevel.Red, explanation: 'Deletes source files irreversibly.' }; + const t = setupRiskGateTool(setup, store); + + // Cancel synchronously while the classifier is running: the gate must abandon the + // assessment and propagate cancellation rather than mask it as a risk-skip result. + const cts = store.add(new CancellationTokenSource()); + setup.riskAssessmentService.onAssess = () => cts.cancel(); + + await assert.rejects(() => t.invoke(cts.token), err => isCancellationError(err)); + assert.deepStrictEqual( + { invoked: t.wasInvoked(), assessCalls: setup.riskAssessmentService.assessCalls.length }, + { invoked: false, assessCalls: 1 }, + ); + }); + + test('autopilot risk gate surfaces an info note to the user when it skips a high-risk tool', async () => { + const setup = createTestToolsService(store); + setup.riskAssessmentService.enabled = true; + setup.riskAssessmentService.assessment = { risk: ToolRiskLevel.Red, explanation: 'Deletes source files irreversibly.' }; + const t = setupRiskGateTool(setup, store); + + // The tool invocation part hides itself after completion, so the reason is surfaced + // as a separate info note appended to the response stream. + const progresses: IChatProgress[] = []; + setup.chatService.appendProgress = (_request, progress) => { progresses.push(progress); }; + + await t.invoke(); + + const info = progresses.find((p): p is IChatInfoMessage => p.kind === 'info'); + assert.deepStrictEqual( + { + hasInfo: !!info, + mentionsRisk: !!info && info.content.value.includes('Deletes source files irreversibly.'), + }, + { hasInfo: true, mentionsRisk: true }, + ); + }); + test('bypass approvals auto-approves terminal tool with confirmation messages', async () => { const { service: testService, chatService: testChatService } = createTestToolsService(store, { configureServices: config => {