feat: xterm logger

src/lib/pty.test.ts

@@ -3,9 +3,10 @@
 /* eslint-disable @typescript-eslint/no-explicit-any */
 
 import * as pty from 'node-pty';
+import { assert } from 'tsafe';
 import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest';
 
-import { PtyManager } from '@/lib/pty';
+import { PtyManager, PtyNotFound } from '@/lib/pty';
 
 vi.mock('node-pty');
 
@@ -109,4 +110,84 @@ describe('PtyManager', () => {
     expect(ids).toContain(entry2.id);
     expect(ids).toHaveLength(2);
   });
+
+  it('should run a command and wait for completion', async () => {
+    const onData = vi.fn();
+    const onExit = vi.fn();
+    const entry = ptyManager.create({ onData, onExit });
+
+    // Setup mock behavior for command execution
+    mockPtyProcess.write.mockImplementation((data: string) => {
+      // Extract the marker from the command
+      const markerMatch = data.match(/__CMD_MARKER_(.+?)__/);
+      if (markerMatch) {
+        const marker = `__CMD_MARKER_${markerMatch[1]}__`;
+        // Simulate output with the marker and exit code 0
+        const callback = mockPtyProcess.onData.mock.calls[0][0];
+        setTimeout(() => {
+          callback(`command output\r\n${marker}:0\r\n`);
+        }, 10);
+      }
+    });
+
+    const result = await ptyManager.runCommand(entry.id, 'echo "test"');
+
+    assert(result !== PtyNotFound);
+    expect(result.exitCode).toBe(0);
+    expect(result.output).toContain('command output');
+    expect(mockPtyProcess.write).toHaveBeenCalled();
+  });
+
+  it('should handle command failure correctly', async () => {
+    const onData = vi.fn();
+    const onExit = vi.fn();
+    const entry = ptyManager.create({ onData, onExit });
+
+    mockPtyProcess.write.mockImplementation((data: string) => {
+      const markerMatch = data.match(/__CMD_MARKER_(.+?)__/);
+      if (markerMatch) {
+        const marker = `__CMD_MARKER_${markerMatch[1]}__`;
+        const callback = mockPtyProcess.onData.mock.calls[0][0];
+        setTimeout(() => {
+          callback(`error output\r\n${marker}:1\r\n`);
+        }, 10);
+      }
+    });
+
+    const result = await ptyManager.runCommand(entry.id, 'invalid_command');
+
+    assert(result !== PtyNotFound);
+    expect(result.exitCode).toBe(1);
+    expect(result.output).toContain('error output');
+  });
+
+  it('should reject commands when PTY is disposed', async () => {
+    const onData = vi.fn();
+    const onExit = vi.fn();
+    const entry = ptyManager.create({ onData, onExit });
+
+    // Start a command but don't let it complete
+    const commandPromise = ptyManager.runCommand(entry.id, 'long_running_command');
+
+    // Dispose the PTY before the command completes
+    ptyManager.dispose(entry.id);
+
+    // The command should reject
+    await expect(commandPromise).rejects.toThrow('PTY was disposed');
+  });
+
+  it('should handle command timeout', async () => {
+    const onData = vi.fn();
+    const onExit = vi.fn();
+    const entry = ptyManager.create({ onData, onExit });
+
+    // Command that won't complete in time
+    mockPtyProcess.write.mockImplementation(() => {
+      // Never send the completion marker
+    });
+
+    const commandPromise = ptyManager.runCommand(entry.id, 'hang', { timeout: 50 });
+
+    await expect(commandPromise).rejects.toThrow('Command timed out');
+  });
 });

src/lib/pty.ts

@@ -27,12 +27,22 @@ type CreatePtyArgs = {
   options?: PtyOptions;
 };
 
