updatedInput for PreToolUse (microsoft#293575)

Author: roblourens

src/vs/workbench/contrib/chat/browser/tools/languageModelToolsService.ts

@@ -23,6 +23,7 @@ import { ThemeIcon } from '../../../../../base/common/themables.js';
 import { localize, localize2 } from '../../../../../nls.js';
 import { IAccessibilityService } from '../../../../../platform/accessibility/common/accessibility.js';
 import { AccessibilitySignal, IAccessibilitySignalService } from '../../../../../platform/accessibilitySignal/browser/accessibilitySignalService.js';
+import { ICommandService } from '../../../../../platform/commands/common/commands.js';
 import { IConfigurationService } from '../../../../../platform/configuration/common/configuration.js';
 import { IContextKey, IContextKeyService } from '../../../../../platform/contextkey/common/contextkey.js';
 import { IDialogService } from '../../../../../platform/dialogs/common/dialogs.js';
@@ -122,6 +123,7 @@ export class LanguageModelToolsService extends Disposable implements ILanguageMo
 		@IStorageService private readonly _storageService: IStorageService,
 		@ILanguageModelToolsConfirmationService private readonly _confirmationService: ILanguageModelToolsConfirmationService,
 		@IHooksExecutionService private readonly _hooksExecutionService: IHooksExecutionService,
+		@ICommandService private readonly _commandService: ICommandService,
 	) {
 		super();
 
@@ -422,6 +424,38 @@ export class LanguageModelToolsService extends Disposable implements ILanguageMo
 		return { hookResult };
 	}
 
+	/**
+	 * Validate updatedInput from a preToolUse hook against the tool's input schema
+	 * using the json.validate command from the JSON extension.
+	 * @returns An error message string if validation fails, or undefined if valid.
+	 */
+	private async _validateUpdatedInput(toolId: string, toolData: IToolData | undefined, updatedInput: object): Promise<string | undefined> {
+		if (!toolData?.inputSchema) {
+			return undefined;
+		}
+
+		type JsonDiagnostic = {
+			message: string;
+			range: { line: number; character: number }[];
+			severity: string;
+			code?: string | number;
+		};
+
+		try {
+			const schemaUri = createToolSchemaUri(toolId);
+			const inputJson = JSON.stringify(updatedInput);
+			const diagnostics = await this._commandService.executeCommand<JsonDiagnostic[]>('json.validate', schemaUri, inputJson) || [];
+			if (diagnostics.length > 0) {
+				return diagnostics.map(d => d.message).join('; ');
+			}
+		} catch (e) {
+			// json extension may not be available; skip validation
+			this._logService.debug(`[LanguageModelToolsService#_validateUpdatedInput] json.validate command failed, skipping validation: ${toErrorMessage(e)}`);
+		}
+
+		return undefined;
+	}
+
 	/**
 	 * Execute the postToolUse hook after tool completion.
 	 * If the hook returns a "block" decision, additional context is appended to the tool result
@@ -517,6 +551,17 @@ export class LanguageModelToolsService extends Disposable implements ILanguageMo
 			return hookDenialResult;
 		}
 
+		// Apply updatedInput from preToolUse hook if provided, after validating against the tool's input schema
+		if (preToolUseHookResult?.updatedInput) {
+			const validationError = await this._validateUpdatedInput(dto.toolId, toolData, preToolUseHookResult.updatedInput);
+			if (validationError) {
+				this._logService.warn(`[LanguageModelToolsService#invokeTool] Tool ${dto.toolId} updatedInput from preToolUse hook failed schema validation: ${validationError}`);
+			} else {
+				this._logService.debug(`[LanguageModelToolsService#invokeTool] Tool ${dto.toolId} input modified by preToolUse hook`);
+				dto.parameters = preToolUseHookResult.updatedInput;
+			}
+		}
+
 		// Fire the event to notify listeners that a tool is being invoked
 		this._onDidInvokeTool.fire({
 			toolId: dto.toolId,

src/vs/workbench/contrib/chat/common/hooks/hooksCommandTypes.ts

@@ -88,6 +88,7 @@ export interface IPreToolUseCommandOutput extends IHookCommandOutput {
 		readonly hookEventName?: string;
 		readonly permissionDecision?: 'allow' | 'deny';
 		readonly permissionDecisionReason?: string;
+		readonly updatedInput?: object;
 		readonly additionalContext?: string;
 	};
 }

src/vs/workbench/contrib/chat/common/hooks/hooksExecutionService.ts

@@ -390,6 +390,7 @@ export class HooksExecutionService extends Disposable implements IHooksExecution
 		let mostRestrictiveDecision: PreToolUsePermissionDecision | undefined;
 		let winningResult: IHookResult | undefined;
 		let winningReason: string | undefined;
+		let lastUpdatedInput: object | undefined;
 
 		for (const result of results) {
 			if (result.success && typeof result.output === 'object' && result.output !== null) {
@@ -408,6 +409,11 @@ export class HooksExecutionService extends Disposable implements IHooksExecution
 							allAdditionalContext.push(hookSpecificOutput.additionalContext);
 						}
 
+						// Track the last updatedInput (later hooks override earlier ones)
+						if (hookSpecificOutput.updatedInput) {
+							lastUpdatedInput = hookSpecificOutput.updatedInput;
+						}
+
 						// Track the most restrictive decision: deny > ask > allow
 						const decision = hookSpecificOutput.permissionDecision;
 						if (decision && this._isMoreRestrictive(decision, mostRestrictiveDecision)) {
@@ -422,14 +428,16 @@ export class HooksExecutionService extends Disposable implements IHooksExecution
 			}
 		}
 
-		if (!mostRestrictiveDecision || !winningResult) {
+		if (!mostRestrictiveDecision && !lastUpdatedInput && allAdditionalContext.length === 0) {
 			return undefined;
 		}
 
+		const baseResult = winningResult ?? results[0];
 		return {
-			...winningResult,
+			...baseResult,
 			permissionDecision: mostRestrictiveDecision,
 			permissionDecisionReason: winningReason,
+			updatedInput: lastUpdatedInput,
 			additionalContext: allAdditionalContext.length > 0 ? allAdditionalContext : undefined,
 		};
 	}

src/vs/workbench/contrib/chat/common/hooks/hooksTypes.ts

@@ -16,7 +16,7 @@
  * External types (in hooksCommandTypes.ts) define the contract with spawned commands.
  */
 
