refactor(agent): harden local agent volume lifecycle

This commit is contained in:
Nicolas Meienberger 2026-05-05 21:40:23 +02:00
parent 2062beac68
commit 27eb5c2f5e
No known key found for this signature in database
15 changed files with 471 additions and 68 deletions

View file

@ -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();
});

View file

@ -10,9 +10,16 @@ import { toMessage } from "../utils/errors";
import { VOLUME_MOUNT_BASE } from "../core/constants"; import { VOLUME_MOUNT_BASE } from "../core/constants";
import { db } from "../db/db"; import { db } from "../db/db";
import { withContext } from "../core/request-context"; import { withContext } from "../core/request-context";
import { config } from "../core/config";
import { LOCAL_AGENT_ID } from "../modules/agents/constants";
export class CleanupDanglingMountsJob extends Job { export class CleanupDanglingMountsJob extends Job {
async run() { 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({}); const organizations = await db.query.organization.findMany({});
if (organizations.length === 0) { if (organizations.length === 0) {
logger.warn("No organizations found; skipping dangling mount cleanup to avoid false positives."); logger.warn("No organizations found; skipping dangling mount cleanup to avoid false positives.");
@ -22,7 +29,7 @@ export class CleanupDanglingMountsJob extends Job {
const allVolumes = []; const allVolumes = [];
for (const org of organizations) { for (const org of organizations) {
const volumes = await withContext({ organizationId: org.id }, async () => volumeService.listVolumes()); 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(); const allSystemMounts = await readMountInfo();

View file

@ -3,14 +3,21 @@ import type { BackupRunPayload, VolumeCommand, VolumeCommandResult } from "@zero
import { Effect } from "effect"; import { Effect } from "effect";
import { config } from "../../core/config"; import { config } from "../../core/config";
import { createAgentManagerRuntime, type AgentManagerEvent } from "./controller/server"; import { createAgentManagerRuntime, type AgentManagerEvent } from "./controller/server";
import { LOCAL_AGENT_ID } from "./constants";
import { spawnLocalAgentProcess, stopLocalAgentProcess } from "./local/process"; import { spawnLocalAgentProcess, stopLocalAgentProcess } from "./local/process";
import type { BackupExecutionProgress, BackupExecutionResult } from "./helpers/runtime-state"; import {
import { createAgentRuntimeState } from "./helpers/runtime-state"; createAgentRuntimeState,
type AgentRuntimeState,
type BackupExecutionProgress,
type BackupExecutionResult,
} from "./helpers/runtime-state";
import { getDevAgentRuntimeState } from "./helpers/runtime-state.dev"; import { getDevAgentRuntimeState } from "./helpers/runtime-state.dev";
export type { BackupExecutionProgress, BackupExecutionResult } from "./helpers/runtime-state"; export type { BackupExecutionProgress, BackupExecutionResult } from "./helpers/runtime-state";
export type { ProcessWithAgentRuntime } from "./helpers/runtime-state.dev"; export type { ProcessWithAgentRuntime } from "./helpers/runtime-state.dev";
const productionRuntimeState = createAgentRuntimeState(); type ProcessWithProductionAgentRuntime = NodeJS.Process & {
__zerobyteProductionAgentRuntime?: AgentRuntimeState;
};
type AgentRunBackupRequest = { type AgentRunBackupRequest = {
scheduleId: number; scheduleId: number;
@ -19,7 +26,18 @@ type AgentRunBackupRequest = {
onProgress: (progress: BackupExecutionProgress) => void; 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 getAgentManagerRuntime = () => getAgentRuntimeState().agentManager;
const getActiveBackupsByScheduleId = () => getAgentRuntimeState().activeBackupsByScheduleId; const getActiveBackupsByScheduleId = () => getAgentRuntimeState().activeBackupsByScheduleId;
const getActiveBackupScheduleIdsByJobId = () => getAgentRuntimeState().activeBackupScheduleIdsByJobId; const getActiveBackupScheduleIdsByJobId = () => getAgentRuntimeState().activeBackupScheduleIdsByJobId;
@ -286,7 +304,16 @@ export const agentManager = {
}; };
export const startLocalAgent = async () => { 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 // fallow-ignore-next-line unused-export

View file

@ -46,6 +46,7 @@ class StopAgentManagerServerError extends Data.TaggedError("StopAgentManagerServ
export function createAgentManagerRuntime(onEvent: (event: AgentManagerEvent) => void) { export function createAgentManagerRuntime(onEvent: (event: AgentManagerEvent) => void) {
let sessions = new Map<string, ControllerAgentSessionHandle>(); let sessions = new Map<string, ControllerAgentSessionHandle>();
let runtimeScope: Scope.CloseableScope | null = null; let runtimeScope: Scope.CloseableScope | null = null;
const sleep = (ms: number) => new Promise((resolve) => setTimeout(resolve, ms));
const closeSession = (sessionHandle: ControllerAgentSessionHandle) => const closeSession = (sessionHandle: ControllerAgentSessionHandle) =>
Effect.gen(function* () { Effect.gen(function* () {
@ -79,6 +80,11 @@ export function createAgentManagerRuntime(onEvent: (event: AgentManagerEvent) =>
const getSessionHandle = (agentId: string) => sessions.get(agentId); const getSessionHandle = (agentId: string) => sessions.get(agentId);
const getSession = (agentId: string) => getSessionHandle(agentId)?.session; 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 handleSessionEvent = (params: { agentId: string; agentName: string; sessionId: string }) => {
const { agentId, agentName } = params; const { agentId, agentName } = params;
@ -290,6 +296,19 @@ export function createAgentManagerRuntime(onEvent: (event: AgentManagerEvent) =>
return { return {
start, 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) => sendBackup: (agentId: string, payload: BackupRunPayload) =>
Effect.gen(function* () { Effect.gen(function* () {
const session = getSession(agentId); const session = getSession(agentId);

View file

@ -1,9 +1,10 @@
import { afterEach, describe, expect, test, vi } from "vitest"; import { afterEach, describe, expect, test, vi } from "vitest";
import { Scheduler } from "../../../core/scheduler"; 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 * as bootstrapModule from "../bootstrap";
import { agentManager } from "../../agents/agents-manager";
import { createTestVolume } from "~/test/helpers/volume"; import { createTestVolume } from "~/test/helpers/volume";
import { config } from "~/server/core/config";
import { db } from "~/server/db/db";
const loadShutdownModule = async () => { const loadShutdownModule = async () => {
const moduleUrl = new URL("../shutdown.ts", import.meta.url); const moduleUrl = new URL("../shutdown.ts", import.meta.url);
@ -12,11 +13,12 @@ const loadShutdownModule = async () => {
}; };
afterEach(() => { afterEach(() => {
config.flags.enableLocalAgent = true;
vi.restoreAllMocks(); vi.restoreAllMocks();
}); });
describe("shutdown", () => { 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 events: string[] = [];
const stopScheduler = vi.fn(async () => { const stopScheduler = vi.fn(async () => {
events.push("scheduler.stop"); events.push("scheduler.stop");
@ -24,35 +26,51 @@ describe("shutdown", () => {
const stopApplicationRuntime = vi.fn(async () => { const stopApplicationRuntime = vi.fn(async () => {
events.push("agents.stop"); events.push("agents.stop");
}); });
const unmountVolume = vi.fn(async () => { const runVolumeCommand = vi.spyOn(agentManager, "runVolumeCommand");
events.push("backend.unmount");
return { status: "unmounted" as const };
});
await createTestVolume({ const volume = await createTestVolume({
name: "Shutdown test volume", name: "Shutdown test volume",
config: { config: {
backend: "directory", backend: "directory",
path: "/Applications", path: "/Applications",
}, },
status: "mounted", status: "mounted",
agentId: "agent-1",
}); });
vi.spyOn(Scheduler, "stop").mockImplementation(stopScheduler); vi.spyOn(Scheduler, "stop").mockImplementation(stopScheduler);
vi.spyOn(bootstrapModule, "stopApplicationRuntime").mockImplementation(stopApplicationRuntime); 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(); const { shutdown } = await loadShutdownModule();
await shutdown(); 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");
}); });
}); });

View file

@ -2,25 +2,28 @@ import { Scheduler } from "../../core/scheduler";
import { db } from "../../db/db"; import { db } from "../../db/db";
import { logger } from "@zerobyte/core/node"; import { logger } from "@zerobyte/core/node";
import { stopApplicationRuntime } from "./bootstrap"; import { stopApplicationRuntime } from "./bootstrap";
import { decryptVolumeConfig } from "../volumes/volume-config-secrets"; import { withContext } from "../../core/request-context";
import { createVolumeBackend } from "../../../../apps/agent/src/volume-host"; 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 () => { export const shutdown = async () => {
await Scheduler.stop(); await Scheduler.stop();
await stopApplicationRuntime();
const volumes = await db.query.volumesTable.findMany({ if (!config.flags.enableLocalAgent) {
where: { status: "mounted" }, const volumes = await db.query.volumesTable.findMany({
}); where: { AND: [{ status: "mounted" }, { agentId: LOCAL_AGENT_ID }] },
for (const volume of volumes) {
const backend = createVolumeBackend({
...volume,
config: await decryptVolumeConfig(volume.config),
provisioningId: volume.provisioningId ?? null,
}); });
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();
}; };

View file

@ -17,6 +17,7 @@ import { backupsService } from "../backups/backups.service";
import { config } from "~/server/core/config"; import { config } from "~/server/core/config";
import { syncProvisionedResources } from "../provisioning/provisioning"; import { syncProvisionedResources } from "../provisioning/provisioning";
import { toMessage } from "~/server/utils/errors"; import { toMessage } from "~/server/utils/errors";
import { LOCAL_AGENT_ID } from "../agents/constants";
const ensureLatestConfigurationSchema = async () => { const ensureLatestConfigurationSchema = async () => {
const volumes = await db.query.volumesTable.findMany({}); 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`); logger.warn(`Removed ${deletedSchedules} orphaned backup schedule(s) during startup`);
} }
const volumes = await db.query.volumesTable.findMany({ if (!config.flags.enableLocalAgent) {
where: { const volumes = await db.query.volumesTable.findMany({
OR: [ where: {
{ status: "mounted" }, AND: [
{ { agentId: LOCAL_AGENT_ID },
AND: [{ autoRemount: true }, { status: "error" }], {
}, 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}`); },
});
}); });
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 await db
@ -102,7 +110,9 @@ export const startup = async () => {
logger.error(`Failed to update stuck backup schedules on startup: ${err.message}`); 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(VolumeHealthCheckJob).schedule("*/30 * * * *");
Scheduler.build(RepositoryHealthCheckJob).schedule("50 12 * * *"); Scheduler.build(RepositoryHealthCheckJob).schedule("50 12 * * *");
Scheduler.build(BackupExecutionJob).schedule("* * * * *"); Scheduler.build(BackupExecutionJob).schedule("* * * * *");

View file

@ -18,7 +18,13 @@ import { getOrganizationId } from "~/server/core/request-context";
import { type ShortId } from "~/server/utils/branded"; import { type ShortId } from "~/server/utils/branded";
import { decryptVolumeConfig, encryptVolumeConfig } from "./volume-config-secrets"; import { decryptVolumeConfig, encryptVolumeConfig } from "./volume-config-secrets";
import type { VolumeCommand, VolumeCommandResult } from "@zerobyte/contracts/agent-protocol"; 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 { import {
browseFilesystem as browseHostFilesystem, browseFilesystem as browseHostFilesystem,
listVolumeFiles, listVolumeFiles,
@ -63,21 +69,25 @@ const volumeForAgent = async (volume: Volume): Promise<Volume> => ({
config: await decryptVolumeConfig(volume.config), config: await decryptVolumeConfig(volume.config),
}); });
const volumeForHost = async (volume: Volume): Promise<Volume> => ({ const volumeForHost = async (volume: Volume): Promise<AgentVolume> => ({
...volume, ...volume,
shortId: volume.shortId, shortId: volume.shortId,
config: await decryptVolumeConfig(volume.config), config: (await decryptVolumeConfig(volume.config)) as HostBackendConfig,
provisioningId: volume.provisioningId ?? null, 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 shouldRunViaAgent = (volume: Volume) => volume.agentId !== LOCAL_AGENT_ID || config.flags.enableLocalAgent;
const shouldUseControllerLocalVolumeFallback = (volume: Volume) => !shouldRunViaAgent(volume);
const runVolumeBackendCommand = async ( const runVolumeBackendCommand = async (
volume: Volume, volume: Volume,
name: "volume.mount" | "volume.unmount" | "volume.checkHealth", name: "volume.mount" | "volume.unmount" | "volume.checkHealth",
) => { ) => {
if (!shouldRunViaAgent(volume)) { if (shouldUseControllerLocalVolumeFallback(volume)) {
const backend = createVolumeBackend(await volumeForHost(volume)); const backend = createVolumeBackend(await volumeForHost(volume));
switch (name) { switch (name) {
case "volume.mount": case "volume.mount":
@ -202,7 +212,7 @@ const getVolume = async (shortId: ShortId) => {
let statfs: Partial<StatFs> = {}; let statfs: Partial<StatFs> = {};
if (volume.status === "mounted") { if (volume.status === "mounted") {
statfs = await withTimeout( statfs = await withTimeout(
shouldRunViaAgent(volume) !shouldUseControllerLocalVolumeFallback(volume)
? runVolumeCommand(volume.agentId, { ? runVolumeCommand(volume.agentId, {
name: "volume.statfs", name: "volume.statfs",
volume: await volumeForAgent(volume), volume: await volumeForAgent(volume),
@ -280,7 +290,7 @@ const updateVolume = async (shortId: ShortId, volumeData: UpdateVolumeBody) => {
const testConnection = async (backendConfig: BackendConfig) => { const testConnection = async (backendConfig: BackendConfig) => {
if (!config.flags.enableLocalAgent) { 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 }); 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 { try {
if (!shouldRunViaAgent(volume)) { if (shouldUseControllerLocalVolumeFallback(volume)) {
return await listVolumeFiles(await volumeForHost(volume), subPath, offset, limit); return await listVolumeFiles(await volumeForHost(volume), subPath, offset, limit);
} }

View file

@ -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<ControllerCommandContext>({
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<VolumeCommandPayload>({
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<VolumeCommandPayload>({
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,
);
});

View file

@ -47,7 +47,8 @@ const executeVolumeCommand = (command: VolumeCommand) =>
} }
case "volume.listFiles": { case "volume.listFiles": {
const result = yield* Effect.tryPromise({ 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 }), catch: (error) => new VolumeCommandError({ cause: error }),
}); });
@ -80,14 +81,14 @@ export const handleVolumeCommand = (context: ControllerCommandContext, payload:
return command; return command;
}).pipe( }).pipe(
Effect.tapError((error) => { Effect.catchAll((error) =>
return context.offerOutbound( context.offerOutbound(
createAgentMessage("volume.commandResult", { createAgentMessage("volume.commandResult", {
commandId: payload.commandId, commandId: payload.commandId,
status: "error", status: "error",
error: toMessage(error?.cause), error: toMessage(error?.cause),
}), }),
); ),
}), ),
); );
}; };

View file

@ -1,14 +1,32 @@
import { logger } from "@zerobyte/core/node"; import { logger } from "@zerobyte/core/node";
import { createControllerSession, type ControllerSession } from "./controller-session"; import { createControllerSession, type ControllerSession } from "./controller-session";
import { cleanupDanglingVolumeMountDirectories } from "./volume-host/cleanup";
const controllerUrl = process.env.ZEROBYTE_CONTROLLER_URL; const controllerUrl = process.env.ZEROBYTE_CONTROLLER_URL;
const agentToken = process.env.ZEROBYTE_AGENT_TOKEN; const agentToken = process.env.ZEROBYTE_AGENT_TOKEN;
const RECONNECT_DELAY_MS = 1000; const RECONNECT_DELAY_MS = 1000;
const VOLUME_CLEANUP_INTERVAL_MS = 60 * 60 * 1000;
export class Agent { export class Agent {
private ws: WebSocket | null = null; private ws: WebSocket | null = null;
private controllerSession: ControllerSession | null = null; private controllerSession: ControllerSession | null = null;
private reconnectTimeout: ReturnType<typeof setTimeout> | null = null; private reconnectTimeout: ReturnType<typeof setTimeout> | null = null;
private volumeCleanupInterval: ReturnType<typeof setInterval> | 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() { private scheduleReconnect() {
if (this.reconnectTimeout) { if (this.reconnectTimeout) {
@ -22,6 +40,8 @@ export class Agent {
} }
connect() { connect() {
this.startVolumeCleanup();
if (this.reconnectTimeout) { if (this.reconnectTimeout) {
clearTimeout(this.reconnectTimeout); clearTimeout(this.reconnectTimeout);
this.reconnectTimeout = null; this.reconnectTimeout = null;

View file

@ -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();
});

View file

@ -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<AgentVolume> => {
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");
});

View file

@ -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)}`);
});
}
};