From 5a16b5d486ba05f90f55640e15bfd310a63e9e33 Mon Sep 17 00:00:00 2001 From: psychedelicious <4822129+psychedelicious@users.noreply.github.com> Date: Fri, 4 Apr 2025 19:02:23 +1100 Subject: [PATCH 1/2] feat: add awaitable command-running api to PtyManager --- src/lib/pty.test.ts | 83 ++++++++++++++++++++++++++- src/lib/pty.ts | 133 +++++++++++++++++++++++++++++++++++++++++++- 2 files changed, 214 insertions(+), 2 deletions(-) diff --git a/src/lib/pty.test.ts b/src/lib/pty.test.ts index 8dd34bc..a52c725 100644 --- a/src/lib/pty.test.ts +++ b/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'); + }); }); diff --git a/src/lib/pty.ts b/src/lib/pty.ts index 9face0d..6e4e175 100644 --- a/src/lib/pty.ts +++ b/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 = new Map(); options: PtyManagerOptions; subscriptions: Set<() => void> = new Set(); + // Add this new property to track active commands + private activeCommands = new Map(); constructor(options?: Partial) { this.options = { ...DEFAULT_PTY_MANAGER_OPTIONS, ...options }; @@ -52,6 +62,9 @@ export class PtyManager { const historyBuffer = new SlidingBuffer(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 }); + } + } + } } From e34a807187e9cf64b91e42c02a3fbd01d834682b Mon Sep 17 00:00:00 2001 From: psychedelicious <4822129+psychedelicious@users.noreply.github.com> Date: Fri, 4 Apr 2025 19:02:42 +1100 Subject: [PATCH 2/2] feat: install manager uses pty to do install (wip) --- src/main/index.ts | 9 +++++---- src/main/install-manager.ts | 32 +++++++++++++++++++++++++++++++- 2 files changed, 36 insertions(+), 5 deletions(-) diff --git a/src/main/index.ts b/src/main/index.ts index 4b510fa..411a596 100644 --- a/src/main/index.ts +++ b/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, }); diff --git a/src/main/install-manager.ts b/src/main/install-manager.ts index c40b618..e6afc4a 100644 --- a/src/main/install-manager.ts +++ b/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) => 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; sendToWindow: (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);