refactor: gate local agent behind feature flag

app/server/app.ts

@@ -58,7 +58,7 @@ export const createApp = () => {
 			limit: 1000,
 			keyGenerator: (c) => c.req.header("x-forwarded-for") ?? "",
 			skip: () => {
-				return config.disableRateLimiting;
+				return config.flags.disableRateLimiting;
 			},
 		}),
 	);

app/server/core/config.ts

@@ -43,6 +43,7 @@ const envSchema = z
 		APP_SECRET_FILE: z.string().optional(),
 		BASE_URL: z.string(),
 		ENABLE_DEV_PANEL: z.string().default("false"),
+		ENABLE_LOCAL_AGENT: z.string().default("false"),
 		PROVISIONING_PATH: z.string().optional(),
 	})
 	.transform((s, ctx) => {
@@ -125,11 +126,14 @@ const envSchema = z
 			appVersion: s.APP_VERSION,
 			trustedOrigins: trustedOrigins,
 			trustProxy: s.TRUST_PROXY === "true",
-			disableRateLimiting: s.DISABLE_RATE_LIMITING === "true" || s.NODE_ENV === "test",
 			appSecret: appSecret ?? "",
 			baseUrl,
 			isSecure: baseUrl.startsWith("https://"),
-			enableDevPanel: s.ENABLE_DEV_PANEL === "true",
+			flags: {
+				disableRateLimiting: s.DISABLE_RATE_LIMITING === "true" || s.NODE_ENV === "test",
+				enableDevPanel: s.ENABLE_DEV_PANEL === "true",
+				enableLocalAgent: s.ENABLE_LOCAL_AGENT === "true",
+			},
 			provisioningPath: s.PROVISIONING_PATH,
 			allowedHosts,
 		};

app/server/lib/auth.ts

@@ -30,13 +30,13 @@ export const auth = betterAuth({
 	},
 	trustedOrigins: config.trustedOrigins,
 	rateLimit: {
-		enabled: !config.disableRateLimiting,
+		enabled: !config.flags.disableRateLimiting,
 	},
 	advanced: {
 		cookiePrefix: "zerobyte",
 		useSecureCookies: config.isSecure,
 		ipAddress: {
-			disableIpTracking: config.disableRateLimiting,
+			disableIpTracking: config.flags.disableRateLimiting,
 		},
 	},
 	onAPIError: {

app/server/modules/agents/controller/server.ts

@@ -47,18 +47,19 @@ export function createAgentManagerRuntime() {
 	let backupHandlers: AgentBackupEventHandlers = {};
 	let runtimeScope: Scope.CloseableScope | null = null;
 
-	const closeSession = (sessionHandle: ControllerAgentSessionHandle) => {
-		Effect.runSync(Fiber.interrupt(sessionHandle.runFiber));
-		Effect.runSync(Scope.close(sessionHandle.scope, Exit.succeed(undefined)));
-	};
+	const closeSession = (sessionHandle: ControllerAgentSessionHandle) =>
+		Effect.gen(function* () {
+			yield* Fiber.interrupt(sessionHandle.runFiber);
+			yield* Scope.close(sessionHandle.scope, Exit.succeed(undefined));
+		});
 
-	const closeAllSessions = () => {
+	const closeAllSessions = Effect.gen(function* () {
 		const currentSessions = sessions;
 		sessions = new Map();
 		for (const sessionHandle of currentSessions.values()) {
-			closeSession(sessionHandle);
+			yield* closeSession(sessionHandle);
 		}
-	};
+	});
 
 	const getSessionHandle = (agentId: string) => sessions.get(agentId);
 
@@ -105,7 +106,9 @@ export function createAgentManagerRuntime() {
 	const setSession = (agentId: string, sessionHandle: ControllerAgentSessionHandle) => {
 		const existingSession = getSessionHandle(agentId);
 		if (existingSession) {
-			closeSession(existingSession);
+			void Effect.runPromise(closeSession(existingSession)).catch((error) => {
+				logger.error(`Failed to close existing agent session for ${agentId}: ${toMessage(error)}`);
+			});
 		}
 
 		sessions.set(agentId, sessionHandle);
@@ -118,7 +121,9 @@ export function createAgentManagerRuntime() {
 		}
 
 		sessions.delete(agentId);
-		closeSession(sessionHandle);
+		void Effect.runPromise(closeSession(sessionHandle)).catch((error) => {
+			logger.error(`Failed to close agent session for ${agentId}: ${toMessage(error)}`);
+		});
 	};
 
 	const acquireServer = Effect.acquireRelease(
@@ -180,7 +185,7 @@ export function createAgentManagerRuntime() {
 			}),
 		),
 		(server) =>
-			Effect.sync(closeAllSessions).pipe(
+			closeAllSessions.pipe(
 				Effect.andThen(
 					Effect.tryPromise({
 						try: () => server.stop(true),

app/server/modules/agents/local/process.ts

@@ -12,6 +12,10 @@ type LocalAgentState = {
 export async function spawnLocalAgentProcess(runtime: LocalAgentState) {
 	await stopLocalAgentProcess(runtime);
 
+	if (!config.flags.enableLocalAgent) {
+		return;
+	}
+
 	const sourceEntryPoint = path.join(process.cwd(), "apps", "agent", "src", "index.ts");
 	const productionEntryPoint = path.join(process.cwd(), ".output", "agent", "index.mjs");
 

app/server/modules/lifecycle/__tests__/shutdown.test.ts

@@ -1,4 +1,4 @@
-import { afterEach, describe, expect, mock, spyOn, test } from "bun:test";
+import { afterEach, describe, expect, test, vi } from "vitest";
 import { Scheduler } from "../../../core/scheduler";
 import * as backendModule from "../../backends/backend";
 import type { VolumeBackend } from "../../backends/backend";
@@ -12,19 +12,19 @@ const loadShutdownModule = async () => {
 };
 
 afterEach(() => {
-	mock.restore();
+	vi.restoreAllMocks();
 });
 
 describe("shutdown", () => {
 	test("stops the agent runtime before unmounting mounted volumes", async () => {
 		const events: string[] = [];
-		const stopScheduler = mock(async () => {
+		const stopScheduler = vi.fn(async () => {
 			events.push("scheduler.stop");
 		});
-		const stopApplicationRuntime = mock(async () => {
+		const stopApplicationRuntime = vi.fn(async () => {
 			events.push("agents.stop");
 		});
-		const unmountVolume = mock(async () => {
+		const unmountVolume = vi.fn(async () => {
 			events.push("backend.unmount");
 			return { status: "unmounted" as const };
 		});
@@ -38,9 +38,9 @@ describe("shutdown", () => {
 			status: "mounted",
 		});
 
-		spyOn(Scheduler, "stop").mockImplementation(stopScheduler);
-		spyOn(bootstrapModule, "stopApplicationRuntime").mockImplementation(stopApplicationRuntime);
-		spyOn(backendModule, "createVolumeBackend").mockImplementation(
+		vi.spyOn(Scheduler, "stop").mockImplementation(stopScheduler);
+		vi.spyOn(bootstrapModule, "stopApplicationRuntime").mockImplementation(stopApplicationRuntime);
+		vi.spyOn(backendModule, "createVolumeBackend").mockImplementation(
 			() =>
 				({
 					mount: async () => ({ status: "mounted" as const }),

app/server/modules/lifecycle/bootstrap.ts

@@ -8,7 +8,7 @@ let bootstrapPromise: Promise<void> | undefined;
 const runBootstrap = async () => {
 	await runDbMigrations();
 	await runMigrations();
-	agentManager.start();
+	await agentManager.start();
 	await spawnLocalAgent();
 	await startup();
 };

app/server/modules/system/system.service.ts

@@ -107,7 +107,7 @@ const setRegistrationEnabled = async (enabled: boolean) => {
 	logger.info(`Registration enabled set to: ${enabled}`);
 };
 
-const isDevPanelEnabled = () => config.enableDevPanel;
+const isDevPanelEnabled = () => config.flags.enableDevPanel;
 
 export const systemService = {
 	getSystemInfo,

apps/agent/package.json

@@ -4,7 +4,8 @@
 	"type": "module",
 	"module": "src/index.ts",
 	"scripts": {
-		"tsc": "tsc --noEmit"
+		"tsc": "tsc --noEmit",
+		"test": "bunx --bun vitest run --config ./vitest.config.ts"
 	},
 	"dependencies": {
 		"@zerobyte/contracts": "workspace:*",

apps/agent/src/__tests__/controller-session.test.ts

@@ -1,4 +1,4 @@
-import { afterEach, expect, mock, spyOn, test } from "bun:test";
+import { afterEach, expect, test, vi } from "vitest";
 import { Effect } from "effect";
 import waitForExpect from "wait-for-expect";
 import { fromPartial } from "@total-typescript/shoehorn";
@@ -7,11 +7,11 @@ import * as resticServer from "@zerobyte/core/restic/server";
 import { createControllerSession } from "../controller-session";
 
 afterEach(() => {
-	mock.restore();
+	vi.restoreAllMocks();
 });
 
 test("emits backup.failed when a backup command hits a restic error", async () => {
-	spyOn(resticServer, "createRestic").mockReturnValue(
+	vi.spyOn(resticServer, "createRestic").mockReturnValue(
 		fromPartial({
 			backup: () => Effect.fail("source path missing"),
 		}),

apps/agent/src/commands/backup-run.test.ts

@@ -1,4 +1,4 @@
-import { afterEach, expect, mock, spyOn, test } from "bun:test";
+import { afterEach, expect, test, vi } from "vitest";
 import { Effect } from "effect";
 import waitForExpect from "wait-for-expect";
 import { fromPartial } from "@total-typescript/shoehorn";
@@ -18,17 +18,19 @@ const createDeferred = <T>() => {
 };
 
 afterEach(() => {
-	mock.restore();
+	vi.restoreAllMocks();
 });
 
 test("waits for running-job registration before returning to the processor loop", async () => {
 	const outboundMessages: string[] = [];
 	const runningJobs = new Map<string, RunningJob>();
 	const setRunningJobGate = createDeferred<void>();
+	const processorLoopGate = createDeferred<void>();
+	const commandCompleted = createDeferred<void>();
 	const backupGate = createDeferred<{ exitCode: number; result: null; warningDetails: null }>();
 	let registeredAbortController: AbortController | undefined;
 
-	spyOn(resticServer, "createRestic").mockReturnValue(
+	vi.spyOn(resticServer, "createRestic").mockReturnValue(
 		fromPartial({
 			backup: () =>
 				Effect.async<{ exitCode: number; result: null; warningDetails: null }, never>((resume) => {
@@ -82,11 +84,21 @@ test("waits for running-job registration before returning to the processor loop"
 		scheduleId: "schedule-1",
 	});
 
-	const runPromise = Effect.runPromise(handleBackupRunCommand(context, runPayload));
+	const processorLoopPromise = Effect.runPromise(
+		Effect.gen(function* () {
+			yield* handleBackupRunCommand(context, runPayload);
+			commandCompleted.resolve(undefined);
+			yield* Effect.async<void, never>((resume) => {
+				void processorLoopGate.promise.then(() => {
+					resume(Effect.void);
+				});
+			});
+		}),
+	);
 
 	try {
 		const returnedBeforeRegistration = await Promise.race([
-			runPromise.then(() => true),
+			commandCompleted.promise.then(() => true),
 			new Promise<false>((resolve) => {
 				setTimeout(() => resolve(false), 0);
 			}),
@@ -95,7 +107,7 @@ test("waits for running-job registration before returning to the processor loop"
 		expect(returnedBeforeRegistration).toBe(false);
 
 		setRunningJobGate.resolve(undefined);
-		await runPromise;
+		await commandCompleted.promise;
 
 		await Effect.runPromise(handleBackupCancelCommand(context, cancelPayload));
 		expect(registeredAbortController?.signal.aborted).toBe(true);
@@ -111,7 +123,9 @@ test("waits for running-job registration before returning to the processor loop"
 			expect(runningJobs.has("job-1")).toBe(false);
 		});
 	} finally {
+		processorLoopGate.resolve(undefined);
 		setRunningJobGate.resolve(undefined);
 		backupGate.resolve({ exitCode: 0, result: null, warningDetails: null });
+		await processorLoopPromise;
 	}
 });

apps/agent/vitest.config.ts

@@ -0,0 +1,13 @@
+import { defineConfig } from "vitest/config";
+
+export default defineConfig({
+	test: {
+		environment: "node",
+		server: {
+			deps: {
+				inline: ["zod"],
+			},
+		},
+		include: ["src/**/*.test.ts", "src/**/*.spec.ts"],
+	},
+});