From 7d61e7d465cb816681be95ee5a3b8362dd423a42 Mon Sep 17 00:00:00 2001
From: Nicolas Meienberger <github@thisprops.com>
Date: Tue, 31 Mar 2026 17:53:10 +0200
Subject: [PATCH] fix: app lifecycle shutdown

---
 .../modules/lifecycle/__tests__/shutdown.test.ts   |  6 +++---
 app/server/modules/lifecycle/bootstrap.ts          | 10 +++++++++-
 app/server/modules/lifecycle/shutdown.ts           |  4 ++--
 app/server/plugins/bootstrap.ts                    | 14 ++------------
 4 files changed, 16 insertions(+), 18 deletions(-)

diff --git a/app/server/modules/lifecycle/__tests__/shutdown.test.ts b/app/server/modules/lifecycle/__tests__/shutdown.test.ts
index f60839c5..7625c890 100644
--- a/app/server/modules/lifecycle/__tests__/shutdown.test.ts
+++ b/app/server/modules/lifecycle/__tests__/shutdown.test.ts
@@ -1,8 +1,8 @@
 import { afterEach, describe, expect, mock, spyOn, test } from "bun:test";
 import { Scheduler } from "../../../core/scheduler";
-import * as agentsManagerModule from "../../agents/agents-manager";
 import * as backendModule from "../../backends/backend";
 import type { VolumeBackend } from "../../backends/backend";
+import * as bootstrapModule from "../bootstrap";
 import { createTestVolume } from "~/test/helpers/volume";
 
 const loadShutdownModule = async () => {
@@ -21,7 +21,7 @@ describe("shutdown", () => {
 		const stopScheduler = mock(async () => {
 			events.push("scheduler.stop");
 		});
-		const stopAgentRuntime = mock(async () => {
+		const stopApplicationRuntime = mock(async () => {
 			events.push("agents.stop");
 		});
 		const unmountVolume = mock(async () => {
@@ -39,7 +39,7 @@ describe("shutdown", () => {
 		});
 
 		spyOn(Scheduler, "stop").mockImplementation(stopScheduler);
-		spyOn(agentsManagerModule, "stopAgentRuntime").mockImplementation(stopAgentRuntime);
+		spyOn(bootstrapModule, "stopApplicationRuntime").mockImplementation(stopApplicationRuntime);
 		spyOn(backendModule, "createVolumeBackend").mockImplementation(
 			() =>
 				({
diff --git a/app/server/modules/lifecycle/bootstrap.ts b/app/server/modules/lifecycle/bootstrap.ts
index 0866be58..791ebce8 100644
--- a/app/server/modules/lifecycle/bootstrap.ts
+++ b/app/server/modules/lifecycle/bootstrap.ts
@@ -1,5 +1,5 @@
 import { runDbMigrations } from "../../db/db";
-import { agentManager, spawnLocalAgent } from "../agents/agents-manager";
+import { agentManager, spawnLocalAgent, stopAgentRuntime } from "../agents/agents-manager";
 import { runMigrations } from "./migrations";
 import { startup } from "./startup";
 
@@ -25,3 +25,11 @@ export const bootstrapApplication = async () => {
 		throw err;
 	}
 };
+
+export const stopApplicationRuntime = async () => {
+	try {
+		await stopAgentRuntime();
+	} finally {
+		bootstrapPromise = undefined;
+	}
+};
diff --git a/app/server/modules/lifecycle/shutdown.ts b/app/server/modules/lifecycle/shutdown.ts
index 78c0ef21..4deff2af 100644
--- a/app/server/modules/lifecycle/shutdown.ts
+++ b/app/server/modules/lifecycle/shutdown.ts
@@ -2,11 +2,11 @@ import { Scheduler } from "../../core/scheduler";
 import { db } from "../../db/db";
 import { logger } from "@zerobyte/core/node";
 import { createVolumeBackend } from "../backends/backend";
-import { stopAgentRuntime } from "../agents/agents-manager";
+import { stopApplicationRuntime } from "./bootstrap";
 
 export const shutdown = async () => {
 	await Scheduler.stop();
-	await stopAgentRuntime();
+	await stopApplicationRuntime();
 
 	const volumes = await db.query.volumesTable.findMany({
 		where: { status: "mounted" },
diff --git a/app/server/plugins/bootstrap.ts b/app/server/plugins/bootstrap.ts
index c661205f..ede70123 100644
--- a/app/server/plugins/bootstrap.ts
+++ b/app/server/plugins/bootstrap.ts
@@ -1,20 +1,10 @@
 import { definePlugin } from "nitro";
-import { bootstrapApplication } from "../modules/lifecycle/bootstrap";
+import { bootstrapApplication, stopApplicationRuntime } from "../modules/lifecycle/bootstrap";
 import { logger } from "@zerobyte/core/node";
 import { toMessage } from "../utils/errors";
-import { stopAgentRuntime } from "../modules/agents/agents-manager";
-
-type ProcessWithAgentCloseHook = NodeJS.Process & {
-	__zerobyteAgentRuntimeCloseHookRegistered?: boolean;
-};
 
 export default definePlugin(async (nitroApp) => {
-	const runtimeProcess = process as ProcessWithAgentCloseHook;
-
-	if (!runtimeProcess.__zerobyteAgentRuntimeCloseHookRegistered) {
-		nitroApp.hooks.hook("close", stopAgentRuntime);
-		runtimeProcess.__zerobyteAgentRuntimeCloseHookRegistered = true;
-	}
+	nitroApp.hooks.hook("close", stopApplicationRuntime);
 
 	await bootstrapApplication().catch((err) => {
 		logger.error(`Bootstrap failed: ${toMessage(err)}`);