From 35d005464a811018baed8baf17bceb0e13b3d540 Mon Sep 17 00:00:00 2001
From: Robin Huang <robin.j.huang@gmail.com>
Date: Wed, 20 Nov 2024 19:09:05 -0800
Subject: [PATCH] Alias pip to uv pip in the integrated terminal. (#313)

* Add..

* Zsh as the default shell on mac.
---
 src/main-process/comfyDesktopApp.ts | 35 ++++++++++++++++-------------
 src/terminal.ts                     | 25 +++++++++++++++++++--
 2 files changed, 42 insertions(+), 18 deletions(-)

diff --git a/src/main-process/comfyDesktopApp.ts b/src/main-process/comfyDesktopApp.ts
index 3cf3dfef..9ffb36a9 100644
--- a/src/main-process/comfyDesktopApp.ts
+++ b/src/main-process/comfyDesktopApp.ts
@@ -20,15 +20,13 @@ import { getAppResourcesPath } from '../install/resourcePaths';
 
 export class ComfyDesktopApp {
   public comfyServer: ComfyServer | null = null;
-  private terminal: Terminal;
+  private terminal: Terminal | null = null; // Only created after server starts.
 
   constructor(
     public basePath: string,
     public comfySettings: ComfySettings,
     public appWindow: AppWindow
-  ) {
-    this.terminal = new Terminal(this.appWindow, getAppResourcesPath());
-  }
+  ) {}
 
   get pythonInstallPath() {
     return app.isPackaged ? this.basePath : path.join(app.getAppPath(), 'assets');
@@ -50,6 +48,22 @@ export class ComfyDesktopApp {
     });
   }
 
+  private initializeTerminal(virtualEnvironment: VirtualEnvironment) {
+    this.terminal = new Terminal(this.appWindow, this.basePath, virtualEnvironment.uvPath);
+
+    ipcMain.handle(IPC_CHANNELS.TERMINAL_WRITE, (_event, command: string) => {
+      this.terminal?.write(command);
+    });
+
+    ipcMain.handle(IPC_CHANNELS.TERMINAL_RESIZE, (_event, cols: number, rows: number) => {
+      this.terminal?.resize(cols, rows);
+    });
+
+    ipcMain.handle(IPC_CHANNELS.TERMINAL_RESTORE, (_event) => {
+      return this.terminal?.restore();
+    });
+  }
+
   async setupGPUContext(): Promise<void> {
     log.debug('Setting up GPU context');
     try {
@@ -112,18 +126,6 @@ export class ComfyDesktopApp {
         return null;
       }
     });
-
-    ipcMain.handle(IPC_CHANNELS.TERMINAL_WRITE, (_event, command: string) => {
-      this.terminal.write(command);
-    });
-
-    ipcMain.handle(IPC_CHANNELS.TERMINAL_RESIZE, (_event, cols: number, rows: number) => {
-      this.terminal.resize(cols, rows);
-    });
-
-    ipcMain.handle(IPC_CHANNELS.TERMINAL_RESTORE, (_event) => {
-      return this.terminal.restore();
-    });
   }
 
   /**
@@ -175,6 +177,7 @@ export class ComfyDesktopApp {
     this.appWindow.sendServerStartProgress(ProgressStatus.STARTING_SERVER);
     this.comfyServer = new ComfyServer(this.basePath, serverArgs, virtualEnvironment, this.appWindow);
     await this.comfyServer.start();
+    this.initializeTerminal(virtualEnvironment);
   }
 
   static async create(appWindow: AppWindow): Promise<ComfyDesktopApp> {
diff --git a/src/terminal.ts b/src/terminal.ts
index a42cc3df..4b27c1cc 100644
--- a/src/terminal.ts
+++ b/src/terminal.ts
@@ -7,6 +7,7 @@ export class Terminal {
   #pty: pty.IPty | undefined;
   #window: AppWindow | undefined;
   #cwd: string | undefined;
+  #uvPath: string | undefined;
 
   readonly sessionBuffer: string[] = [];
   readonly size = { cols: 80, rows: 30 };
@@ -21,9 +22,10 @@ export class Terminal {
     return this.#window;
   }
 
-  constructor(window: AppWindow, cwd: string) {
+  constructor(window: AppWindow, cwd: string, uvPath: string) {
     this.#window = window;
     this.#cwd = cwd;
+    this.#uvPath = uvPath;
   }
 
   write(data: string) {
@@ -46,7 +48,7 @@ export class Terminal {
   #createPty() {
     const window = this.window;
     // TODO: does this want to be a setting?
-    const shell = os.platform() === 'win32' ? 'powershell.exe' : 'bash';
+    const shell = this.#getDefaultShell();
     const instance = pty.spawn(shell, [], {
       handleFlowControl: false,
       conptyInheritCursor: false,
@@ -56,6 +58,14 @@ export class Terminal {
       cwd: this.#cwd,
     });
 
+    if (process.platform === 'win32') {
+      // PowerShell function
+      instance.write(`function pip { & "${this.#uvPath}" pip $args }\r\n`);
+    } else {
+      // Bash/Zsh alias
+      instance.write(`alias pip='"${this.#uvPath}" pip'\r\n`);
+    }
+
     instance.onData((data) => {
       this.sessionBuffer.push(data);
       window.send(IPC_CHANNELS.TERMINAL_ON_OUTPUT, data);
@@ -64,4 +74,15 @@ export class Terminal {
 
     return instance;
   }
+
+  #getDefaultShell(): string {
+    switch (os.platform()) {
+      case 'win32':
+        return 'powershell.exe';
+      case 'darwin':
+        return 'zsh';
+      default: // Linux and others
+        return 'bash';
+    }
+  }
 }