From 403071707c559feda67ae7ce808547b14ce1f92e Mon Sep 17 00:00:00 2001
From: Kartik Raj <karraj@microsoft.com>
Date: Wed, 7 Feb 2024 22:31:52 +0530
Subject: [PATCH] Use terminal data write event to figure out whether shell
 integration is working (#22872)

Closes https://github.com/microsoft/vscode-python/issues/22439

Blocked on https://github.com/microsoft/vscode/issues/204616
---
 .../shellIntegrationService.ts                | 81 +++++++++++++++++--
 src/client/terminals/types.ts                 |  1 +
 2 files changed, 77 insertions(+), 5 deletions(-)

diff --git a/src/client/terminals/envCollectionActivation/shellIntegrationService.ts b/src/client/terminals/envCollectionActivation/shellIntegrationService.ts
index bfa963537ec9..bb5168bc6e09 100644
--- a/src/client/terminals/envCollectionActivation/shellIntegrationService.ts
+++ b/src/client/terminals/envCollectionActivation/shellIntegrationService.ts
@@ -2,10 +2,16 @@
 // Licensed under the MIT License.
 
 import { injectable, inject } from 'inversify';
-import { IApplicationShell, ITerminalManager, IWorkspaceService } from '../../common/application/types';
+import { EventEmitter } from 'vscode';
+import {
+    IApplicationEnvironment,
+    IApplicationShell,
+    ITerminalManager,
+    IWorkspaceService,
+} from '../../common/application/types';
 import { identifyShellFromShellPath } from '../../common/terminal/shellDetectors/baseShellDetector';
 import { TerminalShellType } from '../../common/terminal/types';
-import { IPersistentStateFactory } from '../../common/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';
@@ -34,12 +40,55 @@ export class ShellIntegrationService implements IShellIntegrationService {
      */
     private readonly USE_COMMAND_APPROACH = false;
 
+    private isWorkingForShell = new Set<TerminalShellType>();
+
+    private readonly didChange = new EventEmitter<void>();
+
+    private isDataWriteEventWorking = true;
+
     constructor(
         @inject(ITerminalManager) private readonly terminalManager: ITerminalManager,
         @inject(IApplicationShell) private readonly appShell: IApplicationShell,
         @inject(IWorkspaceService) private readonly workspaceService: IWorkspaceService,
         @inject(IPersistentStateFactory) private readonly persistentStateFactory: IPersistentStateFactory,
-    ) {}
+        @inject(IApplicationEnvironment) private readonly appEnvironment: IApplicationEnvironment,
+        @inject(IDisposableRegistry) private readonly disposables: IDisposableRegistry,
+    ) {
+        try {
+            this.appShell.onDidWriteTerminalData(
+                (e) => {
+                    if (e.data.includes('\x1b]633;A\x07')) {
+                        let { shell } = this.appEnvironment;
+                        if ('shellPath' in e.terminal.creationOptions && e.terminal.creationOptions.shellPath) {
+                            shell = e.terminal.creationOptions.shellPath;
+                        }
+                        const shellType = identifyShellFromShellPath(shell);
+                        const wasWorking = this.isWorkingForShell.has(shellType);
+                        this.isWorkingForShell.add(shellType);
+                        if (!wasWorking) {
+                            // If it wasn't working previously, status has changed.
+                            this.didChange.fire();
+                        }
+                    }
+                },
+                this,
+                this.disposables,
+            );
+            this.appEnvironment.onDidChangeShell(
+                async (shell: string) => {
+                    this.createDummyHiddenTerminal(shell);
+                },
+                this,
+                this.disposables,
+            );
+            this.createDummyHiddenTerminal(this.appEnvironment.shell);
+        } catch (ex) {
+            this.isDataWriteEventWorking = false;
+            traceError('Unable to check if shell integration is active', ex);
+        }
+    }
+
+    public readonly onDidChangeStatus = this.didChange.event;
 
     public async isWorking(shell: string): Promise<boolean> {
         return this._isWorking(shell).catch((ex) => {
@@ -62,8 +111,20 @@ export class ShellIntegrationService implements IShellIntegrationService {
             return false;
         }
         if (!this.USE_COMMAND_APPROACH) {
-            // For now, based on problems with using the command approach, assume it always works.
-            return true;
+            // 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 persistedResult = this.persistentStateFactory.createGlobalPersistentState<boolean>(key);
@@ -76,6 +137,16 @@ export class ShellIntegrationService implements IShellIntegrationService {
         return result;
     }
 
+    /**
+     * Creates a dummy terminal so that we are guaranteed a data write event for this shell type.
+     */
+    private createDummyHiddenTerminal(shell: string) {
+        this.terminalManager.createTerminal({
+            shellPath: shell,
+            hideFromUser: true,
+        });
+    }
+
     private async checkIfWorkingByRunningCommand(shell: string): Promise<boolean> {
         const shellType = identifyShellFromShellPath(shell);
         const deferred = createDeferred<void>();
diff --git a/src/client/terminals/types.ts b/src/client/terminals/types.ts
index 0fb268fe192c..fbdd490c175c 100644
--- a/src/client/terminals/types.ts
+++ b/src/client/terminals/types.ts
@@ -44,6 +44,7 @@ export interface ITerminalEnvVarCollectionService {
 
 export const IShellIntegrationService = Symbol('IShellIntegrationService');
 export interface IShellIntegrationService {
+    onDidChangeStatus: Event<void>;
     isWorking(shell: string): Promise<boolean>;
 }