From a85f00171db69c72c4c34d869f37bcdc6eb00d79 Mon Sep 17 00:00:00 2001 From: Nicolas Meienberger Date: Thu, 9 Apr 2026 20:54:01 +0200 Subject: [PATCH] refactor: gate local agent behind feature flag --- app/server/app.ts | 2 +- app/server/core/config.ts | 8 ++++-- app/server/lib/auth.ts | 4 +-- .../modules/agents/controller/server.ts | 25 +++++++++++------- app/server/modules/agents/local/process.ts | 4 +++ .../lifecycle/__tests__/shutdown.test.ts | 16 ++++++------ app/server/modules/lifecycle/bootstrap.ts | 2 +- app/server/modules/system/system.service.ts | 2 +- apps/agent/package.json | 3 ++- .../src/__tests__/controller-session.test.ts | 6 ++--- apps/agent/src/commands/backup-run.test.ts | 26 ++++++++++++++----- apps/agent/vitest.config.ts | 13 ++++++++++ 12 files changed, 76 insertions(+), 35 deletions(-) create mode 100644 apps/agent/vitest.config.ts diff --git a/app/server/app.ts b/app/server/app.ts index bdf541fe..042f9d6c 100644 --- a/app/server/app.ts +++ b/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; }, }), ); diff --git a/app/server/core/config.ts b/app/server/core/config.ts index ef08b087..7e4e3862 100644 --- a/app/server/core/config.ts +++ b/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, }; diff --git a/app/server/lib/auth.ts b/app/server/lib/auth.ts index 6344f3af..9fd54c43 100644 --- a/app/server/lib/auth.ts +++ b/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: { diff --git a/app/server/modules/agents/controller/server.ts b/app/server/modules/agents/controller/server.ts index 3bb025d7..ef7663b1 100644 --- a/app/server/modules/agents/controller/server.ts +++ b/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), diff --git a/app/server/modules/agents/local/process.ts b/app/server/modules/agents/local/process.ts index c844a65d..6c353782 100644 --- a/app/server/modules/agents/local/process.ts +++ b/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"); diff --git a/app/server/modules/lifecycle/__tests__/shutdown.test.ts b/app/server/modules/lifecycle/__tests__/shutdown.test.ts index 7625c890..6e8adddf 100644 --- a/app/server/modules/lifecycle/__tests__/shutdown.test.ts +++ b/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 }), diff --git a/app/server/modules/lifecycle/bootstrap.ts b/app/server/modules/lifecycle/bootstrap.ts index 791ebce8..ee0535ff 100644 --- a/app/server/modules/lifecycle/bootstrap.ts +++ b/app/server/modules/lifecycle/bootstrap.ts @@ -8,7 +8,7 @@ let bootstrapPromise: Promise | undefined; const runBootstrap = async () => { await runDbMigrations(); await runMigrations(); - agentManager.start(); + await agentManager.start(); await spawnLocalAgent(); await startup(); }; diff --git a/app/server/modules/system/system.service.ts b/app/server/modules/system/system.service.ts index d591f303..64ced883 100644 --- a/app/server/modules/system/system.service.ts +++ b/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, diff --git a/apps/agent/package.json b/apps/agent/package.json index c3bfdbb2..d530bb71 100644 --- a/apps/agent/package.json +++ b/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:*", diff --git a/apps/agent/src/__tests__/controller-session.test.ts b/apps/agent/src/__tests__/controller-session.test.ts index 87d1fdbc..b4afda3b 100644 --- a/apps/agent/src/__tests__/controller-session.test.ts +++ b/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"), }), diff --git a/apps/agent/src/commands/backup-run.test.ts b/apps/agent/src/commands/backup-run.test.ts index 70fc44ce..fb65677b 100644 --- a/apps/agent/src/commands/backup-run.test.ts +++ b/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 = () => { }; 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(); const setRunningJobGate = createDeferred(); + const processorLoopGate = createDeferred(); + const commandCompleted = createDeferred(); 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((resume) => { + void processorLoopGate.promise.then(() => { + resume(Effect.void); + }); + }); + }), + ); try { const returnedBeforeRegistration = await Promise.race([ - runPromise.then(() => true), + commandCompleted.promise.then(() => true), new Promise((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; } }); diff --git a/apps/agent/vitest.config.ts b/apps/agent/vitest.config.ts new file mode 100644 index 00000000..534c3696 --- /dev/null +++ b/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"], + }, +});