fix(terminal): wait for the shell integration to be ready

Author: sheremet-va

src/api/child_process.ts

@@ -90,7 +90,7 @@ export async function createVitestProcess(pkg: VitestPackage) {
     vitest.on('exit', onExit)
     vitest.on('error', onError)
 
-    waitForWsConnection(wss, pkg, false, 'child_process')
+    waitForWsConnection(wss, pkg, false, 'child_process', false)
       .then((resolved) => {
         resolved.handlers.onStdout = (callback: (data: string) => void) => {
           stdoutCallbacks.add(callback)

src/api/terminal.ts

@@ -3,16 +3,18 @@ import type { WebSocket } from 'ws'
 import type { ResolvedMeta } from '../api'
 import type { VitestPackage } from './pkg'
 import type { ExtensionWorkerProcess } from './types'
+import type { WsConnectionMetadata } from './ws'
 import { createServer } from 'node:http'
 import { pathToFileURL } from 'node:url'
+import { stripVTControlCharacters } from 'node:util'
 import getPort from 'get-port'
 import * as vscode from 'vscode'
 import { WebSocketServer } from 'ws'
 import { getConfig } from '../config'
 import { workerPath } from '../constants'
 import { createErrorLogger, log } from '../log'
-import { formatPkg, showVitestError } from '../utils'
-import { waitForWsConnection, WsConnectionMetadata } from './ws'
+import { formatPkg } from '../utils'
+import { waitForWsConnection } from './ws'
 
 export async function createVitestTerminalProcess(pkg: VitestPackage): Promise<ResolvedMeta> {
   const pnpLoader = pkg.loader
@@ -42,27 +44,80 @@ export async function createVitestTerminalProcess(pkg: VitestPackage): Promise<R
       NODE_ENV: env.NODE_ENV ?? process.env.NODE_ENV ?? 'test',
     },
   })
+  // TODO: make sure it is desposed even if it throws, the same for child_process
+
+  const shellIntegration = await new Promise<vscode.TerminalShellIntegration | undefined>((resolve) => {
+    const disposable = vscode.window.onDidChangeTerminalShellIntegration((e) => {
+      const timeout = setTimeout(() => {
+        disposable.dispose()
+        resolve(undefined)
+      }, 3_000)
+
+      if (e.terminal === terminal) {
+        disposable.dispose()
+        clearTimeout(timeout)
+        resolve(e.shellIntegration)
+      }
+    })
+  })
+
+  const processId = await terminal.processId
+  if (terminal.exitStatus && terminal.exitStatus.code != null) {
+    throw new Error(`Terminal was ${getExitReason(terminal.exitStatus.reason)} with code ${terminal.exitStatus.code}`)
+  }
+
   let command = 'node'
   if (pnpLoader && pnp) {
     command += ` --require ${pnp} --experimental-loader ${pathToFileURL(pnpLoader).toString()}`
   }
-  command += ` ${workerPath}`
-  log.info('[API]', `Initiated ws connection via ${wsAddress}`)
-  log.info('[API]', `Starting ${formatPkg(pkg)} in the terminal: ${command}`)
-  terminal.sendText(command, true)
+  command += ` ${workerPath};`
+
+  log.info('[TERMINAL]', `Initiated ws connection via ${wsAddress}`)
+  log.info('[TERMINAL]', `Starting ${formatPkg(pkg)} in the terminal: ${command}`)
+
+  if (shellIntegration) {
+    log.info('[TERMINAL] Shell integration is initiated.')
+    let execution: vscode.TerminalShellExecution
+    const onWriteShell = vscode.window.onDidStartTerminalShellExecution(async (e) => {
+      if (e.execution !== execution) {
+        return
+      }
+
+      log.info('[TERMINAL] Reporting the shell output.')
+      for await (const line of e.execution.read()) {
+        log.worker('info', stripVTControlCharacters(line))
+        onWriteShell.dispose()
+      }
+    })
+
+    const onEndShell = vscode.window.onDidEndTerminalShellExecution((e) => {
+      if (e.execution === execution) {
+        log.info('[TERMINAL] The shell execution was finished.')
+        onWriteShell.dispose()
+        onEndShell.dispose()
+      }
+    })
+    execution = shellIntegration.executeCommand(command)
+  }
+  else {
+    log.info('[TERMINAL] Shell integration is not initiated, fallback to `terminal.sendText`.')
+    terminal.sendText(command, true)
+  }
+
   const meta = await new Promise<WsConnectionMetadata>((resolve, reject) => {
     const timeout = setTimeout(() => {
       terminal.show(false)
       reject(new Error(`The extension could not connect to the terminal in 5 seconds. See the "vitest" terminal output for more details.`))
-    }, 5000)
-    waitForWsConnection(wss, pkg, false, 'terminal').then(resolve, reject).finally(() => {
+    }, 5_000)
+    wss.once('connection', () => {
       clearTimeout(timeout)
     })
+    waitForWsConnection(wss, pkg, false, 'terminal', !!shellIntegration).then(resolve, reject)
   })
-  const processId = (await terminal.processId) ?? Math.random()
+
   log.info('[API]', `${formatPkg(pkg)} terminal process ${processId} created`)
   const vitestProcess = new ExtensionTerminalProcess(
-    processId,
+    processId ?? Math.random(),
     terminal,
     server,
     meta.ws,
@@ -77,6 +132,22 @@ export async function createVitestTerminalProcess(pkg: VitestPackage): Promise<R
   }
 }
 
+function getExitReason(reason: vscode.TerminalExitReason) {
+  switch (reason) {
+    case vscode.TerminalExitReason.Extension:
+      return 'clsoed by extension'
+    case vscode.TerminalExitReason.Process:
+      return 'closed by the process'
+    case vscode.TerminalExitReason.Shutdown:
+      return 'reloaded or closed'
+    case vscode.TerminalExitReason.User:
+      return 'closed by the user'
+    case vscode.TerminalExitReason.Unknown:
+    default:
+      return 'unexpectedly closed'
+  }
+}
+
 export class ExtensionTerminalProcess implements ExtensionWorkerProcess {
   private _onDidExit = new vscode.EventEmitter<number | null>()
 

src/api/ws.ts

@@ -17,6 +17,7 @@ export function waitForWsConnection(
   pkg: VitestPackage,
   debug: boolean,
   shellType: 'terminal' | 'child_process',
+  hasShellIntegration: boolean,
 ) {
   return new Promise<WsConnectionMetadata>((resolve, reject) => {
     wss.once('connection', (ws) => {
@@ -25,6 +26,7 @@ export function waitForWsConnection(
         pkg,
         debug,
         shellType,
+        hasShellIntegration,
         meta => resolve(meta),
         err => reject(err),
       )
@@ -51,6 +53,7 @@ export function onWsConnection(
   pkg: VitestPackage,
   debug: boolean,
   shellType: 'terminal' | 'child_process',
+  hasShellIntegration: boolean,
   onStart: (meta: WsConnectionMetadata) => unknown,
   onFail: (err: Error) => unknown,
 ) {
@@ -116,6 +119,7 @@ export function onWsConnection(
     type: 'init',
     meta: {
       shellType,
+      hasShellIntegration,
       vitestNodePath: pkg.vitestNodePath,
       env: getConfig(pkg.folder).env || undefined,
       configFile: pkg.configFile,

src/debug.ts

@@ -55,7 +55,7 @@ export async function debugTests(
     smartStep: true,
     ...(config.shellType === 'terminal'
       ? {
-          command: `${runtimeExecutable} ${workerPath}`,
+          command: `${runtimeExecutable} ${workerPath};exit`,
         }
       : {
           program: workerPath,
@@ -116,6 +116,7 @@ export async function debugTests(
       pkg,
       true,
       config.shellType,
+      false,
       async (metadata) => {
         try {
           const api = new VitestFolderAPI(pkg, {

src/worker/init.ts

@@ -12,7 +12,7 @@ export async function initVitest(meta: WorkerInitMetadata, options?: UserConfig)
   let stdout: Writable | undefined
   let stderr: Writable | undefined
 
-  if (meta.shellType === 'terminal') {
+  if (meta.shellType === 'terminal' && !meta.hasShellIntegration) {
     stdout = new Writable({
       write(chunk, __, callback) {
         const log = chunk.toString()

src/worker/types.ts

@@ -7,6 +7,7 @@ export interface WorkerInitMetadata {
   workspaceFile?: string
   env: Record<string, any> | undefined
   shellType: 'terminal' | 'child_process'
+  hasShellIntegration: boolean
   pnpApi?: string
   pnpLoader?: string
 }