Implement graceful shutdown for API server

src/cli/commands/dev.ts

@@ -223,13 +223,13 @@ export async function devCommand() {
     s.start(promptMessages[0] + '\n');
 
     // Initialize the API server
-    const { server, store } = await createPromptServer();
+    const { listen, store } = await createPromptServer();
 
     console.log('\n');
     s.stop(color.green('Server started successfully!'));
 
     // Start initial server
-    httpServer = server.listen(port, () => {
+    httpServer = listen(port, () => {
       console.log(`\nDevelopment server started on ${color.cyan(`http://localhost:${port}`)}`);
 
       const routeList = formatRoutes(store.routes.all('next'), port);
@@ -239,22 +239,6 @@ export async function devCommand() {
       console.log('Available tasks:\n' + taskList + '\n')
     });
 
-    const shutdown = () => {
-      if (httpServer) {
-        httpServer.close(async () => {
-          p.outro('Development server stopped');
-          store.tasks.cleanup();
-          process.exit(0);
-        });
-      } else {
-        p.outro('Development server stopped');
-        process.exit(0);
-      }
-    };
-
-    process.on('SIGINT', shutdown);
-    process.on('SIGTERM', shutdown);
-
   } catch (error) {
     console.log({ error })
     s.stop(color.red('Failed to start server'));

src/core/dataprompt.ts

@@ -145,6 +145,6 @@ export async function createPromptServer(options: {
 } = { startTasks: true }) {
   const { config, startTasks } = options;
   const store = await dataprompt(config);
-  const server = await createApiServer({ store, startTasks });
-  return { store, server };
+  const { app, listen } = await createApiServer({ store, startTasks });
+  return { store, app, listen, server: app };
 }

src/plugins/firebase/public.ts

@@ -1,3 +1,3 @@
 export { getFirebaseApp } from './app.js'
 export { firestorePlugin } from './firestore/index.js'
-export { FirebasePluginConfig } from './types.js'
+export type { FirebasePluginConfig } from './types.js'

src/routing/server.ts

@@ -1,4 +1,5 @@
 import express from 'express';
+import http from 'node:http';
 import { DatapromptStore } from '../core/dataprompt.js';
 import { events } from '../core/events.js';
 import { getLogManager } from '../utils/logging.js';
@@ -88,8 +89,31 @@ export async function createApiServer({ store, startTasks }: {
     );
   }
 
-  // TODO: Add support for graceful shutdown (handling SIGTERM, SIGINT) to close server and database connections properly.
-  // TODO: Consider allowing the caller to specify the port or return a `listen` helper.
+  const app = server;
 
-  return server;
+  const listen = (port: number, callback?: () => void) => {
+    const httpServer = app.listen(port, callback);
+
+    const shutdown = () => {
+      console.log('\nGracefully shutting down...');
+      httpServer.close(() => {
+        console.log('HTTP server closed.');
+        store.tasks.cleanup();
+        process.exit(0);
+      });
+
+      // Force close after 10s
+      setTimeout(() => {
+        console.error('Could not close connections in time, forcefully shutting down');
+        process.exit(1);
+      }, 10000);
+    };
+
+    process.on('SIGTERM', shutdown);
+    process.on('SIGINT', shutdown);
+
+    return httpServer;
+  };
+
+  return { app, listen };
 }