Preserve registered shell init revisions

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

Copilot-Session: bef9b9c0-98f2-4b26-8b8b-18119efe7b13

Author: anthonykim1

src/vs/platform/agentHost/node/copilot/copilotAgentSession.ts

@@ -999,6 +999,8 @@ export class CopilotAgentSession extends Disposable {
 	private _shellInitScriptRegistered = false;
 	private _shellInitScriptMaterializationAttempted = false;
 	private _shellInitSandboxGrantApplied = false;
+	private _shellInitScriptRevision = 0;
+	private _registeredShellInitScriptResource: URI | undefined;
 	private readonly _shellInitScriptSequencer = new Sequencer();
 	private _shellInitScriptDisposing = false;
 	/**
@@ -3916,49 +3918,68 @@ export class CopilotAgentSession extends Disposable {
 			// The SDK requires init scripts to be readable when registered.
 			await this._applyEffectiveSandboxConfig(true);
 			this._shellInitSandboxGrantApplied = true;
-			const ref = await this._materializeShellInitScript(snippets[0]);
-			if (!ref) {
+			const materialized = await this._materializeShellInitScript(snippets[0]);
+			if (!materialized) {
 				// The write failed. Leave the cache unchanged so the next turn
 				// retries rather than permanently treating the script as applied.
 				return;
 			}
-			const result = await this._wrapper.session.rpc.options.update({ shell: { initScripts: [ref] } });
+			const result = await this._wrapper.session.rpc.options.update({ shell: { initScripts: [materialized.ref] } });
 			if (!result.success) {
+				await this._deleteShellInitScriptFile(materialized.resource);
 				throw new Error('Copilot SDK rejected shell init script update');
 			}
+			const previousResource = this._registeredShellInitScriptResource;
 			this._shellInitScriptRegistered = true;
+			this._registeredShellInitScriptResource = materialized.resource;
 			this._lastAppliedShellInitScripts = serialized;
+			if (previousResource) {
+				await this._deleteShellInitScriptFile(previousResource);
+			}
 			this._logService.trace(`[Copilot:${this.sessionId}] Applied shell init script`);
 		} catch (err) {
 			this._logService.warn(`[Copilot:${this.sessionId}] Failed to update shell init scripts`, err);
 		}
 	}
 
-	private async _materializeShellInitScript(script: IShellInitScript): Promise<{ shell: IShellInitScript['shell']; path: string } | undefined> {
-		const resource = URI.joinPath(this._shellInitScriptInstanceDirectory(), script.shell === 'powershell' ? 'init.ps1' : 'init.sh');
+	private async _materializeShellInitScript(script: IShellInitScript): Promise<{ ref: { shell: IShellInitScript['shell']; path: string }; resource: URI } | undefined> {
+		const extension = script.shell === 'powershell' ? 'ps1' : 'sh';
+		const resource = URI.joinPath(this._shellInitScriptInstanceDirectory(), `init-${++this._shellInitScriptRevision}.${extension}`);
 		const atomic = this._fileService.hasCapability(resource, FileSystemProviderCapabilities.FileAtomicWrite)
 			? { postfix: '.vsctmp' }
 			: false;
 		this._shellInitScriptMaterializationAttempted = true;
 		try {
 			await this._fileService.writeFile(resource, VSBuffer.fromString(script.script), { atomic });
-			return { shell: script.shell, path: resource.fsPath };
+			return { ref: { shell: script.shell, path: resource.fsPath }, resource };
 		} catch (error) {
 			this._logService.warn(`[Copilot:${this.sessionId}] Failed to write shell init script: ${getErrorMessage(error)}`);
 			return undefined;
 		}
 	}
 
+	private async _deleteShellInitScriptFile(resource: URI): Promise<void> {
+		try {
+			await this._fileService.del(resource);
+		} catch (error) {
+			if (!(error instanceof Error) || toFileOperationResult(error) !== FileOperationResult.FILE_NOT_FOUND) {
+				this._logService.warn(`[Copilot:${this.sessionId}] Failed to remove replaced shell init script: ${getErrorMessage(error)}`);
+			}
+		}
+	}
+
 	private async _clearShellInitScript(): Promise<void> {
 		if (!this._shellInitScriptMaterializationAttempted) {
 			return;
 		}
 		try {
 			await this._fileService.del(this._shellInitScriptInstanceDirectory(), { recursive: true });
 			this._shellInitScriptMaterializationAttempted = false;
+			this._registeredShellInitScriptResource = undefined;
 		} catch (error) {
 			if (error instanceof Error && toFileOperationResult(error) === FileOperationResult.FILE_NOT_FOUND) {
 				this._shellInitScriptMaterializationAttempted = false;
+				this._registeredShellInitScriptResource = undefined;
 			} else {
 				this._logService.warn(`[Copilot:${this.sessionId}] Failed to remove shell init script: ${getErrorMessage(error)}`);
 			}

src/vs/platform/agentHost/test/node/copilotAgentSession.test.ts

@@ -12421,7 +12421,7 @@ Use the attached image as context.
 
 			await session.send('go', undefined, 'turn-1', 'interactive');
 			const scriptPath = (mockSession.shellInitScriptUpdates.at(-1) as Array<{ path: string }>)?.[0]?.path;
-			assert.ok(scriptPath?.startsWith(TEST_SHELL_INIT_DIR) && scriptPath.endsWith('init.sh'), String(scriptPath));
+			assert.ok(scriptPath?.startsWith(TEST_SHELL_INIT_DIR) && /init-\d+\.sh$/.test(scriptPath), String(scriptPath));
 
 			const afterFirst = mockSession.shellInitScriptUpdates.length;
 			await session.send('go', undefined, 'turn-2', 'interactive');
@@ -12430,12 +12430,14 @@ Use the attached image as context.
 			setConfigValue(SessionConfigKey.ShellInitSnippets, [{ ...initScript, script: 'changed' }]);
 			fireSessionConfigChange({ [SessionConfigKey.ShellInitSnippets]: [{ ...initScript, script: 'changed' }] });
 			await timeout(0);
+			const changedScriptPath = (mockSession.shellInitScriptUpdates.at(-1) as Array<{ path: string }>)?.[0]?.path;
+			assert.ok(changedScriptPath && changedScriptPath !== scriptPath, String(changedScriptPath));
 
 			setConfigValue(SessionConfigKey.ShellInitSnippets, []);
 			await session.send('go', undefined, 'turn-3', 'interactive');
 			assert.deepStrictEqual(mockSession.shellInitScriptUpdates, [
 				[{ shell: 'bash', path: scriptPath }],
-				[{ shell: 'bash', path: scriptPath }],
+				[{ shell: 'bash', path: changedScriptPath }],
 				[],
 			]);
 		});
@@ -12455,6 +12457,32 @@ Use the attached image as context.
 			assert.deepStrictEqual(mockSession.shellInitScriptUpdates, [registered]);
 		});
 
+		test('keeps the registered script file unchanged when a replacement is rejected', async () => {
+			const { session, mockSession, storedFileContents, setConfigValue } = await createAgentSession(disposables);
+			setConfigValue(SessionConfigKey.ShellInitSnippets, [initScript]);
+			await session.send('go', undefined, 'turn-1', 'interactive');
+			const registeredPath = (mockSession.shellInitScriptUpdates.at(-1) as Array<{ path: string }>)[0].path;
+
+			mockSession.shellInitScriptUpdateSuccess = false;
+			setConfigValue(SessionConfigKey.ShellInitSnippets, [{ ...initScript, script: 'rejected' }]);
+			await session.send('go', undefined, 'turn-2', 'interactive');
+			const updatesAfterRejection = mockSession.shellInitScriptUpdates.length;
+
+			mockSession.shellInitScriptUpdateSuccess = true;
+			setConfigValue(SessionConfigKey.ShellInitSnippets, [initScript]);
+			await session.send('go', undefined, 'turn-3', 'interactive');
+
+			assert.deepStrictEqual({
+				registeredContent: storedFileContents.get(URI.file(registeredPath).toString()),
+				registeredFileCount: [...storedFileContents.keys()].filter(key => key.includes('/agentHost/shellInit/')).length,
+				updateCount: mockSession.shellInitScriptUpdates.length,
+			}, {
+				registeredContent: initScript.script,
+				registeredFileCount: 1,
+				updateCount: updatesAfterRejection,
+			});
+		});
+
 		test('each instance owns a distinct directory so a stale cleanup cannot delete a successor script', async () => {
 			// dispose() queues the deletion without awaiting it; a resumed
 			// replacement for the same SDK session may register its script first.
@@ -12489,7 +12517,7 @@ Use the attached image as context.
 			});
 			setConfigValue(SessionConfigKey.ShellInitSnippets, [initScript]);
 			await session.send('go', undefined, 'turn-1', 'interactive');
-			assert.strictEqual([...storedFileContents.keys()].filter(key => key.includes('/test-session-1/') && key.endsWith('init.sh')).length, 2);
+			assert.strictEqual([...storedFileContents.keys()].filter(key => key.includes('/test-session-1/') && key.endsWith('.sh')).length, 2);
 
 			// The session-directory prune must not take the successor's script
 			// with it; only this instance's directory may be removed.