-import { vEnum, vObj, vOptionalProp, vString } from '../../../../../base/common/validation.js';
+import { vEnum, vObj, vObjAny, vOptionalProp, vString } from '../../../../../base/common/validation.js';
 
 //#region Common Hook Types
 
@@ -67,8 +67,9 @@ export interface IPreToolUseCallerInput {
 export const preToolUseOutputValidator = vObj({
 	hookSpecificOutput: vOptionalProp(vObj({
 		hookEventName: vOptionalProp(vString()),
-		permissionDecision: vEnum('allow', 'deny', 'ask'),
+		permissionDecision: vOptionalProp(vEnum('allow', 'deny', 'ask')),
 		permissionDecisionReason: vOptionalProp(vString()),
+		updatedInput: vOptionalProp(vObjAny()),
 		additionalContext: vOptionalProp(vString()),
 	})),
 });
@@ -88,6 +89,12 @@ export type PreToolUsePermissionDecision = 'allow' | 'deny' | 'ask';
 export interface IPreToolUseHookResult extends IHookResult {
 	readonly permissionDecision?: PreToolUsePermissionDecision;
 	readonly permissionDecisionReason?: string;
+	/**
+	 * Modified tool input parameters from the hook.
+	 * When set, replaces the original tool input before execution.
+	 * Combine with 'allow' to auto-approve, or 'ask' to show modified input to the user.
+	 */
+	readonly updatedInput?: object;
 	readonly additionalContext?: string[];
 }
 

src/vs/workbench/contrib/chat/test/browser/tools/languageModelToolsService.test.ts

@@ -14,6 +14,7 @@ import { ensureNoDisposablesAreLeakedInTestSuite } from '../../../../../../base/
 import { IAccessibilityService } from '../../../../../../platform/accessibility/common/accessibility.js';
 import { TestAccessibilityService } from '../../../../../../platform/accessibility/test/common/testAccessibilityService.js';
 import { AccessibilitySignal, IAccessibilitySignalService } from '../../../../../../platform/accessibilitySignal/browser/accessibilitySignalService.js';
+import { ICommandService } from '../../../../../../platform/commands/common/commands.js';
 import { ConfigurationTarget, IConfigurationChangeEvent } from '../../../../../../platform/configuration/common/configuration.js';
 import { TestConfigurationService } from '../../../../../../platform/configuration/test/common/testConfigurationService.js';
 import { ContextKeyService } from '../../../../../../platform/contextkey/browser/contextKeyService.js';
@@ -147,6 +148,7 @@ interface TestToolsServiceOptions {
 	accessibilitySignalService?: Partial<IAccessibilitySignalService>;
 	telemetryService?: Partial<ITelemetryService>;
 	hooksExecutionService?: MockHooksExecutionService;
+	commandService?: Partial<ICommandService>;
 	/** Called after configurationService is created but before the service is instantiated */
 	configureServices?: (config: TestConfigurationService) => void;
 }
@@ -181,6 +183,9 @@ function createTestToolsService(store: ReturnType<typeof ensureNoDisposablesAreL
 	if (options?.telemetryService) {
 		instaService.stub(ITelemetryService, options.telemetryService);
 	}
+	if (options?.commandService) {
+		instaService.stub(ICommandService, options.commandService as ICommandService);
+	}
 
 	const service = store.add(instaService.createInstance(LanguageModelToolsService));
 	return { configurationService, chatService, service, contextKeyService };
@@ -4037,6 +4042,152 @@ suite('LanguageModelToolsService', () => {
 			assert.strictEqual(result.content[0].kind, 'text');
 			assert.strictEqual((result.content[0] as IToolResultTextPart).value, 'success');
 		});
+
+		test('when hook returns updatedInput, tool is invoked with replaced parameters', async () => {
+			let receivedParameters: Record<string, any> | undefined;
+			mockHooksService.preToolUseHookResult = {
+				output: undefined,
+				success: true,
+				permissionDecision: 'allow',
+				updatedInput: { safeCommand: 'echo hello' },
+			};
+
+			const tool = registerToolForTest(hookService, store, 'hookUpdatedInputTool', {
+				invoke: async (dto) => {
+					receivedParameters = dto.parameters;
+					return { content: [{ kind: 'text', value: 'done' }] };
+				},
+				prepareToolInvocation: async () => ({
+					confirmationMessages: {
+						title: 'Confirm?',
+						message: 'Confirm action',
+						allowAutoConfirm: true
+					}
+				})
+			});
+
+			stubGetSession(hookChatService, 'hook-test-updated-input', { requestId: 'req1' });
+
+			await hookService.invokeTool(
+				tool.makeDto({ originalCommand: 'rm -rf /' }, { sessionId: 'hook-test-updated-input' }),
+				async () => 0,
+				CancellationToken.None
+			);
+
+			assert.deepStrictEqual(receivedParameters, { safeCommand: 'echo hello' });
+		});
+
+		test('when hook returns updatedInput that fails schema validation, original parameters are kept', async () => {
+			const mockCommandService = {
+				executeCommand: async (commandId: string) => {
+					if (commandId === 'json.validate') {
+						return [{ message: 'Missing required property "command"', range: [{ line: 0, character: 0 }, { line: 0, character: 1 }], severity: 'Error' }];
+					}
+					return undefined;
+				}
+			};
+
+			const mockHooks = new MockHooksExecutionService();
+			const setup = createTestToolsService(store, {
+				hooksExecutionService: mockHooks,
+				commandService: mockCommandService as ICommandService,
+			});
+
+			let receivedParameters: Record<string, any> | undefined;
+			mockHooks.preToolUseHookResult = {
+				output: undefined,
+				success: true,
+				permissionDecision: 'allow',
+				updatedInput: { invalidField: 'wrong' },
+			};
+
+			const tool = registerToolForTest(setup.service, store, 'hookValidationFailTool', {
+				invoke: async (dto) => {
+					receivedParameters = dto.parameters;
+					return { content: [{ kind: 'text', value: 'done' }] };
+				},
+				prepareToolInvocation: async () => ({
+					confirmationMessages: {
+						title: 'Confirm?',
+						message: 'Confirm action',
+						allowAutoConfirm: true
+					}
+				})
+			}, {
+				inputSchema: {
+					type: 'object',
+					properties: { command: { type: 'string' } },
+					required: ['command'],
+				}
+			});
+
+			stubGetSession(setup.chatService, 'hook-test-validation-fail', { requestId: 'req1' });
+
+			await setup.service.invokeTool(
+				tool.makeDto({ command: 'original' }, { sessionId: 'hook-test-validation-fail' }),
+				async () => 0,
+				CancellationToken.None
+			);
+
+			// Original parameters should be kept since validation failed
+			assert.deepStrictEqual(receivedParameters, { command: 'original' });
+		});
+
+		test('when hook returns updatedInput that passes schema validation, parameters are replaced', async () => {
+			const mockCommandService = {
+				executeCommand: async (commandId: string) => {
+					if (commandId === 'json.validate') {
+						return []; // no diagnostics = valid
+					}
+					return undefined;
+				}
+			};
+
+			const mockHooks = new MockHooksExecutionService();
+			const setup = createTestToolsService(store, {
+				hooksExecutionService: mockHooks,
+				commandService: mockCommandService as ICommandService,
+			});
+
+			let receivedParameters: Record<string, any> | undefined;
+			mockHooks.preToolUseHookResult = {
+				output: undefined,
+				success: true,
+				permissionDecision: 'allow',
+				updatedInput: { command: 'safe-command' },
+			};
+
+			const tool = registerToolForTest(setup.service, store, 'hookValidationPassTool', {
+				invoke: async (dto) => {
+					receivedParameters = dto.parameters;
+					return { content: [{ kind: 'text', value: 'done' }] };
+				},
+				prepareToolInvocation: async () => ({
+					confirmationMessages: {
+						title: 'Confirm?',
+						message: 'Confirm action',
+						allowAutoConfirm: true
+					}
+				})
+			}, {
+				inputSchema: {
+					type: 'object',
+					properties: { command: { type: 'string' } },
+					required: ['command'],
+				}
+			});
+
+			stubGetSession(setup.chatService, 'hook-test-validation-pass', { requestId: 'req1' });
+
+			await setup.service.invokeTool(
+				tool.makeDto({ command: 'original' }, { sessionId: 'hook-test-validation-pass' }),
+				async () => 0,
+				CancellationToken.None
+			);
+
+			// Updated parameters should be applied since validation passed
+			assert.deepStrictEqual(receivedParameters, { command: 'safe-command' });
+		});
 	});
 
 	suite('postToolUse hooks', () => {