From 27eb5c2f5e07de454648fce969c7acbf37c1fdb1 Mon Sep 17 00:00:00 2001 From: Nicolas Meienberger Date: Tue, 5 May 2026 21:40:23 +0200 Subject: [PATCH] refactor(agent): harden local agent volume lifecycle --- .../migration.sql | 2 +- .../jobs/__tests__/cleanup-dangling.test.ts | 18 +++ app/server/jobs/cleanup-dangling.ts | 9 +- app/server/modules/agents/agents-manager.ts | 37 ++++- .../modules/agents/controller/server.ts | 19 +++ .../lifecycle/__tests__/shutdown.test.ts | 52 ++++--- app/server/modules/lifecycle/shutdown.ts | 31 +++-- app/server/modules/lifecycle/startup.ts | 44 +++--- app/server/modules/volumes/volume.service.ts | 26 ++-- apps/agent/src/commands/volume.test.ts | 127 ++++++++++++++++++ apps/agent/src/commands/volume.ts | 11 +- apps/agent/src/index.ts | 20 +++ .../src/volume-host/__tests__/cleanup.test.ts | 49 +++++++ .../volume-host/__tests__/operations.test.ts | 66 +++++++++ apps/agent/src/volume-host/cleanup.ts | 28 ++++ 15 files changed, 471 insertions(+), 68 deletions(-) create mode 100644 app/server/jobs/__tests__/cleanup-dangling.test.ts create mode 100644 apps/agent/src/commands/volume.test.ts create mode 100644 apps/agent/src/volume-host/__tests__/cleanup.test.ts create mode 100644 apps/agent/src/volume-host/__tests__/operations.test.ts create mode 100644 apps/agent/src/volume-host/cleanup.ts diff --git a/app/drizzle/20260505165117_early_purple_man/migration.sql b/app/drizzle/20260505165117_early_purple_man/migration.sql index db99a6c0..8273a520 100644 --- a/app/drizzle/20260505165117_early_purple_man/migration.sql +++ b/app/drizzle/20260505165117_early_purple_man/migration.sql @@ -13,4 +13,4 @@ CREATE TABLE `agents_table` ( ); --> statement-breakpoint CREATE INDEX `agents_table_organization_id_idx` ON `agents_table` (`organization_id`);--> statement-breakpoint -CREATE INDEX `agents_table_status_idx` ON `agents_table` (`status`); \ No newline at end of file +CREATE INDEX `agents_table_status_idx` ON `agents_table` (`status`); diff --git a/app/server/jobs/__tests__/cleanup-dangling.test.ts b/app/server/jobs/__tests__/cleanup-dangling.test.ts new file mode 100644 index 00000000..11e52f02 --- /dev/null +++ b/app/server/jobs/__tests__/cleanup-dangling.test.ts @@ -0,0 +1,18 @@ +import { afterEach, expect, test, vi } from "vitest"; +import { config } from "../../core/config"; +import { CleanupDanglingMountsJob } from "../cleanup-dangling"; +import * as mountinfo from "../../utils/mountinfo"; + +afterEach(() => { + config.flags.enableLocalAgent = true; + vi.restoreAllMocks(); +}); + +test("skips controller-local mount inspection when local volume execution is agent-owned", async () => { + config.flags.enableLocalAgent = true; + const readMountInfo = vi.spyOn(mountinfo, "readMountInfo"); + + await new CleanupDanglingMountsJob().run(); + + expect(readMountInfo).not.toHaveBeenCalled(); +}); diff --git a/app/server/jobs/cleanup-dangling.ts b/app/server/jobs/cleanup-dangling.ts index a1c7eef7..ac0b618f 100644 --- a/app/server/jobs/cleanup-dangling.ts +++ b/app/server/jobs/cleanup-dangling.ts @@ -10,9 +10,16 @@ import { toMessage } from "../utils/errors"; import { VOLUME_MOUNT_BASE } from "../core/constants"; import { db } from "../db/db"; import { withContext } from "../core/request-context"; +import { config } from "../core/config"; +import { LOCAL_AGENT_ID } from "../modules/agents/constants"; export class CleanupDanglingMountsJob extends Job { async run() { + if (config.flags.enableLocalAgent) { + logger.debug("Skipping controller-local dangling mount cleanup because local volume execution is agent-owned."); + return { done: true, timestamp: new Date() }; + } + const organizations = await db.query.organization.findMany({}); if (organizations.length === 0) { logger.warn("No organizations found; skipping dangling mount cleanup to avoid false positives."); @@ -22,7 +29,7 @@ export class CleanupDanglingMountsJob extends Job { const allVolumes = []; for (const org of organizations) { const volumes = await withContext({ organizationId: org.id }, async () => volumeService.listVolumes()); - allVolumes.push(...volumes); + allVolumes.push(...volumes.filter((volume) => volume.agentId === LOCAL_AGENT_ID)); } const allSystemMounts = await readMountInfo(); diff --git a/app/server/modules/agents/agents-manager.ts b/app/server/modules/agents/agents-manager.ts index 962a4e03..9e63d9bd 100644 --- a/app/server/modules/agents/agents-manager.ts +++ b/app/server/modules/agents/agents-manager.ts @@ -3,14 +3,21 @@ import type { BackupRunPayload, VolumeCommand, VolumeCommandResult } from "@zero import { Effect } from "effect"; import { config } from "../../core/config"; import { createAgentManagerRuntime, type AgentManagerEvent } from "./controller/server"; +import { LOCAL_AGENT_ID } from "./constants"; import { spawnLocalAgentProcess, stopLocalAgentProcess } from "./local/process"; -import type { BackupExecutionProgress, BackupExecutionResult } from "./helpers/runtime-state"; -import { createAgentRuntimeState } from "./helpers/runtime-state"; +import { + createAgentRuntimeState, + type AgentRuntimeState, + type BackupExecutionProgress, + type BackupExecutionResult, +} from "./helpers/runtime-state"; import { getDevAgentRuntimeState } from "./helpers/runtime-state.dev"; export type { BackupExecutionProgress, BackupExecutionResult } from "./helpers/runtime-state"; export type { ProcessWithAgentRuntime } from "./helpers/runtime-state.dev"; -const productionRuntimeState = createAgentRuntimeState(); +type ProcessWithProductionAgentRuntime = NodeJS.Process & { + __zerobyteProductionAgentRuntime?: AgentRuntimeState; +}; type AgentRunBackupRequest = { scheduleId: number; @@ -19,7 +26,18 @@ type AgentRunBackupRequest = { onProgress: (progress: BackupExecutionProgress) => void; }; -const getAgentRuntimeState = () => (config.__prod__ ? productionRuntimeState : getDevAgentRuntimeState()); +const getProductionAgentRuntimeState = () => { + // Nitro production builds can bundle startup plugins and API handlers into separate chunks. + // Keep the live controller on process so both chunks see the same agent sessions. + const runtimeProcess = process as ProcessWithProductionAgentRuntime; + if (!runtimeProcess.__zerobyteProductionAgentRuntime) { + runtimeProcess.__zerobyteProductionAgentRuntime = createAgentRuntimeState(); + } + + return runtimeProcess.__zerobyteProductionAgentRuntime; +}; + +const getAgentRuntimeState = () => (config.__prod__ ? getProductionAgentRuntimeState() : getDevAgentRuntimeState()); const getAgentManagerRuntime = () => getAgentRuntimeState().agentManager; const getActiveBackupsByScheduleId = () => getAgentRuntimeState().activeBackupsByScheduleId; const getActiveBackupScheduleIdsByJobId = () => getAgentRuntimeState().activeBackupScheduleIdsByJobId; @@ -286,7 +304,16 @@ export const agentManager = { }; export const startLocalAgent = async () => { - await spawnLocalAgentProcess(getAgentRuntimeState()); + const runtime = getAgentRuntimeState(); + await spawnLocalAgentProcess(runtime); + + if (!runtime.agentManager) { + return; + } + + if (!(await runtime.agentManager.waitForAgentReady(LOCAL_AGENT_ID))) { + throw new Error("Local agent did not become ready before startup"); + } }; // fallow-ignore-next-line unused-export diff --git a/app/server/modules/agents/controller/server.ts b/app/server/modules/agents/controller/server.ts index 995ba237..f1c55e3b 100644 --- a/app/server/modules/agents/controller/server.ts +++ b/app/server/modules/agents/controller/server.ts @@ -46,6 +46,7 @@ class StopAgentManagerServerError extends Data.TaggedError("StopAgentManagerServ export function createAgentManagerRuntime(onEvent: (event: AgentManagerEvent) => void) { let sessions = new Map(); let runtimeScope: Scope.CloseableScope | null = null; + const sleep = (ms: number) => new Promise((resolve) => setTimeout(resolve, ms)); const closeSession = (sessionHandle: ControllerAgentSessionHandle) => Effect.gen(function* () { @@ -79,6 +80,11 @@ export function createAgentManagerRuntime(onEvent: (event: AgentManagerEvent) => const getSessionHandle = (agentId: string) => sessions.get(agentId); const getSession = (agentId: string) => getSessionHandle(agentId)?.session; + const isAgentReady = (agentId: string) => { + const session = getSession(agentId); + return !!session && Effect.runSync(session.isReady()); + }; + const handleSessionEvent = (params: { agentId: string; agentName: string; sessionId: string }) => { const { agentId, agentName } = params; @@ -290,6 +296,19 @@ export function createAgentManagerRuntime(onEvent: (event: AgentManagerEvent) => return { start, + waitForAgentReady: async (agentId: string, timeoutMs = 10_000) => { + const deadline = Date.now() + timeoutMs; + + while (Date.now() < deadline) { + if (isAgentReady(agentId)) { + return true; + } + + await sleep(50); + } + + return isAgentReady(agentId); + }, sendBackup: (agentId: string, payload: BackupRunPayload) => Effect.gen(function* () { const session = getSession(agentId); diff --git a/app/server/modules/lifecycle/__tests__/shutdown.test.ts b/app/server/modules/lifecycle/__tests__/shutdown.test.ts index da7aa379..4f5e1727 100644 --- a/app/server/modules/lifecycle/__tests__/shutdown.test.ts +++ b/app/server/modules/lifecycle/__tests__/shutdown.test.ts @@ -1,9 +1,10 @@ import { afterEach, describe, expect, test, vi } from "vitest"; import { Scheduler } from "../../../core/scheduler"; -import * as volumeHostModule from "../../../../../apps/agent/src/volume-host"; -import type { VolumeBackend } from "../../../../../apps/agent/src/volume-host"; import * as bootstrapModule from "../bootstrap"; +import { agentManager } from "../../agents/agents-manager"; import { createTestVolume } from "~/test/helpers/volume"; +import { config } from "~/server/core/config"; +import { db } from "~/server/db/db"; const loadShutdownModule = async () => { const moduleUrl = new URL("../shutdown.ts", import.meta.url); @@ -12,11 +13,12 @@ const loadShutdownModule = async () => { }; afterEach(() => { + config.flags.enableLocalAgent = true; vi.restoreAllMocks(); }); describe("shutdown", () => { - test("stops the agent runtime before unmounting mounted volumes", async () => { + test("does not unmount agent-owned volumes during controller shutdown", async () => { const events: string[] = []; const stopScheduler = vi.fn(async () => { events.push("scheduler.stop"); @@ -24,35 +26,51 @@ describe("shutdown", () => { const stopApplicationRuntime = vi.fn(async () => { events.push("agents.stop"); }); - const unmountVolume = vi.fn(async () => { - events.push("backend.unmount"); - return { status: "unmounted" as const }; - }); + const runVolumeCommand = vi.spyOn(agentManager, "runVolumeCommand"); - await createTestVolume({ + const volume = await createTestVolume({ name: "Shutdown test volume", config: { backend: "directory", path: "/Applications", }, status: "mounted", + agentId: "agent-1", }); vi.spyOn(Scheduler, "stop").mockImplementation(stopScheduler); vi.spyOn(bootstrapModule, "stopApplicationRuntime").mockImplementation(stopApplicationRuntime); - vi.spyOn(volumeHostModule, "createVolumeBackend").mockImplementation( - () => - ({ - mount: async () => ({ status: "mounted" as const }), - unmount: unmountVolume, - checkHealth: async () => ({ status: "mounted" as const }), - }) satisfies VolumeBackend, - ); const { shutdown } = await loadShutdownModule(); await shutdown(); - expect(events).toEqual(["scheduler.stop", "agents.stop", "backend.unmount"]); + expect(events).toEqual(["scheduler.stop", "agents.stop"]); + expect(runVolumeCommand).not.toHaveBeenCalled(); + }); + + test("keeps legacy controller-local fallback unmount on shutdown", async () => { + config.flags.enableLocalAgent = false; + const events: string[] = []; + vi.spyOn(Scheduler, "stop").mockImplementation(async () => { + events.push("scheduler.stop"); + }); + vi.spyOn(bootstrapModule, "stopApplicationRuntime").mockImplementation(async () => { + events.push("agents.stop"); + }); + + const volume = await createTestVolume({ + name: "Fallback shutdown test volume", + config: { backend: "directory", path: "/Applications" }, + status: "mounted", + }); + + const { shutdown } = await loadShutdownModule(); + + await shutdown(); + + const updated = await db.query.volumesTable.findFirst({ where: { id: volume.id } }); + expect(events).toEqual(["scheduler.stop", "agents.stop"]); + expect(updated?.status).toBe("unmounted"); }); }); diff --git a/app/server/modules/lifecycle/shutdown.ts b/app/server/modules/lifecycle/shutdown.ts index 66ec7bf2..53c14356 100644 --- a/app/server/modules/lifecycle/shutdown.ts +++ b/app/server/modules/lifecycle/shutdown.ts @@ -2,25 +2,28 @@ import { Scheduler } from "../../core/scheduler"; import { db } from "../../db/db"; import { logger } from "@zerobyte/core/node"; import { stopApplicationRuntime } from "./bootstrap"; -import { decryptVolumeConfig } from "../volumes/volume-config-secrets"; -import { createVolumeBackend } from "../../../../apps/agent/src/volume-host"; +import { withContext } from "../../core/request-context"; +import { volumeService } from "../volumes/volume.service"; +import { toMessage } from "../../utils/errors"; +import { config } from "../../core/config"; +import { LOCAL_AGENT_ID } from "../agents/constants"; export const shutdown = async () => { await Scheduler.stop(); - await stopApplicationRuntime(); - const volumes = await db.query.volumesTable.findMany({ - where: { status: "mounted" }, - }); - - for (const volume of volumes) { - const backend = createVolumeBackend({ - ...volume, - config: await decryptVolumeConfig(volume.config), - provisioningId: volume.provisioningId ?? null, + if (!config.flags.enableLocalAgent) { + const volumes = await db.query.volumesTable.findMany({ + where: { AND: [{ status: "mounted" }, { agentId: LOCAL_AGENT_ID }] }, }); - const { status, error } = await backend.unmount(); - logger.info(`Volume ${volume.name} unmount status: ${status}${error ? `, error: ${error}` : ""}`); + for (const volume of volumes) { + const { status, error } = await withContext({ organizationId: volume.organizationId }, () => + volumeService.unmountVolume(volume.shortId), + ).catch((error) => ({ status: "error" as const, error: toMessage(error) })); + + logger.info(`Volume ${volume.name} unmount status: ${status}${error ? `, error: ${error}` : ""}`); + } } + + await stopApplicationRuntime(); }; diff --git a/app/server/modules/lifecycle/startup.ts b/app/server/modules/lifecycle/startup.ts index bd55caac..05fe79e8 100644 --- a/app/server/modules/lifecycle/startup.ts +++ b/app/server/modules/lifecycle/startup.ts @@ -17,6 +17,7 @@ import { backupsService } from "../backups/backups.service"; import { config } from "~/server/core/config"; import { syncProvisionedResources } from "../provisioning/provisioning"; import { toMessage } from "~/server/utils/errors"; +import { LOCAL_AGENT_ID } from "../agents/constants"; const ensureLatestConfigurationSchema = async () => { const volumes = await db.query.volumesTable.findMany({}); @@ -71,23 +72,30 @@ export const startup = async () => { logger.warn(`Removed ${deletedSchedules} orphaned backup schedule(s) during startup`); } - const volumes = await db.query.volumesTable.findMany({ - where: { - OR: [ - { status: "mounted" }, - { - AND: [{ autoRemount: true }, { status: "error" }], - }, - ], - }, - }); - - for (const volume of volumes) { - await withContext({ organizationId: volume.organizationId }, async () => { - await volumeService.mountVolume(volume.shortId).catch((err) => { - logger.error(`Error auto-remounting volume ${volume.name} on startup: ${err.message}`); - }); + if (!config.flags.enableLocalAgent) { + const volumes = await db.query.volumesTable.findMany({ + where: { + AND: [ + { agentId: LOCAL_AGENT_ID }, + { + OR: [ + { status: "mounted" }, + { + AND: [{ autoRemount: true }, { status: "error" }], + }, + ], + }, + ], + }, }); + + for (const volume of volumes) { + await withContext({ organizationId: volume.organizationId }, async () => { + await volumeService.mountVolume(volume.shortId).catch((err) => { + logger.error(`Error auto-remounting volume ${volume.name} on startup: ${err.message}`); + }); + }); + } } await db @@ -102,7 +110,9 @@ export const startup = async () => { logger.error(`Failed to update stuck backup schedules on startup: ${err.message}`); }); - Scheduler.build(CleanupDanglingMountsJob).schedule("0 * * * *"); + if (!config.flags.enableLocalAgent) { + Scheduler.build(CleanupDanglingMountsJob).schedule("0 * * * *"); + } Scheduler.build(VolumeHealthCheckJob).schedule("*/30 * * * *"); Scheduler.build(RepositoryHealthCheckJob).schedule("50 12 * * *"); Scheduler.build(BackupExecutionJob).schedule("* * * * *"); diff --git a/app/server/modules/volumes/volume.service.ts b/app/server/modules/volumes/volume.service.ts index d53cff21..44cc3e60 100644 --- a/app/server/modules/volumes/volume.service.ts +++ b/app/server/modules/volumes/volume.service.ts @@ -18,7 +18,13 @@ import { getOrganizationId } from "~/server/core/request-context"; import { type ShortId } from "~/server/utils/branded"; import { decryptVolumeConfig, encryptVolumeConfig } from "./volume-config-secrets"; import type { VolumeCommand, VolumeCommandResult } from "@zerobyte/contracts/agent-protocol"; -import { createVolumeBackend, getStatFs, getVolumePath } from "../../../../apps/agent/src/volume-host"; +import { + createVolumeBackend, + getStatFs, + getVolumePath, + type AgentVolume, + type BackendConfig as HostBackendConfig, +} from "../../../../apps/agent/src/volume-host"; import { browseFilesystem as browseHostFilesystem, listVolumeFiles, @@ -63,21 +69,25 @@ const volumeForAgent = async (volume: Volume): Promise => ({ config: await decryptVolumeConfig(volume.config), }); -const volumeForHost = async (volume: Volume): Promise => ({ +const volumeForHost = async (volume: Volume): Promise => ({ ...volume, shortId: volume.shortId, - config: await decryptVolumeConfig(volume.config), + config: (await decryptVolumeConfig(volume.config)) as HostBackendConfig, provisioningId: volume.provisioningId ?? null, }); -// TODO(agent-rollout): Remove the local host execution branch once all installs run volume operations through agents. +// Transitional fallback: older controller-only installs do not run the supervised local agent. +// Keep all controller-local host execution behind this predicate so the fallback is easy to delete +// once volume operations always go through the local agent. const shouldRunViaAgent = (volume: Volume) => volume.agentId !== LOCAL_AGENT_ID || config.flags.enableLocalAgent; +const shouldUseControllerLocalVolumeFallback = (volume: Volume) => !shouldRunViaAgent(volume); + const runVolumeBackendCommand = async ( volume: Volume, name: "volume.mount" | "volume.unmount" | "volume.checkHealth", ) => { - if (!shouldRunViaAgent(volume)) { + if (shouldUseControllerLocalVolumeFallback(volume)) { const backend = createVolumeBackend(await volumeForHost(volume)); switch (name) { case "volume.mount": @@ -202,7 +212,7 @@ const getVolume = async (shortId: ShortId) => { let statfs: Partial = {}; if (volume.status === "mounted") { statfs = await withTimeout( - shouldRunViaAgent(volume) + !shouldUseControllerLocalVolumeFallback(volume) ? runVolumeCommand(volume.agentId, { name: "volume.statfs", volume: await volumeForAgent(volume), @@ -280,7 +290,7 @@ const updateVolume = async (shortId: ShortId, volumeData: UpdateVolumeBody) => { const testConnection = async (backendConfig: BackendConfig) => { if (!config.flags.enableLocalAgent) { - return Effect.runPromise(testVolumeConnection(backendConfig)); + return Effect.runPromise(testVolumeConnection(backendConfig as HostBackendConfig)); } const command = await runVolumeCommand(LOCAL_AGENT_ID, { name: "volume.testConnection", backendConfig }); @@ -376,7 +386,7 @@ const listFiles = async (shortId: ShortId, subPath?: string, offset: number = 0, } try { - if (!shouldRunViaAgent(volume)) { + if (shouldUseControllerLocalVolumeFallback(volume)) { return await listVolumeFiles(await volumeForHost(volume), subPath, offset, limit); } diff --git a/apps/agent/src/commands/volume.test.ts b/apps/agent/src/commands/volume.test.ts new file mode 100644 index 00000000..938a11aa --- /dev/null +++ b/apps/agent/src/commands/volume.test.ts @@ -0,0 +1,127 @@ +import { afterEach, expect, test, vi } from "vitest"; +import { Effect } from "effect"; +import { fromPartial } from "@total-typescript/shoehorn"; +import { + parseAgentMessage, + type AgentWireMessage, + type VolumeCommandPayload, +} from "@zerobyte/contracts/agent-protocol"; +import type { ControllerCommandContext } from "../context"; + +const volumeHostMock = vi.hoisted(() => ({ + createVolumeBackend: vi.fn(), + getStatFs: vi.fn(), + getVolumePath: vi.fn(), +})); + +const operationsMock = vi.hoisted(() => ({ + browseFilesystem: vi.fn(), + listVolumeFiles: vi.fn(), + testVolumeConnection: vi.fn(), +})); + +vi.mock("../volume-host", () => volumeHostMock); +vi.mock("../volume-host/operations", () => operationsMock); + +import { handleVolumeCommand } from "./volume"; + +afterEach(() => { + vi.restoreAllMocks(); + volumeHostMock.createVolumeBackend.mockReset(); + volumeHostMock.getStatFs.mockReset(); + volumeHostMock.getVolumePath.mockReset(); + operationsMock.browseFilesystem.mockReset(); + operationsMock.listVolumeFiles.mockReset(); + operationsMock.testVolumeConnection.mockReset(); +}); + +const runVolumeCommand = async (payload: VolumeCommandPayload) => { + const outboundMessages: string[] = []; + const context = fromPartial({ + offerOutbound: (message: AgentWireMessage) => + Effect.sync(() => { + outboundMessages.push(message); + return true; + }), + }); + + await Effect.runPromise(handleVolumeCommand(context, payload)); + return outboundMessages.map((message) => parseAgentMessage(message)); +}; + +test("runs backend-backed volume commands on the agent host", async () => { + const mount = vi.fn(async () => ({ status: "mounted" as const })); + volumeHostMock.createVolumeBackend.mockReturnValue({ mount }); + + const messages = await runVolumeCommand( + fromPartial({ + commandId: "command-1", + command: { + name: "volume.mount", + volume: { id: 1, config: { backend: "directory", path: "/tmp/source" }, provisioningId: undefined }, + }, + }), + ); + + expect(volumeHostMock.createVolumeBackend).toHaveBeenCalledWith( + expect.objectContaining({ id: 1, config: { backend: "directory", path: "/tmp/source" }, provisioningId: null }), + ); + expect(mount).toHaveBeenCalledOnce(); + expect(messages[0]?.success).toBe(true); + if (messages[0]?.success && messages[0].data.type === "volume.commandResult") { + expect(messages[0].data.payload).toEqual({ + commandId: "command-1", + status: "success", + command: { name: "volume.mount", result: { status: "mounted" } }, + }); + } +}); + +test("returns command errors without throwing", async () => { + operationsMock.browseFilesystem.mockRejectedValue(new Error("permission denied")); + + const messages = await runVolumeCommand({ + commandId: "command-2", + command: { name: "filesystem.browse", path: "/root" }, + }); + + expect(messages[0]?.success).toBe(true); + if (messages[0]?.success && messages[0].data.type === "volume.commandResult") { + expect(messages[0].data.payload).toEqual({ + commandId: "command-2", + status: "error", + error: "permission denied", + }); + } +}); + +test("routes file listing commands to host operations", async () => { + operationsMock.listVolumeFiles.mockResolvedValue({ + files: [], + path: "/logs", + offset: 0, + limit: 10, + total: 0, + hasMore: false, + }); + + await runVolumeCommand( + fromPartial({ + commandId: "command-3", + command: { + name: "volume.listFiles", + volume: { id: 1, config: { backend: "directory", path: "/tmp/source" } }, + subPath: "/logs", + offset: 0, + limit: 10, + }, + }), + ); + + expect(operationsMock.listVolumeFiles).toHaveBeenCalledWith( + expect.objectContaining({ id: 1, config: { backend: "directory", path: "/tmp/source" } }), + "/logs", + 0, + 10, + ); +}); diff --git a/apps/agent/src/commands/volume.ts b/apps/agent/src/commands/volume.ts index 55b3c1a5..55b384dc 100644 --- a/apps/agent/src/commands/volume.ts +++ b/apps/agent/src/commands/volume.ts @@ -47,7 +47,8 @@ const executeVolumeCommand = (command: VolumeCommand) => } case "volume.listFiles": { const result = yield* Effect.tryPromise({ - try: () => listVolumeFiles(asVolume(command.volume), command.subPath, command.offset, command.limit), + try: () => + listVolumeFiles(asVolume(command.volume), command.subPath, command.offset, command.limit), catch: (error) => new VolumeCommandError({ cause: error }), }); @@ -80,14 +81,14 @@ export const handleVolumeCommand = (context: ControllerCommandContext, payload: return command; }).pipe( - Effect.tapError((error) => { - return context.offerOutbound( + Effect.catchAll((error) => + context.offerOutbound( createAgentMessage("volume.commandResult", { commandId: payload.commandId, status: "error", error: toMessage(error?.cause), }), - ); - }), + ), + ), ); }; diff --git a/apps/agent/src/index.ts b/apps/agent/src/index.ts index 2593792c..4f6b210e 100644 --- a/apps/agent/src/index.ts +++ b/apps/agent/src/index.ts @@ -1,14 +1,32 @@ import { logger } from "@zerobyte/core/node"; import { createControllerSession, type ControllerSession } from "./controller-session"; +import { cleanupDanglingVolumeMountDirectories } from "./volume-host/cleanup"; const controllerUrl = process.env.ZEROBYTE_CONTROLLER_URL; const agentToken = process.env.ZEROBYTE_AGENT_TOKEN; const RECONNECT_DELAY_MS = 1000; +const VOLUME_CLEANUP_INTERVAL_MS = 60 * 60 * 1000; export class Agent { private ws: WebSocket | null = null; private controllerSession: ControllerSession | null = null; private reconnectTimeout: ReturnType | null = null; + private volumeCleanupInterval: ReturnType | null = null; + + private runVolumeCleanup() { + void cleanupDanglingVolumeMountDirectories().catch((error) => { + logger.warn(`Agent volume cleanup failed: ${error instanceof Error ? error.message : String(error)}`); + }); + } + + private startVolumeCleanup() { + if (this.volumeCleanupInterval) { + return; + } + + this.runVolumeCleanup(); + this.volumeCleanupInterval = setInterval(() => this.runVolumeCleanup(), VOLUME_CLEANUP_INTERVAL_MS); + } private scheduleReconnect() { if (this.reconnectTimeout) { @@ -22,6 +40,8 @@ export class Agent { } connect() { + this.startVolumeCleanup(); + if (this.reconnectTimeout) { clearTimeout(this.reconnectTimeout); this.reconnectTimeout = null; diff --git a/apps/agent/src/volume-host/__tests__/cleanup.test.ts b/apps/agent/src/volume-host/__tests__/cleanup.test.ts new file mode 100644 index 00000000..1b191a5f --- /dev/null +++ b/apps/agent/src/volume-host/__tests__/cleanup.test.ts @@ -0,0 +1,49 @@ +import * as fs from "node:fs/promises"; +import * as os from "node:os"; +import * as path from "node:path"; +import { afterEach, expect, test, vi } from "vitest"; + +let tempRoot: string | undefined; +let mockMountPoints: string[] = []; + +afterEach(async () => { + vi.doUnmock("./constants"); + vi.doUnmock("./fs"); + vi.resetModules(); + mockMountPoints = []; + if (tempRoot) { + await fs.rm(tempRoot, { recursive: true, force: true }); + tempRoot = undefined; + } +}); + +const loadCleanup = async () => { + vi.resetModules(); + tempRoot = await fs.mkdtemp(path.join(os.tmpdir(), "zerobyte-agent-cleanup-")); + vi.doMock("./constants", () => ({ VOLUME_MOUNT_BASE: tempRoot })); + vi.doMock("./fs", () => ({ + readMountInfo: async () => mockMountPoints.map((mountPoint) => ({ mountPoint, fstype: "fuse.sshfs" })), + })); + + return import("../cleanup"); +}; + +test("removes stale volume directories that are not mounted on the agent host", async () => { + const { cleanupDanglingVolumeMountDirectories } = await loadCleanup(); + await fs.mkdir(path.join(tempRoot!, "stale-volume", "_data"), { recursive: true }); + + await cleanupDanglingVolumeMountDirectories(); + + await expect(fs.access(path.join(tempRoot!, "stale-volume"))).rejects.toThrow(); +}); + +test("keeps volume directories that are still mounted on the agent host", async () => { + const { cleanupDanglingVolumeMountDirectories } = await loadCleanup(); + const localVolumeDir = path.join(tempRoot!, "mounted-volume"); + mockMountPoints = [path.join(localVolumeDir, "_data")]; + await fs.mkdir(path.join(localVolumeDir, "_data"), { recursive: true }); + + await cleanupDanglingVolumeMountDirectories(); + + await expect(fs.access(localVolumeDir)).resolves.toBeNull(); +}); diff --git a/apps/agent/src/volume-host/__tests__/operations.test.ts b/apps/agent/src/volume-host/__tests__/operations.test.ts new file mode 100644 index 00000000..b0e60dbd --- /dev/null +++ b/apps/agent/src/volume-host/__tests__/operations.test.ts @@ -0,0 +1,66 @@ +import * as fs from "node:fs/promises"; +import * as os from "node:os"; +import * as path from "node:path"; +import { afterEach, expect, test } from "vitest"; +import { listVolumeFiles } from "../operations"; +import type { AgentVolume } from "../types"; + +let tempRoot: string | undefined; + +afterEach(async () => { + if (tempRoot) { + await fs.rm(tempRoot, { recursive: true, force: true }); + tempRoot = undefined; + } +}); + +const createDirectoryVolume = async (): Promise => { + tempRoot = await fs.mkdtemp(path.join(os.tmpdir(), "zerobyte-volume-ops-")); + return { + id: 1, + shortId: "volume-1", + name: "Test volume", + config: { backend: "directory", path: tempRoot }, + createdAt: Date.now(), + updatedAt: Date.now(), + lastHealthCheck: Date.now(), + type: "directory", + status: "mounted", + lastError: null, + provisioningId: null, + autoRemount: true, + agentId: "local", + organizationId: "org-1", + }; +}; + +test("listVolumeFiles returns sorted paginated entries inside the volume", async () => { + const volume = await createDirectoryVolume(); + await fs.mkdir(path.join(tempRoot!, "z-dir")); + await fs.mkdir(path.join(tempRoot!, "a-dir")); + await fs.writeFile(path.join(tempRoot!, "b-file.txt"), "hello"); + + const result = await listVolumeFiles(volume, undefined, 1, 2); + + expect(result).toMatchObject({ + path: "/", + offset: 1, + limit: 2, + total: 3, + hasMore: false, + }); + expect(result.files.map((entry) => entry.name)).toEqual(["z-dir", "b-file.txt"]); + expect(result.files[1]).toMatchObject({ path: "/b-file.txt", type: "file", size: 5 }); +}); + +test("listVolumeFiles rejects traversal outside the volume", async () => { + const volume = await createDirectoryVolume(); + + await expect(listVolumeFiles(volume, "../outside", 0, 10)).rejects.toThrow("Invalid path"); +}); + +test("listVolumeFiles reports missing directories consistently", async () => { + const volume = await createDirectoryVolume(); + + await expect(listVolumeFiles(volume, "missing", 0, 10)).rejects.toThrow("Directory not found"); +}); diff --git a/apps/agent/src/volume-host/cleanup.ts b/apps/agent/src/volume-host/cleanup.ts new file mode 100644 index 00000000..6c6a3a4c --- /dev/null +++ b/apps/agent/src/volume-host/cleanup.ts @@ -0,0 +1,28 @@ +import * as fs from "node:fs/promises"; +import * as path from "node:path"; +import { logger } from "@zerobyte/core/node"; +import { toMessage } from "@zerobyte/core/utils"; +import { VOLUME_MOUNT_BASE } from "./constants"; +import { readMountInfo } from "./fs"; + +export const cleanupDanglingVolumeMountDirectories = async () => { + const mounts = await readMountInfo().catch((error) => { + logger.warn(`Failed to read mount info for volume cleanup: ${toMessage(error)}`); + return []; + }); + const mountedPaths = new Set(mounts.map((mount) => mount.mountPoint)); + const volumeDirs = await fs.readdir(VOLUME_MOUNT_BASE).catch(() => []); + + for (const dir of volumeDirs) { + const mountPath = path.join(VOLUME_MOUNT_BASE, dir, "_data"); + if (mountedPaths.has(mountPath)) { + continue; + } + + const fullPath = path.join(VOLUME_MOUNT_BASE, dir); + logger.info(`Removing stale volume mount directory at ${fullPath}`); + await fs.rm(fullPath, { recursive: true, force: true }).catch((error) => { + logger.warn(`Failed to remove stale volume mount directory ${fullPath}: ${toMessage(error)}`); + }); + } +};