diff --git a/.env.test b/.env.test index 131beb69..d6805656 100644 --- a/.env.test +++ b/.env.test @@ -1,3 +1,4 @@ ZEROBYTE_DATABASE_URL=:memory: APP_SECRET=8b9acd4456dd5db0a4a3c4f4e1240b2c3ae08bb59690167197425e4a25dd9a69 BASE_URL=http://localhost:4096 +ENABLE_LOCAL_AGENT=true diff --git a/app/server/modules/agents/__tests__/agents-manager.backups.test.ts b/app/server/modules/agents/__tests__/agents-manager.backups.test.ts new file mode 100644 index 00000000..279f9b2e --- /dev/null +++ b/app/server/modules/agents/__tests__/agents-manager.backups.test.ts @@ -0,0 +1,47 @@ +import { afterEach, expect, test, vi } from "vitest"; +import waitForExpect from "wait-for-expect"; +import { fromAny, fromPartial } from "@total-typescript/shoehorn"; +import { agentManager, type ProcessWithAgentRuntime } from "../agents-manager"; +import type { AgentManagerRuntime } from "../controller/server"; +import type { BackupRunPayload } from "@zerobyte/contracts/agent-protocol"; + +const setAgentRuntime = (agentManagerRuntime: Partial | null) => { + (process as ProcessWithAgentRuntime).__zerobyteAgentRuntime = { + agentManager: fromAny(agentManagerRuntime), + localAgent: null, + isStoppingLocalAgent: false, + localAgentRestartTimeout: null, + }; +}; + +afterEach(() => { + delete (process as ProcessWithAgentRuntime).__zerobyteAgentRuntime; + vi.restoreAllMocks(); +}); + +test("cancelBackup resolves a running backup when the cancel command cannot be delivered", async () => { + const sendBackup = vi.fn().mockResolvedValue(true); + const cancelBackup = vi.fn().mockResolvedValue(false); + setAgentRuntime({ sendBackup, cancelBackup }); + + const resultPromise = agentManager.runBackup("local", { + scheduleId: 42, + payload: fromPartial({ + jobId: "job-1", + scheduleId: "schedule-1", + }), + signal: new AbortController().signal, + onProgress: vi.fn(), + }); + + await waitForExpect(() => { + expect(sendBackup).toHaveBeenCalledTimes(1); + }); + + await expect(agentManager.cancelBackup("local", 42)).resolves.toBe(true); + await expect(resultPromise).resolves.toEqual({ status: "cancelled" }); + expect(cancelBackup).toHaveBeenCalledWith("local", { + jobId: "job-1", + scheduleId: "schedule-1", + }); +}); diff --git a/app/server/modules/agents/__tests__/agents-manager.test.ts b/app/server/modules/agents/__tests__/agents-manager.test.ts new file mode 100644 index 00000000..c55d222d --- /dev/null +++ b/app/server/modules/agents/__tests__/agents-manager.test.ts @@ -0,0 +1,94 @@ +import { EventEmitter } from "node:events"; +import { PassThrough } from "node:stream"; +import { afterEach, beforeEach, expect, test, vi } from "vitest"; +import type { ProcessWithAgentRuntime } from "../helpers/runtime-state.dev"; + +const spawnMock = vi.fn(); + +vi.mock("node:child_process", async () => { + return { spawn: spawnMock }; +}); + +let spawnLocalAgent: (typeof import("../agents-manager"))["spawnLocalAgent"]; +let stopLocalAgent: (typeof import("../agents-manager"))["stopLocalAgent"]; + +const processWithAgentRuntime = process as ProcessWithAgentRuntime; + +const setAgentRuntime = () => { + processWithAgentRuntime.__zerobyteAgentRuntime = { + agentManager: null, + localAgent: null, + isStoppingLocalAgent: false, + localAgentRestartTimeout: null, + }; +}; + +type FakeChildProcess = EventEmitter & { + stdout: PassThrough; + stderr: PassThrough; + exitCode: number | null; + signalCode: NodeJS.Signals | null; + kill: ReturnType; +}; + +const createFakeChild = () => { + const child = new EventEmitter() as FakeChildProcess; + + child.stdout = new PassThrough(); + child.stderr = new PassThrough(); + child.exitCode = null; + child.signalCode = null; + child.kill = vi.fn(() => { + child.exitCode = 0; + child.emit("exit", 0, null); + return true; + }); + + return child; +}; + +beforeEach(async () => { + vi.resetModules(); + setAgentRuntime(); + ({ spawnLocalAgent, stopLocalAgent } = await import("../agents-manager")); +}); + +afterEach(async () => { + await stopLocalAgent(); + delete processWithAgentRuntime.__zerobyteAgentRuntime; + spawnMock.mockReset(); + vi.restoreAllMocks(); + vi.useRealTimers(); +}); + +test("respawns the local agent after an unexpected exit", async () => { + vi.useFakeTimers(); + + const firstChild = createFakeChild(); + const secondChild = createFakeChild(); + spawnMock.mockReturnValueOnce(firstChild).mockReturnValueOnce(secondChild); + + await spawnLocalAgent(); + + firstChild.exitCode = 1; + firstChild.emit("exit", 1, null); + + await vi.advanceTimersByTimeAsync(1_000); + + expect(spawnMock).toHaveBeenCalledTimes(2); +}); + +test("does not respawn the local agent after an intentional stop", async () => { + vi.useFakeTimers(); + + const child = createFakeChild(); + spawnMock.mockReturnValue(child); + + await spawnLocalAgent(); + await stopLocalAgent(); + + await vi.advanceTimersByTimeAsync(1_000); + + expect(spawnMock).toHaveBeenCalledTimes(1); + expect(child.kill).toHaveBeenCalledTimes(1); +}); diff --git a/app/server/modules/agents/__tests__/session.test.ts b/app/server/modules/agents/__tests__/session.test.ts index b787e1ea..de184c01 100644 --- a/app/server/modules/agents/__tests__/session.test.ts +++ b/app/server/modules/agents/__tests__/session.test.ts @@ -1,27 +1,37 @@ import { Effect, Exit, Scope } from "effect"; import { expect, test, vi } from "vitest"; +import waitForExpect from "wait-for-expect"; import { fromPartial } from "@total-typescript/shoehorn"; import { createAgentMessage } from "@zerobyte/contracts/agent-protocol"; import { createControllerAgentSession } from "../controller/session"; -const createSocket = () => { - return fromPartial[0]>({ +const createSocket = (overrides: Partial[0]> = {}) => { + return { data: { id: "connection-1", agentId: "local", organizationId: null, agentName: "Local Agent" }, - send: vi.fn(), - }); + send: vi.fn(() => 1), + close: vi.fn(), + ...overrides, + }; }; -const createSession = (handlers: Parameters[1] = {}) => { +const createSession = (handlers: Parameters[1] = {}, socket = createSocket()) => { const scope = Effect.runSync(Scope.make()); try { - const session = Effect.runSync(Scope.extend(createControllerAgentSession(createSocket(), handlers), scope)); + const session = Effect.runSync(Scope.extend(createControllerAgentSession(fromPartial(socket), handlers), scope)); return { session, + run: () => { + Effect.runFork(Scope.extend(session.run, scope)); + }, + socket, close: () => { Effect.runSync(Scope.close(scope, Exit.succeed(undefined))); }, + closeAsync: () => { + return Effect.runPromise(Scope.close(scope, Exit.succeed(undefined))); + }, }; } catch (error) { Effect.runSync(Scope.close(scope, Exit.fail(error))); @@ -47,11 +57,12 @@ test("close emits a synthetic backup.cancelled for a started backup", () => { close(); expect(onBackupCancelled).toHaveBeenCalledTimes(1); - expect(onBackupCancelled).toHaveBeenCalledWith({ - jobId: "job-1", - scheduleId: "schedule-1", - message: "The connection to the backup agent was lost. Restart the backup to ensure it completes.", - }); + expect(onBackupCancelled).toHaveBeenCalledWith( + expect.objectContaining({ + jobId: "job-1", + scheduleId: "schedule-1", + }), + ); }); test.each([ @@ -139,9 +150,49 @@ test("close emits a synthetic backup.cancelled for a queued backup", () => { close(); expect(onBackupCancelled).toHaveBeenCalledTimes(1); - expect(onBackupCancelled).toHaveBeenCalledWith({ - jobId: "job-queued", - scheduleId: "schedule-queued", - message: "The connection to the backup agent was lost. Restart the backup to ensure it completes.", - }); + expect(onBackupCancelled).toHaveBeenLastCalledWith( + expect.objectContaining({ + jobId: "job-queued", + scheduleId: "schedule-queued", + }), + ); +}); + +test("a dropped backup.cancel closes the session and emits a synthetic backup.cancelled", async () => { + const send = vi.fn(() => 0); + const socket = createSocket({ send, close: vi.fn() }); + const onBackupCancelled = vi.fn(); + const { session, run, closeAsync } = createSession({ onBackupCancelled }, socket); + + try { + run(); + Effect.runSync( + session.handleMessage( + createAgentMessage("backup.started", { + jobId: "job-1", + scheduleId: "schedule-1", + }), + ), + ); + Effect.runSync( + session.sendBackupCancel({ + jobId: "job-1", + scheduleId: "schedule-1", + }), + ); + + await waitForExpect(() => { + expect(send).toHaveBeenCalledTimes(1); + expect(socket.close).toHaveBeenCalledTimes(1); + expect(onBackupCancelled).toHaveBeenCalledTimes(1); + expect(onBackupCancelled).toHaveBeenCalledWith( + expect.objectContaining({ + jobId: "job-1", + scheduleId: "schedule-1", + }), + ); + }); + } finally { + await closeAsync(); + } }); diff --git a/app/server/modules/agents/agents-manager.ts b/app/server/modules/agents/agents-manager.ts index c1cd9635..6782dcad 100644 --- a/app/server/modules/agents/agents-manager.ts +++ b/app/server/modules/agents/agents-manager.ts @@ -1,49 +1,216 @@ -import type { ChildProcess } from "node:child_process"; -import type { AgentManagerRuntime } from "./controller/server"; +import { logger } from "@zerobyte/core/node"; +import type { BackupRunPayload } from "@zerobyte/contracts/agent-protocol"; +import { config } from "../../core/config"; +import type { AgentBackupEventHandlers } from "./controller/server"; import { spawnLocalAgentProcess, stopLocalAgentProcess } from "./local/process"; +import type { BackupExecutionProgress, BackupExecutionResult } from "./helpers/runtime-state"; +import { createAgentRuntimeState } 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"; -export type { AgentBackupEventHandlers } from "./controller/server"; +const productionRuntimeState = createAgentRuntimeState(); -type AgentRuntimeState = { - agentManager: AgentManagerRuntime | null; - localAgent: ChildProcess | null; +export type AgentRunBackupRequest = { + scheduleId: number; + payload: BackupRunPayload; + signal: AbortSignal; + onProgress: (progress: BackupExecutionProgress) => void; }; -type ProcessWithAgentRuntime = NodeJS.Process & { - __zerobyteAgentRuntime?: AgentRuntimeState; -}; +const getAgentRuntimeState = () => (config.__prod__ ? productionRuntimeState : getDevAgentRuntimeState()); +const getAgentManagerRuntime = () => getAgentRuntimeState().agentManager; +const getActiveBackupsByScheduleId = () => getAgentRuntimeState().activeBackupsByScheduleId; +const getActiveBackupScheduleIdsByJobId = () => getAgentRuntimeState().activeBackupScheduleIdsByJobId; -const getAgentRuntimeState = () => { - const runtimeProcess = process as ProcessWithAgentRuntime; - const existingRuntime = runtimeProcess.__zerobyteAgentRuntime; - - if (existingRuntime) { - return existingRuntime; +const clearActiveBackupRun = (scheduleId: number) => { + const activeBackupsByScheduleId = getActiveBackupsByScheduleId(); + const activeBackupScheduleIdsByJobId = getActiveBackupScheduleIdsByJobId(); + const activeBackupRun = activeBackupsByScheduleId.get(scheduleId); + if (!activeBackupRun) { + return null; } - const runtime = { - agentManager: null, - localAgent: null, - }; - - runtimeProcess.__zerobyteAgentRuntime = runtime; - return runtime; + activeBackupsByScheduleId.delete(scheduleId); + activeBackupScheduleIdsByJobId.delete(activeBackupRun.jobId); + return activeBackupRun; }; -const getAgentManagerRuntime = () => getAgentRuntimeState().agentManager; +const resolveActiveBackupRun = (scheduleId: number, result: BackupExecutionResult) => { + const activeBackupRun = clearActiveBackupRun(scheduleId); + if (!activeBackupRun) { + return false; + } + + activeBackupRun.resolve(result); + return true; +}; + +const getActiveBackupRun = (jobId: string, scheduleId: string, eventName: string, agentId: string) => { + const trackedScheduleId = getActiveBackupScheduleIdsByJobId().get(jobId); + if (trackedScheduleId === undefined) { + logger.warn(`Received ${eventName} for unknown job ${jobId} from agent ${agentId}`); + return null; + } + + const activeBackupRun = getActiveBackupsByScheduleId().get(trackedScheduleId); + if (!activeBackupRun) { + logger.warn(`Received ${eventName} for inactive job ${jobId} from agent ${agentId}`); + return null; + } + + if (activeBackupRun.scheduleShortId !== scheduleId) { + logger.warn(`Ignoring ${eventName} for job ${jobId} due to schedule mismatch ${scheduleId} from agent ${agentId}`); + return null; + } + + return activeBackupRun; +}; + +const requestBackupCancellation = async (agentId: string, scheduleId: number) => { + const activeBackupRun = getActiveBackupsByScheduleId().get(scheduleId); + if (!activeBackupRun) { + return false; + } + + if (activeBackupRun.cancellationRequested) { + return true; + } + + activeBackupRun.cancellationRequested = true; + + const runtime = getAgentManagerRuntime(); + if (!runtime) { + resolveActiveBackupRun(scheduleId, { status: "cancelled" }); + return true; + } + + if ( + await runtime.cancelBackup(agentId, { + jobId: activeBackupRun.jobId, + scheduleId: activeBackupRun.scheduleShortId, + }) + ) { + return true; + } + + resolveActiveBackupRun(scheduleId, { status: "cancelled" }); + return true; +}; + +const backupEventHandlers: AgentBackupEventHandlers = { + onBackupStarted: ({ agentId, payload }) => { + getActiveBackupRun(payload.jobId, payload.scheduleId, "backup.started", agentId); + }, + onBackupProgress: ({ agentId, payload }) => { + const activeBackupRun = getActiveBackupRun(payload.jobId, payload.scheduleId, "backup.progress", agentId); + if (!activeBackupRun) { + return; + } + + activeBackupRun.onProgress(payload.progress); + }, + onBackupCompleted: ({ agentId, payload }) => { + const activeBackupRun = getActiveBackupRun(payload.jobId, payload.scheduleId, "backup.completed", agentId); + if (!activeBackupRun) { + return; + } + + resolveActiveBackupRun(activeBackupRun.scheduleId, { + status: "completed", + exitCode: payload.exitCode, + result: payload.result, + warningDetails: payload.warningDetails ?? null, + }); + }, + onBackupFailed: ({ agentId, payload }) => { + const activeBackupRun = getActiveBackupRun(payload.jobId, payload.scheduleId, "backup.failed", agentId); + if (!activeBackupRun) { + return; + } + + resolveActiveBackupRun(activeBackupRun.scheduleId, { + status: "failed", + error: payload.errorDetails ?? payload.error, + }); + }, + onBackupCancelled: ({ agentId, payload }) => { + const activeBackupRun = getActiveBackupRun(payload.jobId, payload.scheduleId, "backup.cancelled", agentId); + if (!activeBackupRun) { + return; + } + + resolveActiveBackupRun(activeBackupRun.scheduleId, { + status: "cancelled", + message: activeBackupRun.cancellationRequested ? undefined : payload.message, + }); + }, +}; export const startAgentRuntime = async () => { const runtime = getAgentRuntimeState(); if (runtime.agentManager) { await runtime.agentManager.stop(); + runtime.agentManager = null; } const { createAgentManagerRuntime } = await import("./controller/server"); - const agentManager = createAgentManagerRuntime(); + const nextAgentManager = createAgentManagerRuntime(); + nextAgentManager.setBackupEventHandlers(backupEventHandlers); - await agentManager.start(); - runtime.agentManager = agentManager; + await nextAgentManager.start(); + runtime.agentManager = nextAgentManager; +}; + +export const agentManager = { + runBackup: async (agentId: string, request: AgentRunBackupRequest) => { + const runtime = getAgentManagerRuntime(); + if (!runtime) { + return { + status: "unavailable", + error: new Error(`Backup agent ${agentId} is not connected`), + } satisfies BackupExecutionResult; + } + + if (request.signal.aborted) { + throw request.signal.reason || new Error("Operation aborted"); + } + + const completion = new Promise((resolve) => { + getActiveBackupsByScheduleId().set(request.scheduleId, { + scheduleId: request.scheduleId, + jobId: request.payload.jobId, + scheduleShortId: request.payload.scheduleId, + onProgress: request.onProgress, + resolve, + cancellationRequested: false, + }); + getActiveBackupScheduleIdsByJobId().set(request.payload.jobId, request.scheduleId); + }); + + try { + if (!(await runtime.sendBackup(agentId, request.payload))) { + clearActiveBackupRun(request.scheduleId); + return { + status: "unavailable", + error: new Error(`Failed to send backup command to agent ${agentId}`), + } satisfies BackupExecutionResult; + } + + if (request.signal.aborted) { + await requestBackupCancellation(agentId, request.scheduleId); + } + + return completion; + } catch (error) { + clearActiveBackupRun(request.scheduleId); + throw error; + } + }, + cancelBackup: async (agentId: string, scheduleId: number) => { + return requestBackupCancellation(agentId, scheduleId); + }, }; export const spawnLocalAgent = async () => { diff --git a/app/server/modules/agents/controller/session.ts b/app/server/modules/agents/controller/session.ts index cae39436..999cec41 100644 --- a/app/server/modules/agents/controller/session.ts +++ b/app/server/modules/agents/controller/session.ts @@ -57,6 +57,7 @@ export const createControllerAgentSession = ( handlers: ControllerAgentSessionHandlers = {}, ): Effect.Effect => Effect.gen(function* () { + let isClosed = false; const outboundQueue = yield* Queue.bounded(64); const trackedBackupJobs = yield* Ref.make>(new Map()); const state = yield* Ref.make({ @@ -102,7 +103,7 @@ export const createControllerAgentSession = ( yield* updateState((current) => ({ ...current, isReady: false })); const trackedJobs = yield* takeTrackedBackupJobs; for (const [jobId, trackedJob] of trackedJobs) { - let message = "The connection to the backup agent was lost. Restart the backup to ensure it completes."; + const message = "The connection to the backup agent was lost. Restart the backup to ensure it completes."; yield* Effect.sync(() => { handlers.onBackupCancelled?.({ jobId, scheduleId: trackedJob.scheduleId, message }); @@ -112,7 +113,31 @@ export const createControllerAgentSession = ( yield* Queue.shutdown(outboundQueue); }); - yield* Effect.addFinalizer(() => releaseSession); + const closeSession = () => + Effect.suspend(() => { + if (isClosed) { + return Effect.sync(() => undefined); + } + + isClosed = true; + return releaseSession; + }); + + yield* Effect.addFinalizer(() => closeSession()); + + const handleSendFailure = (reason: string) => { + logger.error( + `Closing session for agent ${socket.data.agentId} on ${socket.data.id} after an outbound websocket send failed: ${reason}`, + ); + + socket.close(); + + void Effect.runPromise(closeSession()).catch((error) => { + logger.error( + `Failed to close session for agent ${socket.data.agentId} on ${socket.data.id}: ${toMessage(error)}`, + ); + }); + }; const run = Effect.gen(function* () { yield* Effect.forkScoped( @@ -121,11 +146,12 @@ export const createControllerAgentSession = ( const message = yield* Queue.take(outboundQueue); yield* Effect.sync(() => { try { - socket.send(message); + const sendResult = socket.send(message); + if (sendResult === 0) { + handleSendFailure("connection issue"); + } } catch (error) { - logger.error( - `Failed to send message to agent ${socket.data.agentId} on ${socket.data.id}: ${toMessage(error)}`, - ); + handleSendFailure(toMessage(error)); } }); }), diff --git a/app/server/modules/agents/helpers/runtime-state.dev.ts b/app/server/modules/agents/helpers/runtime-state.dev.ts new file mode 100644 index 00000000..daccadbb --- /dev/null +++ b/app/server/modules/agents/helpers/runtime-state.dev.ts @@ -0,0 +1,37 @@ +import { createAgentRuntimeState, type AgentRuntimeState } from "./runtime-state"; + +type LegacyAgentRuntimeState = Omit & + Partial>; + +export type ProcessWithAgentRuntime = NodeJS.Process & { + __zerobyteAgentRuntime?: LegacyAgentRuntimeState; +}; + +const hasActiveBackupMaps = (runtime: LegacyAgentRuntimeState): runtime is AgentRuntimeState => { + return runtime.activeBackupsByScheduleId instanceof Map && runtime.activeBackupScheduleIdsByJobId instanceof Map; +}; + +const hydrateAgentRuntimeState = (runtime: LegacyAgentRuntimeState): AgentRuntimeState => ({ + ...runtime, + activeBackupsByScheduleId: runtime.activeBackupsByScheduleId ?? new Map(), + activeBackupScheduleIdsByJobId: runtime.activeBackupScheduleIdsByJobId ?? new Map(), +}); + +export const getDevAgentRuntimeState = (): AgentRuntimeState => { + // Bun reloads modules in place during development, so keep the live runtime on process. + const runtimeProcess = process as ProcessWithAgentRuntime; + const existingRuntime = runtimeProcess.__zerobyteAgentRuntime; + if (!existingRuntime) { + const runtime = createAgentRuntimeState(); + runtimeProcess.__zerobyteAgentRuntime = runtime; + return runtime; + } + + if (hasActiveBackupMaps(existingRuntime)) { + return existingRuntime; + } + + const runtime = hydrateAgentRuntimeState(existingRuntime); + runtimeProcess.__zerobyteAgentRuntime = runtime; + return runtime; +}; diff --git a/app/server/modules/agents/helpers/runtime-state.ts b/app/server/modules/agents/helpers/runtime-state.ts new file mode 100644 index 00000000..51de3fb4 --- /dev/null +++ b/app/server/modules/agents/helpers/runtime-state.ts @@ -0,0 +1,38 @@ +import type { ChildProcess } from "node:child_process"; +import type { ResticBackupOutputDto } from "@zerobyte/core/restic"; +import type { BackupProgressPayload } from "@zerobyte/contracts/agent-protocol"; +import type { AgentManagerRuntime } from "../controller/server"; + +export type BackupExecutionProgress = BackupProgressPayload["progress"]; +export type BackupExecutionResult = + | { status: "unavailable"; error: Error } + | { status: "completed"; exitCode: number; result: ResticBackupOutputDto | null; warningDetails: string | null } + | { status: "failed"; error: string } + | { status: "cancelled"; message?: string }; + +type ActiveBackupRun = { + scheduleId: number; + jobId: string; + scheduleShortId: string; + onProgress: (progress: BackupExecutionProgress) => void; + resolve: (result: BackupExecutionResult) => void; + cancellationRequested: boolean; +}; + +export type AgentRuntimeState = { + agentManager: AgentManagerRuntime | null; + localAgent: ChildProcess | null; + isStoppingLocalAgent: boolean; + localAgentRestartTimeout: ReturnType | null; + activeBackupsByScheduleId: Map; + activeBackupScheduleIdsByJobId: Map; +}; + +export const createAgentRuntimeState = (): AgentRuntimeState => ({ + agentManager: null, + localAgent: null, + isStoppingLocalAgent: false, + localAgentRestartTimeout: null, + activeBackupsByScheduleId: new Map(), + activeBackupScheduleIdsByJobId: new Map(), +}); diff --git a/app/server/modules/agents/local/process.ts b/app/server/modules/agents/local/process.ts index 6c353782..3852b36a 100644 --- a/app/server/modules/agents/local/process.ts +++ b/app/server/modules/agents/local/process.ts @@ -7,6 +7,8 @@ import { deriveLocalAgentToken } from "../helpers/tokens"; type LocalAgentState = { localAgent: ChildProcess | null; + isStoppingLocalAgent: boolean; + localAgentRestartTimeout: ReturnType | null; }; export async function spawnLocalAgentProcess(runtime: LocalAgentState) { @@ -48,27 +50,47 @@ export async function spawnLocalAgentProcess(runtime: LocalAgentState) { }); agentProcess.on("exit", (code, signal) => { + const shouldRestart = runtime.localAgent === agentProcess && !runtime.isStoppingLocalAgent; if (runtime.localAgent === agentProcess) { runtime.localAgent = null; } logger.info(`Agent process exited with code ${code} and signal ${signal}`); + + if (!shouldRestart) { + return; + } + + runtime.localAgentRestartTimeout = setTimeout(() => { + runtime.localAgentRestartTimeout = null; + void spawnLocalAgentProcess(runtime).catch((error) => { + logger.error(`Failed to restart local agent: ${error instanceof Error ? error.message : String(error)}`); + }); + }, 1_000); }); } export async function stopLocalAgentProcess(runtime: LocalAgentState) { + if (runtime.localAgentRestartTimeout) { + clearTimeout(runtime.localAgentRestartTimeout); + runtime.localAgentRestartTimeout = null; + } + if (!runtime.localAgent) { return; } const agentProcess = runtime.localAgent; runtime.localAgent = null; + runtime.isStoppingLocalAgent = true; if (agentProcess.exitCode !== null || agentProcess.signalCode !== null) { + runtime.isStoppingLocalAgent = false; return; } const exited = new Promise((resolve) => { agentProcess.once("exit", () => { + runtime.isStoppingLocalAgent = false; resolve(); }); }); diff --git a/app/server/modules/backups/__tests__/backups.execution.test.ts b/app/server/modules/backups/__tests__/backups.service.execution.test.ts similarity index 65% rename from app/server/modules/backups/__tests__/backups.execution.test.ts rename to app/server/modules/backups/__tests__/backups.service.execution.test.ts index fbb20003..cdc4bd4d 100644 --- a/app/server/modules/backups/__tests__/backups.execution.test.ts +++ b/app/server/modules/backups/__tests__/backups.service.execution.test.ts @@ -1,7 +1,6 @@ import waitForExpect from "wait-for-expect"; import { afterEach, describe, expect, test, vi } from "vitest"; import { backupsService } from "../backups.service"; -import { backupsExecutionService } from "../backups.execution"; import { createTestVolume } from "~/test/helpers/volume"; import { createTestBackupSchedule } from "~/test/helpers/backup"; import { createTestRepository } from "~/test/helpers/repository"; @@ -14,9 +13,14 @@ import type { SafeSpawnParams } from "@zerobyte/core/node"; import { logger } from "@zerobyte/core/node"; import { restic } from "~/server/core/restic"; import { NotFoundError, BadRequestError } from "http-errors-enhanced"; +import { fromAny } from "@total-typescript/shoehorn"; +import { scheduleQueries } from "../backups.queries"; import { repositoriesService } from "~/server/modules/repositories/repositories.service"; import { repoMutex } from "~/server/core/repository-mutex"; import { notificationsService } from "~/server/modules/notifications/notifications.service"; +import { agentManager } from "~/server/modules/agents/agents-manager"; +import { createAgentBackupMocks } from "~/test/helpers/agent-mock"; +import { getScheduleByIdOrShortId } from "../helpers/backup-schedule-lookups"; const setup = () => { const resticBackupMock = vi.fn((_: SafeSpawnParams) => @@ -24,6 +28,7 @@ const setup = () => { ); const resticForgetMock = vi.fn(() => Promise.resolve({ success: true, data: null })); const resticCopyMock = vi.fn(() => Promise.resolve({ success: true, output: "" })); + const { runBackupMock, cancelBackupMock } = createAgentBackupMocks(resticBackupMock); const refreshStatsMock = vi.fn(() => Promise.resolve({ total_size: 0, @@ -39,12 +44,16 @@ const setup = () => { vi.spyOn(restic, "forget").mockImplementation(resticForgetMock); vi.spyOn(restic, "copy").mockImplementation(resticCopyMock); vi.spyOn(repositoriesService, "refreshRepositoryStats").mockImplementation(refreshStatsMock); + vi.spyOn(agentManager, "runBackup").mockImplementation(runBackupMock); + vi.spyOn(agentManager, "cancelBackup").mockImplementation(cancelBackupMock); vi.spyOn(context, "getOrganizationId").mockReturnValue(TEST_ORG_ID); return { resticBackupMock, resticForgetMock, resticCopyMock, + runBackupMock, + cancelBackupMock, refreshStatsMock, }; }; @@ -65,7 +74,7 @@ describe("backup execution - validation failures", () => { }); // act - const result = await backupsExecutionService.validateBackupExecution(schedule.id); + const result = await backupsService.validateBackupExecution(schedule.id); // assert expect(result.type).toBe("failure"); @@ -76,10 +85,70 @@ describe("backup execution - validation failures", () => { expect(resticBackupMock).not.toHaveBeenCalled(); }); + test("should fail backup when volume does not exist", async () => { + // arrange + setup(); + const volume = await createTestVolume(); + const repository = await createTestRepository(); + const schedule = await createTestBackupSchedule({ + volumeId: volume.id, + repositoryId: repository.id, + }); + + const hydratedSchedule = await scheduleQueries.findById(schedule.id, TEST_ORG_ID); + expect(hydratedSchedule).toBeDefined(); + const scheduleWithoutVolume = { + ...hydratedSchedule, + volume: null, + }; + vi.spyOn(scheduleQueries, "findById").mockResolvedValueOnce(fromAny(scheduleWithoutVolume)); + + // act + const result = await backupsService.validateBackupExecution(schedule.id); + + // assert + expect(result.type).toBe("failure"); + if (result.type === "failure") { + expect(result.error).toBeInstanceOf(NotFoundError); + expect(result.error.message).toBe("Volume not found"); + expect(result.partialContext?.schedule).toBeDefined(); + } + }); + + test("should fail backup when repository does not exist", async () => { + // arrange + setup(); + const volume = await createTestVolume(); + const repository = await createTestRepository(); + const schedule = await createTestBackupSchedule({ + volumeId: volume.id, + repositoryId: repository.id, + }); + + const hydratedSchedule = await scheduleQueries.findById(schedule.id, TEST_ORG_ID); + expect(hydratedSchedule).toBeDefined(); + const scheduleWithoutRepository = { + ...hydratedSchedule, + repository: null, + }; + vi.spyOn(scheduleQueries, "findById").mockResolvedValueOnce(fromAny(scheduleWithoutRepository)); + + // act + const result = await backupsService.validateBackupExecution(schedule.id); + + // assert + expect(result.type).toBe("failure"); + if (result.type === "failure") { + expect(result.error).toBeInstanceOf(NotFoundError); + expect(result.error.message).toBe("Repository not found"); + expect(result.partialContext?.schedule).toBeDefined(); + expect(result.partialContext?.volume).toBeDefined(); + } + }); test("should fail backup when schedule does not exist", async () => { setup(); // act - const result = await backupsExecutionService.validateBackupExecution(99999); + const result = await backupsService.validateBackupExecution(99999); // assert expect(result.type).toBe("failure"); @@ -106,7 +175,7 @@ describe("backup execution - validation failures", () => { Promise.resolve({ exitCode: 1, summary: generateBackupOutput(), error: "failed" }), ); - await backupsExecutionService.executeBackup(schedule.id); + await backupsService.executeBackup(schedule.id); expect(notificationSpy).toHaveBeenCalled(); expect(notificationSpy.mock.calls.at(-1)?.[2]?.error).toBe("failed"); @@ -128,12 +197,33 @@ describe("backup execution - validation failures", () => { Promise.resolve({ exitCode: 1, summary: generateBackupOutput(), error: "manual failure" }), ); - await backupsExecutionService.executeBackup(schedule.id, true); + await backupsService.executeBackup(schedule.id, true); expect( errorSpy.mock.calls.some(([message]) => String(message).includes('Failed to parse cron expression ""')), ).toBe(false); }); + + test("should fail backup when the local agent is unavailable", async () => { + const { runBackupMock } = setup(); + const volume = await createTestVolume(); + const repository = await createTestRepository(); + const schedule = await createTestBackupSchedule({ + volumeId: volume.id, + repositoryId: repository.id, + }); + + runBackupMock.mockResolvedValueOnce({ + status: "unavailable", + error: new Error("Local backup agent is not connected"), + }); + + await backupsService.executeBackup(schedule.id); + + const updatedSchedule = await getScheduleByIdOrShortId(schedule.id); + expect(updatedSchedule.lastBackupStatus).toBe("error"); + expect(updatedSchedule.lastBackupError).toBe("Local backup agent is not connected"); + }); }); describe("stop backup", () => { @@ -156,9 +246,9 @@ describe("stop backup", () => { }); }); - await backupsExecutionService.executeBackup(schedule.id); + await backupsService.executeBackup(schedule.id); - const updatedSchedule = await backupsService.getScheduleById(schedule.id); + const updatedSchedule = await getScheduleByIdOrShortId(schedule.id); expect(updatedSchedule.lastBackupStatus).toBe("warning"); expect(updatedSchedule.lastBackupError).toBe("error: open /mnt/data/private.db: permission denied"); }); @@ -184,15 +274,86 @@ describe("stop backup", () => { }); }); - await backupsExecutionService.executeBackup(schedule.id); + await backupsService.executeBackup(schedule.id); - const updatedSchedule = await backupsService.getScheduleById(schedule.id); + const updatedSchedule = await getScheduleByIdOrShortId(schedule.id); expect(updatedSchedule.lastBackupStatus).toBe("error"); expect(updatedSchedule.lastBackupError).toBe( "Permissions 0755 for '/tmp/zerobyte-ssh-key' are too open.\nThis private key will be ignored.", ); }); + test("should settle and mark the backup as failed when the backup process throws", async () => { + const { resticBackupMock } = setup(); + const volume = await createTestVolume(); + const repository = await createTestRepository(); + const schedule = await createTestBackupSchedule({ + volumeId: volume.id, + repositoryId: repository.id, + }); + + resticBackupMock.mockImplementationOnce(() => Promise.reject(new Error("restic crashed"))); + + const result = await Promise.race([ + backupsService.executeBackup(schedule.id).then(() => "settled"), + new Promise((resolve) => setTimeout(() => resolve("timed-out"), 100)), + ]); + + expect(result).toBe("settled"); + + const updatedSchedule = await getScheduleByIdOrShortId(schedule.id); + expect(updatedSchedule.lastBackupStatus).toBe("error"); + expect(updatedSchedule.lastBackupError).toBe("Error: restic crashed"); + }); + + test("should block forget on the same repository until the active backup completes", async () => { + const { resticBackupMock, resticForgetMock, runBackupMock } = setup(); + const volume = await createTestVolume(); + const repository = await createTestRepository(); + const schedule = await createTestBackupSchedule({ + volumeId: volume.id, + repositoryId: repository.id, + retentionPolicy: { keepHourly: 24 }, + }); + + let completeBackup: (() => void) | undefined; + resticBackupMock.mockImplementationOnce( + () => + new Promise((resolve) => { + completeBackup = () => resolve({ exitCode: 0, summary: generateBackupOutput(), error: "" }); + }), + ); + + const backupPromise = backupsService.executeBackup(schedule.id); + + await waitForExpect(() => { + expect(runBackupMock).toHaveBeenCalledTimes(1); + }); + + let forgetFinished = false; + const forgetPromise = backupsService.runForget(schedule.id).finally(() => { + forgetFinished = true; + }); + + await new Promise((resolve) => setTimeout(resolve, 50)); + + expect(resticForgetMock).not.toHaveBeenCalled(); + expect(forgetFinished).toBe(false); + + expect(completeBackup).toBeDefined(); + completeBackup?.(); + + await backupPromise; + await forgetPromise; + + expect(resticForgetMock).toHaveBeenCalled(); + expect(resticForgetMock).toHaveBeenCalledWith( + repository.config, + expect.objectContaining({ keepHourly: 24 }), + expect.objectContaining({ tag: schedule.shortId, organizationId: TEST_ORG_ID }), + ); + }); + test("should stop a running backup", async () => { // arrange const { resticBackupMock } = setup(); @@ -220,19 +381,46 @@ describe("stop backup", () => { }); }); - const executePromise = backupsExecutionService.executeBackup(schedule.id); + const executePromise = backupsService.executeBackup(schedule.id); await waitForExpect(async () => { - const runningSchedule = await backupsService.getScheduleById(schedule.id); + const runningSchedule = await getScheduleByIdOrShortId(schedule.id); expect(runningSchedule.lastBackupStatus).toBe("in_progress"); }); // act - await backupsExecutionService.stopBackup(schedule.id); + await backupsService.stopBackup(schedule.id); await executePromise; // assert - const updatedSchedule = await backupsService.getScheduleById(schedule.id); + const updatedSchedule = await getScheduleByIdOrShortId(schedule.id); + expect(updatedSchedule.lastBackupStatus).toBe("warning"); + expect(updatedSchedule.lastBackupError).toBe("Backup was stopped by the user"); + }); + + test("should stop a running backup when the cancel command cannot be delivered", async () => { + const { resticBackupMock, cancelBackupMock } = setup(); + const volume = await createTestVolume(); + const repository = await createTestRepository(); + const schedule = await createTestBackupSchedule({ + volumeId: volume.id, + repositoryId: repository.id, + }); + + resticBackupMock.mockImplementation(() => new Promise(() => {})); + cancelBackupMock.mockResolvedValueOnce(false); + + const executePromise = backupsService.executeBackup(schedule.id); + + await waitForExpect(async () => { + const runningSchedule = await getScheduleByIdOrShortId(schedule.id); + expect(runningSchedule.lastBackupStatus).toBe("in_progress"); + }); + + await backupsService.stopBackup(schedule.id); + await executePromise; + + const updatedSchedule = await getScheduleByIdOrShortId(schedule.id); expect(updatedSchedule.lastBackupStatus).toBe("warning"); expect(updatedSchedule.lastBackupError).toBe("Backup was stopped by the user"); }); @@ -246,36 +434,25 @@ describe("stop backup", () => { repositoryId: repository.id, }); - vi.spyOn(repoMutex, "acquireShared").mockImplementation((_repositoryId, _operation, signal) => { - return new Promise((_, reject) => { - if (signal?.aborted) { - reject(signal.reason instanceof Error ? signal.reason : new Error("Operation aborted")); - return; - } + const releaseLock = await repoMutex.acquireExclusive(repository.id, "test"); + const executePromise = backupsService.executeBackup(schedule.id); - signal?.addEventListener( - "abort", - () => { - reject(signal.reason instanceof Error ? signal.reason : new Error("Operation aborted")); - }, - { once: true }, - ); + try { + await waitForExpect(async () => { + const queuedSchedule = await getScheduleByIdOrShortId(schedule.id); + expect(queuedSchedule.lastBackupStatus).toBe("in_progress"); }); - }); - const executePromise = backupsExecutionService.executeBackup(schedule.id); + expect(resticBackupMock).not.toHaveBeenCalled(); - await waitForExpect(async () => { - const queuedSchedule = await backupsService.getScheduleById(schedule.id); - expect(queuedSchedule.lastBackupStatus).toBe("in_progress"); - }); + await backupsService.stopBackup(schedule.id); + } finally { + releaseLock(); + } - expect(resticBackupMock).not.toHaveBeenCalled(); - - await backupsExecutionService.stopBackup(schedule.id); await executePromise; - const updatedSchedule = await backupsService.getScheduleById(schedule.id); + const updatedSchedule = await getScheduleByIdOrShortId(schedule.id); expect(updatedSchedule.lastBackupStatus).toBe("warning"); expect(updatedSchedule.lastBackupError).toBe("Backup was stopped by the user"); expect(resticBackupMock).not.toHaveBeenCalled(); @@ -297,9 +474,9 @@ describe("stop backup", () => { Promise.resolve({ exitCode: 1, summary: generateBackupOutput(), error: "retry me" }), ); - await backupsExecutionService.executeBackup(schedule.id); + await backupsService.executeBackup(schedule.id); - const failedSchedule = await backupsService.getScheduleById(schedule.id); + const failedSchedule = await getScheduleByIdOrShortId(schedule.id); expect(failedSchedule.failureRetryCount).toBe(1); resticBackupMock.mockImplementationOnce(({ signal }: SafeSpawnParams) => { @@ -319,17 +496,17 @@ describe("stop backup", () => { }); }); - const executePromise = backupsExecutionService.executeBackup(schedule.id); + const executePromise = backupsService.executeBackup(schedule.id); await waitForExpect(async () => { - const retryingSchedule = await backupsService.getScheduleById(schedule.id); + const retryingSchedule = await getScheduleByIdOrShortId(schedule.id); expect(retryingSchedule.lastBackupStatus).toBe("in_progress"); }); - await backupsExecutionService.stopBackup(schedule.id); + await backupsService.stopBackup(schedule.id); await executePromise; - const cancelledSchedule = await backupsService.getScheduleById(schedule.id); + const cancelledSchedule = await getScheduleByIdOrShortId(schedule.id); expect(cancelledSchedule.lastBackupStatus).toBe("warning"); expect(cancelledSchedule.failureRetryCount).toBe(0); }); @@ -348,20 +525,40 @@ describe("stop backup", () => { }); // act & assert - await expect(backupsExecutionService.stopBackup(schedule.id)).rejects.toThrow( + await expect(backupsService.stopBackup(schedule.id)).rejects.toThrow( "No backup is currently running for this schedule", ); - const updatedSchedule = await backupsService.getScheduleById(schedule.id); + const updatedSchedule = await getScheduleByIdOrShortId(schedule.id); expect(updatedSchedule.lastBackupAt).toBe(previousLastBackupAt); expect(updatedSchedule.lastBackupStatus).toBe("warning"); expect(updatedSchedule.lastBackupError).toBe("Backup was stopped by the user"); }); + test("should reset a stuck in_progress status even when no backup is running", async () => { + // arrange + setup(); + const volume = await createTestVolume(); + const repository = await createTestRepository(); + const schedule = await createTestBackupSchedule({ + volumeId: volume.id, + repositoryId: repository.id, + lastBackupStatus: "in_progress", + }); + + // act + await backupsService.stopBackup(schedule.id).catch(() => {}); + + // assert + const updatedSchedule = await getScheduleByIdOrShortId(schedule.id); + expect(updatedSchedule.lastBackupStatus).toBe("warning"); + expect(updatedSchedule.lastBackupError).toBe("Backup was stopped by the user"); + }); + test("should throw NotFoundError when schedule does not exist", async () => { setup(); // act & assert - await expect(backupsExecutionService.stopBackup(99999)).rejects.toThrow("Backup schedule not found"); + await expect(backupsService.stopBackup(99999)).rejects.toThrow("Backup schedule not found"); }); }); @@ -382,7 +579,7 @@ describe("retention policy - runForget", () => { }); // act - await backupsExecutionService.runForget(schedule.id); + await backupsService.runForget(schedule.id); // assert expect(resticForgetMock).toHaveBeenCalledWith( @@ -411,7 +608,7 @@ describe("retention policy - runForget", () => { }); // act & assert - await expect(backupsExecutionService.runForget(schedule.id)).rejects.toThrow( + await expect(backupsService.runForget(schedule.id)).rejects.toThrow( "No retention policy configured for this schedule", ); }); @@ -419,7 +616,7 @@ describe("retention policy - runForget", () => { test("should throw NotFoundError when schedule does not exist", async () => { setup(); // act & assert - await expect(backupsExecutionService.runForget(99999)).rejects.toThrow("Backup schedule not found"); + await expect(backupsService.runForget(99999)).rejects.toThrow("Backup schedule not found"); }); test("should throw NotFoundError when repository does not exist", async () => { @@ -432,9 +629,7 @@ describe("retention policy - runForget", () => { }); // act & assert - await expect(backupsExecutionService.runForget(schedule.id, "non-existent-repo")).rejects.toThrow( - "Repository not found", - ); + await expect(backupsService.runForget(schedule.id, "non-existent-repo")).rejects.toThrow("Repository not found"); }); }); @@ -453,7 +648,7 @@ describe("mirror operations", () => { await createTestBackupScheduleMirror(schedule.id, mirrorRepository.id); // act - await backupsExecutionService.copyToMirrors(schedule.id, sourceRepository, null); + await backupsService.copyToMirrors(schedule.id, sourceRepository, null); // assert expect(resticCopyMock).toHaveBeenCalledWith( @@ -480,7 +675,7 @@ describe("mirror operations", () => { await createTestBackupScheduleMirror(schedule.id, mirrorRepository.id, { enabled: false }); // act - await backupsExecutionService.copyToMirrors(schedule.id, sourceRepository, null); + await backupsService.copyToMirrors(schedule.id, sourceRepository, null); // assert expect(resticCopyMock).not.toHaveBeenCalled(); @@ -500,7 +695,7 @@ describe("mirror operations", () => { const mirror = await createTestBackupScheduleMirror(schedule.id, mirrorRepository.id); // act - await backupsExecutionService.copyToMirrors(schedule.id, sourceRepository, null); + await backupsService.copyToMirrors(schedule.id, sourceRepository, null); // assert const mirrors = await backupsService.getMirrors(schedule.id); @@ -531,7 +726,7 @@ describe("mirror operations", () => { }); // act - await backupsExecutionService.copyToMirrors(schedule.id, sourceRepository, null); + await backupsService.copyToMirrors(schedule.id, sourceRepository, null); // assert const mirrors = await backupsService.getMirrors(schedule.id); @@ -558,7 +753,7 @@ describe("mirror operations", () => { resticCopyMock.mockImplementationOnce(() => Promise.reject(new Error("Copy failed"))); // act - await backupsExecutionService.copyToMirrors(schedule.id, sourceRepository, null); + await backupsService.copyToMirrors(schedule.id, sourceRepository, null); // assert const mirrors = await backupsService.getMirrors(schedule.id); @@ -586,7 +781,7 @@ describe("mirror operations", () => { resticCopyMock.mockImplementation(() => Promise.resolve({ success: true, output: "" })); // act - await backupsExecutionService.copyToMirrors(schedule.id, sourceRepository, schedule.retentionPolicy); + await backupsService.copyToMirrors(schedule.id, sourceRepository, schedule.retentionPolicy); await waitForExpect(() => { expect(resticCopyMock).toHaveBeenCalled(); @@ -617,7 +812,7 @@ describe("mirror operations", () => { resticForgetMock.mockClear(); // act - await backupsExecutionService.copyToMirrors(schedule.id, sourceRepository, schedule.retentionPolicy); + await backupsService.copyToMirrors(schedule.id, sourceRepository, schedule.retentionPolicy); await waitForExpect(() => { expect(resticCopyMock).toHaveBeenCalled(); @@ -660,10 +855,10 @@ describe("mirror operations", () => { ); resticCopyMock.mockImplementation(() => Promise.resolve({ success: true, output: "" })); - const firstCopyPromise = backupsExecutionService.copyToMirrors(firstSchedule.id, sourceRepository, null); + const firstCopyPromise = backupsService.copyToMirrors(firstSchedule.id, sourceRepository, null); await firstCopyStarted; - const secondCopyPromise = backupsExecutionService.copyToMirrors(secondSchedule.id, sourceRepository, null); + const secondCopyPromise = backupsService.copyToMirrors(secondSchedule.id, sourceRepository, null); try { const secondCopyState = await Promise.race<"resolved" | "timeout">([ diff --git a/app/server/modules/backups/__tests__/backups.service.test.ts b/app/server/modules/backups/__tests__/backups.service.test.ts index e0d82282..6d6c0185 100644 --- a/app/server/modules/backups/__tests__/backups.service.test.ts +++ b/app/server/modules/backups/__tests__/backups.service.test.ts @@ -8,15 +8,19 @@ import { createTestRepository } from "~/test/helpers/repository"; import { generateBackupOutput } from "~/test/helpers/restic"; import { faker } from "@faker-js/faker"; import * as spawnModule from "@zerobyte/core/node"; +import type { SafeSpawnParams } from "@zerobyte/core/node"; import { db } from "~/server/db/db"; import { backupScheduleMirrorsTable, repositoriesTable, volumesTable } from "~/server/db/schema"; import { TEST_ORG_ID } from "~/test/helpers/organization"; import * as context from "~/server/core/request-context"; -import { backupsExecutionService } from "../backups.execution"; import { repositoriesService } from "~/server/modules/repositories/repositories.service"; +import { agentManager } from "~/server/modules/agents/agents-manager"; +import { createAgentBackupMocks } from "~/test/helpers/agent-mock"; +import { getScheduleByIdOrShortId } from "../helpers/backup-schedule-lookups"; const setup = () => { - const resticBackupMock = vi.fn(() => Promise.resolve({ exitCode: 0, summary: "", error: "" })); + const resticBackupMock = vi.fn((_: SafeSpawnParams) => Promise.resolve({ exitCode: 0, summary: "", error: "" })); + const { runBackupMock, cancelBackupMock } = createAgentBackupMocks(resticBackupMock); const refreshStatsMock = vi.fn(() => Promise.resolve({ total_size: 0, @@ -29,10 +33,14 @@ const setup = () => { ); vi.spyOn(spawnModule, "safeSpawn").mockImplementation(resticBackupMock); vi.spyOn(repositoriesService, "refreshRepositoryStats").mockImplementation(refreshStatsMock); + vi.spyOn(agentManager, "runBackup").mockImplementation(runBackupMock); + vi.spyOn(agentManager, "cancelBackup").mockImplementation(cancelBackupMock); vi.spyOn(context, "getOrganizationId").mockReturnValue(TEST_ORG_ID); return { resticBackupMock, + runBackupMock, + cancelBackupMock, refreshStatsMock, }; }; @@ -59,10 +67,10 @@ describe("execute backup", () => { ); // act - await backupsExecutionService.executeBackup(schedule.id); + await backupsService.executeBackup(schedule.id); // assert - const updatedSchedule = await backupsService.getScheduleById(schedule.id); + const updatedSchedule = await getScheduleByIdOrShortId(schedule.id); expect(updatedSchedule.nextBackupAt).not.toBeNull(); const nextBackupAt = new Date(updatedSchedule.nextBackupAt ?? 0); @@ -84,7 +92,7 @@ describe("execute backup", () => { }); // act - await backupsExecutionService.executeBackup(schedule.id); + await backupsService.executeBackup(schedule.id); // assert expect(resticBackupMock).not.toHaveBeenCalled(); @@ -106,11 +114,11 @@ describe("execute backup", () => { ); // act - await backupsExecutionService.executeBackup(schedule.id, true); + await backupsService.executeBackup(schedule.id, true); // assert expect(resticBackupMock).toHaveBeenCalled(); - const updatedSchedule = await backupsService.getScheduleById(schedule.id); + const updatedSchedule = await getScheduleByIdOrShortId(schedule.id); expect(updatedSchedule.lastBackupStatus).toBe("success"); expect(updatedSchedule.lastBackupAt).not.toBeNull(); }); @@ -132,10 +140,10 @@ describe("execute backup", () => { ); // act - await backupsExecutionService.executeBackup(schedule.id, true); + await backupsService.executeBackup(schedule.id, true); // assert - const updatedSchedule = await backupsService.getScheduleById(schedule.id); + const updatedSchedule = await getScheduleByIdOrShortId(schedule.id); expect(updatedSchedule.nextBackupAt).toBeNull(); }); @@ -155,13 +163,13 @@ describe("execute backup", () => { }); // act - void backupsExecutionService.executeBackup(schedule.id); + void backupsService.executeBackup(schedule.id); await waitForExpect(() => { expect(resticBackupMock).toHaveBeenCalledTimes(1); }); - await backupsExecutionService.executeBackup(schedule.id); + await backupsService.executeBackup(schedule.id); // assert expect(resticBackupMock).toHaveBeenCalledTimes(1); @@ -182,10 +190,10 @@ describe("execute backup", () => { ); // act - await backupsExecutionService.executeBackup(schedule.id); + await backupsService.executeBackup(schedule.id); // assert - const updatedSchedule = await backupsService.getScheduleById(schedule.id); + const updatedSchedule = await getScheduleByIdOrShortId(schedule.id); expect(updatedSchedule.lastBackupStatus).toBe("warning"); }); @@ -204,10 +212,10 @@ describe("execute backup", () => { ); // act - await backupsExecutionService.executeBackup(schedule.id); + await backupsService.executeBackup(schedule.id); // assert - const updatedSchedule = await backupsService.getScheduleById(schedule.id); + const updatedSchedule = await getScheduleByIdOrShortId(schedule.id); expect(updatedSchedule.lastBackupStatus).toBe("error"); }); }); @@ -229,7 +237,7 @@ describe("getSchedulesToExecute", () => { }); // act - const schedulesToExecute = await backupsExecutionService.getSchedulesToExecute(); + const schedulesToExecute = await backupsService.getSchedulesToExecute(); // assert expect(schedulesToExecute).toContain(schedule.id); @@ -246,7 +254,7 @@ describe("getScheduleByIdOrShortId", () => { repositoryId: repository.id, }); - const found = await backupsService.getScheduleByIdOrShortId(String(schedule.id)); + const found = await getScheduleByIdOrShortId(String(schedule.id)); expect(found.id).toBe(schedule.id); expect(found.shortId).toBe(schedule.shortId); @@ -261,7 +269,7 @@ describe("getScheduleByIdOrShortId", () => { repositoryId: repository.id, }); - const found = await backupsService.getScheduleByIdOrShortId(schedule.shortId); + const found = await getScheduleByIdOrShortId(schedule.shortId); expect(found.id).toBe(schedule.id); expect(found.shortId).toBe(schedule.shortId); @@ -274,10 +282,8 @@ describe("getScheduleByIdOrShortId", () => { organizationId: otherOrgId, }); - await expect(backupsService.getScheduleByIdOrShortId(schedule.shortId)).rejects.toThrow( - "Backup schedule not found", - ); - await expect(backupsService.getScheduleByIdOrShortId(schedule.id)).rejects.toThrow("Backup schedule not found"); + await expect(getScheduleByIdOrShortId(schedule.shortId)).rejects.toThrow("Backup schedule not found"); + await expect(getScheduleByIdOrShortId(schedule.id)).rejects.toThrow("Backup schedule not found"); }); }); diff --git a/app/server/modules/backups/backup-executor.ts b/app/server/modules/backups/backup-executor.ts index 60c9918c..7a074e5e 100644 --- a/app/server/modules/backups/backup-executor.ts +++ b/app/server/modules/backups/backup-executor.ts @@ -1,9 +1,15 @@ import { Effect } from "effect"; -import { restic } from "../../core/restic"; -import type { BackupSchedule, Repository, Volume } from "../../db/schema"; -import type { ResticBackupOutputDto, ResticBackupProgressDto } from "@zerobyte/core/restic"; -import { createBackupOptions } from "./backup.helpers"; +import type { BackupSchedule, Volume, Repository } from "../../db/schema"; +import { config } from "../../core/config"; +import { restic, resticDeps } from "../../core/restic"; +import type { BackupRunPayload } from "@zerobyte/contracts/agent-protocol"; +import { agentManager, type BackupExecutionProgress } from "../agents/agents-manager"; import { getVolumePath } from "../volumes/helpers"; +import { decryptRepositoryConfig } from "../repositories/repository-config-secrets"; +import { createBackupOptions } from "./backup.helpers"; +import { toErrorDetails } from "../../utils/errors"; + +const LOCAL_AGENT_ID = "local"; type BackupExecutionRequest = { scheduleId: number; @@ -15,30 +21,85 @@ type BackupExecutionRequest = { onProgress: (progress: BackupExecutionProgress) => void; }; -export type BackupExecutionProgress = ResticBackupProgressDto; - -export type BackupExecutionResult = - | { - status: "unavailable"; - error: Error; - } - | { - status: "completed"; - exitCode: number; - result: ResticBackupOutputDto | null; - warningDetails: string | null; - } - | { - status: "failed"; - error: unknown; - } - | { - status: "cancelled"; - message?: string; - }; +export type { BackupExecutionResult } from "../agents/agents-manager"; const activeControllersByScheduleId = new Map(); +const createBackupRunPayload = async ({ + jobId, + schedule, + volume, + repository, + organizationId, +}: BackupExecutionRequest & { jobId: string }): Promise => { + const sourcePath = getVolumePath(volume); + const { signal: _, ...options } = createBackupOptions(schedule, sourcePath); + const repositoryConfig = await decryptRepositoryConfig(repository.config); + const encryptedResticPassword = await resticDeps.getOrganizationResticPassword(organizationId); + const resticPassword = await resticDeps.resolveSecret(encryptedResticPassword); + + return { + jobId, + scheduleId: schedule.shortId, + organizationId, + sourcePath, + repositoryConfig, + options: { + ...options, + compressionMode: repository.compressionMode ?? "auto", + }, + runtime: { + password: resticPassword, + cacheDir: resticDeps.resticCacheDir, + passFile: resticDeps.resticPassFile, + defaultExcludes: resticDeps.defaultExcludes, + rcloneConfigFile: resticDeps.rcloneConfigFile, + hostname: resticDeps.hostname, + }, + }; +}; + +const executeBackupWithoutAgent = async ( + payload: BackupRunPayload, + { signal, onProgress }: Pick, +) => { + try { + const execution = await Effect.runPromise( + restic + .backup(payload.repositoryConfig, payload.sourcePath, { + ...payload.options, + organizationId: payload.organizationId, + signal, + onProgress, + }) + .pipe( + Effect.map((result) => ({ success: true as const, result })), + Effect.catchAll((error) => Effect.succeed({ success: false as const, error })), + ), + ); + + if (!execution.success) { + return { + status: "failed" as const, + error: toErrorDetails(execution.error), + }; + } + + const { exitCode, result, warningDetails } = execution.result; + return { + status: "completed" as const, + exitCode, + result, + warningDetails, + }; + } catch (error) { + return { + status: "failed" as const, + error: toErrorDetails(error), + }; + } +}; + export const backupExecutor = { track: (scheduleId: number) => { const abortController = new AbortController(); @@ -50,43 +111,45 @@ export const backupExecutor = { activeControllersByScheduleId.delete(scheduleId); } }, - execute: async (params: BackupExecutionRequest): Promise => { - const { schedule, volume, repository, organizationId, signal, onProgress } = params; - try { - const volumePath = getVolumePath(volume); - const backupOptions = createBackupOptions(schedule, volumePath, signal); - - const execution = await Effect.runPromise( - restic - .backup(repository.config, volumePath, { - ...backupOptions, - compressionMode: repository.compressionMode ?? "auto", - organizationId, - onProgress, - }) - .pipe( - Effect.map((result) => ({ success: true as const, result })), - Effect.catchAll((error) => Effect.succeed({ success: false as const, error })), - ), - ); - - if (!execution.success) { - throw execution.error; - } - - const { exitCode, result, warningDetails } = execution.result; - return { status: "completed", exitCode, result, warningDetails }; - } catch (error) { - return { status: "failed", error }; + execute: async (request: Omit) => { + const trackedAbortController = activeControllersByScheduleId.get(request.scheduleId); + if (!trackedAbortController || trackedAbortController.signal !== request.signal) { + throw new Error(`Backup execution for schedule ${request.scheduleId} was not tracked`); } + + if (request.signal.aborted) { + throw request.signal.reason || new Error("Operation aborted"); + } + + const jobId = Bun.randomUUIDv7(); + + const payload = await createBackupRunPayload({ ...request, jobId }); + + if (request.signal.aborted) { + throw request.signal.reason || new Error("Operation aborted"); + } + + const executionResult = await agentManager.runBackup(LOCAL_AGENT_ID, { + scheduleId: request.scheduleId, + payload, + signal: request.signal, + onProgress: request.onProgress, + }); + + if (executionResult.status === "unavailable" && !config.flags.enableLocalAgent) { + return executeBackupWithoutAgent(payload, request); + } + + return executionResult; }, - cancel: (scheduleId: number) => { + cancel: async (scheduleId: number) => { const abortController = activeControllersByScheduleId.get(scheduleId); if (!abortController) { return false; } abortController.abort(); + await agentManager.cancelBackup(LOCAL_AGENT_ID, scheduleId); return true; }, }; diff --git a/app/server/modules/backups/backups.execution.ts b/app/server/modules/backups/backups.execution.ts deleted file mode 100644 index f85388bf..00000000 --- a/app/server/modules/backups/backups.execution.ts +++ /dev/null @@ -1,11 +0,0 @@ -import { backupsService } from "./backups.service"; - -export const backupsExecutionService = { - executeBackup: backupsService.executeBackup, - validateBackupExecution: backupsService.validateBackupExecution, - getSchedulesToExecute: backupsService.getSchedulesToExecute, - stopBackup: backupsService.stopBackup, - runForget: backupsService.runForget, - copyToMirrors: backupsService.copyToMirrors, - getBackupProgress: backupsService.getBackupProgress, -}; diff --git a/app/server/modules/backups/backups.service.ts b/app/server/modules/backups/backups.service.ts index a01047b3..d39479f0 100644 --- a/app/server/modules/backups/backups.service.ts +++ b/app/server/modules/backups/backups.service.ts @@ -37,14 +37,6 @@ const listSchedules = async () => { return schedules.filter((schedule) => schedule.volume && schedule.repository); }; -const getScheduleById = async (scheduleId: number) => { - return getScheduleByIdOrShortId(scheduleId); -}; - -const getScheduleByShortId = async (shortId: ShortId) => { - return getScheduleByIdOrShortId(shortId); -}; - const createSchedule = async (data: CreateBackupScheduleBody) => { const organizationId = getOrganizationId(); if (data.cronExpression && !isValidCron(data.cronExpression)) { @@ -467,7 +459,7 @@ const stopBackup = async (scheduleId: number) => { } try { - if (!backupExecutor.cancel(scheduleId)) { + if (!(await backupExecutor.cancel(scheduleId))) { throw new ConflictError("No backup is currently running for this schedule"); } @@ -479,9 +471,6 @@ const stopBackup = async (scheduleId: number) => { export const backupsService = { listSchedules, - getScheduleById, - getScheduleByShortId, - getScheduleByIdOrShortId, createSchedule, updateSchedule, deleteSchedule, diff --git a/app/server/modules/backups/helpers/backup-lifecycle.ts b/app/server/modules/backups/helpers/backup-lifecycle.ts index 4bb33b0f..acc887f1 100644 --- a/app/server/modules/backups/helpers/backup-lifecycle.ts +++ b/app/server/modules/backups/helpers/backup-lifecycle.ts @@ -10,7 +10,7 @@ import { getOrganizationId } from "~/server/core/request-context"; import type { BackupProgressEventDto } from "~/schemas/events-dto"; import { calculateNextRun } from "../backup.helpers"; import { scheduleQueries } from "../backups.queries"; -import type { BackupExecutionProgress } from "../backup-executor"; +import type { BackupExecutionProgress } from "../../agents/agents-manager"; import { repositoriesService } from "../../repositories/repositories.service"; import { copyToMirrors, runForget } from "./backup-maintenance"; diff --git a/app/test/helpers/agent-mock.ts b/app/test/helpers/agent-mock.ts new file mode 100644 index 00000000..4fbf20e5 --- /dev/null +++ b/app/test/helpers/agent-mock.ts @@ -0,0 +1,98 @@ +import { vi } from "vitest"; +import { fromAny, fromPartial } from "@total-typescript/shoehorn"; +import type { SafeSpawnParams } from "@zerobyte/core/node"; +import type { BackupExecutionResult } from "~/server/modules/agents/agents-manager"; + +export const createAgentBackupMocks = ( + resticBackupMock: (params: SafeSpawnParams) => Promise<{ + exitCode: number; + summary: string; + error: string; + stderr?: string; + }>, +) => { + const runningBackups = new Map void; cancelled: boolean }>(); + + const runBackupMock = vi.fn( + async (_agentId: string, request: { scheduleId: number; payload: { jobId: string }; signal: AbortSignal }) => { + return new Promise((resolve) => { + runningBackups.set(request.scheduleId, { resolve, cancelled: false }); + + request.signal.addEventListener( + "abort", + () => { + const running = runningBackups.get(request.scheduleId); + if (!running || running.cancelled) { + return; + } + + running.cancelled = true; + runningBackups.delete(request.scheduleId); + resolve({ status: "cancelled" }); + }, + { once: true }, + ); + + void (async () => { + const stderrLines: string[] = []; + const result = await resticBackupMock( + fromPartial({ + signal: request.signal, + onStderr: (line: string) => { + stderrLines.push(line); + }, + }), + ); + const running = runningBackups.get(request.scheduleId); + if (!running || running.cancelled) { + return; + } + + runningBackups.delete(request.scheduleId); + + if (result.exitCode === 0 || result.exitCode === 3) { + let parsedResult: Record | null = null; + if (result.summary) { + try { + parsedResult = JSON.parse(result.summary) as Record; + } catch { + parsedResult = null; + } + } + + resolve({ + status: "completed", + exitCode: result.exitCode, + result: fromAny(parsedResult), + warningDetails: stderrLines.join("\n") || null, + }); + return; + } + + const resultWithStderr = result as typeof result & { stderr?: string }; + resolve({ + status: "failed", + error: stderrLines.join("\n") || resultWithStderr.stderr || result.error, + }); + })().catch((err) => { + runningBackups.delete(request.scheduleId); + resolve({ status: "failed", error: String(err) }); + }); + }); + }, + ); + + const cancelBackupMock = vi.fn(async (_agentId: string, scheduleId: number) => { + const running = runningBackups.get(scheduleId); + if (!running) { + return false; + } + + running.cancelled = true; + runningBackups.delete(scheduleId); + running.resolve({ status: "cancelled" }); + return true; + }); + + return { runBackupMock, cancelBackupMock }; +}; diff --git a/apps/agent/src/__tests__/controller-session.test.ts b/apps/agent/src/__tests__/controller-session.test.ts index e9ff4e83..4f238b8f 100644 --- a/apps/agent/src/__tests__/controller-session.test.ts +++ b/apps/agent/src/__tests__/controller-session.test.ts @@ -70,3 +70,25 @@ test("emits backup.failed when a backup command hits a restic error", async () = session.close(); } }); + +test("closes the websocket when an outbound send throws", async () => { + const close = vi.fn(() => undefined); + const session = createControllerSession( + fromPartial({ + send: () => { + throw new Error("socket write failed"); + }, + close, + }), + ); + + try { + session.onOpen(); + + await waitForExpect(() => { + expect(close).toHaveBeenCalledTimes(1); + }); + } finally { + session.close(); + } +}); diff --git a/apps/agent/src/controller-session.ts b/apps/agent/src/controller-session.ts index 9948d00e..d8f82e24 100644 --- a/apps/agent/src/controller-session.ts +++ b/apps/agent/src/controller-session.ts @@ -17,6 +17,7 @@ export type ControllerSession = { }; export const createControllerSession = (ws: WebSocket): ControllerSession => { + let isClosed = false; const outboundQueue = Effect.runSync(Queue.bounded(64)); const inboundQueue = Effect.runSync(Queue.bounded(64)); const runningJobsRef = Effect.runSync(Ref.make>(new Map())); @@ -63,6 +64,31 @@ export const createControllerSession = (ws: WebSocket): ControllerSession => { offerOutbound, }; + const closeSession = () => { + if (isClosed) { + return; + } + + isClosed = true; + void Effect.runPromise(abortRunningJobs).catch(() => {}); + void Effect.runPromise(Fiber.interrupt(writerFiber)).catch(() => {}); + void Effect.runPromise(Fiber.interrupt(processorFiber)).catch(() => {}); + void Effect.runPromise(Queue.shutdown(outboundQueue)).catch(() => {}); + void Effect.runPromise(Queue.shutdown(inboundQueue)).catch(() => {}); + }; + + const handleSendFailure = (reason: string) => { + logger.error(`Closing agent session after an outbound websocket send failed: ${reason}`); + + try { + ws.close(); + } catch (error) { + logger.error(`Failed to close controller websocket after send failure: ${toMessage(error)}`); + } + + closeSession(); + }; + const writerFiber = Effect.runFork( Effect.forever( Effect.gen(function* () { @@ -71,7 +97,7 @@ export const createControllerSession = (ws: WebSocket): ControllerSession => { try { ws.send(message); } catch (error) { - logger.error(`Failed to send controller message: ${toMessage(error)}`); + handleSendFailure(toMessage(error)); } }); }), @@ -106,16 +132,15 @@ export const createControllerSession = (ws: WebSocket): ControllerSession => { }); }, onMessage: (data) => { + if (typeof data !== "string") { + logger.warn("Agent received a non-text message"); + return; + } + void Effect.runPromise(offerInbound(data as ControllerWireMessage)).catch((error) => { logger.error(`Failed to queue inbound message: ${toMessage(error)}`); }); }, - close: () => { - void Effect.runPromise(abortRunningJobs).catch(() => {}); - void Effect.runPromise(Fiber.interrupt(writerFiber)).catch(() => {}); - void Effect.runPromise(Fiber.interrupt(processorFiber)).catch(() => {}); - void Effect.runPromise(Queue.shutdown(outboundQueue)).catch(() => {}); - void Effect.runPromise(Queue.shutdown(inboundQueue)).catch(() => {}); - }, + close: closeSession, }; }; diff --git a/apps/agent/src/index.ts b/apps/agent/src/index.ts index f9e08b6d..2593792c 100644 --- a/apps/agent/src/index.ts +++ b/apps/agent/src/index.ts @@ -5,7 +5,7 @@ const controllerUrl = process.env.ZEROBYTE_CONTROLLER_URL; const agentToken = process.env.ZEROBYTE_AGENT_TOKEN; const RECONNECT_DELAY_MS = 1000; -class Agent { +export class Agent { private ws: WebSocket | null = null; private controllerSession: ControllerSession | null = null; private reconnectTimeout: ReturnType | null = null; @@ -22,6 +22,10 @@ class Agent { } connect() { + if (this.reconnectTimeout) { + clearTimeout(this.reconnectTimeout); + this.reconnectTimeout = null; + } if (this.ws) { return; } diff --git a/apps/agent/tsconfig.json b/apps/agent/tsconfig.json index e30ff81c..4ac5f0a9 100644 --- a/apps/agent/tsconfig.json +++ b/apps/agent/tsconfig.json @@ -7,7 +7,7 @@ "module": "Preserve", "moduleDetection": "force", "jsx": "react-jsx", - "types": ["node", "bun"], + "types": ["bun", "node"], "allowJs": true, // Bundler mode