Improve shell integration reliability for zsh (#22891)

Kartik Raj · web-flow · commit 5174d5c1d552 · 2024-02-11T00:08:44.000+05:30
Closes #22881 If status changes, re-run activation. Also persist once we know shell integration works for a shell.
diff --git a/src/client/terminals/envCollectionActivation/service.ts b/src/client/terminals/envCollectionActivation/service.ts
@@ -26,7 +26,7 @@ import {
     IPathUtils,
 } from '../../common/types';
 import { Interpreters } from '../../common/utils/localize';
-import { traceError, traceVerbose, traceWarn } from '../../logging';
+import { traceError, traceInfo, traceVerbose, traceWarn } from '../../logging';
 import { IInterpreterService } from '../../interpreter/contracts';
 import { defaultShells } from '../../interpreter/activation/service';
 import { IEnvironmentActivationService } from '../../interpreter/activation/types';
@@ -111,6 +111,14 @@ export class TerminalEnvVarCollectionService implements IExtensionActivationServ
                     this,
                     this.disposables,
                 );
+                this.shellIntegrationService.onDidChangeStatus(
+                    async () => {
+                        traceInfo("Shell integration status changed, can confirm it's working.");
+                        await this._applyCollection(undefined).ignoreErrors();
+                    },
+                    this,
+                    this.disposables,
+                );
                 this.applicationEnvironment.onDidChangeShell(
                     async (shell: string) => {
                         this.processEnvVars = undefined;
@@ -125,7 +133,9 @@ export class TerminalEnvVarCollectionService implements IExtensionActivationServ
                 const isActive = await this.shellIntegrationService.isWorking(shell);
                 const shellType = identifyShellFromShellPath(shell);
                 if (!isActive && shellType !== TerminalShellType.commandPrompt) {
-                    traceWarn(`Shell integration is not active, environment activated maybe overriden by the shell.`);
+                    traceWarn(
+                        `Shell integration may not be active, environment activated may be overridden by the shell.`,
+                    );
                 }
                 this.registeredOnce = true;
             }
diff --git a/src/client/terminals/envCollectionActivation/shellIntegrationService.ts b/src/client/terminals/envCollectionActivation/shellIntegrationService.ts
@@ -12,9 +12,8 @@ import {
 import { identifyShellFromShellPath } from '../../common/terminal/shellDetectors/baseShellDetector';
 import { TerminalShellType } from '../../common/terminal/types';
 import { IDisposableRegistry, IPersistentStateFactory } from '../../common/types';
-import { createDeferred, sleep } from '../../common/utils/async';
-import { cache } from '../../common/utils/decorators';
-import { traceError, traceInfo, traceVerbose } from '../../logging';
+import { sleep } from '../../common/utils/async';
+import { traceError, traceVerbose } from '../../logging';
 import { IShellIntegrationService } from '../types';
 
 /**
@@ -29,17 +28,12 @@ const ShellIntegrationShells = [
     TerminalShellType.fish,
 ];
 
-export const isShellIntegrationWorkingKey = 'SHELL_INTEGRATION_WORKING_KEY';
+export enum isShellIntegrationWorking {
+    key = 'SHELL_INTEGRATION_WORKING_KEY',
+}
 
 @injectable()
 export class ShellIntegrationService implements IShellIntegrationService {
-    /**
-     * It seems to have a couple of issues:
-     * * Ends up cluterring terminal history
-     * * Does not work for hidden terminals: https://github.com/microsoft/vscode/issues/199611
-     */
-    private readonly USE_COMMAND_APPROACH = false;
-
     private isWorkingForShell = new Set<TerminalShellType>();
 
     private readonly didChange = new EventEmitter<void>();
@@ -55,6 +49,12 @@ export class ShellIntegrationService implements IShellIntegrationService {
         @inject(IDisposableRegistry) private readonly disposables: IDisposableRegistry,
     ) {
         try {
+            const activeShellType = identifyShellFromShellPath(this.appEnvironment.shell);
+            const key = getKeyForShell(activeShellType);
+            const persistedResult = this.persistentStateFactory.createGlobalPersistentState<boolean>(key);
+            if (persistedResult.value) {
+                this.isWorkingForShell.add(activeShellType);
+            }
             this.appShell.onDidWriteTerminalData(
                 (e) => {
                     if (e.data.includes('\x1b]633;A\x07')) {
@@ -63,6 +63,7 @@ export class ShellIntegrationService implements IShellIntegrationService {
                             shell = e.terminal.creationOptions.shellPath;
                         }
                         const shellType = identifyShellFromShellPath(shell);
+                        traceVerbose('Received shell integration sequence for', shellType);
                         const wasWorking = this.isWorkingForShell.has(shellType);
                         this.isWorkingForShell.add(shellType);
                         if (!wasWorking) {
@@ -86,6 +87,12 @@ export class ShellIntegrationService implements IShellIntegrationService {
             this.isDataWriteEventWorking = false;
             traceError('Unable to check if shell integration is active', ex);
         }
+        const isEnabled = !!this.workspaceService
+            .getConfiguration('terminal')
+            .get<boolean>('integrated.shellIntegration.enabled');
+        if (!isEnabled) {
+            traceVerbose('Shell integration is disabled in user settings.');
+        }
     }
 
     public readonly onDidChangeStatus = this.didChange.event;
@@ -97,46 +104,48 @@ export class ShellIntegrationService implements IShellIntegrationService {
         });
     }
 
-    @cache(-1, true)
     public async _isWorking(shell: string): Promise<boolean> {
-        const isEnabled = this.workspaceService
-            .getConfiguration('terminal')
-            .get<boolean>('integrated.shellIntegration.enabled')!;
-        if (!isEnabled) {
-            traceVerbose('Shell integrated is disabled in user settings.');
-        }
         const shellType = identifyShellFromShellPath(shell);
-        const isSupposedToWork = isEnabled && ShellIntegrationShells.includes(shellType);
+        const isSupposedToWork = ShellIntegrationShells.includes(shellType);
         if (!isSupposedToWork) {
             return false;
         }
-        if (!this.USE_COMMAND_APPROACH) {
-            // For now, based on problems with using the command approach, use terminal data write event.
-            if (!this.isDataWriteEventWorking) {
-                // Assume shell integration is working, if data write event isn't working.
-                return true;
-            }
-            if (shellType === TerminalShellType.powershell || shellType === TerminalShellType.powershellCore) {
-                // Due to upstream bug: https://github.com/microsoft/vscode/issues/204616, assume shell integration is working for now.
-                return true;
-            }
-            if (!this.isWorkingForShell.has(shellType)) {
-                // Maybe data write event has not been processed yet, wait a bit.
-                await sleep(1000);
-            }
-            return this.isWorkingForShell.has(shellType);
-        }
-        const key = `${isShellIntegrationWorkingKey}_${shellType}`;
+        const key = getKeyForShell(shellType);
         const persistedResult = this.persistentStateFactory.createGlobalPersistentState<boolean>(key);
         if (persistedResult.value !== undefined) {
             return persistedResult.value;
         }
-        const result = await this.checkIfWorkingByRunningCommand(shell);
-        // Persist result to storage to avoid running commands unncecessary.
-        await persistedResult.updateValue(result);
+        const result = await this.useDataWriteApproach(shellType);
+        if (result) {
+            // Once we know that shell integration is working for a shell, persist it so we need not do this check every session.
+            await persistedResult.updateValue(result);
+        }
         return result;
     }
 
+    private async useDataWriteApproach(shellType: TerminalShellType) {
+        // For now, based on problems with using the command approach, use terminal data write event.
+        if (!this.isDataWriteEventWorking) {
+            // Assume shell integration is working, if data write event isn't working.
+            return true;
+        }
+        if (shellType === TerminalShellType.powershell || shellType === TerminalShellType.powershellCore) {
+            // Due to upstream bug: https://github.com/microsoft/vscode/issues/204616, assume shell integration is working for now.
+            return true;
+        }
+        if (!this.isWorkingForShell.has(shellType)) {
+            // Maybe data write event has not been processed yet, wait a bit.
+            await sleep(1000);
+        }
+        traceVerbose(
+            'Did we determine shell integration to be working for',
+            shellType,
+            '?',
+            this.isWorkingForShell.has(shellType),
+        );
+        return this.isWorkingForShell.has(shellType);
+    }
+
     /**
      * Creates a dummy terminal so that we are guaranteed a data write event for this shell type.
      */
@@ -146,39 +155,8 @@ export class ShellIntegrationService implements IShellIntegrationService {
             hideFromUser: true,
         });
     }
+}
 
-    private async checkIfWorkingByRunningCommand(shell: string): Promise<boolean> {
-        const shellType = identifyShellFromShellPath(shell);
-        const deferred = createDeferred<void>();
-        const timestamp = new Date().getTime();
-        const name = `Python ${timestamp}`;
-        const onDidExecuteTerminalCommand = this.appShell.onDidExecuteTerminalCommand?.bind(this.appShell);
-        if (!onDidExecuteTerminalCommand) {
-            // Proposed API is not available, assume shell integration is working at this point.
-            return true;
-        }
-        try {
-            const disposable = onDidExecuteTerminalCommand((e) => {
-                if (e.terminal.name === name) {
-                    deferred.resolve();
-                }
-            });
-            const terminal = this.terminalManager.createTerminal({
-                name,
-                shellPath: shell,
-                hideFromUser: true,
-            });
-            terminal.sendText(`echo ${shell}`);
-            const success = await Promise.race([sleep(3000).then(() => false), deferred.promise.then(() => true)]);
-            disposable.dispose();
-            if (!success) {
-                traceInfo(`Shell integration is not working for ${shellType}`);
-            }
-            return success;
-        } catch (ex) {
-            traceVerbose(`Proposed API is not available, failed to subscribe to onDidExecuteTerminalCommand`, ex);
-            // Proposed API is not available, assume shell integration is working at this point.
-            return true;
-        }
-    }
+function getKeyForShell(shellType: TerminalShellType) {
+    return `${isShellIntegrationWorking.key}_${shellType}`;
 }
