mcp: shut down idle session daemons (#96)

Author: benvinegar

src/mcp/daemonState.ts

@@ -163,6 +163,10 @@ export class HunkDaemonState {
     return comments.filter((comment) => comment.filePath === filter.filePath);
   }
 
+  getSessionCount() {
+    return this.sessions.size;
+  }
+
   getPendingCommandCount() {
     return this.pendingCommands.size;
   }

src/mcp/server.ts

@@ -11,8 +11,9 @@ import {
   type SessionDaemonResponse,
 } from "../session/protocol";
 
-const STALE_SESSION_TTL_MS = 45_000;
-const STALE_SESSION_SWEEP_INTERVAL_MS = 15_000;
+const DEFAULT_STALE_SESSION_TTL_MS = 45_000;
+const DEFAULT_STALE_SESSION_SWEEP_INTERVAL_MS = 15_000;
+const DEFAULT_IDLE_TIMEOUT_MS = 60_000;
 
 const SUPPORTED_SESSION_ACTIONS: SessionDaemonAction[] = [
   "list",
@@ -26,6 +27,12 @@ const SUPPORTED_SESSION_ACTIONS: SessionDaemonAction[] = [
   "comment-clear",
 ];
 
+export interface ServeHunkMcpServerOptions {
+  idleTimeoutMs?: number;
+  staleSessionTtlMs?: number;
+  staleSessionSweepIntervalMs?: number;
+}
+
 function formatDaemonServeError(error: unknown, host: string, port: number) {
   const message = error instanceof Error ? error.message : String(error);
   const normalized = message.toLowerCase();
@@ -154,18 +161,90 @@ async function handleSessionApiRequest(state: HunkDaemonState, request: Request)
 }
 
 /** Serve the local Hunk session daemon and websocket session broker. */
-export function serveHunkMcpServer() {
+export function serveHunkMcpServer(options: ServeHunkMcpServerOptions = {}) {
   const config = resolveHunkMcpConfig();
+  const idleTimeoutMs = options.idleTimeoutMs ?? DEFAULT_IDLE_TIMEOUT_MS;
+  const staleSessionTtlMs = options.staleSessionTtlMs ?? DEFAULT_STALE_SESSION_TTL_MS;
+  const staleSessionSweepIntervalMs =
+    options.staleSessionSweepIntervalMs ?? DEFAULT_STALE_SESSION_SWEEP_INTERVAL_MS;
   const state = new HunkDaemonState();
   const startedAt = Date.now();
+  let lastActivityAt = startedAt;
   let shuttingDown = false;
+  let sweepTimer: Timer | null = null;
+  let idleTimer: Timer | null = null;
+  let server: ReturnType<typeof Bun.serve<{}>> | null = null;
+
+  const hasActiveWork = () => state.getSessionCount() > 0 || state.getPendingCommandCount() > 0;
+
+  const clearIdleShutdownTimer = () => {
+    if (!idleTimer) {
+      return;
+    }
+
+    clearTimeout(idleTimer);
+    idleTimer = null;
+  };
+
+  const shutdown = () => {
+    if (shuttingDown) {
+      return;
+    }
+
+    shuttingDown = true;
+    if (sweepTimer) {
+      clearInterval(sweepTimer);
+      sweepTimer = null;
+    }
+
+    clearIdleShutdownTimer();
+    process.off("SIGINT", shutdown);
+    process.off("SIGTERM", shutdown);
+
+    state.shutdown();
+    server?.stop(true);
+  };
+
+  const refreshIdleShutdownTimer = () => {
+    clearIdleShutdownTimer();
+
+    if (shuttingDown || idleTimeoutMs <= 0 || hasActiveWork()) {
+      return;
+    }
+
+    const idleForMs = Date.now() - lastActivityAt;
+    const remainingMs = Math.max(0, idleTimeoutMs - idleForMs);
 
-  const sweepTimer = setInterval(() => {
-    state.pruneStaleSessions({ ttlMs: STALE_SESSION_TTL_MS });
-  }, STALE_SESSION_SWEEP_INTERVAL_MS);
+    idleTimer = setTimeout(() => {
+      idleTimer = null;
+
+      if (shuttingDown || hasActiveWork()) {
+        return;
+      }
+
+      if (Date.now() - lastActivityAt < idleTimeoutMs) {
+        refreshIdleShutdownTimer();
+        return;
+      }
+
+      shutdown();
+    }, remainingMs);
+    idleTimer.unref?.();
+  };
+
+  const noteActivity = () => {
+    lastActivityAt = Date.now();
+    refreshIdleShutdownTimer();
+  };
+
+  sweepTimer = setInterval(() => {
+    const removed = state.pruneStaleSessions({ ttlMs: staleSessionTtlMs });
+    if (removed > 0) {
+      noteActivity();
+    }
+  }, staleSessionSweepIntervalMs);
   sweepTimer.unref?.();
 
-  let server: ReturnType<typeof Bun.serve<{}>>;
   try {
     server = Bun.serve<{}>({
       hostname: config.host,
@@ -174,7 +253,11 @@ export function serveHunkMcpServer() {
         const url = new URL(request.url);
 
         if (url.pathname === "/health") {
-          state.pruneStaleSessions({ ttlMs: STALE_SESSION_TTL_MS });
+          const removed = state.pruneStaleSessions({ ttlMs: staleSessionTtlMs });
+          if (removed > 0) {
+            noteActivity();
+          }
+
           return Response.json({
             ok: true,
             pid: process.pid,
@@ -183,17 +266,19 @@ export function serveHunkMcpServer() {
             sessionApi: `${config.httpOrigin}${HUNK_SESSION_API_PATH}`,
             sessionCapabilities: `${config.httpOrigin}${HUNK_SESSION_CAPABILITIES_PATH}`,
             sessionSocket: `${config.wsOrigin}${HUNK_SESSION_SOCKET_PATH}`,
-            sessions: state.listSessions().length,
+            sessions: state.getSessionCount(),
             pendingCommands: state.getPendingCommandCount(),
-            staleSessionTtlMs: STALE_SESSION_TTL_MS,
+            staleSessionTtlMs,
           });
         }
 
         if (url.pathname === HUNK_SESSION_CAPABILITIES_PATH) {
+          noteActivity();
           return Response.json(sessionCapabilities());
         }
 
         if (url.pathname === HUNK_SESSION_API_PATH) {
+          noteActivity();
           return handleSessionApiRequest(state, request);
         }
 
@@ -230,44 +315,41 @@ export function serveHunkMcpServer() {
           switch (parsed.type) {
             case "register":
               state.registerSession(socket, parsed.registration, parsed.snapshot);
+              noteActivity();
               break;
             case "snapshot":
               state.updateSnapshot(parsed.sessionId, parsed.snapshot);
+              noteActivity();
               break;
             case "heartbeat":
               state.markSessionSeen(parsed.sessionId);
+              noteActivity();
               break;
             case "command-result":
               state.handleCommandResult(parsed);
+              noteActivity();
               break;
           }
         },
         close: (socket) => {
           state.unregisterSocket(socket);
+          noteActivity();
         },
       },
     });
   } catch (error) {
-    clearInterval(sweepTimer);
-    throw formatDaemonServeError(error, config.host, config.port);
-  }
-
-  const shutdown = () => {
-    if (shuttingDown) {
-      return;
+    if (sweepTimer) {
+      clearInterval(sweepTimer);
+      sweepTimer = null;
     }
 
-    shuttingDown = true;
-    clearInterval(sweepTimer);
-    process.off("SIGINT", shutdown);
-    process.off("SIGTERM", shutdown);
-
-    state.shutdown();
-    server.stop(true);
-  };
+    clearIdleShutdownTimer();
+    throw formatDaemonServeError(error, config.host, config.port);
+  }
 
   process.once("SIGINT", shutdown);
   process.once("SIGTERM", shutdown);
+  refreshIdleShutdownTimer();
 
   console.log(`Hunk session daemon listening on ${config.httpOrigin}${HUNK_SESSION_API_PATH}`);
   console.log(`Hunk session websocket listening on ${config.wsOrigin}${HUNK_SESSION_SOCKET_PATH}`);