Merge pull request microsoft#273131 from microsoft/tyriar/272704

Re-write && to ; in pwsh

Author: Tyriar

src/vs/workbench/contrib/terminalContrib/chatAgentTools/browser/commandSimplifier.ts

@@ -9,22 +9,32 @@ import { IWorkspaceContextService } from '../../../../../platform/workspace/comm
 import type { ITerminalInstance } from '../../../terminal/browser/terminal.js';
 import { isPowerShell } from './runInTerminalHelpers.js';
 import type { IRunInTerminalInputParams } from './tools/runInTerminalTool.js';
+import { TreeSitterCommandParserLanguage, type TreeSitterCommandParser } from './treeSitterCommandParser.js';
 
 export class CommandSimplifier {
 	constructor(
 		private readonly _osBackend: Promise<OperatingSystem>,
+		private readonly _treeSitterCommandParser: TreeSitterCommandParser,
 		@IWorkspaceContextService private readonly _workspaceContextService: IWorkspaceContextService,
 	) {
 	}
 
 	async rewriteIfNeeded(args: IRunInTerminalInputParams, instance: Pick<ITerminalInstance, 'getCwdResource'> | undefined, shell: string): Promise<string> {
-		const commandLine = args.command;
 		const os = await this._osBackend;
 
+		let commandLine = args.command;
+		commandLine = await this._removeRedundantCdPrefix(instance, commandLine, os, shell);
+		commandLine = await this._rewriteUnsupportedPwshChainOperators(commandLine, os, shell);
+
+		return commandLine;
+	}
+
+	private async _removeRedundantCdPrefix(instance: Pick<ITerminalInstance, 'getCwdResource'> | undefined, commandLine: string, os: OperatingSystem, shell: string): Promise<string> {
+		const isPwsh = isPowerShell(shell, os);
+
 		// Re-write the command if it starts with `cd <dir> && <suffix>` or `cd <dir>; <suffix>`
 		// to just `<suffix>` if the directory matches the current terminal's cwd. This simplifies
 		// the result in the chat by removing redundancies that some models like to add.
-		const isPwsh = isPowerShell(shell, os);
 		const cdPrefixMatch = commandLine.match(
 			isPwsh
 				? /^(?:cd(?: \/d)?|Set-Location(?: -Path)?) (?<dir>[^\s]+) ?(?:&&|;)\s+(?<suffix>.+)$/i
@@ -68,7 +78,25 @@ export class CommandSimplifier {
 				}
 			}
 		}
+		return commandLine;
+	}
 
+	private async _rewriteUnsupportedPwshChainOperators(commandLine: string, os: OperatingSystem, shell: string) {
+		// TODO: This should just be Windows PowerShell in the future when the powershell grammar
+		// supports chain operators https://github.com/airbus-cert/tree-sitter-powershell/issues/27
+		if (isPowerShell(shell, os)) {
+			const doubleAmpersandCaptures = await this._treeSitterCommandParser.queryTree(TreeSitterCommandParserLanguage.PowerShell, commandLine, [
+				'(',
+				'  (command',
+				'    (command_elements',
+				'      (generic_token) @double.ampersand',
+				'        (#eq? @double.ampersand "&&")))',
+				')',
+			].join('\n'));
+			for (const capture of doubleAmpersandCaptures.reverse()) {
+				commandLine = `${commandLine.substring(0, capture.node.startIndex)};${commandLine.substring(capture.node.endIndex)}`;
+			}
+		}
 		return commandLine;
 	}
 }

src/vs/workbench/contrib/terminalContrib/chatAgentTools/browser/tools/runInTerminalTool.ts

