Ensure socket messages are always sent for all npm installs (#419)

* feat: hide the installation toast when a server driven npm install finishes

* feat: ensure the first and second npm install when creating an ai app are properly sequenced, and that the process gets deleted at the end

* feat: send more websocket messages to make sure the client always knows about in flight npm installs

* feat: send a socket message on connection that relays the current npm install status

* feat: always listen for deps:status messages and keep the usePackageJson context up to date

* feat: add mechanism to automatically broadcast deps:* messages whenever dependencies are installed

Previously, events were sent bespoke in each implementation. Many missed
events and that led to cases where the client often had no idea the
state of the server.

Now, always send all events always.

* fix: add missed `onExit` handler

* fix: run npm run format

* fix: address linter warning

* feat: wrap all returns coming out of npmInstall in waitForProcessToComplete

* fix: modify wss.onJoin to take the standard set of parameters as other handler functions

packages/api/apps/app.mts

@@ -1,10 +1,10 @@
 import { randomid, type AppType } from '@srcbook/shared';
 import { db } from '../db/index.mjs';
 import { type App as DBAppType, apps as appsTable } from '../db/schema.mjs';
-import { applyPlan, createViteApp, deleteViteApp, pathToApp, getFlatFilesForApp } from './disk.mjs';
+import { applyPlan, createViteApp, deleteViteApp, getFlatFilesForApp } from './disk.mjs';
 import { CreateAppSchemaType, CreateAppWithAiSchemaType } from './schemas.mjs';
 import { asc, desc, eq } from 'drizzle-orm';
-import { npmInstall } from '../exec.mjs';
+import { npmInstall } from './processes.mjs';
 import { generateApp } from '../ai/generate.mjs';
 import { toValidPackageName } from '../apps/utils.mjs';
 import { getPackagesToInstall, parsePlan } from '../ai/plan-parser.mjs';
@@ -40,8 +40,7 @@ export async function createAppWithAi(data: CreateAppWithAiSchemaType): Promise<
 
   // Note: we don't surface issues or retries and this is "running in the background".
   // In this case it works in our favor because we'll kickoff generation while it happens
-  npmInstall({
-    cwd: pathToApp(app.externalId),
+  const firstNpmInstallProcess = npmInstall(app.externalId, {
     stdout(data) {
       console.log(data.toString('utf8'));
     },
@@ -61,9 +60,10 @@ export async function createAppWithAi(data: CreateAppWithAiSchemaType): Promise<
   const packagesToInstall = getPackagesToInstall(plan);
 
   if (packagesToInstall.length > 0) {
+    await firstNpmInstallProcess;
+
     console.log('installing packages', packagesToInstall);
-    npmInstall({
-      cwd: pathToApp(app.externalId),
+    npmInstall(app.externalId, {
       packages: packagesToInstall,
       stdout(data) {
         console.log(data.toString('utf8'));
@@ -92,8 +92,7 @@ export async function createApp(data: CreateAppSchemaType): Promise<DBAppType> {
   // TODO: handle this better.
   // This should be done somewhere else and surface issues or retries.
   // Not awaiting here because it's "happening in the background".
-  npmInstall({
-    cwd: pathToApp(app.externalId),
+  npmInstall(app.externalId, {
     stdout(data) {
       console.log(data.toString('utf8'));
     },

packages/api/apps/processes.mts

@@ -1,6 +1,7 @@
 import { ChildProcess } from 'node:child_process';
 import { pathToApp } from './disk.mjs';
 import { npmInstall as execNpmInstall, vite as execVite } from '../exec.mjs';
+import { wss } from '../index.mjs';
 
 export type ProcessType = 'npm:install' | 'vite:server';
 
@@ -62,23 +63,88 @@ export function deleteAppProcess(appId: string, process: ProcessType) {
   processes.del(appId, process);
 }
 
+async function waitForProcessToComplete(process: AppProcessType) {
+  if (process.process.exitCode !== null) {
+    return process;
+  }
+
+  return new Promise<AppProcessType>((resolve, reject) => {
+    process.process.once('exit', () => {
+      resolve(process);
+    });
+    process.process.once('error', (err) => {
+      reject(err);
+    });
+  });
+}
+
 /**
  * Runs npm install for the given app.
  *
  * If there's already a process running npm install, it will return that process.
  */
 export function npmInstall(
   appId: string,
-  options: Omit<Parameters<typeof execNpmInstall>[0], 'cwd'>,
+  options: Omit<Partial<Parameters<typeof execNpmInstall>[0]>, 'cwd'> & { onStart?: () => void },
 ) {
-  if (!processes.has(appId, 'npm:install')) {
-    processes.set(appId, {
-      type: 'npm:install',
-      process: execNpmInstall({ cwd: pathToApp(appId), ...options }),
-    });
+  const runningProcess = processes.get(appId, 'npm:install');
+  if (runningProcess) {
+    return waitForProcessToComplete(runningProcess);
+  }
+
+  wss.broadcast(`app:${appId}`, 'deps:install:status', { status: 'installing' });
+  if (options.onStart) {
+    options.onStart();
   }
 
-  return processes.get(appId, 'npm:install');
+  const newlyStartedProcess: NpmInstallProcessType = {
+    type: 'npm:install',
+    process: execNpmInstall({
+      ...options,
+
+      cwd: pathToApp(appId),
+      stdout: (data) => {
+        wss.broadcast(`app:${appId}`, 'deps:install:log', {
+          log: { type: 'stdout', data: data.toString('utf8') },
+        });
+
+        if (options.stdout) {
+          options.stdout(data);
+        }
+      },
+      stderr: (data) => {
+        wss.broadcast(`app:${appId}`, 'deps:install:log', {
+          log: { type: 'stderr', data: data.toString('utf8') },
+        });
+
+        if (options.stderr) {
+          options.stderr(data);
+        }
+      },
+      onExit: (code, signal) => {
+        // We must clean up this process so that we can run npm install again
+        deleteAppProcess(appId, 'npm:install');
+
+        wss.broadcast(`app:${appId}`, 'deps:install:status', {
+          status: code === 0 ? 'complete' : 'failed',
+          code,
+        });
+
+        if (code === 0) {
+          wss.broadcast(`app:${appId}`, 'deps:status:response', {
+            nodeModulesExists: true,
+          });
+        }
+
+        if (options.onExit) {
+          options.onExit(code, signal);
+        }
+      },
+    }),
+  };
+  processes.set(appId, newlyStartedProcess);
+
+  return waitForProcessToComplete(newlyStartedProcess);
 }
 
 /**

packages/api/server/channels/app.mts

@@ -164,45 +164,15 @@ async function previewStop(
   result.process.kill('SIGTERM');
 }
 
-async function dependenciesInstall(
-  payload: DepsInstallPayloadType,
-  context: AppContextType,
-  wss: WebSocketServer,
-) {
+async function dependenciesInstall(payload: DepsInstallPayloadType, context: AppContextType) {
   const app = await loadApp(context.params.appId);
 
   if (!app) {
     return;
   }
 
   npmInstall(app.externalId, {
-    args: [],
     packages: payload.packages ?? undefined,
-    stdout: (data) => {
-      wss.broadcast(`app:${app.externalId}`, 'deps:install:log', {
-        log: { type: 'stdout', data: data.toString('utf8') },
-      });
-    },
-    stderr: (data) => {
-      wss.broadcast(`app:${app.externalId}`, 'deps:install:log', {
-        log: { type: 'stderr', data: data.toString('utf8') },
-      });
-    },
-    onExit: (code) => {
-      // We must clean up this process so that we can run npm install again
-      deleteAppProcess(app.externalId, 'npm:install');
-
-      wss.broadcast(`app:${app.externalId}`, 'deps:install:status', {
-        status: code === 0 ? 'complete' : 'failed',
-        code,
-      });
-
-      if (code === 0) {
-        wss.broadcast(`app:${app.externalId}`, 'deps:status:response', {
-          nodeModulesExists: true,
-        });
-      }
-    },
   });
 }
 
@@ -261,10 +231,17 @@ export function register(wss: WebSocketServer) {
       previewStart(payload, context, wss),
     )
     .on('preview:stop', PreviewStopPayloadSchema, previewStop)
-    .on('deps:install', DepsInstallPayloadSchema, (payload, context) => {
-      dependenciesInstall(payload, context, wss);
-    })
+    .on('deps:install', DepsInstallPayloadSchema, dependenciesInstall)
     .on('deps:clear', DepsInstallPayloadSchema, clearNodeModules)
     .on('deps:status', DepsStatusPayloadSchema, dependenciesStatus)
-    .on('file:updated', FileUpdatedPayloadSchema, onFileUpdated);
+    .on('file:updated', FileUpdatedPayloadSchema, onFileUpdated)
+    .onJoin((_payload, context, conn) => {
+      const appExternalId = (context as AppContextType).params.appId;
+
+      // When connecting, send back info about an in flight npm install if one exists
+      const npmInstallProcess = getAppProcess(appExternalId, 'npm:install');
+      if (npmInstallProcess) {
+        conn.reply(`app:${appExternalId}`, 'deps:install:status', { status: 'installing' });
+      }
+    });
 }

packages/api/server/ws-client.mts

@@ -17,6 +17,12 @@ export interface ConnectionContextType {
   reply: (topic: string, event: string, payload: Record<string, any>) => void;
 }
 
+type HandlerType = (
+  payload: Record<string, any>,
+  context: MessageContextType,
+  conn: ConnectionContextType,
+) => void;
+
 /**
  * Channel is responsible for dispatching incoming messages for a given topic.
  *
@@ -52,15 +58,11 @@ export class Channel {
     string,
     {
       schema: z.ZodTypeAny;
-      handler: (
-        payload: Record<string, any>,
-        context: MessageContextType,
-        conn: ConnectionContextType,
-      ) => void;
+      handler: HandlerType;
     }
   > = {};
 
-  onJoinCallback: (topic: string, ws: WebSocket) => void = () => {};
+  onJoinCallback: HandlerType = () => {};
 
   constructor(topic: string) {
     this.topic = topic;
@@ -130,7 +132,7 @@ export class Channel {
     return this;
   }
 
-  onJoin(callback: (topic: string, ws: WebSocket) => void) {
+  onJoin(callback: HandlerType) {
     this.onJoinCallback = callback;
     return this;
   }
@@ -214,7 +216,11 @@ export default class WebSocketServer {
     if (event === 'subscribe') {
       conn.subscriptions.push(topic);
       conn.reply(topic, 'subscribed', { id: payload.id });
-      channel.onJoinCallback(topic, conn.socket);
+      channel.onJoinCallback(
+        payload,
+        { topic: match.topic, event: event, params: match.params },
+        conn,
+      );
       return;
     }
 

packages/shared/src/schemas/websockets.mts

@@ -189,10 +189,13 @@ export const DepsInstallLogPayloadSchema = z.object({
   ]),
 });
 
-export const DepsInstallStatusPayloadSchema = z.object({
-  status: z.enum(['complete', 'failed']),
-  code: z.number().int(),
-});
+export const DepsInstallStatusPayloadSchema = z.union([
+  z.object({ status: z.literal('installing') }),
+  z.object({
+    status: z.enum(['complete', 'failed']),
+    code: z.number().int(),
+  }),
+]);
 
 export const DepsClearPayloadSchema = z.object({});
 export const DepsStatusPayloadSchema = z.object({});

packages/web/src/components/apps/package-install-toast.tsx

@@ -33,6 +33,8 @@ const PackageInstallToast: React.FunctionComponent = () => {
   useEffect(() => {
     if (nodeModulesExists === false && (status === 'idle' || status === 'complete')) {
       setShowToast(true);
+    } else if (nodeModulesExists === true) {
+      setShowToast(false);
     }
   }, [nodeModulesExists, status]);
 

packages/web/src/components/apps/use-package-json.tsx

@@ -52,6 +52,17 @@ export function PackageJsonProvider({ channel, children }: ProviderPropsType) {
     };
   }, [channel]);
 
+  useEffect(() => {
+    const callback = (payload: DepsInstallStatusPayloadType) => {
+      setStatus(payload.status);
+    };
+    channel.on('deps:install:status', callback);
+
+    return () => {
+      channel.off('deps:install:status', callback);
+    };
+  }, [channel]);
+
   const npmInstall = useCallback(
     async (packages?: Array<string>) => {
       addLog(
@@ -71,21 +82,27 @@ export function PackageJsonProvider({ channel, children }: ProviderPropsType) {
         };
         channel.on('deps:install:log', logCallback);
 
-        const statusCallback = ({ status, code }: DepsInstallStatusPayloadType) => {
-          channel.off('deps:install:log', logCallback);
-          channel.off('deps:install:status', statusCallback);
-          setStatus(status);
-
-          addLog(
-            'info',
-            'srcbook',
-            `${!packages ? 'npm install' : `npm install ${packages.join(' ')}`} exited with status code ${code}`,
-          );
-
-          if (status === 'complete') {
-            resolve();
-          } else {
-            reject(new Error(`Error running npm install: ${contents}`));
+        const statusCallback = (payload: DepsInstallStatusPayloadType) => {
+          switch (payload.status) {
+            case 'installing':
+              break;
+            case 'failed':
+            case 'complete':
+              channel.off('deps:install:log', logCallback);
+              channel.off('deps:install:status', statusCallback);
+
+              addLog(
+                'info',
+                'srcbook',
+                `${!packages ? 'npm install' : `npm install ${packages.join(' ')}`} exited with status code ${payload.code}`,
+              );
+
+              if (payload.status === 'complete') {
+                resolve();
+              } else {
+                reject(new Error(`Error running npm install: ${contents}`));
+              }
+              break;
           }
         };
         channel.on('deps:install:status', statusCallback);