diff --git a/app/schemas/events-dto.ts b/app/schemas/events-dto.ts index a671205c..eeaffa30 100644 --- a/app/schemas/events-dto.ts +++ b/app/schemas/events-dto.ts @@ -2,7 +2,7 @@ import { z } from "zod"; import { resticBackupProgressMetricsSchema, resticBackupRunSummarySchema } from "@zerobyte/core/restic"; const backupEventStatusSchema = z.enum(["success", "error", "stopped", "warning"]); -const restoreEventStatusSchema = z.enum(["success", "error"]); +const restoreEventStatusSchema = z.enum(["success", "error", "cancelled"]); const backupEventBaseSchema = z.object({ scheduleId: z.string(), @@ -15,6 +15,7 @@ const organizationScopedSchema = z.object({ }); const restoreEventBaseSchema = z.object({ + restoreId: z.string(), repositoryId: z.string(), snapshotId: z.string(), }); @@ -51,6 +52,8 @@ const restoreProgressEventSchema = restoreEventBaseSchema.extend(restoreProgress const restoreCompletedEventSchema = restoreEventBaseSchema.extend({ status: restoreEventStatusSchema, error: z.string().optional(), + filesRestored: z.number().optional(), + filesSkipped: z.number().optional(), }); const serverBackupStartedEventSchema = organizationScopedSchema.extend(backupStartedEventSchema.shape); diff --git a/app/server/modules/agents/__tests__/agents-manager.events.test.ts b/app/server/modules/agents/__tests__/agents-manager.events.test.ts index 2b12ef5f..fd27d505 100644 --- a/app/server/modules/agents/__tests__/agents-manager.events.test.ts +++ b/app/server/modules/agents/__tests__/agents-manager.events.test.ts @@ -1,7 +1,7 @@ import { afterEach, expect, test, vi } from "vitest"; import { Effect } from "effect"; import { fromAny, fromPartial } from "@total-typescript/shoehorn"; -import type { BackupRunPayload } from "@zerobyte/contracts/agent-protocol"; +import type { BackupRunPayload, RestoreRunPayload } from "@zerobyte/contracts/agent-protocol"; import type { AgentManagerEvent } from "../controller/server"; import type { ProcessWithAgentRuntime } from "../agents-manager"; @@ -9,6 +9,8 @@ const controllerMock = vi.hoisted(() => ({ onEvent: null as null | ((event: AgentManagerEvent) => void), sendBackup: vi.fn(), cancelBackup: vi.fn(), + sendRestore: vi.fn(), + cancelRestore: vi.fn(), stop: vi.fn(), })); @@ -22,6 +24,8 @@ vi.mock("../controller/server", async () => { stop: Effect.sync(controllerMock.stop), sendBackup: controllerMock.sendBackup, cancelBackup: controllerMock.cancelBackup, + sendRestore: controllerMock.sendRestore, + cancelRestore: controllerMock.cancelRestore, }; }), }; @@ -37,6 +41,7 @@ const resetAgentRuntime = () => { localAgentRestartTimeout: null, activeBackupsByScheduleId: new Map(), activeBackupScheduleIdsByJobId: new Map(), + activeRestoresByRestoreId: new Map(), }; }; @@ -44,12 +49,17 @@ const backupPayload = fromPartial({ jobId: "job-1", scheduleId: "schedule-1", }); +const restorePayload = fromPartial({ + restoreId: "restore-1", +}); afterEach(() => { delete processWithAgentRuntime.__zerobyteAgentRuntime; controllerMock.onEvent = null; controllerMock.sendBackup.mockReset(); controllerMock.cancelBackup.mockReset(); + controllerMock.sendRestore.mockReset(); + controllerMock.cancelRestore.mockReset(); controllerMock.stop.mockReset(); vi.resetModules(); vi.restoreAllMocks(); @@ -73,7 +83,11 @@ test("backup progress is delivered to the running backup callback", async () => type: "backup.progress", agentId: "local", agentName: "Local Agent", - payload: fromAny({ jobId: "job-1", scheduleId: "schedule-1", progress: { percentDone: 0.5 } }), + payload: fromAny({ + jobId: "job-1", + scheduleId: "schedule-1", + progress: { percentDone: 0.5 }, + }), }); controllerMock.onEvent?.({ type: "backup.completed", @@ -108,7 +122,12 @@ test("backup failed and cancelled events resolve the matching running backup", a type: "backup.failed", agentId: "local", agentName: "Local Agent", - payload: { jobId: "job-1", scheduleId: "schedule-1", error: "failed", errorDetails: "restic failed" }, + payload: { + jobId: "job-1", + scheduleId: "schedule-1", + error: "failed", + errorDetails: "restic failed", + }, }); await expect(failedPromise).resolves.toEqual({ status: "failed", error: "restic failed" }); @@ -124,7 +143,10 @@ test("backup failed and cancelled events resolve the matching running backup", a agentName: "Local Agent", payload: { jobId: "job-2", scheduleId: "schedule-2", message: "cancelled remotely" }, }); - await expect(cancelledPromise).resolves.toEqual({ status: "cancelled", message: "cancelled remotely" }); + await expect(cancelledPromise).resolves.toEqual({ + status: "cancelled", + message: "cancelled remotely", + }); await stopAgentController(); }); @@ -147,7 +169,11 @@ test("agent disconnect cancels only backups owned by that agent", async () => { onProgress: vi.fn(), }); - controllerMock.onEvent?.({ type: "agent.disconnected", agentId: "local", agentName: "Local Agent" }); + controllerMock.onEvent?.({ + type: "agent.disconnected", + agentId: "local", + agentName: "Local Agent", + }); controllerMock.onEvent?.({ type: "backup.completed", agentId: "remote", @@ -230,6 +256,117 @@ test("runBackup requests cancellation when the abort signal fires while sending" }); expect(result).toEqual({ status: "cancelled" }); - expect(controllerMock.cancelBackup).toHaveBeenCalledWith("local", { jobId: "job-1", scheduleId: "schedule-1" }); + expect(controllerMock.cancelBackup).toHaveBeenCalledWith("local", { + jobId: "job-1", + scheduleId: "schedule-1", + }); + await stopAgentController(); +}); + +test("restore events are delivered to the running restore callbacks", async () => { + resetAgentRuntime(); + controllerMock.sendRestore.mockImplementation(() => Effect.succeed(true)); + const { agentManager, startAgentController, stopAgentController } = await import("../agents-manager"); + const onStarted = vi.fn(); + const onProgress = vi.fn(); + + await startAgentController(); + const started = await agentManager.startRestore("local", { + payload: restorePayload, + signal: new AbortController().signal, + onStarted, + onProgress, + }); + if (started.status !== "started") { + throw new Error("Expected restore to start"); + } + + controllerMock.onEvent?.({ + type: "restore.started", + agentId: "local", + agentName: "Local Agent", + payload: { + restoreId: "restore-1", + organizationId: "org-1", + repositoryId: "repo-1", + snapshotId: "snapshot-1", + }, + }); + controllerMock.onEvent?.({ + type: "restore.progress", + agentId: "local", + agentName: "Local Agent", + payload: fromAny({ restoreId: "restore-1", progress: { percent_done: 0.5 } }), + }); + controllerMock.onEvent?.({ + type: "restore.completed", + agentId: "local", + agentName: "Local Agent", + payload: { + restoreId: "restore-1", + organizationId: "org-1", + repositoryId: "repo-1", + snapshotId: "snapshot-1", + result: { message_type: "summary", files_restored: 2, files_skipped: 1 }, + }, + }); + + await expect(started.result).resolves.toEqual({ + status: "completed", + result: { message_type: "summary", files_restored: 2, files_skipped: 1 }, + }); + expect(onStarted).toHaveBeenCalledOnce(); + expect(onProgress).toHaveBeenCalledWith({ percent_done: 0.5 }); + await stopAgentController(); +}); + +test("agent disconnect cancels only restores owned by that agent", async () => { + resetAgentRuntime(); + controllerMock.sendRestore.mockImplementation(() => Effect.succeed(true)); + const { agentManager, startAgentController, stopAgentController } = await import("../agents-manager"); + + await startAgentController(); + const localStarted = await agentManager.startRestore("local", { + payload: restorePayload, + signal: new AbortController().signal, + onStarted: vi.fn(), + onProgress: vi.fn(), + }); + const remoteStarted = await agentManager.startRestore("remote", { + payload: fromPartial({ restoreId: "restore-2" }), + signal: new AbortController().signal, + onStarted: vi.fn(), + onProgress: vi.fn(), + }); + if (localStarted.status !== "started" || remoteStarted.status !== "started") { + throw new Error("Expected restores to start"); + } + + controllerMock.onEvent?.({ + type: "agent.disconnected", + agentId: "local", + agentName: "Local Agent", + }); + controllerMock.onEvent?.({ + type: "restore.completed", + agentId: "remote", + agentName: "Remote Agent", + payload: { + restoreId: "restore-2", + organizationId: "org-1", + repositoryId: "repo-1", + snapshotId: "snapshot-1", + result: { message_type: "summary", files_restored: 1, files_skipped: 0 }, + }, + }); + + await expect(localStarted.result).resolves.toEqual({ + status: "cancelled", + message: "The connection to the restore agent was lost. Restart the restore to ensure it completes.", + }); + await expect(remoteStarted.result).resolves.toEqual({ + status: "completed", + result: { message_type: "summary", files_restored: 1, files_skipped: 0 }, + }); await stopAgentController(); }); diff --git a/app/server/modules/agents/__tests__/controller-runtime.test.ts b/app/server/modules/agents/__tests__/controller-runtime.test.ts index 7c3ef8b3..801646af 100644 --- a/app/server/modules/agents/__tests__/controller-runtime.test.ts +++ b/app/server/modules/agents/__tests__/controller-runtime.test.ts @@ -215,6 +215,37 @@ test("websocket lifecycle updates agent connection status", async () => { expect(stop).toHaveBeenCalledWith(true); }); +test("websocket restore events are forwarded with agent metadata", async () => { + const serve = vi + .spyOn(Bun, "serve") + .mockReturnValue(fromPartial({ port: 3001, stop: vi.fn(() => Promise.resolve()) })); + const { runtime, onEvent } = await startRuntime(vi.fn()); + const websocket = serve.mock.calls[0]?.[0].websocket; + const socket = createSocket("connection-1"); + + await websocket?.open?.(fromPartial(socket)); + await websocket?.message?.( + fromPartial(socket), + createAgentMessage("restore.completed", { + restoreId: "restore-1", + organizationId: "org-1", + repositoryId: "repo-1", + snapshotId: "snapshot-1", + result: { message_type: "summary", files_restored: 2, files_skipped: 0 }, + }), + ); + await Effect.runPromise(runtime.stop); + + expect(onEvent).toHaveBeenCalledWith( + expect.objectContaining({ + type: "restore.completed", + agentId: LOCAL_AGENT_ID, + agentName: LOCAL_AGENT_NAME, + payload: expect.objectContaining({ restoreId: "restore-1" }), + }), + ); +}); + test("websocket open failure closes the upgraded socket", async () => { agentsServiceMocks.markAgentConnecting.mockRejectedValueOnce(new Error("db unavailable")); const serve = vi diff --git a/app/server/modules/agents/agents-manager.ts b/app/server/modules/agents/agents-manager.ts index 7a5b9c90..0f1624a8 100644 --- a/app/server/modules/agents/agents-manager.ts +++ b/app/server/modules/agents/agents-manager.ts @@ -1,5 +1,10 @@ import { logger } from "@zerobyte/core/node"; -import type { BackupRunPayload, VolumeCommand, VolumeCommandResult } from "@zerobyte/contracts/agent-protocol"; +import type { + BackupRunPayload, + RestoreRunPayload, + VolumeCommand, + VolumeCommandResult, +} from "@zerobyte/contracts/agent-protocol"; import { Effect } from "effect"; import { config } from "../../core/config"; import { createAgentManagerRuntime, type AgentManagerEvent } from "./controller/server"; @@ -10,9 +15,16 @@ import { type AgentRuntimeState, type BackupExecutionProgress, type BackupExecutionResult, + type RestoreExecutionProgress, + type RestoreExecutionResult, } from "./helpers/runtime-state"; import { getDevAgentRuntimeState } from "./helpers/runtime-state.dev"; -export type { BackupExecutionProgress, BackupExecutionResult } from "./helpers/runtime-state"; +export type { + BackupExecutionProgress, + BackupExecutionResult, + RestoreExecutionProgress, + RestoreExecutionResult, +} from "./helpers/runtime-state"; export type { ProcessWithAgentRuntime } from "./helpers/runtime-state.dev"; type ProcessWithProductionAgentRuntime = NodeJS.Process & { @@ -26,6 +38,17 @@ type AgentRunBackupRequest = { onProgress: (progress: BackupExecutionProgress) => void; }; +type AgentStartRestoreRequest = { + payload: RestoreRunPayload; + signal: AbortSignal; + onStarted: () => void; + onProgress: (progress: RestoreExecutionProgress) => void; +}; + +type AgentRestoreStartResult = + | { status: "started"; result: Promise } + | { status: "unavailable"; error: Error }; + 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. @@ -41,6 +64,7 @@ const getAgentRuntimeState = () => (config.__prod__ ? getProductionAgentRuntimeS const getAgentManagerRuntime = () => getAgentRuntimeState().agentManager; const getActiveBackupsByScheduleId = () => getAgentRuntimeState().activeBackupsByScheduleId; const getActiveBackupScheduleIdsByJobId = () => getAgentRuntimeState().activeBackupScheduleIdsByJobId; +const getActiveRestoresByRestoreId = () => getAgentRuntimeState().activeRestoresByRestoreId; const clearActiveBackupRun = (scheduleId: number) => { const activeBackupsByScheduleId = getActiveBackupsByScheduleId(); @@ -65,6 +89,27 @@ const resolveActiveBackupRun = (scheduleId: number, result: BackupExecutionResul return true; }; +const clearActiveRestoreRun = (restoreId: string) => { + const activeRestoresByRestoreId = getActiveRestoresByRestoreId(); + const activeRestoreRun = activeRestoresByRestoreId.get(restoreId); + if (!activeRestoreRun) { + return null; + } + + activeRestoresByRestoreId.delete(restoreId); + return activeRestoreRun; +}; + +const resolveActiveRestoreRun = (restoreId: string, result: RestoreExecutionResult) => { + const activeRestoreRun = clearActiveRestoreRun(restoreId); + if (!activeRestoreRun) { + return false; + } + + activeRestoreRun.resolve(result); + return true; +}; + const cancelActiveBackupRunsForAgent = (agentId: string, message: string) => { const activeBackupsByScheduleId = getActiveBackupsByScheduleId(); const matchingScheduleIds = [...activeBackupsByScheduleId.values()] @@ -76,6 +121,17 @@ const cancelActiveBackupRunsForAgent = (agentId: string, message: string) => { } }; +const cancelActiveRestoreRunsForAgent = (agentId: string, message: string) => { + const activeRestoresByRestoreId = getActiveRestoresByRestoreId(); + const matchingRestoreIds = [...activeRestoresByRestoreId.values()] + .filter((activeRestoreRun) => activeRestoreRun.agentId === agentId) + .map((activeRestoreRun) => activeRestoreRun.restoreId); + + for (const restoreId of matchingRestoreIds) { + resolveActiveRestoreRun(restoreId, { status: "cancelled", message }); + } +}; + const getActiveBackupRun = (jobId: string, scheduleId: string, eventName: string, agentId: string) => { const trackedScheduleId = getActiveBackupScheduleIdsByJobId().get(jobId); if (trackedScheduleId === undefined) { @@ -99,6 +155,21 @@ const getActiveBackupRun = (jobId: string, scheduleId: string, eventName: string return activeBackupRun; }; +const getActiveRestoreRun = (restoreId: string, eventName: string, agentId: string) => { + const activeRestoreRun = getActiveRestoresByRestoreId().get(restoreId); + if (!activeRestoreRun) { + logger.warn(`Received ${eventName} for unknown restore ${restoreId} from agent ${agentId}`); + return null; + } + + if (activeRestoreRun.agentId !== agentId) { + logger.warn(`Ignoring ${eventName} for restore ${restoreId} from unexpected agent ${agentId}`); + return null; + } + + return activeRestoreRun; +}; + const requestBackupCancellation = async (agentId: string, scheduleId: number) => { const activeBackupRun = getActiveBackupsByScheduleId().get(scheduleId); if (!activeBackupRun) { @@ -132,6 +203,32 @@ const requestBackupCancellation = async (agentId: string, scheduleId: number) => return true; }; +const requestRestoreCancellation = async (agentId: string, restoreId: string) => { + const activeRestoreRun = getActiveRestoresByRestoreId().get(restoreId); + if (!activeRestoreRun) { + return false; + } + + if (activeRestoreRun.cancellationRequested) { + return true; + } + + activeRestoreRun.cancellationRequested = true; + + const runtime = getAgentManagerRuntime(); + if (!runtime) { + resolveActiveRestoreRun(restoreId, { status: "cancelled" }); + return true; + } + + if (await Effect.runPromise(runtime.cancelRestore(agentId, { restoreId }))) { + return true; + } + + resolveActiveRestoreRun(restoreId, { status: "cancelled" }); + return true; +}; + const handleAgentManagerEvent = (event: AgentManagerEvent) => { switch (event.type) { case "agent.disconnected": { @@ -139,6 +236,10 @@ const handleAgentManagerEvent = (event: AgentManagerEvent) => { event.agentId, "The connection to the backup agent was lost. Restart the backup to ensure it completes.", ); + cancelActiveRestoreRunsForAgent( + event.agentId, + "The connection to the restore agent was lost. Restart the restore to ensure it completes.", + ); break; } case "backup.started": { @@ -212,6 +313,60 @@ const handleAgentManagerEvent = (event: AgentManagerEvent) => { }); break; } + case "restore.started": { + const activeRestoreRun = getActiveRestoreRun(event.payload.restoreId, event.type, event.agentId); + if (!activeRestoreRun) { + break; + } + + activeRestoreRun.onStarted(); + break; + } + case "restore.progress": { + const activeRestoreRun = getActiveRestoreRun(event.payload.restoreId, event.type, event.agentId); + if (!activeRestoreRun) { + break; + } + + activeRestoreRun.onProgress(event.payload.progress); + break; + } + case "restore.completed": { + const activeRestoreRun = getActiveRestoreRun(event.payload.restoreId, event.type, event.agentId); + if (!activeRestoreRun) { + break; + } + + resolveActiveRestoreRun(activeRestoreRun.restoreId, { + status: "completed", + result: event.payload.result, + }); + break; + } + case "restore.failed": { + const activeRestoreRun = getActiveRestoreRun(event.payload.restoreId, event.type, event.agentId); + if (!activeRestoreRun) { + break; + } + + resolveActiveRestoreRun(activeRestoreRun.restoreId, { + status: "failed", + error: event.payload.errorDetails ?? event.payload.error, + }); + break; + } + case "restore.cancelled": { + const activeRestoreRun = getActiveRestoreRun(event.payload.restoreId, event.type, event.agentId); + if (!activeRestoreRun) { + break; + } + + resolveActiveRestoreRun(activeRestoreRun.restoreId, { + status: "cancelled", + message: activeRestoreRun.cancellationRequested ? undefined : event.payload.message, + }); + break; + } } }; @@ -307,6 +462,52 @@ export const agentManager = { return response.command; }, + startRestore: async (agentId: string, request: AgentStartRestoreRequest): Promise => { + const runtime = getAgentManagerRuntime(); + if (!runtime) { + return { + status: "unavailable", + error: new Error(`Restore agent ${agentId} is not connected`), + }; + } + + if (request.signal.aborted) { + throw request.signal.reason || new Error("Operation aborted"); + } + + const completion = new Promise((resolve) => { + getActiveRestoresByRestoreId().set(request.payload.restoreId, { + agentId, + restoreId: request.payload.restoreId, + onStarted: request.onStarted, + onProgress: request.onProgress, + resolve, + cancellationRequested: false, + }); + }); + + try { + if (!(await Effect.runPromise(runtime.sendRestore(agentId, request.payload)))) { + clearActiveRestoreRun(request.payload.restoreId); + return { + status: "unavailable", + error: new Error(`Failed to send restore command to agent ${agentId}`), + }; + } + + if (request.signal.aborted) { + await requestRestoreCancellation(agentId, request.payload.restoreId); + } + + return { status: "started", result: completion }; + } catch (error) { + clearActiveRestoreRun(request.payload.restoreId); + throw error; + } + }, + cancelRestore: async (agentId: string, restoreId: string) => { + return requestRestoreCancellation(agentId, restoreId); + }, }; export const startLocalAgent = async () => { diff --git a/app/server/modules/agents/controller/server.ts b/app/server/modules/agents/controller/server.ts index 9d7670b8..e5c431b0 100644 --- a/app/server/modules/agents/controller/server.ts +++ b/app/server/modules/agents/controller/server.ts @@ -5,6 +5,8 @@ import type { AgentMessage, BackupCancelPayload, BackupRunPayload, + RestoreCancelPayload, + RestoreRunPayload, VolumeCommand, VolumeCommandResponsePayload, } from "@zerobyte/contracts/agent-protocol"; @@ -22,16 +24,9 @@ type AgentEventContext = { agentName: string; }; -type AgentBackupMessage = Extract< - AgentMessage, - { - type: "backup.started" | "backup.progress" | "backup.completed" | "backup.failed" | "backup.cancelled"; - } ->; - export type AgentManagerEvent = | (AgentEventContext & { type: "agent.disconnected" }) - | (AgentEventContext & AgentBackupMessage); + | (AgentEventContext & AgentMessage); type ControllerAgentSessionHandle = { agentId: string; @@ -86,7 +81,7 @@ export function createAgentManagerRuntime(onEvent: (event: AgentManagerEvent) => return !!session && Effect.runSync(session.isReady()); }; - const handleSessionEvent = (params: { agentId: string; agentName: string; sessionId: string }) => { + const handleSessionEvent = (params: { agentId: string; agentName: string }) => { const { agentId, agentName } = params; return (event: ControllerAgentSessionEvent) => { @@ -109,15 +104,6 @@ export function createAgentManagerRuntime(onEvent: (event: AgentManagerEvent) => case "agent.disconnected": { return Effect.sync(() => onEvent({ type: "agent.disconnected", agentId, agentName })); } - case "restore.cancelled": - case "restore.completed": - case "restore.failed": - case "restore.progress": - case "restore.started": { - // TODO: once we implement the app side - return Effect.void; - } - default: { return Effect.sync(() => onEvent({ ...event, agentId, agentName })); } @@ -135,7 +121,6 @@ export function createAgentManagerRuntime(onEvent: (event: AgentManagerEvent) => handleSessionEvent({ agentId: ws.data.agentId, agentName: ws.data.agentName, - sessionId: ws.data.id, }), ), scope, @@ -368,6 +353,47 @@ export function createAgentManagerRuntime(onEvent: (event: AgentManagerEvent) => logger.info(`Sent backup cancel for command ${payload.jobId} to agent ${agentId}`); return true; }), + sendRestore: (agentId: string, payload: RestoreRunPayload) => + Effect.gen(function* () { + const session = getSession(agentId); + + if (!session) { + logger.warn(`Cannot send restore command. Agent ${agentId} is not connected.`); + return false; + } + + if (!(yield* session.isReady())) { + logger.warn(`Cannot send restore command. Agent ${agentId} is not ready.`); + return false; + } + + if (!(yield* session.sendRestore(payload))) { + logger.warn(`Cannot send restore command. Agent ${agentId} is no longer accepting commands.`); + return false; + } + + logger.info( + `Sent restore command ${payload.restoreId} to agent ${agentId} for snapshot ${payload.snapshotId}`, + ); + return true; + }), + cancelRestore: (agentId: string, payload: RestoreCancelPayload) => + Effect.gen(function* () { + const session = getSession(agentId); + + if (!session) { + logger.warn(`Cannot cancel restore command. Agent ${agentId} is not connected.`); + return false; + } + + if (!(yield* session.sendRestoreCancel(payload))) { + logger.warn(`Cannot cancel restore command. Agent ${agentId} is no longer accepting commands.`); + return false; + } + + logger.info(`Sent restore cancel for command ${payload.restoreId} to agent ${agentId}`); + return true; + }), runVolumeCommand: ( agentId: string, command: VolumeCommand, diff --git a/app/server/modules/agents/controller/session.ts b/app/server/modules/agents/controller/session.ts index 5a68402e..c17fafff 100644 --- a/app/server/modules/agents/controller/session.ts +++ b/app/server/modules/agents/controller/session.ts @@ -7,6 +7,8 @@ import { type BackupCancelPayload, type BackupRunPayload, type ControllerWireMessage, + type RestoreCancelPayload, + type RestoreRunPayload, type VolumeCommand, type VolumeCommandResponsePayload, } from "@zerobyte/contracts/agent-protocol"; @@ -29,22 +31,19 @@ type SessionState = { lastPongAt: number | null; }; -type PendingCommand = { - deferred: Deferred.Deferred; - description: string; -}; +type PendingCommand = { deferred: Deferred.Deferred; description: string }; export type ControllerAgentSessionEvent = | Exclude - | { - type: "agent.disconnected"; - }; + | { type: "agent.disconnected" }; export type ControllerAgentSession = { readonly connectionId: string; handleMessage: (data: string) => Effect.Effect; sendBackup: (payload: BackupRunPayload) => Effect.Effect; sendBackupCancel: (payload: BackupCancelPayload) => Effect.Effect; + sendRestore: (payload: RestoreRunPayload) => Effect.Effect; + sendRestoreCancel: (payload: RestoreCancelPayload) => Effect.Effect; runVolumeCommand: (command: VolumeCommand) => Effect.Effect; isReady: () => Effect.Effect; run: Effect.Effect; @@ -68,7 +67,9 @@ export const createControllerAgentSession = ( Queue.offer(outboundQueue, message).pipe( Effect.catchAllCause((cause) => Effect.sync(() => { - logger.error(`Failed to queue outbound message for agent ${socket.data.agentId}: ${toMessage(cause)}`); + logger.error( + `Failed to queue outbound message for agent ${socket.data.agentId}: ${toMessage(cause)}`, + ); return false; }), ), @@ -189,7 +190,11 @@ export const createControllerAgentSession = ( } case "heartbeat.pong": { const seenAt = Date.now(); - yield* updateState((current) => ({ ...current, lastSeenAt: seenAt, lastPongAt: message.payload.sentAt })); + yield* updateState((current) => ({ + ...current, + lastSeenAt: seenAt, + lastPongAt: message.payload.sentAt, + })); yield* onEvent(message); break; } @@ -216,7 +221,9 @@ export const createControllerAgentSession = ( } if (!parsed.success) { - yield* logger.effect.warn(`Invalid agent message from ${socket.data.agentId}: ${parsed.error.message}`); + yield* logger.effect.warn( + `Invalid agent message from ${socket.data.agentId}: ${parsed.error.message}`, + ); return; } @@ -225,6 +232,8 @@ export const createControllerAgentSession = ( }, sendBackup: (payload) => offerOutbound(createControllerMessage("backup.run", payload)), sendBackupCancel: (payload) => offerOutbound(createControllerMessage("backup.cancel", payload)), + sendRestore: (payload) => offerOutbound(createControllerMessage("restore.run", payload)), + sendRestoreCancel: (payload) => offerOutbound(createControllerMessage("restore.cancel", payload)), runVolumeCommand: (command) => Effect.gen(function* () { const commandId = Bun.randomUUIDv7(); @@ -232,7 +241,9 @@ export const createControllerAgentSession = ( const deferred = yield* Deferred.make(); yield* setPendingCommand(commandId, { deferred, description }); - const queued = yield* offerOutbound(createControllerMessage("volume.command", { commandId, command })); + const queued = yield* offerOutbound( + createControllerMessage("volume.command", { commandId, command }), + ); if (!queued) { yield* removePendingCommand(commandId); return yield* Effect.fail(new Error(`Failed to queue volume command ${command.name}`)); diff --git a/app/server/modules/agents/helpers/runtime-state.dev.ts b/app/server/modules/agents/helpers/runtime-state.dev.ts index daccadbb..89f41f8f 100644 --- a/app/server/modules/agents/helpers/runtime-state.dev.ts +++ b/app/server/modules/agents/helpers/runtime-state.dev.ts @@ -1,20 +1,26 @@ import { createAgentRuntimeState, type AgentRuntimeState } from "./runtime-state"; -type LegacyAgentRuntimeState = Omit & - Partial>; +type RuntimeMapKey = "activeBackupsByScheduleId" | "activeBackupScheduleIdsByJobId" | "activeRestoresByRestoreId"; + +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 hasActiveRuntimeMaps = (runtime: LegacyAgentRuntimeState): runtime is AgentRuntimeState => { + return ( + runtime.activeBackupsByScheduleId instanceof Map && + runtime.activeBackupScheduleIdsByJobId instanceof Map && + runtime.activeRestoresByRestoreId instanceof Map + ); }; const hydrateAgentRuntimeState = (runtime: LegacyAgentRuntimeState): AgentRuntimeState => ({ ...runtime, activeBackupsByScheduleId: runtime.activeBackupsByScheduleId ?? new Map(), activeBackupScheduleIdsByJobId: runtime.activeBackupScheduleIdsByJobId ?? new Map(), + activeRestoresByRestoreId: runtime.activeRestoresByRestoreId ?? new Map(), }); export const getDevAgentRuntimeState = (): AgentRuntimeState => { @@ -27,7 +33,7 @@ export const getDevAgentRuntimeState = (): AgentRuntimeState => { return runtime; } - if (hasActiveBackupMaps(existingRuntime)) { + if (hasActiveRuntimeMaps(existingRuntime)) { return existingRuntime; } diff --git a/app/server/modules/agents/helpers/runtime-state.ts b/app/server/modules/agents/helpers/runtime-state.ts index 13fc2a15..e650a731 100644 --- a/app/server/modules/agents/helpers/runtime-state.ts +++ b/app/server/modules/agents/helpers/runtime-state.ts @@ -1,12 +1,27 @@ import type { ChildProcess } from "node:child_process"; import type { ResticBackupOutputDto } from "@zerobyte/core/restic"; -import type { BackupProgressPayload } from "@zerobyte/contracts/agent-protocol"; +import type { + BackupProgressPayload, + RestoreCompletedPayload, + RestoreProgressPayload, +} 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: "completed"; + exitCode: number; + result: ResticBackupOutputDto | null; + warningDetails: string | null; + } + | { status: "failed"; error: string } + | { status: "cancelled"; message?: string }; +export type RestoreExecutionProgress = RestoreProgressPayload["progress"]; +export type RestoreExecutionResult = + | { status: "unavailable"; error: Error } + | { status: "completed"; result: RestoreCompletedPayload["result"] } | { status: "failed"; error: string } | { status: "cancelled"; message?: string }; @@ -20,6 +35,15 @@ type ActiveBackupRun = { cancellationRequested: boolean; }; +type ActiveRestoreRun = { + agentId: string; + restoreId: string; + onStarted: () => void; + onProgress: (progress: RestoreExecutionProgress) => void; + resolve: (result: RestoreExecutionResult) => void; + cancellationRequested: boolean; +}; + export type AgentRuntimeState = { agentManager: AgentManagerRuntime | null; localAgent: ChildProcess | null; @@ -27,6 +51,7 @@ export type AgentRuntimeState = { localAgentRestartTimeout: ReturnType | null; activeBackupsByScheduleId: Map; activeBackupScheduleIdsByJobId: Map; + activeRestoresByRestoreId: Map; }; export const createAgentRuntimeState = (): AgentRuntimeState => ({ @@ -36,4 +61,5 @@ export const createAgentRuntimeState = (): AgentRuntimeState => ({ localAgentRestartTimeout: null, activeBackupsByScheduleId: new Map(), activeBackupScheduleIdsByJobId: new Map(), + activeRestoresByRestoreId: new Map(), }); diff --git a/app/server/modules/repositories/__tests__/repositories.service.test.ts b/app/server/modules/repositories/__tests__/repositories.service.test.ts index 14ab2450..1381543e 100644 --- a/app/server/modules/repositories/__tests__/repositories.service.test.ts +++ b/app/server/modules/repositories/__tests__/repositories.service.test.ts @@ -4,22 +4,24 @@ import * as os from "node:os"; import nodePath from "node:path"; import { randomUUID } from "node:crypto"; import { Readable } from "node:stream"; +import { Effect } from "effect"; import { afterEach, beforeAll, beforeEach, describe, expect, test, vi } from "vitest"; import type { RepositoryConfig } from "@zerobyte/core/restic"; import { REPOSITORY_BASE } from "~/server/core/constants"; +import { config } from "~/server/core/config"; import { withContext } from "~/server/core/request-context"; import { db } from "~/server/db/db"; -import { repositoriesTable } from "~/server/db/schema"; +import { agentsTable, repositoriesTable, type RepositoryInsert } from "~/server/db/schema"; import { generateShortId } from "~/server/utils/id"; import { restic } from "~/server/core/restic"; +import { agentManager } from "~/server/modules/agents/agents-manager"; import { createTestSession } from "~/test/helpers/auth"; import { createTestBackupSchedule } from "~/test/helpers/backup"; import { cache, cacheKeys } from "~/server/utils/cache"; import { ResticError } from "@zerobyte/core/restic/server"; import { repositoriesService } from "../repositories.service"; -import { Effect } from "effect"; -const createTestRepository = async (organizationId: string) => { +const createTestRepository = async (organizationId: string, overrides: Partial = {}) => { const id = randomUUID(); const shortId = generateShortId(); const [repository] = await db @@ -33,6 +35,7 @@ const createTestRepository = async (organizationId: string) => { compressionMode: "auto", status: "healthy", organizationId, + ...overrides, }) .returning(); return repository; @@ -115,7 +118,11 @@ describe("repositoriesService.createRepository", () => { test("keeps an explicit local repository path unchanged when importing existing repository", async () => { // arrange const explicitPath = `${REPOSITORY_BASE}/custom-${randomUUID()}`; - const config: RepositoryConfig = { backend: "local", path: explicitPath, isExistingRepository: true }; + const config: RepositoryConfig = { + backend: "local", + path: explicitPath, + isExistingRepository: true, + }; vi.spyOn(restic, "snapshots").mockImplementation(() => Effect.succeed([])); @@ -184,7 +191,9 @@ describe("repositoriesService repository stats", () => { expect(refreshed).toEqual(expectedStats); expect(statsSpy).toHaveBeenCalledTimes(1); - const persistedRepository = await db.query.repositoriesTable.findFirst({ where: { id: repository.id } }); + const persistedRepository = await db.query.repositoriesTable.findFirst({ + where: { id: repository.id }, + }); expect(persistedRepository?.stats).toEqual(expectedStats); expect(typeof persistedRepository?.statsUpdatedAt).toBe("number"); @@ -404,7 +413,19 @@ describe("repositoriesService.dumpSnapshot", () => { }); describe("repositoriesService.restoreSnapshot", () => { + let originalEnableLocalAgent: boolean; + const createPendingRestoreStart = () => ({ + status: "started" as const, + result: new Promise(() => {}), + }); + + beforeEach(() => { + originalEnableLocalAgent = config.flags.enableLocalAgent; + config.flags.enableLocalAgent = true; + }); + afterEach(() => { + config.flags.enableLocalAgent = originalEnableLocalAgent; vi.restoreAllMocks(); }); @@ -424,19 +445,8 @@ describe("repositoriesService.restoreSnapshot", () => { ]), ); - const restoreMock = vi.fn(() => - Effect.succeed({ - message_type: "summary" as const, - seconds_elapsed: 1, - percent_done: 100, - files_skipped: 0, - total_files: 1, - files_restored: 1, - total_bytes: 1, - bytes_restored: 1, - }), - ); - vi.spyOn(restic, "restore").mockImplementation(restoreMock); + const restoreMock = vi.fn(() => Promise.resolve(createPendingRestoreStart())); + vi.spyOn(agentManager, "startRestore").mockImplementation(restoreMock); return { organizationId, @@ -446,13 +456,15 @@ describe("repositoriesService.restoreSnapshot", () => { }; }; - test("rejects restore targets inside protected roots", async () => { + test("rejects protected targets even when the local agent is enabled", async () => { const { organizationId, userId, repositoryShortId, restoreMock } = await setupRestoreSnapshotScenario(); const targetPath = nodePath.join(os.tmpdir(), "zerobyte-restore-target"); await expect( withContext({ organizationId, userId }, () => - repositoriesService.restoreSnapshot(repositoryShortId, "snapshot-restore", { targetPath }), + repositoriesService.restoreSnapshot(repositoryShortId, "snapshot-restore", { + targetPath, + }), ), ).rejects.toThrow("Restore target path is not allowed"); @@ -465,25 +477,197 @@ describe("repositoriesService.restoreSnapshot", () => { try { await withContext({ organizationId, userId }, () => - repositoriesService.restoreSnapshot(repositoryShortId, "snapshot-restore", { targetPath }), + repositoriesService.restoreSnapshot(repositoryShortId, "snapshot-restore", { + targetPath, + }), ); } finally { await fs.rm(targetPath, { recursive: true, force: true }); } expect(restoreMock).toHaveBeenCalledWith( + "local", expect.objectContaining({ - backend: "local", - }), - "snapshot-restore", - targetPath, - expect.objectContaining({ - organizationId, - basePath: "/var/lib/zerobyte/volumes/vol123/_data", + payload: expect.objectContaining({ + snapshotId: "snapshot-restore", + target: targetPath, + repositoryConfig: expect.objectContaining({ backend: "local" }), + options: expect.objectContaining({ + organizationId, + basePath: "/var/lib/zerobyte/volumes/vol123/_data", + }), + }), }), ); }); + test("rejects starting a second active restore for the same snapshot", async () => { + const { organizationId, userId, repositoryShortId } = await setupRestoreSnapshotScenario(); + const targetPath = await fs.mkdtemp(nodePath.join(process.cwd(), "restore-target-")); + + try { + await withContext({ organizationId, userId }, () => + repositoriesService.restoreSnapshot(repositoryShortId, "snapshot-restore", { + targetPath, + }), + ); + + await expect( + withContext({ organizationId, userId }, () => + repositoriesService.restoreSnapshot(repositoryShortId, "snapshot-restore", { + targetPath, + }), + ), + ).rejects.toThrow("A restore is already running for this snapshot"); + } finally { + await fs.rm(targetPath, { recursive: true, force: true }); + } + }); + + test("routes restore to the requested target agent", async () => { + const organizationId = session.organizationId; + const agentId = `agent-${randomUUID()}`; + const repository = await createTestRepository(organizationId, { + type: "s3", + config: { + backend: "s3", + endpoint: "https://s3.example.com", + bucket: "bucket", + accessKeyId: "access-key", + secretAccessKey: "secret-key", + }, + }); + await db.insert(agentsTable).values({ + id: agentId, + organizationId, + name: "Remote Agent", + kind: "remote", + status: "online", + capabilities: {}, + updatedAt: Date.now(), + }); + vi.spyOn(restic, "snapshots").mockReturnValue( + Effect.succeed([ + { + id: "snapshot-restore", + short_id: "snapshot-restore", + time: new Date().toISOString(), + paths: ["/var/lib/zerobyte/volumes/vol123/_data"], + hostname: "host", + }, + ]), + ); + const restoreMock = vi.fn(() => Promise.resolve(createPendingRestoreStart())); + vi.spyOn(agentManager, "startRestore").mockImplementation(restoreMock); + const targetPath = await fs.mkdtemp(nodePath.join(process.cwd(), "restore-target-")); + + try { + await withContext({ organizationId, userId: session.user.id }, () => + repositoriesService.restoreSnapshot(repository.shortId, "snapshot-restore", { + targetPath, + targetAgentId: agentId, + }), + ); + } finally { + await fs.rm(targetPath, { recursive: true, force: true }); + } + + expect(restoreMock).toHaveBeenCalledWith( + agentId, + expect.objectContaining({ + payload: expect.objectContaining({ + target: targetPath, + }), + }), + ); + }); + + test("rejects a target agent outside the current organization", async () => { + const organizationId = session.organizationId; + const otherSession = await createTestSession(); + const otherAgentId = `agent-${randomUUID()}`; + const repository = await createTestRepository(organizationId, { + type: "s3", + config: { + backend: "s3", + endpoint: "https://s3.example.com", + bucket: "bucket", + accessKeyId: "access-key", + secretAccessKey: "secret-key", + }, + }); + + await db.insert(agentsTable).values({ + id: otherAgentId, + organizationId: otherSession.organizationId, + name: "Other Org Agent", + kind: "remote", + status: "online", + capabilities: {}, + updatedAt: Date.now(), + }); + vi.spyOn(restic, "snapshots").mockReturnValue( + Effect.succeed([ + { + id: "snapshot-restore", + short_id: "snapshot-restore", + time: new Date().toISOString(), + paths: ["/var/lib/zerobyte/volumes/vol123/_data"], + hostname: "host", + }, + ]), + ); + const restoreMock = vi.fn(() => Promise.resolve(createPendingRestoreStart())); + vi.spyOn(agentManager, "startRestore").mockImplementation(restoreMock); + + await expect( + withContext({ organizationId, userId: session.user.id }, () => + repositoriesService.restoreSnapshot(repository.shortId, "snapshot-restore", { + targetAgentId: otherAgentId, + }), + ), + ).rejects.toThrow("Restore target agent not found"); + + expect(restoreMock).not.toHaveBeenCalled(); + }); + + test("uses controller-local restore fallback when local agent supervision is disabled", async () => { + config.flags.enableLocalAgent = false; + const { organizationId, userId, repositoryShortId, restoreMock } = await setupRestoreSnapshotScenario(); + const resticRestoreMock = vi.spyOn(restic, "restore").mockReturnValue( + Effect.succeed({ + message_type: "summary" as const, + files_skipped: 0, + files_restored: 1, + }), + ); + const targetPath = await fs.mkdtemp(nodePath.join(process.cwd(), "restore-target-")); + + try { + await withContext({ organizationId, userId }, () => + repositoriesService.restoreSnapshot(repositoryShortId, "snapshot-restore", { + targetPath, + }), + ); + } finally { + await fs.rm(targetPath, { recursive: true, force: true }); + } + + expect(restoreMock).not.toHaveBeenCalled(); + await waitForExpect(() => { + expect(resticRestoreMock).toHaveBeenCalledWith( + expect.objectContaining({ backend: "local" }), + "snapshot-restore", + targetPath, + expect.objectContaining({ + organizationId, + basePath: "/var/lib/zerobyte/volumes/vol123/_data", + signal: expect.any(AbortSignal), + }), + ); + }); + }); + test("rejects original-location restore for snapshots with non-posix source paths", async () => { const { organizationId, userId, repositoryShortId, restoreMock } = await setupRestoreSnapshotScenario([ "d:\\some\\path", @@ -509,21 +693,26 @@ describe("repositoriesService.restoreSnapshot", () => { try { await withContext({ organizationId, userId }, () => - repositoriesService.restoreSnapshot(repositoryShortId, "snapshot-restore", { targetPath }), + repositoriesService.restoreSnapshot(repositoryShortId, "snapshot-restore", { + targetPath, + }), ); } finally { await fs.rm(targetPath, { recursive: true, force: true }); } expect(restoreMock).toHaveBeenCalledWith( + "local", expect.objectContaining({ - backend: "local", - }), - "snapshot-restore", - targetPath, - expect.objectContaining({ - organizationId, - basePath: "/", + payload: expect.objectContaining({ + snapshotId: "snapshot-restore", + target: targetPath, + repositoryConfig: expect.objectContaining({ backend: "local" }), + options: expect.objectContaining({ + organizationId, + basePath: "/", + }), + }), }), ); }); @@ -536,9 +725,14 @@ describe("repositoriesService.getRetentionCategories", () => { test("recomputes retention categories after repository cache invalidation", async () => { const organizationId = session.organizationId; - const schedule = await createTestBackupSchedule({ organizationId, retentionPolicy: { keepLast: 1 } }); + const schedule = await createTestBackupSchedule({ + organizationId, + retentionPolicy: { keepLast: 1 }, + }); - const repository = await db.query.repositoriesTable.findFirst({ where: { id: schedule.repositoryId } }); + const repository = await db.query.repositoriesTable.findFirst({ + where: { id: schedule.repositoryId }, + }); expect(repository).toBeTruthy(); if (!repository) { @@ -622,7 +816,9 @@ describe("repositoriesService.deleteSnapshot", () => { expect(statsSpy).toHaveBeenCalledTimes(1); }); - const updatedRepository = await db.query.repositoriesTable.findFirst({ where: { id: repository.id } }); + const updatedRepository = await db.query.repositoriesTable.findFirst({ + where: { id: repository.id }, + }); expect(updatedRepository?.stats).toEqual(expectedStats); expect(typeof updatedRepository?.statsUpdatedAt).toBe("number"); }); diff --git a/app/server/modules/repositories/repositories.controller.ts b/app/server/modules/repositories/repositories.controller.ts index 00ccb83a..2988c393 100644 --- a/app/server/modules/repositories/repositories.controller.ts +++ b/app/server/modules/repositories/repositories.controller.ts @@ -182,7 +182,10 @@ export const repositoriesController = new Hono() const offset = Math.max(0, query.offset ?? 0); const limit = Math.min(1000, Math.max(1, query.limit ?? 500)); - const result = await repositoriesService.listSnapshotFiles(shortId, snapshotId, decodedPath, { offset, limit }); + const result = await repositoriesService.listSnapshotFiles(shortId, snapshotId, decodedPath, { + offset, + limit, + }); c.header("Cache-Control", "max-age=300, stale-while-revalidate=600"); @@ -235,7 +238,7 @@ export const repositoriesController = new Hono() const { snapshotId, ...options } = c.req.valid("json"); const result = await repositoriesService.restoreSnapshot(shortId, snapshotId, options); - return c.json(result, 200); + return c.json(result, 202); }) .post("/:shortId/doctor", startDoctorDto, async (c) => { const shortId = asShortId(c.req.param("shortId")); diff --git a/app/server/modules/repositories/repositories.dto.ts b/app/server/modules/repositories/repositories.dto.ts index 79ca2d00..d0704ccc 100644 --- a/app/server/modules/repositories/repositories.dto.ts +++ b/app/server/modules/repositories/repositories.dto.ts @@ -343,14 +343,13 @@ export const restoreSnapshotBody = z.object({ excludeXattr: z.array(z.string()).optional(), delete: z.boolean().optional(), targetPath: z.string().optional(), + targetAgentId: z.string().optional(), overwrite: overwriteModeSchema.optional(), }); const restoreSnapshotResponse = z.object({ - success: z.boolean(), - message: z.string(), - filesRestored: z.number(), - filesSkipped: z.number(), + restoreId: z.string(), + status: z.literal("started"), }); export type RestoreSnapshotDto = z.infer; @@ -360,8 +359,8 @@ export const restoreSnapshotDto = describeRoute({ tags: ["Repositories"], operationId: "restoreSnapshot", responses: { - 200: { - description: "Snapshot restored successfully", + 202: { + description: "Snapshot restore started", content: { "application/json": { schema: resolver(restoreSnapshotResponse), diff --git a/app/server/modules/repositories/repositories.service.ts b/app/server/modules/repositories/repositories.service.ts index 61eeaeb4..fb0c1566 100644 --- a/app/server/modules/repositories/repositories.service.ts +++ b/app/server/modules/repositories/repositories.service.ts @@ -19,7 +19,7 @@ import { logger } from "@zerobyte/core/node"; import { parseRetentionCategories, type RetentionCategory } from "~/server/utils/retention-categories"; import { repoMutex } from "../../core/repository-mutex"; import { db } from "../../db/db"; -import { repositoriesTable, type Repository } from "../../db/schema"; +import { repositoriesTable, type Repository, type RepositoryInsert } from "../../db/schema"; import { cache, cacheKeys } from "../../utils/cache"; import { runEffectPromise, toMessage } from "../../utils/errors"; import { generateShortId } from "../../utils/id"; @@ -30,12 +30,22 @@ import type { DumpPathKind, UpdateRepositoryBody } from "./repositories.dto"; import { findCommonAncestor } from "@zerobyte/core/utils"; import { prepareSnapshotDump } from "./helpers/dump"; import { executeDoctor } from "./helpers/doctor"; +import { restoreExecutor } from "./restore-executor"; import type { ShortId } from "~/server/utils/branded"; import { decryptRepositoryConfig, encryptRepositoryConfig } from "./repository-config-secrets"; import { getScheduleByIdOrShortId } from "../backups/helpers/backup-schedule-lookups"; +import type { RestoreExecutionProgress, RestoreExecutionResult } from "../agents/agents-manager"; +import { agentsService } from "../agents/agents.service"; +import { LOCAL_AGENT_ID } from "../agents/constants"; +import { taskStore } from "../tasks/tasks.store"; +import type { ParsedTask, TaskInput } from "../tasks/tasks.schemas"; import { Effect } from "effect"; const runningDoctors = new Map(); +const RESTORE_TASK_RESOURCE_TYPE = "repository"; + +type RestoreTaskInput = Extract; +type RestoreTask = ParsedTask & { kind: "restore"; input: RestoreTaskInput }; const emptyRepositoryStats: ResticStatsDto = { total_size: 0, @@ -56,6 +66,175 @@ const getBlockedRestoreTargets = () => { ].filter((e) => e !== undefined); }; +const assertAllowedControllerLocalRestoreTarget = (target: string) => { + const resolvedTarget = nodePath.resolve(target); + + for (const blockedTarget of getBlockedRestoreTargets()) { + if (isPathWithin(blockedTarget, resolvedTarget)) { + throw new BadRequestError( + "Restore target path is not allowed. Restoring to this path could overwrite critical system files or application data.", + ); + } + } +}; + +const isRestoreTask = (task: ParsedTask): task is RestoreTask => + task.kind === "restore" && task.input.kind === "restore"; + +const asRestoreTask = (task: ParsedTask, restoreId: string, eventName: string) => { + if (!isRestoreTask(task)) { + logger.warn(`Received ${eventName} for non-restore task ${restoreId}`); + return null; + } + + return task; +}; + +const updateActiveRestoreTask = (restoreId: string, eventName: string, update: () => ParsedTask) => { + try { + return asRestoreTask(update(), restoreId, eventName); + } catch (error) { + logger.warn(`Received ${eventName} for inactive restore ${restoreId}: ${toMessage(error)}`); + return null; + } +}; + +const findActiveRestoreTask = ( + organizationId: string, + repositoryShortId: string, + snapshotId: string, +): RestoreTask | null => { + return ( + taskStore + .listActiveByResource({ + organizationId, + kind: "restore", + resourceType: RESTORE_TASK_RESOURCE_TYPE, + resourceId: repositoryShortId, + }) + .find((task): task is RestoreTask => isRestoreTask(task) && task.input.snapshotId === snapshotId) ?? null + ); +}; + +const emitRestoreStarted = (task: RestoreTask) => { + serverEvents.emit("restore:started", { + restoreId: task.id, + organizationId: task.organizationId, + repositoryId: task.input.repositoryId, + snapshotId: task.input.snapshotId, + }); +}; + +const emitRestoreProgress = (task: RestoreTask, progress: RestoreExecutionProgress) => { + serverEvents.emit("restore:progress", { + restoreId: task.id, + organizationId: task.organizationId, + repositoryId: task.input.repositoryId, + snapshotId: task.input.snapshotId, + ...progress, + }); +}; + +const emitRestoreCompleted = ( + task: RestoreTask, + payload: { + status: "success" | "error" | "cancelled"; + error?: string; + filesRestored?: number; + filesSkipped?: number; + }, +) => { + serverEvents.emit("restore:completed", { + restoreId: task.id, + organizationId: task.organizationId, + repositoryId: task.input.repositoryId, + snapshotId: task.input.snapshotId, + ...payload, + }); +}; + +const markRestoreStarted = (restoreId: string) => { + const task = updateActiveRestoreTask(restoreId, "restore.started", () => taskStore.markRunning(restoreId)); + if (!task) return; + + emitRestoreStarted(task); +}; + +const updateRestoreProgress = (restoreId: string, progress: RestoreExecutionProgress) => { + const task = updateActiveRestoreTask(restoreId, "restore.progress", () => + taskStore.updateProgress(restoreId, { kind: "restore", progress }), + ); + if (!task) return; + + emitRestoreProgress(task, progress); +}; + +const completeRestoreTask = ( + restoreId: string, + result: Extract["result"], +) => { + const task = updateActiveRestoreTask(restoreId, "restore.completed", () => + taskStore.complete(restoreId, { kind: "restore", result }), + ); + if (!task) return; + + emitRestoreCompleted(task, { + status: "success", + filesRestored: result.files_restored, + filesSkipped: result.files_skipped, + }); +}; + +const failRestoreTask = (restoreId: string, error: string) => { + const task = updateActiveRestoreTask(restoreId, "restore.failed", () => taskStore.fail(restoreId, error)); + if (!task) return; + + emitRestoreCompleted(task, { status: "error", error }); +}; + +const cancelRestoreTask = (restoreId: string, message?: string) => { + const task = updateActiveRestoreTask(restoreId, "restore.cancelled", () => + taskStore.cancel(restoreId, message ?? null), + ); + if (!task) return; + + emitRestoreCompleted(task, { status: "cancelled", error: task.cancellationRequested ? undefined : message }); +}; + +const finishRestoreExecution = async (restoreId: string, resultPromise: Promise) => { + try { + const result = await resultPromise; + + switch (result.status) { + case "completed": + completeRestoreTask(restoreId, result.result); + return; + case "failed": + failRestoreTask(restoreId, result.error); + return; + case "cancelled": + cancelRestoreTask(restoreId, result.message); + return; + case "unavailable": + failRestoreTask(restoreId, result.error.message); + return; + } + } catch (error) { + failRestoreTask(restoreId, toMessage(error)); + } +}; + +const assertAllowedRestoreAgent = async (agentId: string, organizationId: string) => { + if (agentId === LOCAL_AGENT_ID) { + return; + } + + const agent = await agentsService.getAgent(agentId); + if (!agent || agent.organizationId !== organizationId) { + throw new NotFoundError("Restore target agent not found"); + } +}; + const findRepository = async (shortId: ShortId) => { const organizationId = getOrganizationId(); return await db.query.repositoriesTable.findFirst({ @@ -268,7 +447,13 @@ const listSnapshotFiles = async ( const cacheKey = cacheKeys.repository.ls(repository.id, snapshotId, path, offset, limit); type LsResult = { - snapshot: { id: string; short_id: string; time: string; hostname: string; paths: string[] } | null; + snapshot: { + id: string; + short_id: string; + time: string; + hostname: string; + paths: string[]; + } | null; nodes: { name: string; type: string; path: string; size?: number; mode?: number }[]; pagination: { offset: number; limit: number; total: number; hasMore: boolean }; }; @@ -327,6 +512,7 @@ const restoreSnapshot = async ( excludeXattr?: string[]; delete?: boolean; targetPath?: string; + targetAgentId?: string; overwrite?: OverwriteMode; }, ) => { @@ -337,77 +523,71 @@ const restoreSnapshot = async ( throw new NotFoundError("Repository not found"); } - const target = options?.targetPath || "/"; - const resolvedTarget = nodePath.resolve(target); - const blockedTargets = getBlockedRestoreTargets(); + const { targetAgentId, targetPath, ...restoreExecutionOptions } = options ?? {}; + const target = targetPath || "/"; - for (const blockedTarget of blockedTargets) { - if (isPathWithin(blockedTarget, resolvedTarget)) { - throw new BadRequestError( - "Restore target path is not allowed. Restoring to this path could overwrite critical system files or application data.", - ); - } - } + const snapshot = await getSnapshotDetails(repository.shortId, snapshotId); + const hasNonPosixSnapshotPaths = snapshot.paths.some((path) => !path.startsWith("/")); - const { paths } = await getSnapshotDetails(repository.shortId, snapshotId); - const hasNonPosixSnapshotPaths = paths.some((path) => !path.startsWith("/")); - - if (hasNonPosixSnapshotPaths && !options?.targetPath) { + if (hasNonPosixSnapshotPaths && !targetPath) { throw new BadRequestError( "Original location restore is unavailable for this snapshot. Restore it to a custom location instead.", ); } - const basePath = hasNonPosixSnapshotPaths ? "/" : findCommonAncestor(paths); + const basePath = hasNonPosixSnapshotPaths ? "/" : findCommonAncestor(snapshot.paths); + const executionAgentId = targetAgentId ?? LOCAL_AGENT_ID; + const useControllerLocalRestoreFallback = executionAgentId === LOCAL_AGENT_ID && !appConfig.flags.enableLocalAgent; + await assertAllowedRestoreAgent(executionAgentId, organizationId); - const releaseLock = await repoMutex.acquireShared(repository.id, `restore:${snapshotId}`); - try { - serverEvents.emit("restore:started", { - organizationId, - repositoryId: repository.shortId, - snapshotId, - }); - - const result = await runEffectPromise( - restic.restore(repository.config, snapshotId, target, { - basePath, - ...options, - organizationId, - onProgress: (progress) => { - serverEvents.emit("restore:progress", { - organizationId, - repositoryId: repository.shortId, - snapshotId, - ...progress, - }); - }, - }), + if (!useControllerLocalRestoreFallback && repository.type === "local" && executionAgentId !== LOCAL_AGENT_ID) { + throw new BadRequestError( + "Local repository restores must run on the agent that can access the repository path.", ); + } - serverEvents.emit("restore:completed", { + if (executionAgentId === LOCAL_AGENT_ID) { + assertAllowedControllerLocalRestoreTarget(target); + } + + const activeRestore = findActiveRestoreTask(organizationId, repository.shortId, snapshotId); + if (activeRestore) { + throw new ConflictError("A restore is already running for this snapshot"); + } + + const task = taskStore.create({ + organizationId, + resourceType: RESTORE_TASK_RESOURCE_TYPE, + resourceId: repository.shortId, + targetAgentId: useControllerLocalRestoreFallback ? null : executionAgentId, + input: { kind: "restore", repositoryId: repository.shortId, snapshotId, target }, + }); + const restoreId = task.id; + try { + const repositoryConfig = await decryptRepositoryConfig(repository.config); + const execution = await restoreExecutor.start({ + restoreId, organizationId, - repositoryId: repository.shortId, + repositoryId: repository.id, + repositoryShortId: repository.shortId, + repositoryConfig, snapshotId, - status: "success", + target, + executionAgentId, + options: { + basePath, + ...restoreExecutionOptions, + }, + onStarted: () => markRestoreStarted(restoreId), + onProgress: (progress) => updateRestoreProgress(restoreId, progress), }); - return { - success: true, - message: "Snapshot restored successfully", - filesRestored: result.files_restored, - filesSkipped: result.files_skipped, - }; + void finishRestoreExecution(restoreId, execution.result); + + return { restoreId, status: "started" as const }; } catch (error) { - serverEvents.emit("restore:completed", { - organizationId, - repositoryId: repository.shortId, - snapshotId, - status: "error", - error: toMessage(error), - }); + failRestoreTask(restoreId, toMessage(error)); throw error; - } finally { - releaseLock(); } }; @@ -424,7 +604,11 @@ const dumpSnapshot = async (shortId: ShortId, snapshotId: string, path?: string, try { const snapshot = await getSnapshotDetails(repository.shortId, snapshotId); - const preparedDump = prepareSnapshotDump({ snapshotId, snapshotPaths: snapshot.paths, requestedPath: path }); + const preparedDump = prepareSnapshotDump({ + snapshotId, + snapshotPaths: snapshot.paths, + requestedPath: path, + }); const dumpOptions: Parameters[2] = { organizationId, path: preparedDump.path, @@ -738,23 +922,22 @@ const updateRepository = async (shortId: ShortId, updates: UpdateRepositoryBody) const configChanged = updates.config && JSON.stringify(decryptedExisting) !== JSON.stringify(parsedConfig); const encryptedConfig = updates.config ? await encryptRepositoryConfig(parsedConfig) : existingConfig; const updatedAt = Date.now(); - const updatePayload = { + const updatePayload: Partial = { name: newName, compressionMode: updates.compressionMode ?? existing.compressionMode, updatedAt, config: encryptedConfig, - ...(configChanged - ? { - status: "unknown" as const, - lastChecked: null, - lastError: null, - doctorResult: null, - stats: null, - statsUpdatedAt: null, - } - : {}), }; + if (configChanged) { + updatePayload.status = "unknown"; + updatePayload.lastChecked = null; + updatePayload.lastError = null; + updatePayload.doctorResult = null; + updatePayload.stats = null; + updatePayload.statsUpdatedAt = null; + } + const [updated] = await db .update(repositoriesTable) .set(updatePayload) @@ -816,7 +999,14 @@ const execResticCommand = async ( addCommonArgs(resticArgs, env, repository.config); try { - const result = await safeSpawn({ command: "restic", args: resticArgs, env, signal, onStdout, onStderr }); + const result = await safeSpawn({ + command: "restic", + args: resticArgs, + env, + signal, + onStdout, + onStderr, + }); return { exitCode: result.exitCode }; } finally { await cleanupTemporaryKeys(env, resticDeps); diff --git a/app/server/modules/repositories/restore-executor.ts b/app/server/modules/repositories/restore-executor.ts new file mode 100644 index 00000000..ae8b7114 --- /dev/null +++ b/app/server/modules/repositories/restore-executor.ts @@ -0,0 +1,122 @@ +import type { RepositoryConfig } from "@zerobyte/core/restic"; +import type { RestoreRunPayload } from "@zerobyte/contracts/agent-protocol"; +import { config as appConfig } from "~/server/core/config"; +import { repoMutex } from "../../core/repository-mutex"; +import { restic, resticDeps } from "../../core/restic"; +import { runEffectPromise, toMessage } from "../../utils/errors"; +import { agentManager, type RestoreExecutionProgress, type RestoreExecutionResult } from "../agents/agents-manager"; +import { LOCAL_AGENT_ID } from "../agents/constants"; + +type RestoreExecutionOptions = Omit[3], "organizationId" | "signal" | "onProgress">; + +type RestoreExecutionRequest = { + restoreId: string; + organizationId: string; + repositoryId: string; + repositoryShortId: string; + repositoryConfig: RepositoryConfig; + snapshotId: string; + target: string; + executionAgentId: string; + options: RestoreExecutionOptions; + onStarted: () => void; + onProgress: (progress: RestoreExecutionProgress) => void; +}; + +type RestoreExecutionHandle = { + result: Promise; +}; + +const shouldRunInController = (agentId: string) => agentId === LOCAL_AGENT_ID && !appConfig.flags.enableLocalAgent; + +const createRestoreRunPayload = async (request: RestoreExecutionRequest): Promise => { + const encryptedResticPassword = await resticDeps.getOrganizationResticPassword(request.organizationId); + const resticPassword = await resticDeps.resolveSecret(encryptedResticPassword); + + return { + restoreId: request.restoreId, + organizationId: request.organizationId, + repositoryId: request.repositoryShortId, + snapshotId: request.snapshotId, + target: request.target, + repositoryConfig: request.repositoryConfig, + runtime: { password: resticPassword }, + options: { + ...request.options, + organizationId: request.organizationId, + }, + }; +}; + +const executeControllerRestore = async ( + request: RestoreExecutionRequest, + signal: AbortSignal, +): Promise => { + if (signal.aborted) { + return { status: "cancelled", message: "Restore was cancelled" }; + } + + request.onStarted(); + + try { + const result = await runEffectPromise( + restic.restore(request.repositoryConfig, request.snapshotId, request.target, { + ...request.options, + organizationId: request.organizationId, + signal, + onProgress: request.onProgress, + }), + ); + + return { status: "completed", result }; + } catch (error) { + if (signal.aborted) { + return { status: "cancelled", message: "Restore was cancelled" }; + } + + return { status: "failed", error: toMessage(error) }; + } +}; + +export const restoreExecutor = { + start: async (request: RestoreExecutionRequest): Promise => { + const abortController = new AbortController(); + + let releaseLock: (() => void) | null = null; + try { + releaseLock = await repoMutex.acquireShared( + request.repositoryId, + `restore:${request.restoreId}`, + abortController.signal, + ); + + let result: Promise; + if (shouldRunInController(request.executionAgentId)) { + result = executeControllerRestore(request, abortController.signal); + } else { + const payload = await createRestoreRunPayload(request); + const started = await agentManager.startRestore(request.executionAgentId, { + payload, + signal: abortController.signal, + onStarted: request.onStarted, + onProgress: request.onProgress, + }); + + if (started.status === "unavailable") { + throw started.error; + } + + result = started.result; + } + + return { + result: result.finally(() => { + releaseLock?.(); + }), + }; + } catch (error) { + releaseLock?.(); + throw error; + } + }, +}; diff --git a/app/server/modules/tasks/__tests__/tasks.store.test.ts b/app/server/modules/tasks/__tests__/tasks.store.test.ts index 04057ff5..3a6ed237 100644 --- a/app/server/modules/tasks/__tests__/tasks.store.test.ts +++ b/app/server/modules/tasks/__tests__/tasks.store.test.ts @@ -6,14 +6,18 @@ import { generateBackupOutput } from "~/test/helpers/restic"; import { taskStore } from "../tasks.store"; import type { TaskInput, TaskProgress, TaskResult } from "../tasks.schemas"; -const backupInput = (scheduleId = 1): TaskInput => ({ +type BackupTaskInput = Extract; +type BackupTaskProgress = Extract; +type BackupTaskResult = Extract; + +const backupInput = (scheduleId = 1): BackupTaskInput => ({ kind: "backup", scheduleId, scheduleShortId: `schedule-${scheduleId}`, manual: false, }); -const backupProgress = (percentDone = 0.5): TaskProgress => ({ +const backupProgress = (percentDone = 0.5): BackupTaskProgress => ({ kind: "backup", progress: { message_type: "status", @@ -28,7 +32,7 @@ const backupProgress = (percentDone = 0.5): TaskProgress => ({ }, }); -const backupResult = (): TaskResult => ({ +const backupResult = (): BackupTaskResult => ({ kind: "backup", exitCode: 0, result: JSON.parse(generateBackupOutput()), @@ -45,6 +49,16 @@ const createBackupTask = (overrides: Partial ...overrides, }); +const createRestoreTask = (overrides: Partial[0]> = {}) => + taskStore.create({ + organizationId: TEST_ORG_ID, + resourceType: "repository", + resourceId: "repo-short", + targetAgentId: "local", + input: { kind: "restore", repositoryId: "repo-short", snapshotId: "snapshot-1", target: "/tmp/restore" }, + ...overrides, + }); + beforeEach(async () => { await ensureTestOrganization(); await db.delete(tasksTable); @@ -97,6 +111,10 @@ test("moves an active task through running, progress, cancellation request, and const completed = taskStore.complete(task.id, backupResult()); expect(completed.status).toBe("succeeded"); + expect(completed.result?.kind).toBe("backup"); + if (completed.result?.kind !== "backup") { + throw new Error("Expected backup result"); + } expect(completed.result?.result?.snapshot_id).toBe("abcd1234"); expect(completed.finishedAt).toEqual(expect.any(Number)); expect(completed.cancellationRequested).toBe(true); @@ -117,6 +135,44 @@ test("records failed and cancelled terminal task states", () => { expect(cancelled.cancellationRequested).toBe(true); }); +test("moves restore tasks through progress and success", () => { + const task = createRestoreTask({ id: "restore-task" }); + + const running = taskStore.markRunning(task.id); + expect(running.kind).toBe("restore"); + expect(running.input).toMatchObject({ kind: "restore", snapshotId: "snapshot-1" }); + + const progressed = taskStore.updateProgress(task.id, { + kind: "restore", + progress: { + message_type: "status", + seconds_elapsed: 2, + percent_done: 0.25, + total_files: 4, + files_restored: 1, + total_bytes: 400, + bytes_restored: 100, + }, + }); + expect(progressed.progress?.progress.percent_done).toBe(0.25); + + const completed = taskStore.complete(task.id, { + kind: "restore", + result: { + message_type: "summary", + total_files: 4, + files_restored: 4, + files_skipped: 0, + }, + }); + expect(completed.status).toBe("succeeded"); + expect(completed.result?.kind).toBe("restore"); + if (completed.result?.kind !== "restore") { + throw new Error("Expected restore result"); + } + expect(completed.result?.result.files_restored).toBe(4); +}); + test("finds the newest active task for a resource and marks only matching active tasks stale", async () => { createBackupTask({ id: "task-a-old", resourceId: "shared" }); const newest = createBackupTask({ id: "task-z-new", resourceId: "shared" }); diff --git a/app/server/modules/tasks/tasks.schemas.ts b/app/server/modules/tasks/tasks.schemas.ts index d707a119..e9902cba 100644 --- a/app/server/modules/tasks/tasks.schemas.ts +++ b/app/server/modules/tasks/tasks.schemas.ts @@ -1,4 +1,9 @@ -import { resticBackupOutputSchema, resticBackupProgressSchema } from "@zerobyte/core/restic"; +import { + resticBackupOutputSchema, + resticBackupProgressSchema, + resticRestoreOutputSchema, + restoreProgressSchema, +} from "@zerobyte/core/restic"; import { z } from "zod"; export const taskStatuses = ["queued", "running", "cancelling", "cancelled", "succeeded", "failed", "stale"] as const; @@ -6,7 +11,7 @@ export const activeTaskStatuses = ["queued", "running", "cancelling"] as const; export const taskStatusSchema = z.enum(taskStatuses); export const activeTaskStatusSchema = z.enum(activeTaskStatuses); -export const taskKindSchema = z.enum(["backup"]); +export const taskKindSchema = z.enum(["backup", "restore"]); export const taskInputSchema = z.discriminatedUnion("kind", [ z.object({ @@ -15,6 +20,12 @@ export const taskInputSchema = z.discriminatedUnion("kind", [ scheduleShortId: z.string(), manual: z.boolean(), }), + z.object({ + kind: z.literal("restore"), + repositoryId: z.string(), + snapshotId: z.string(), + target: z.string(), + }), ]); export const taskProgressSchema = z.discriminatedUnion("kind", [ @@ -22,6 +33,10 @@ export const taskProgressSchema = z.discriminatedUnion("kind", [ kind: z.literal("backup"), progress: resticBackupProgressSchema, }), + z.object({ + kind: z.literal("restore"), + progress: restoreProgressSchema, + }), ]); export const taskResultSchema = z.discriminatedUnion("kind", [ @@ -31,6 +46,10 @@ export const taskResultSchema = z.discriminatedUnion("kind", [ result: resticBackupOutputSchema.nullable(), warningDetails: z.string().nullable(), }), + z.object({ + kind: z.literal("restore"), + result: resticRestoreOutputSchema, + }), ]); export const taskSchema = z diff --git a/app/server/modules/tasks/tasks.store.ts b/app/server/modules/tasks/tasks.store.ts index 1727259b..f7bd3b73 100644 --- a/app/server/modules/tasks/tasks.store.ts +++ b/app/server/modules/tasks/tasks.store.ts @@ -30,9 +30,7 @@ type CreateTaskParams = { input: TaskInput; }; -type MarkActiveStaleParams = Partial & { - error?: string; -}; +type MarkActiveStaleParams = Partial & { error?: string }; const parseTask = (row: unknown): ParsedTask => taskSchema.parse(row); @@ -188,6 +186,17 @@ export const taskStore = { return row ? parseTask(row) : null; }, + listActiveByResource: (params: TaskResource): ParsedTask[] => { + const rows = db + .select() + .from(tasksTable) + .where(and(...buildActiveConditions(params))) + .orderBy(desc(tasksTable.createdAt), desc(tasksTable.id)) + .all(); + + return rows.map(parseTask); + }, + markActiveStale: (params: MarkActiveStaleParams = {}): ParsedTask[] => { const now = Date.now(); const rows = db