@@ -56,8 +56,9 @@ function createPowerShellModelDescription(shell: string): string {
 		`This tool allows you to execute ${isWinPwsh ? 'Windows PowerShell 5.1' : 'PowerShell'} commands in a persistent terminal session, preserving environment variables, working directory, and other context across multiple commands.`,
 		'',
 		'Command Execution:',
-		// Even for pwsh 7+ we want to use `;` to chain commands since the tree sitter grammar
-		// doesn't parse `&&`. See https://github.com/airbus-cert/tree-sitter-powershell/issues/27
+		// TODO: Even for pwsh 7+ we want to use `;` to chain commands since the tree sitter grammar
+		// doesn't parse `&&`. We want to change this to avoid `&&` only in Windows PowerShell when
+		// the grammar supports it https://github.com/airbus-cert/tree-sitter-powershell/issues/27
 		'- Use semicolons ; to chain commands on one line, NEVER use && even when asked explicitly',
 		'- Prefer pipelines | for object-based data flow',
 		'- Never create a sub-shell (eg. powershell -c "command") unless explicitly asked',
@@ -293,8 +294,8 @@ export class RunInTerminalTool extends Disposable implements IToolImpl {
 		this._osBackend = this._remoteAgentService.getEnvironment().then(remoteEnv => remoteEnv?.os ?? OS);
 
 		this._terminalToolCreator = _instantiationService.createInstance(ToolTerminalCreator);
-		this._commandSimplifier = _instantiationService.createInstance(CommandSimplifier, this._osBackend);
 		this._treeSitterCommandParser = this._instantiationService.createInstance(TreeSitterCommandParser);
+		this._commandSimplifier = _instantiationService.createInstance(CommandSimplifier, this._osBackend, this._treeSitterCommandParser);
 		this._telemetry = _instantiationService.createInstance(RunInTerminalToolTelemetry);
 		this._commandLineAutoApprover = this._register(_instantiationService.createInstance(CommandLineAutoApprover));
 		this._profileFetcher = _instantiationService.createInstance(TerminalProfileFetcher);

src/vs/workbench/contrib/terminalContrib/chatAgentTools/browser/treeSitterCommandParser.ts

@@ -6,7 +6,7 @@
 import { BugIndicatingError } from '../../../../../base/common/errors.js';
 import { derived, waitForState } from '../../../../../base/common/observable.js';
 import { ITreeSitterLibraryService } from '../../../../../editor/common/services/treeSitter/treeSitterLibraryService.js';
-import type { Language, Parser, Query } from '@vscode/tree-sitter-wasm';
+import type { Language, Parser, Query, QueryCapture } from '@vscode/tree-sitter-wasm';
 
 export const enum TreeSitterCommandParserLanguage {
 	Bash = 'bash',
@@ -42,9 +42,32 @@ export class TreeSitterCommandParser {
 
 		const captures = query.captures(tree.rootNode);
 		const subCommands = captures.map(e => e.node.text);
+
 		return subCommands;
 	}
 
+	async queryTree(languageId: TreeSitterCommandParserLanguage, commandLine: string, querySource: string): Promise<QueryCapture[]> {
+		const parser = await this._parser;
+		const language = await waitForState(derived(reader => {
+			return this._treeSitterLibraryService.getLanguage(languageId, true, reader);
+		}));
+		parser.setLanguage(language);
+
+		const tree = parser.parse(commandLine);
+		if (!tree) {
+			throw new BugIndicatingError('Failed to parse tree');
+		}
+
+		const query = await this._treeSitterLibraryService.createQuery(language, querySource);
+		if (!query) {
+			throw new BugIndicatingError('Failed to create tree sitter query');
+		}
+
+		const captures = query.captures(tree.rootNode);
+
+		return captures;
+	}
+
 	private async _getQuery(language: Language): Promise<Query> {
 		let query = this._queries.get(language);
 		if (!query) {

…s/test/browser/commandSimplifier.test.ts …ectron-browser/commandSimplifier.test.tssrc/vs/workbench/contrib/terminalContrib/chatAgentTools/test/browser/commandSimplifier.test.ts renamed to src/vs/workbench/contrib/terminalContrib/chatAgentTools/test/electron-browser/commandSimplifier.test.ts src/vs/workbench/contrib/terminalContrib/chatAgentTools/test/browser/commandSimplifier.test.ts renamed to src/vs/workbench/contrib/terminalContrib/chatAgentTools/test/electron-browser/commandSimplifier.test.ts

@@ -15,13 +15,21 @@ import type { ITerminalInstance } from '../../../../terminal/browser/terminal.js
 import { URI } from '../../../../../../base/common/uri.js';
 import type { IRunInTerminalInputParams } from '../../browser/tools/runInTerminalTool.js';
 import { Workspace } from '../../../../../../platform/workspace/test/common/testWorkspace.js';
+import { TreeSitterCommandParser } from '../../browser/treeSitterCommandParser.js';
+import { ITreeSitterLibraryService } from '../../../../../../editor/common/services/treeSitter/treeSitterLibraryService.js';
+import { TestIPCFileSystemProvider } from '../../../../../test/electron-browser/workbenchTestServices.js';
+import { NullLogService } from '../../../../../../platform/log/common/log.js';
+import { FileService } from '../../../../../../platform/files/common/fileService.js';
+import { Schemas } from '../../../../../../base/common/network.js';
+import { TreeSitterLibraryService } from '../../../../../services/treeSitter/browser/treeSitterLibraryService.js';
 
 suite('command re-writing', () => {
 	const store = ensureNoDisposablesAreLeakedInTestSuite();
 
 	let instantiationService: TestInstantiationService;
 	let workspaceService: TestContextService;
 
+	let parser: TreeSitterCommandParser;
 	let commandSimplifier: CommandSimplifier;
 
 	function createRewriteParams(command: string, chatSessionId?: string): IRunInTerminalInputParams {
@@ -45,14 +53,27 @@ suite('command re-writing', () => {
 	}
 
 	setup(() => {
-		instantiationService = workbenchInstantiationService(undefined, store);
+		const fileService = store.add(new FileService(new NullLogService()));
+		const fileSystemProvider = new TestIPCFileSystemProvider();
+		store.add(fileService.registerProvider(Schemas.file, fileSystemProvider));
+
+		instantiationService = workbenchInstantiationService({
+			fileService: () => fileService,
+		}, store);
+
+		const treeSitterLibraryService = store.add(instantiationService.createInstance(TreeSitterLibraryService));
+		treeSitterLibraryService.isTest = true;
+		instantiationService.stub(ITreeSitterLibraryService, treeSitterLibraryService);
+
+		parser = instantiationService.createInstance(TreeSitterCommandParser);
+
 		workspaceService = instantiationService.get(IWorkspaceContextService) as TestContextService;
 	});
 
 	suite('cd <cwd> && <suffix> -> <suffix>', () => {
 		(!isWindows ? suite : suite.skip)('Posix', () => {
 			setup(() => {
-				commandSimplifier = instantiationService.createInstance(CommandSimplifier, Promise.resolve(OperatingSystem.Linux));
+				commandSimplifier = instantiationService.createInstance(CommandSimplifier, Promise.resolve(OperatingSystem.Linux), parser);
 			});
 
 			test('should return original command when no cd prefix pattern matches', async () => {
@@ -168,7 +189,7 @@ suite('command re-writing', () => {
 
 		(isWindows ? suite : suite.skip)('Windows', () => {
 			setup(() => {
-				commandSimplifier = instantiationService.createInstance(CommandSimplifier, Promise.resolve(OperatingSystem.Windows));
+				commandSimplifier = instantiationService.createInstance(CommandSimplifier, Promise.resolve(OperatingSystem.Windows), parser);
 			});
 
 			test('should ignore any trailing back slash', async () => {
@@ -310,7 +331,7 @@ suite('command re-writing', () => {
 			test('should not rewrite cd /d when directory does not match cwd', async () => {
 				const testDir = 'C:\\test\\workspace';
 				const differentDir = 'C:\\different\\path';
-				const command = `cd /d ${differentDir} && echo hello`;
+				const command = `cd /d ${differentDir} ; echo hello`;
 				const options = createRewriteParams(command, 'session-1');
 				setWorkspaceFolders([URI.file(testDir)]);
 
@@ -344,4 +365,22 @@ suite('command re-writing', () => {
 			});
 		});
 	});
+
+	suite('PowerShell: && -> ;', () => {
+		async function t(originalCommandLine: string, expectedResult: string) {
+			const parameters = createRewriteParams(originalCommandLine);
+			const result = await commandSimplifier.rewriteIfNeeded(parameters, undefined, 'pwsh');
+			strictEqual(result, expectedResult);
+		}
+		setup(() => {
+			commandSimplifier = instantiationService.createInstance(CommandSimplifier, Promise.resolve(OperatingSystem.Windows), parser);
+		});
+
+		test('should rewrite && to ; in PowerShell commands', () => t('echo hello && echo world', 'echo hello ; echo world'));
+		test('should rewrite multiple && to ; in PowerShell commands', () => t('echo first && echo second && echo third', 'echo first ; echo second ; echo third'));
+		test('should handle complex commands with && operators', () => t('npm install && npm test && echo "build complete"', 'npm install ; npm test ; echo "build complete"'));
+		test('should work with Windows PowerShell shell identifier', () => t('Get-Process && Stop-Process', 'Get-Process ; Stop-Process'));
+		test('should preserve existing semicolons', () => t('echo hello; echo world && echo final', 'echo hello; echo world ; echo final'));
+		test('should not rewrite strings', () => t('echo "&&" && Write-Host "&& &&" && "&&"', 'echo "&&" ; Write-Host "&& &&" ; "&&"'));
+	});
 });

src/vs/workbench/contrib/terminalContrib/chatAgentTools/test/electron-browser/treeSitterCommandParser.test.ts

@@ -22,8 +22,7 @@ suite('TreeSitterCommandParser', () => {
 	let parser: TreeSitterCommandParser;
 
 	setup(() => {
-		const logService = new NullLogService();
-		const fileService = store.add(new FileService(logService));
+		const fileService = store.add(new FileService(new NullLogService()));
 		const fileSystemProvider = new TestIPCFileSystemProvider();
 		store.add(fileService.registerProvider(Schemas.file, fileSystemProvider));