Preserve $TMPDIR when retrying terminal commands outside the sandbox (#304601)

* Changes to include environment variables for retry without sandboxing

* adding e2e tests

* fixing code formatting

extensions/vscode-api-tests/src/singlefolder-tests/chat.runInTerminal.test.ts

@@ -71,7 +71,7 @@ function extractTextContent(result: vscode.LanguageModelToolResult): string {
 		participantRegistered = false;
 		pendingResult = undefined;
 		pendingCommand = undefined;
-		pendingTimeout = undefined;
+		pendingOptions = undefined;
 
 		const chatToolsConfig = vscode.workspace.getConfiguration('chat.tools.global');
 		await chatToolsConfig.update('autoApprove', undefined, vscode.ConfigurationTarget.Global);
@@ -82,10 +82,16 @@ function extractTextContent(result: vscode.LanguageModelToolResult): string {
 	 * Helper: invokes run_in_terminal via a chat participant and returns the tool result text.
 	 * Each call creates a new chat session to avoid participant re-registration issues.
 	 */
+	interface RunInTerminalOptions {
+		timeout?: number;
+		requestUnsandboxedExecution?: boolean;
+		requestUnsandboxedExecutionReason?: string;
+	}
+
 	let participantRegistered = false;
 	let pendingResult: DeferredPromise<vscode.LanguageModelToolResult> | undefined;
 	let pendingCommand: string | undefined;
-	let pendingTimeout: number | undefined;
+	let pendingOptions: RunInTerminalOptions | undefined;
 
 	function setupParticipant() {
 		if (participantRegistered) {
@@ -98,10 +104,10 @@ function extractTextContent(result: vscode.LanguageModelToolResult): string {
 			}
 			const currentResult = pendingResult;
 			const currentCommand = pendingCommand;
-			const currentTimeout = pendingTimeout ?? 15000;
+			const currentOptions = pendingOptions ?? {};
 			pendingResult = undefined;
 			pendingCommand = undefined;
-			pendingTimeout = undefined;
+			pendingOptions = undefined;
 			try {
 				const result = await vscode.lm.invokeTool('run_in_terminal', {
 					input: {
@@ -109,7 +115,11 @@ function extractTextContent(result: vscode.LanguageModelToolResult): string {
 						explanation: 'Integration test command',
 						goal: 'Test run_in_terminal output',
 						isBackground: false,
-						timeout: currentTimeout
+						timeout: currentOptions.timeout ?? 15000,
+						...currentOptions.requestUnsandboxedExecution ? {
+							requestUnsandboxedExecution: true,
+							requestUnsandboxedExecutionReason: currentOptions.requestUnsandboxedExecutionReason,
+						} : {},
 					},
 					toolInvocationToken: request.toolInvocationToken,
 				});
@@ -122,13 +132,18 @@ function extractTextContent(result: vscode.LanguageModelToolResult): string {
 		disposables.push(participant);
 	}
 
-	async function invokeRunInTerminal(command: string, timeout = 15000): Promise<string> {
+	async function invokeRunInTerminal(command: string, options?: RunInTerminalOptions): Promise<string>;
+	async function invokeRunInTerminal(command: string, timeout?: number): Promise<string>;
+	async function invokeRunInTerminal(command: string, optionsOrTimeout?: RunInTerminalOptions | number): Promise<string> {
 		setupParticipant();
 
+		const opts: RunInTerminalOptions = typeof optionsOrTimeout === 'number'
+			? { timeout: optionsOrTimeout }
+			: optionsOrTimeout ?? {};
 		const resultPromise = new DeferredPromise<vscode.LanguageModelToolResult>();
 		pendingResult = resultPromise;
 		pendingCommand = command;
-		pendingTimeout = timeout;
+		pendingOptions = opts;
 
 		await vscode.commands.executeCommand('workbench.action.chat.newChat');
 		vscode.commands.executeCommand('workbench.action.chat.open', { query: '@participant test' });
@@ -326,6 +341,27 @@ function extractTextContent(result: vscode.LanguageModelToolResult): string {
 				assert.ok(acceptable.includes(output.trim()), `Unexpected output: ${JSON.stringify(output.trim())}`);
 			});
 
+			test('requestUnsandboxedExecution preserves sandbox $TMPDIR', async function () {
+				this.timeout(60000);
+
+				const marker = `SANDBOX_UNSANDBOX_${Date.now()}`;
+				const sentinelName = `sentinel-${marker}.txt`;
+
+				// Step 1: Write a sentinel file into the sandbox-provided $TMPDIR.
+				const writeOutput = await invokeRunInTerminal(`echo ${marker} > "$TMPDIR/${sentinelName}" && echo ${marker}`);
+				assert.strictEqual(writeOutput.trim(), marker);
+
+				// Step 2: Retry with requestUnsandboxedExecution=true while sandbox
+				// stays enabled. The tool should preserve $TMPDIR from the sandbox so
+				// the sentinel file created in step 1 is still accessible.
+				const retryOutput = await invokeRunInTerminal(`cat "$TMPDIR/${sentinelName}"`, {
+					timeout: 30000,
+					requestUnsandboxedExecution: true,
+					requestUnsandboxedExecutionReason: 'Need to verify $TMPDIR persists on unsandboxed retry',
+				});
+				assert.strictEqual(retryOutput.trim(), marker);
+			});
+
 			test('cannot write to /tmp', async function () {
 				this.timeout(60000);
 

src/vs/workbench/contrib/terminalContrib/chatAgentTools/browser/tools/commandLineRewriter/commandLineSandboxRewriter.ts

@@ -15,9 +15,6 @@ export class CommandLineSandboxRewriter extends Disposable implements ICommandLi
 	}
 
 	async rewrite(options: ICommandLineRewriterOptions): Promise<ICommandLineRewriterResult | undefined> {
-		if (options.requestUnsandboxedExecution) {
-			return undefined;
-		}
 
 		if (!(await this._sandboxService.isEnabled())) {
 			return undefined;
@@ -30,7 +27,7 @@ export class CommandLineSandboxRewriter extends Disposable implements ICommandLi
 			return undefined;
 		}
 
-		const wrappedCommand = this._sandboxService.wrapCommand(options.commandLine);
+		const wrappedCommand = this._sandboxService.wrapCommand(options.commandLine, options.requestUnsandboxedExecution);
 		return {
 			rewritten: wrappedCommand,
 			reasoning: 'Wrapped command for sandbox execution',

src/vs/workbench/contrib/terminalContrib/chatAgentTools/common/terminalSandboxService.ts

@@ -36,7 +36,7 @@ export interface ITerminalSandboxService {
 	readonly _serviceBrand: undefined;
 	isEnabled(): Promise<boolean>;
 	getOS(): Promise<OperatingSystem>;
-	wrapCommand(command: string): string;
+	wrapCommand(command: string, requestUnsandboxedExecution?: boolean): string;
 	getSandboxConfigPath(forceRefresh?: boolean): Promise<string | undefined>;
 	getTempDir(): URI | undefined;
 	setNeedsForceUpdateConfigFile(): void;
@@ -127,10 +127,15 @@ export class TerminalSandboxService extends Disposable implements ITerminalSandb
 		return this._os;
 	}
 
-	public wrapCommand(command: string): string {
+	public wrapCommand(command: string, requestUnsandboxedExecution?: boolean): string {
 		if (!this._sandboxConfigPath || !this._tempDir) {
 			throw new Error('Sandbox config path or temp dir not initialized');
 		}
+		// If requestUnsandboxedExecution is true, need to ensure env variables set during sandbox still apply.
+		if (requestUnsandboxedExecution) {
+			return this._tempDir?.path ? `(TMPDIR="${this._tempDir.path}"; export TMPDIR; ${command})` : command;
+		}
+
 		if (!this._execPath) {
 			throw new Error('Executable path not set to run sandbox commands');
 		}

src/vs/workbench/contrib/terminalContrib/chatAgentTools/test/browser/terminalSandboxService.test.ts

@@ -388,6 +388,20 @@ suite('TerminalSandboxService - allowTrustedDomains', () => {
 		);
 	});
 
+	test('should preserve TMPDIR when unsandboxed execution is requested', async () => {
+		const sandboxService = store.add(instantiationService.createInstance(TerminalSandboxService));
+		await sandboxService.getSandboxConfigPath();
+
+		strictEqual(sandboxService.wrapCommand('echo test', true), `(TMPDIR="${sandboxService.getTempDir()?.path}"; export TMPDIR; echo test)`);
+	});
+
+	test('should preserve TMPDIR for piped unsandboxed commands', async () => {
+		const sandboxService = store.add(instantiationService.createInstance(TerminalSandboxService));
+		await sandboxService.getSandboxConfigPath();
+
+		strictEqual(sandboxService.wrapCommand('echo test | cat', true), `(TMPDIR="${sandboxService.getTempDir()?.path}"; export TMPDIR; echo test | cat)`);
+	});
+
 	test('should pass wrapped command as a single quoted argument', async () => {
 		const sandboxService = store.add(instantiationService.createInstance(TerminalSandboxService));
 		await sandboxService.getSandboxConfigPath();

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

@@ -22,7 +22,7 @@ suite('CommandLineSandboxRewriter', () => {
 		instantiationService.stub(ITerminalSandboxService, {
 			_serviceBrand: undefined,
 			isEnabled: async () => false,
-			wrapCommand: command => command,
+			wrapCommand: (command, _requestUnsandboxedExecution) => command,
 			getSandboxConfigPath: async () => '/tmp/sandbox.json',
 			getTempDir: () => undefined,
 			setNeedsForceUpdateConfigFile: () => { },
@@ -62,7 +62,7 @@ suite('CommandLineSandboxRewriter', () => {
 		const calls: string[] = [];
 		stubSandboxService({
 			isEnabled: async () => true,
-			wrapCommand: command => {
+			wrapCommand: (command, _requestUnsandboxedExecution) => {
 				calls.push('wrapCommand');
 				return `wrapped:${command}`;
 			},
@@ -79,12 +79,12 @@ suite('CommandLineSandboxRewriter', () => {
 		deepStrictEqual(calls, ['getSandboxConfigPath', 'wrapCommand']);
 	});
 
-	test('does not wrap command when sandbox bypass was explicitly requested', async () => {
+	test('wraps command and forwards sandbox bypass flag when explicitly requested', async () => {
 		const calls: string[] = [];
 		stubSandboxService({
 			isEnabled: async () => true,
-			wrapCommand: command => {
-				calls.push(`wrap:${command}`);
+			wrapCommand: (command, requestUnsandboxedExecution) => {
+				calls.push(`wrap:${command}:${String(requestUnsandboxedExecution)}`);
 				return `wrapped:${command}`;
 			},
 			getSandboxConfigPath: async () => {
@@ -99,7 +99,8 @@ suite('CommandLineSandboxRewriter', () => {
 			requestUnsandboxedExecution: true,
 		});
 
-		strictEqual(result, undefined);
-		deepStrictEqual(calls, []);
+		strictEqual(result?.rewritten, 'wrapped:echo hello');
+		strictEqual(result?.reasoning, 'Wrapped command for sandbox execution');
+		deepStrictEqual(calls, ['config', 'wrap:echo hello:true']);
 	});
 });

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

@@ -121,7 +121,7 @@ suite('RunInTerminalTool', () => {
 		terminalSandboxService = {
 			_serviceBrand: undefined,
 			isEnabled: async () => sandboxEnabled,
-			wrapCommand: (command: string) => `sandbox:${command}`,
+			wrapCommand: (command: string, requestUnsandboxedExecution?: boolean) => requestUnsandboxedExecution ? `unsandboxed:${command}` : `sandbox:${command}`,
 			getSandboxConfigPath: async () => sandboxEnabled ? '/tmp/sandbox.json' : undefined,
 			getTempDir: () => undefined,
 			setNeedsForceUpdateConfigFile: () => { },
@@ -552,7 +552,7 @@ suite('RunInTerminalTool', () => {
 			const terminalData = result?.toolSpecificData as IChatTerminalToolInvocationData;
 			strictEqual(terminalData.requestUnsandboxedExecution, true);
 			strictEqual(terminalData.requestUnsandboxedExecutionReason, 'Needs network access outside the sandbox');
-			strictEqual(terminalData.commandLine.toolEdited, undefined);
+			strictEqual(terminalData.commandLine.toolEdited, 'unsandboxed:echo hello');
 
 			const confirmationMessage = result?.confirmationMessages?.message;
 			ok(confirmationMessage && typeof confirmationMessage !== 'string');
@@ -566,7 +566,7 @@ suite('RunInTerminalTool', () => {
 			ok(actions, 'Expected custom actions to be defined');
 			strictEqual(actions.length, 11);
 			ok(!isSeparator(actions[0]));
-			strictEqual(actions[0].label, 'Allow `echo …` in this Session');
+			strictEqual(actions[0].label, 'Allow `unsandboxed:echo …` in this Session');
 			ok(!isSeparator(actions[4]));
 			strictEqual(actions[4].label, 'Allow Exact Command Line in this Session');
 			ok(!isSeparator(actions[10]));