-const PtyNotFound = Symbol('PtyNotFound');
+type CommandExecution = {
+  id: string; // PTY ID
+  marker: string;
+  resolve: (value: { exitCode: number; output: string }) => void;
+  reject: (reason: Error) => void;
+  output: string[];
+};
+
+export const PtyNotFound = Symbol('PtyNotFound');
 
 export class PtyManager {
   ptys: Map<string, PtyEntry> = new Map();
   options: PtyManagerOptions;
   subscriptions: Set<() => void> = new Set();
+  // Add this new property to track active commands
+  private activeCommands = new Map<string, CommandExecution>();
 
   constructor(options?: Partial<PtyManagerOptions>) {
     this.options = { ...DEFAULT_PTY_MANAGER_OPTIONS, ...options };
@@ -52,6 +62,9 @@ export class PtyManager {
     const historyBuffer = new SlidingBuffer<string>(this.options.maxHistorySize);
 
     ptyProcess.onData((data) => {
+      // Check if this data contains a command completion marker
+      this.checkForCommandCompletion(id, data);
+
       const result = ansiSequenceBuffer.append(data);
       if (!result.hasIncomplete) {
         historyBuffer.push(result.complete);
@@ -60,6 +73,14 @@ export class PtyManager {
     });
 
     ptyProcess.onExit(({ exitCode }) => {
+      // Reject any pending commands for this PTY
+      for (const [marker, command] of this.activeCommands.entries()) {
+        if (command.id === id) {
+          command.reject(new Error(`Process exited with code ${exitCode} before command completed`));
+          this.activeCommands.delete(marker);
+        }
+      }
+
       ansiSequenceBuffer.clear();
       historyBuffer.clear();
       this.ptys.delete(id);
@@ -95,6 +116,14 @@ export class PtyManager {
   };
 
   dispose = (id: string): void => {
+    // Reject any pending commands for this PTY
+    for (const [marker, command] of this.activeCommands.entries()) {
+      if (command.id === id) {
+        command.reject(new Error('PTY was disposed'));
+        this.activeCommands.delete(marker);
+      }
+    }
+
     this.do(id, (entry) => {
       entry.process.kill();
       entry.ansiSequenceBuffer.clear();
@@ -104,6 +133,12 @@ export class PtyManager {
   };
 
   teardown = () => {
+    // Reject all pending commands
+    for (const [marker, command] of this.activeCommands.entries()) {
+      command.reject(new Error('PTY manager was torn down'));
+      this.activeCommands.delete(marker);
+    }
+
     const ids = this.ptys.keys();
     for (const id of ids) {
       this.dispose(id);
@@ -124,4 +159,100 @@ export class PtyManager {
     }
     return callback(entry);
   };
+
+  /**
+   * Run a command in a PTY and wait for it to complete.
+   * @returns Promise that resolves with the exit code and output when the command completes
+   */
+  runCommand = (
+    id: string,
+    command: string,
+    options?: { timeout?: number }
+  ): Promise<{ exitCode: number; output: string }> | typeof PtyNotFound => {
+    return this.do(id, (entry) => {
+      // Generate a unique marker for this command
+      const marker = `__CMD_MARKER_${nanoid(8)}__`;
+
+      // Create a promise that will resolve when the command completes
+      const commandPromise = new Promise<{ exitCode: number; output: string }>((resolve, reject) => {
+        // Store command tracking info
+        this.activeCommands.set(marker, {
+          id: id,
+          marker,
+          resolve,
+          reject,
+          output: [],
+        });
+
+        // Create the shell command that includes our marker
+        const wrappedCommand = this.wrapCommandForShell(command, marker);
+
+        // Send the command to the PTY
+        entry.process.write(`${wrappedCommand}\r`);
+      });
+
+      if (!options?.timeout) {
+        return commandPromise;
+      }
+
+      // Add timeout handling if specified
+      const timeoutPromise = new Promise<{ exitCode: number; output: string }>((_, reject) => {
+        setTimeout(() => {
+          if (this.activeCommands.has(marker)) {
+            this.activeCommands.delete(marker);
+            reject(new Error(`Command timed out after ${options.timeout}ms`));
+          }
+        }, options.timeout);
+      });
+
+      return Promise.race([commandPromise, timeoutPromise]);
+    });
+  };
+
+  /**
+   * Wrap a command so that it outputs a marker and exit code when done.
+   * Different shells need different syntax.
+   */
+  private wrapCommandForShell(command: string, marker: string): string {
+    const shell = getShell().toLowerCase();
+    const isWindows = process.platform === 'win32';
+
+    if (isWindows && (shell.includes('powershell') || shell.includes('pwsh'))) {
+      // PowerShell syntax
+      return `& { ${command}; Write-Host "${marker}:$LASTEXITCODE" }`;
+    } else if (isWindows && shell.includes('cmd')) {
+      // Windows CMD syntax
+      return `${command} & echo ${marker}:%ERRORLEVEL%`;
+    } else {
+      // Bash/sh/zsh syntax (default for macOS/Linux)
+      return `{ ${command}; } && echo "${marker}:$?"; true`;
+    }
+  }
+
+  /**
+   * Check if the data from a PTY contains a command completion marker.
+   */
+  private checkForCommandCompletion(ptyId: string, data: string): void {
+    // Check each active command
+    for (const [marker, command] of this.activeCommands.entries()) {
+      if (command.id !== ptyId) {
+        continue;
+      }
+
+      // Collect the output
+      command.output.push(data);
+
+      // Look for the marker pattern: MARKER:EXIT_CODE
+      const pattern = new RegExp(`${marker}:(\\d+)`);
+      const match = data.match(pattern);
+
+      if (match?.[1]) {
+        const exitCode = parseInt(match[1], 10);
+        const output = command.output.join('');
+
+        this.activeCommands.delete(marker);
+        command.resolve({ exitCode, output });
+      }
+    }
+  }
 }

src/main/index.ts

@@ -23,16 +23,17 @@ if (require('electron-squirrel-startup')) {
 
 const main = new MainProcessManager({ store });
 
-const [install, cleanupInstall] = createInstallManager({
+const [ptyManager, cleanupPty] = createPtyManager({
   ipc: main.ipc,
   sendToWindow: main.sendToWindow,
 });
-const [invoke, cleanupInvoke] = createInvokeManager({
-  store,
+const [install, cleanupInstall] = createInstallManager({
   ipc: main.ipc,
   sendToWindow: main.sendToWindow,
+  ptyManager,
 });
-const [_pty, cleanupPty] = createPtyManager({
+const [invoke, cleanupInvoke] = createInvokeManager({
+  store,
   ipc: main.ipc,
   sendToWindow: main.sendToWindow,
 });

src/main/install-manager.ts

@@ -8,6 +8,7 @@ import path, { join } from 'path';
 import { serializeError } from 'serialize-error';
 import { assert } from 'tsafe';
 
+import type { PtyManager } from '@/lib/pty';
 import { withResultAsync } from '@/lib/result';
 import { SimpleLogger } from '@/lib/simple-logger';
 import { FIRST_RUN_MARKER_FILENAME } from '@/main/constants';
@@ -28,8 +29,13 @@ export class InstallManager {
   private onStatusChange: (status: WithTimestamp<InstallProcessStatus>) => void;
   private log: SimpleLogger;
   private abortController: AbortController | null;
+  private ptyManager: PtyManager;
 
-  constructor(arg: { ipcLogger: InstallManager['ipcLogger']; onStatusChange: InstallManager['onStatusChange'] }) {
+  constructor(arg: {
+    ipcLogger: InstallManager['ipcLogger'];
+    onStatusChange: InstallManager['onStatusChange'];
+    ptyManager: InstallManager['ptyManager'];
+  }) {
     this.ipcLogger = arg.ipcLogger;
     this.onStatusChange = arg.onStatusChange;
     this.status = { type: 'uninitialized', timestamp: Date.now() };
@@ -38,6 +44,7 @@ export class InstallManager {
       console[entry.level](entry.message);
     });
     this.abortController = null;
+    this.ptyManager = arg.ptyManager;
   }
 
   logRepairModeMessages = (): void => {
@@ -65,6 +72,27 @@ export class InstallManager {
      * - Delete any existing virtual environment.
      */
 
+    // Create a new PTY for the installation process
+    const installationPty = this.ptyManager.create({
+      onData: (id, data) => {
+        // This just handles the visual output in the UI
+      },
+      onExit: (id, exitCode) => {
+        // Handle unexpected PTY exit
+        if (this.status.type !== 'completed' && this.status.type !== 'canceled') {
+          this.updateStatus({
+            type: 'error',
+            error: {
+              message: `Installation process exited unexpectedly with code ${exitCode}`,
+            },
+            ptyId: id,
+          });
+        }
+      },
+    });
+
+    const ptyId = installationPty.id;
+
     this.updateStatus({ type: 'starting' });
     this.log.info('Starting up...\r\n');
 
@@ -375,6 +403,7 @@ export class InstallManager {
 export const createInstallManager = (arg: {
   ipc: IpcListener<IpcEvents>;
   sendToWindow: <T extends keyof IpcRendererEvents>(channel: T, ...args: IpcRendererEvents[T]) => void;
+  ptyManager: PtyManager;
 }) => {
   const { ipc, sendToWindow } = arg;
 
@@ -385,6 +414,7 @@ export const createInstallManager = (arg: {
     onStatusChange: (status) => {
       sendToWindow('install-process:status', status);
     },
+    ptyManager,
   });
   ipc.handle('install-process:start-install', (_, installationPath, gpuType, version, repair) => {
     installManager.startInstall(installationPath, gpuType, version, repair);