refactor(restore): move to task stream system

app/client/components/__test__/restore-form.test.tsx

@@ -1,6 +1,9 @@
 import { afterEach, beforeEach, describe, expect, test, vi } from "vitest";
+import type { ListTasksResponse } from "~/client/api-client";
+import { restoreTasksOptions } from "~/client/modules/repositories/restore-tasks";
 import { HttpResponse, http, server } from "~/test/msw/server";
-import { cleanup, render, screen, userEvent, waitFor, within } from "~/test/test-utils";
+import { cleanup, createTestQueryClient, render, screen, userEvent, waitFor, within } from "~/test/test-utils";
+import { taskChangedEventName } from "~/schemas/task-events";
 import { fromAny } from "@total-typescript/shoehorn";
 
 vi.mock("@tanstack/react-router", async (importOriginal) => {
@@ -15,25 +18,273 @@ vi.mock("@tanstack/react-router", async (importOriginal) => {
 import { RestoreForm } from "../restore-form";
 
 class MockEventSource {
-	addEventListener() {}
-	close() {}
-	onerror: ((event: Event) => void) | null = null;
+	static instances: MockEventSource[] = [];
 
-	constructor(public url: string) {}
+	onerror: ((event: Event) => void) | null = null;
+	private listeners = new Map<string, Set<(event: Event) => void>>();
+
+	constructor(public url: string) {
+		MockEventSource.instances.push(this);
+	}
+
+	addEventListener(type: string, listener: EventListenerOrEventListenerObject) {
+		const listeners = this.listeners.get(type) ?? new Set<(event: Event) => void>();
+		const callback = typeof listener === "function" ? listener : (event: Event) => listener.handleEvent(event);
+		listeners.add(callback);
+		this.listeners.set(type, listeners);
+	}
+
+	emit(type: string, data: unknown) {
+		const event = new MessageEvent(type, {
+			data: JSON.stringify(data),
+		});
+
+		for (const listener of this.listeners.get(type) ?? []) {
+			listener(event);
+		}
+	}
+
+	close() {}
+
+	static reset() {
+		MockEventSource.instances = [];
+	}
 }
 
 const originalEventSource = globalThis.EventSource;
+const repositoryId = "repo-1";
+const snapshotId = "snap-1";
+type TaskResponse = ListTasksResponse[number];
+type RestoreProgress = Extract<NonNullable<TaskResponse["progress"]>, { kind: "restore" }>["progress"];
+type RestoreResult = Extract<NonNullable<TaskResponse["result"]>, { kind: "restore" }>;
+
+const createRestoreTask = (
+	options: {
+		id?: string;
+		snapshotId?: string;
+		status?: TaskResponse["status"];
+		progress?: RestoreProgress | null;
+		result?: RestoreResult | null;
+		error?: string | null;
+		updatedAt?: number;
+	} = {},
+): TaskResponse => {
+	const status = options.status ?? "running";
+	const updatedAt = options.updatedAt ?? 1711411200000;
+	const isTerminal = status === "cancelled" || status === "succeeded" || status === "failed" || status === "stale";
+
+	return {
+		id: options.id ?? "task-restore",
+		kind: "restore",
+		status,
+		resourceType: "repository",
+		resourceId: repositoryId,
+		targetAgentId: null,
+		input: {
+			kind: "restore",
+			repositoryId,
+			snapshotId: options.snapshotId ?? snapshotId,
+			target: "/",
+		},
+		progress: options.progress ? { kind: "restore", progress: options.progress } : null,
+		result: options.result ?? null,
+		error: options.error ?? null,
+		cancellationRequested: status === "cancelled",
+		createdAt: 1711411200000,
+		startedAt: 1711411200000,
+		updatedAt,
+		finishedAt: isTerminal ? updatedAt : null,
+	};
+};
+
+const snapshotFilesHandler = http.get("/api/v1/repositories/:shortId/snapshots/:snapshotId/files", () => {
+	return HttpResponse.json({
+		files: [
+			{ name: "project", path: "/mnt/project", type: "dir" },
+			{ name: "a.txt", path: "/mnt/project/a.txt", type: "file" },
+		],
+	});
+});
+
+const renderRestoreForm = (queryClient = createTestQueryClient()) => {
+	return render(
+		<RestoreForm
+			repository={fromAny({ shortId: repositoryId, name: "Repo 1" })}
+			snapshotId={snapshotId}
+			returnPath={`/repositories/${repositoryId}/${snapshotId}`}
+			queryBasePath="/mnt/project"
+			displayBasePath="/mnt"
+		/>,
+		{ queryClient },
+	);
+};
+
+const renderRestoreFormWithPrefetchedTasks = async (tasks: ListTasksResponse) => {
+	server.use(
+		http.get("/api/v1/tasks", () => HttpResponse.json(tasks)),
+		snapshotFilesHandler,
+	);
+	const queryClient = createTestQueryClient();
+	await queryClient.ensureQueryData(restoreTasksOptions(repositoryId));
+	renderRestoreForm(queryClient);
+
+	await waitFor(() => {
+		expect(MockEventSource.instances).toHaveLength(1);
+	});
+
+	return MockEventSource.instances[0];
+};
 
 beforeEach(() => {
+	MockEventSource.reset();
 	globalThis.EventSource = MockEventSource as unknown as typeof EventSource;
 });
 
 afterEach(() => {
 	globalThis.EventSource = originalEventSource;
 	cleanup();
+	MockEventSource.reset();
 });
 
 describe("RestoreForm", () => {
+	test("recovers the matching active restore from the prefetched task snapshot", async () => {
+		const matchingRestore = createRestoreTask({ updatedAt: 1711411200000 });
+		const newerRestoreForAnotherSnapshot = createRestoreTask({
+			id: "task-other-snapshot",
+			snapshotId: "snap-2",
+			updatedAt: 1711411201000,
+		});
+
+		const eventSource = await renderRestoreFormWithPrefetchedTasks([
+			newerRestoreForAnotherSnapshot,
+			matchingRestore,
+		]);
+
+		const restoreButton = await screen.findByRole("button", { name: "Restoring..." });
+		expect(restoreButton.hasAttribute("disabled")).toBe(true);
+		expect(screen.getByText("Restore in progress")).toBeTruthy();
+		expect(eventSource?.url).toBe("/api/v1/tasks/events?kind=restore&resourceType=repository&resourceId=repo-1");
+	});
+
+	test("renders canonical task progress without creating another event stream", async () => {
+		const activeRestore = createRestoreTask();
+		const eventSource = await renderRestoreFormWithPrefetchedTasks([activeRestore]);
+
+		eventSource?.emit(
+			taskChangedEventName,
+			createRestoreTask({
+				progress: {
+					message_type: "status",
+					seconds_elapsed: 2,
+					percent_done: 0.25,
+					total_files: 4,
+					files_restored: 1,
+					total_bytes: 400,
+					bytes_restored: 100,
+				},
+				updatedAt: activeRestore.updatedAt + 1,
+			}),
+		);
+
+		expect(await screen.findByText("25%")).toBeTruthy();
+		expect(screen.getByText("1 / 4")).toBeTruthy();
+		expect(screen.getByText("2s")).toBeTruthy();
+		expect(screen.getByText("50 B/s")).toBeTruthy();
+		expect(MockEventSource.instances).toHaveLength(1);
+	});
+
+	test("keeps restoring feedback until the started task is reconciled", async () => {
+		let resolveRestoreResponse: (response: Response) => void = () => {};
+		const pendingRestoreResponse = new Promise<Response>((resolve) => {
+			resolveRestoreResponse = resolve;
+		});
+		server.use(
+			snapshotFilesHandler,
+			http.post("/api/v1/repositories/:shortId/restore", () => pendingRestoreResponse),
+		);
+		const { queryClient } = renderRestoreForm();
+
+		await userEvent.click(screen.getByRole("button", { name: "Restore All" }));
+
+		const restoreButton = await screen.findByRole("button", { name: "Restoring..." });
+		expect(restoreButton.hasAttribute("disabled")).toBe(true);
+		expect(screen.getByText("Restore in progress")).toBeTruthy();
+
+		resolveRestoreResponse(HttpResponse.json({ restoreId: "task-restore", status: "started" }, { status: 202 }));
+		await waitFor(() => {
+			expect(queryClient.getMutationCache().getAll()[0]?.state.status).toBe("success");
+		});
+
+		expect(screen.getByRole("button", { name: "Restoring..." }).hasAttribute("disabled")).toBe(true);
+		expect(screen.queryByRole("button", { name: "Restore All" })).toBeNull();
+
+		MockEventSource.instances[0]?.emit(taskChangedEventName, createRestoreTask());
+
+		expect(await screen.findByRole("button", { name: "Restoring..." })).toBeTruthy();
+	});
+
+	test.each([
+		{
+			status: "succeeded" as const,
+			result: {
+				kind: "restore" as const,
+				result: {
+					message_type: "summary" as const,
+					files_restored: 4,
+					files_skipped: 0,
+				},
+			},
+			error: null,
+			title: "Restore completed",
+			description: "Snapshot snap-1 was restored successfully.",
+		},
+		{
+			status: "failed" as const,
+			result: null,
+			error: "Restic restore failed",
+			title: "Restore failed",
+			description: "Restic restore failed",
+		},
+		{
+			status: "cancelled" as const,
+			result: null,
+			error: null,
+			title: "Restore failed",
+			description: "Snapshot snap-1 could not be restored.",
+		},
+		{
+			status: "stale" as const,
+			result: null,
+			error: "Zerobyte was restarted before this task completed",
+			title: "Restore failed",
+			description: "Zerobyte was restarted before this task completed",
+		},
+	])(
+		"clears restoring state and shows $status terminal feedback",
+		async ({ status, result, error, title, description }) => {
+			const activeRestore = createRestoreTask();
+			const eventSource = await renderRestoreFormWithPrefetchedTasks([activeRestore]);
+
+			eventSource?.emit(
+				taskChangedEventName,
+				createRestoreTask({
+					status,
+					result,
+					error,
+					updatedAt: activeRestore.updatedAt + 1,
+				}),
+			);
+
+			expect(await screen.findByText(title)).toBeTruthy();
+			expect(screen.getByText(description)).toBeTruthy();
+			expect(screen.queryByText("Restore in progress")).toBeNull();
+
+			await userEvent.click(screen.getByRole("button", { name: "OK" }));
+			const restoreButton = await screen.findByRole("button", { name: "Restore All" });
+			expect(restoreButton.hasAttribute("disabled")).toBe(false);
+		},
+	);
+
 	test("restores the selected ancestor folder path from a broader display root", async () => {
 		let restoreRequestBody: unknown;
 

app/client/components/restore-form.tsx

@@ -1,4 +1,4 @@
-import { useCallback, useEffect, useRef, useState } from "react";
+import { useCallback, useEffect, useState } from "react";
 import { useMutation } from "@tanstack/react-query";
 import { AlertTriangle, ChevronDown, Download, FolderOpen, RotateCcw } from "lucide-react";
 import { Button } from "~/client/components/ui/button";
@@ -21,7 +21,7 @@ import { FolderSelector } from "~/client/components/folder-selector";
 import { SnapshotTreeBrowser } from "~/client/components/file-browsers/snapshot-tree-browser";
 import { RestoreProgress } from "~/client/components/restore-progress";
 import { restoreSnapshotMutation } from "~/client/api-client/@tanstack/react-query.gen";
-import { type RestoreCompletedEvent, useServerEvents } from "~/client/hooks/use-server-events";
+import { useRestoreTask } from "~/client/modules/repositories/restore-tasks";
 import { OVERWRITE_MODES, type OverwriteMode } from "@zerobyte/core/restic";
 import { isPathWithin } from "@zerobyte/core/utils";
 import type { Repository } from "~/client/lib/types";
@@ -49,7 +49,6 @@ export function RestoreForm({
 	hasNonPosixSnapshotPaths = false,
 }: RestoreFormProps) {
 	const navigate = useNavigate();
-	const { addEventListener } = useServerEvents();
 
 	const snapshotBasePath = queryBasePath ?? "/";
 	const hasMismatchedDisplayBasePath = displayBasePath && !isPathWithin(displayBasePath, snapshotBasePath);
@@ -62,10 +61,6 @@ export function RestoreForm({
 	const [overwriteMode, setOverwriteMode] = useState<OverwriteMode>("always");
 	const [showAdvanced, setShowAdvanced] = useState(false);
 	const [excludeXattr, setExcludeXattr] = useState("");
-	const [isRestoreActive, setIsRestoreActive] = useState(false);
-	const [restoreResult, setRestoreResult] = useState<RestoreCompletedEvent | null>(null);
-	const [showRestoreResultAlert, setShowRestoreResultAlert] = useState(false);
-	const restoreCompletedRef = useRef(false);
 
 	const [selectedPaths, setSelectedPaths] = useState<Set<string>>(new Set());
 	const [selectedPathKind, setSelectedPathKind] = useState<"file" | "dir" | null>(null);
@@ -79,63 +74,25 @@ export function RestoreForm({
 		}
 	}, [restoreRequiresCustomTarget]);
 
-	useEffect(() => {
-		const abortController = new AbortController();
-		const signal = abortController.signal;
-
-		addEventListener(
-			"restore:started",
-			(startedData) => {
-				if (startedData.repositoryId === repository.shortId && startedData.snapshotId === snapshotId) {
-					restoreCompletedRef.current = false;
-					setIsRestoreActive(true);
-					setRestoreResult(null);
-					setShowRestoreResultAlert(false);
-				}
-			},
-			{ signal },
-		);
-
-		addEventListener(
-			"restore:progress",
-			(progressData) => {
-				if (progressData.repositoryId === repository.shortId && progressData.snapshotId === snapshotId) {
-					if (restoreCompletedRef.current) {
-						return;
-					}
-					setIsRestoreActive(true);
-				}
-			},
-			{ signal },
-		);
-
-		addEventListener(
-			"restore:completed",
-			(completedData) => {
-				if (completedData.repositoryId === repository.shortId && completedData.snapshotId === snapshotId) {
-					restoreCompletedRef.current = true;
-					setIsRestoreActive(false);
-					setRestoreResult(completedData);
-					setShowRestoreResultAlert(true);
-				}
-			},
-			{ signal },
-		);
-
-		return () => {
-			abortController.abort();
-		};
-	}, [addEventListener, repository.shortId, snapshotId]);
-
-	const { mutate: restoreSnapshot, isPending: isRestoring } = useMutation({
+	const {
+		data: restoreStart,
+		mutate: restoreSnapshot,
+		isPending: isRestoring,
+		reset: resetRestoreMutation,
+	} = useMutation({
 		...restoreSnapshotMutation(),
 		onError: (error) => {
-			restoreCompletedRef.current = true;
-			setIsRestoreActive(false);
 			handleRepositoryError("Restore failed", error, repository.shortId);
 		},
 	});
 
+	const {
+		restoreProgress,
+		finishedRestoreTask,
+		clearFinishedRestoreTask,
+		isRestoreRunning: isRestoreTaskRunning,
+	} = useRestoreTask(repository.shortId, snapshotId, restoreStart?.restoreId);
+
 	const handleRestore = useCallback(() => {
 		const excludeXattrValues = excludeXattr
 			.split(",")
@@ -147,10 +104,8 @@ export function RestoreForm({
 
 		const includePaths = Array.from(selectedPaths);
 
-		restoreCompletedRef.current = false;
-		setIsRestoreActive(true);
-		setRestoreResult(null);
-		setShowRestoreResultAlert(false);
+		clearFinishedRestoreTask();
+		resetRestoreMutation();
 
 		restoreSnapshot({
 			path: { shortId: repository.shortId },
@@ -173,6 +128,8 @@ export function RestoreForm({
 		selectedPaths,
 		selectedPathKind,
 		overwriteMode,
+		clearFinishedRestoreTask,
+		resetRestoreMutation,
 		restoreSnapshot,
 	]);
 
@@ -196,21 +153,15 @@ export function RestoreForm({
 	}, [repository.shortId, snapshotId, selectedPathKind, selectedPaths]);
 
 	const acknowledgeRestoreResult = useCallback(() => {
-		setShowRestoreResultAlert(false);
-		setRestoreResult(null);
-	}, []);
-
-	const handleResultAlertOpenChange = useCallback((open: boolean) => {
-		if (open) {
-			setShowRestoreResultAlert(true);
-		}
-	}, []);
+		clearFinishedRestoreTask();
+		resetRestoreMutation();
+	}, [clearFinishedRestoreTask, resetRestoreMutation]);
 
 	const canRestore = restoreRequiresCustomTarget
 		? hasCustomTargetPath
 		: restoreLocation === "original" || hasCustomTargetPath;
 	const canDownload = selectedPathCount <= 1;
-	const isRestoreRunning = isRestoring || isRestoreActive;
+	const isRestoreRunning = isRestoring || isRestoreTaskRunning;
 
 	function getDownloadButtonText(): string {
 		if (selectedPathCount > 0) {
@@ -267,7 +218,7 @@ export function RestoreForm({
 
 			<div className="grid grid-cols-1 lg:grid-cols-3 gap-6">
 				<div className="space-y-6">
-					{isRestoreRunning && <RestoreProgress repositoryId={repository.shortId} snapshotId={snapshotId} />}
+					{isRestoreRunning && <RestoreProgress progress={restoreProgress} />}
 
 					{restoreRequiresCustomTarget && (
 						<Alert variant="warning">
@@ -415,16 +366,16 @@ export function RestoreForm({
 				</Card>
 			</div>
 
-			<AlertDialog open={showRestoreResultAlert} onOpenChange={handleResultAlertOpenChange}>
+			<AlertDialog open={finishedRestoreTask !== null}>
 				<AlertDialogContent>
 					<AlertDialogHeader>
 						<AlertDialogTitle>
-							{restoreResult?.status === "success" ? "Restore completed" : "Restore failed"}
+							{finishedRestoreTask?.status === "succeeded" ? "Restore completed" : "Restore failed"}
 						</AlertDialogTitle>
 						<AlertDialogDescription>
-							{restoreResult?.status === "success"
+							{finishedRestoreTask?.status === "succeeded"
 								? `Snapshot ${snapshotId} was restored successfully.`
-								: restoreResult?.error || `Snapshot ${snapshotId} could not be restored.`}
+								: finishedRestoreTask?.error || `Snapshot ${snapshotId} could not be restored.`}
 						</AlertDialogDescription>
 					</AlertDialogHeader>
 					<AlertDialogFooter>

app/client/components/restore-progress.tsx

@@ -1,48 +1,20 @@
-import { useEffect, useState } from "react";
+import type { ListTasksResponse } from "~/client/api-client";
 import { ByteSize } from "~/client/components/bytes-size";
 import { useFormatBytes } from "~/client/hooks/use-format-bytes";
 import { useRootLoaderData } from "~/client/hooks/use-root-loader-data";
 import { Card } from "~/client/components/ui/card";
 import { Progress } from "~/client/components/ui/progress";
-import { type RestoreProgressEvent, useServerEvents } from "~/client/hooks/use-server-events";
 import { formatDuration } from "~/utils/utils";
 
+type RestoreTaskProgress = Extract<NonNullable<ListTasksResponse[number]["progress"]>, { kind: "restore" }>["progress"];
+
 type Props = {
-	repositoryId: string;
-	snapshotId: string;
+	progress: RestoreTaskProgress | null;
 };
 
-export const RestoreProgress = ({ repositoryId, snapshotId }: Props) => {
+export const RestoreProgress = ({ progress }: Props) => {
 	const formatBytes = useFormatBytes();
 	const { locale } = useRootLoaderData();
-	const { addEventListener } = useServerEvents();
-	const [progress, setProgress] = useState<RestoreProgressEvent | null>(null);
-
-	useEffect(() => {
-		const abortController = new AbortController();
-
-		addEventListener(
-			"restore:progress",
-			(progressData) => {
-				if (progressData.repositoryId === repositoryId && progressData.snapshotId === snapshotId) {
-					setProgress(progressData);
-				}
-			},
-			{ signal: abortController.signal },
-		);
-
-		addEventListener(
-			"restore:completed",
-			(completedData) => {
-				if (completedData.repositoryId === repositoryId && completedData.snapshotId === snapshotId) {
-					setProgress(null);
-				}
-			},
-			{ signal: abortController.signal },
-		);
-
-		return () => abortController.abort();
-	}, [addEventListener, repositoryId, snapshotId]);
 
 	if (!progress) {
 		return (
@@ -55,8 +27,13 @@ export const RestoreProgress = ({ repositoryId, snapshotId }: Props) => {
 		);
 	}
 
-	const percentDone = Math.round(progress.percent_done * 100);
-	const speed = progress.seconds_elapsed > 0 ? formatBytes(progress.bytes_restored / progress.seconds_elapsed) : null;
+	const secondsElapsed = progress.seconds_elapsed ?? 0;
+	const percentDone = Math.round((progress.percent_done ?? 0) * 100);
+	const totalFiles = progress.total_files ?? 0;
+	const filesRestored = progress.files_restored ?? 0;
+	const totalBytes = progress.total_bytes ?? 0;
+	const bytesRestored = progress.bytes_restored ?? 0;
+	const speed = secondsElapsed > 0 ? formatBytes(bytesRestored / secondsElapsed) : null;
 
 	return (
 		<Card className="p-4">
@@ -74,22 +51,22 @@ export const RestoreProgress = ({ repositoryId, snapshotId }: Props) => {
 				<div>
 					<p className="text-xs uppercase text-muted-foreground">Files</p>
 					<p className="font-medium">
-						{progress.files_restored.toLocaleString(locale)}
+						{filesRestored.toLocaleString(locale)}
 						&nbsp;/&nbsp;
-						{progress.total_files.toLocaleString(locale)}
+						{totalFiles.toLocaleString(locale)}
 					</p>
 				</div>
 				<div>
 					<p className="text-xs uppercase text-muted-foreground">Data</p>
 					<p className="font-medium">
-						<ByteSize bytes={progress.bytes_restored} base={1024} />
+						<ByteSize bytes={bytesRestored} base={1024} />
 						&nbsp;/&nbsp;
-						<ByteSize bytes={progress.total_bytes} base={1024} />
+						<ByteSize bytes={totalBytes} base={1024} />
 					</p>
 				</div>
 				<div>
 					<p className="text-xs uppercase text-muted-foreground">Elapsed</p>
-					<p className="font-medium">{formatDuration(progress.seconds_elapsed)}</p>
+					<p className="font-medium">{formatDuration(secondsElapsed)}</p>
 				</div>
 				<div>
 					<p className="text-xs uppercase text-muted-foreground">Speed</p>

app/client/hooks/use-server-events.ts

@@ -25,9 +25,6 @@ type SharedServerEventsState = {
 	subscribers: number;
 };
 
-export type RestoreProgressEvent = ServerEventsPayloadMap["restore:progress"];
-export type RestoreCompletedEvent = ServerEventsPayloadMap["restore:completed"];
-
 const sharedState: SharedServerEventsState = {
 	eventSource: null,
 	handlers: {},

app/client/modules/repositories/restore-tasks.ts

@@ -0,0 +1,71 @@
+import { useCallback, useEffect, useState } from "react";
+import { useQuery } from "@tanstack/react-query";
+import type { ListTasksResponse } from "~/client/api-client";
+import { getTaskOptions } from "~/client/api-client/@tanstack/react-query.gen";
+import { taskEventsOptions, useTaskEvents, type TaskEventsQuery } from "~/client/hooks/use-task-events";
+import { activeTaskStatuses } from "~/schemas/tasks";
+
+type TaskResponse = ListTasksResponse[number];
+
+const restoreTasksFilter = (repositoryId: string): TaskEventsQuery => ({
+	kind: "restore",
+	resourceType: "repository",
+	resourceId: repositoryId,
+});
+
+const matchesRestore = (task: TaskResponse, snapshotId: string) => {
+	return task.kind === "restore" && task.input.kind === "restore" && task.input.snapshotId === snapshotId;
+};
+
+const isActive = (task: TaskResponse) => {
+	return activeTaskStatuses.some((status) => status === task.status);
+};
+
+export const restoreTasksOptions = (repositoryId: string) => {
+	return taskEventsOptions(restoreTasksFilter(repositoryId));
+};
+
+export const useRestoreTask = (repositoryId: string, snapshotId: string, startedTaskId?: string) => {
+	const [streamedFinishedTask, setStreamedFinishedTask] = useState<TaskResponse | null>(null);
+	const { data: streamedTasks } = useTaskEvents(restoreTasksFilter(repositoryId), {
+		onTaskFinished: (task) => {
+			if (matchesRestore(task, snapshotId)) {
+				setStreamedFinishedTask(task);
+			}
+		},
+	});
+
+	const streamedActiveTask = streamedTasks?.find((task) => matchesRestore(task, snapshotId)) ?? null;
+	const taskFromStream = streamedActiveTask ?? streamedFinishedTask;
+	const needsFallback = startedTaskId !== undefined && taskFromStream?.id !== startedTaskId;
+
+	const { data: fallbackTask } = useQuery({
+		...getTaskOptions({ path: { taskId: startedTaskId ?? "" } }),
+		enabled: needsFallback,
+		refetchInterval: ({ state }) => (state.data && !isActive(state.data) ? false : 1000),
+	});
+
+	const matchingFallbackTask = fallbackTask && matchesRestore(fallbackTask, snapshotId) ? fallbackTask : null;
+	const restoreTask = needsFallback ? matchingFallbackTask : taskFromStream;
+	const restoreIsActive = restoreTask ? isActive(restoreTask) : false;
+
+	useEffect(() => {
+		if (streamedActiveTask) {
+			setStreamedFinishedTask(null);
+		}
+	}, [streamedActiveTask]);
+
+	const clearFinishedRestoreTask = useCallback(() => {
+		setStreamedFinishedTask(null);
+	}, []);
+
+	const restoreProgress =
+		restoreIsActive && restoreTask?.progress?.kind === "restore" ? restoreTask.progress.progress : null;
+
+	return {
+		restoreProgress,
+		finishedRestoreTask: restoreIsActive ? null : restoreTask,
+		clearFinishedRestoreTask,
+		isRestoreRunning: restoreTask ? restoreIsActive : startedTaskId !== undefined,
+	};
+};

app/routes/(dashboard)/backups/$backupId/$snapshotId.restore.tsx

@@ -1,6 +1,7 @@
 import { createFileRoute } from "@tanstack/react-router";
 import { getBackupSchedule } from "~/client/api-client";
 import { getRepositoryOptions, getSnapshotDetailsOptions } from "~/client/api-client/@tanstack/react-query.gen";
+import { restoreTasksOptions } from "~/client/modules/repositories/restore-tasks";
 import { RestoreSnapshotPage } from "~/client/modules/repositories/routes/restore-snapshot";
 import { getVolumeMountPath } from "~/client/lib/volume-path";
 import { findCommonAncestor } from "@zerobyte/core/utils";
@@ -15,6 +16,7 @@ export const Route = createFileRoute("/(dashboard)/backups/$backupId/$snapshotId
 			throw new Response("Not Found", { status: 404 });
 		}
 
+		const restoreTaskOptions = restoreTasksOptions(schedule.data.repository.shortId);
 		const [snapshot, repository] = await Promise.all([
 			context.queryClient.ensureQueryData({
 				...getSnapshotDetailsOptions({
@@ -24,6 +26,7 @@ export const Route = createFileRoute("/(dashboard)/backups/$backupId/$snapshotId
 			context.queryClient.ensureQueryData({
 				...getRepositoryOptions({ path: { shortId: schedule.data.repository.shortId } }),
 			}),
+			context.queryClient.ensureQueryData(restoreTaskOptions),
 		]);
 
 		const hasNonPosixSnapshotPaths = snapshot.paths.some((path) => !path.startsWith("/"));

app/routes/(dashboard)/repositories/$repositoryId/$snapshotId/restore.tsx

@@ -1,6 +1,7 @@
 import { createFileRoute } from "@tanstack/react-router";
 import { getBackupSchedule } from "~/client/api-client";
 import { getRepositoryOptions, getSnapshotDetailsOptions } from "~/client/api-client/@tanstack/react-query.gen";
+import { restoreTasksOptions } from "~/client/modules/repositories/restore-tasks";
 import { RestoreSnapshotPage } from "~/client/modules/repositories/routes/restore-snapshot";
 import { getVolumeMountPath } from "~/client/lib/volume-path";
 import { findCommonAncestor } from "@zerobyte/core/utils";
@@ -9,11 +10,15 @@ export const Route = createFileRoute("/(dashboard)/repositories/$repositoryId/$s
 	component: RouteComponent,
 	errorComponent: (e) => <div>{e.error.message}</div>,
 	loader: async ({ params, context }) => {
+		const restoreTaskOptions = restoreTasksOptions(params.repositoryId);
 		const [snapshot, repository] = await Promise.all([
 			context.queryClient.ensureQueryData({
 				...getSnapshotDetailsOptions({ path: { shortId: params.repositoryId, snapshotId: params.snapshotId } }),
 			}),
-			context.queryClient.ensureQueryData({ ...getRepositoryOptions({ path: { shortId: params.repositoryId } }) }),
+			context.queryClient.ensureQueryData({
+				...getRepositoryOptions({ path: { shortId: params.repositoryId } }),
+			}),
+			context.queryClient.ensureQueryData(restoreTaskOptions),
 		]);
 
 		let displayBasePath: string | undefined;
@@ -38,8 +43,14 @@ export const Route = createFileRoute("/(dashboard)/repositories/$repositoryId/$s
 	staticData: {
 		breadcrumb: (match) => [
 			{ label: "Repositories", href: "/repositories" },
-			{ label: match.loaderData?.repository?.name || "Repository", href: `/repositories/${match.params.repositoryId}` },
-			{ label: match.params.snapshotId, href: `/repositories/${match.params.repositoryId}/${match.params.snapshotId}` },
+			{
+				label: match.loaderData?.repository?.name || "Repository",
+				href: `/repositories/${match.params.repositoryId}`,
+			},
+			{
+				label: match.params.snapshotId,
+				href: `/repositories/${match.params.repositoryId}/${match.params.snapshotId}`,
+			},
 			{ label: "Restore" },
 		],
 	},

app/schemas/events-dto.ts

@@ -2,7 +2,6 @@ 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", "cancelled"]);
 
 const backupEventBaseSchema = z.object({
 	scheduleId: z.string(),
@@ -14,28 +13,12 @@ const organizationScopedSchema = z.object({
 	organizationId: z.string(),
 });
 
-const restoreEventBaseSchema = z.object({
-	restoreId: z.string(),
-	repositoryId: z.string(),
-	snapshotId: z.string(),
-});
-
 const dumpStartedEventSchema = z.object({
 	repositoryId: z.string(),
 	snapshotId: z.string(),
 	path: z.string(),
 	filename: z.string(),
 });
-
-const restoreProgressMetricsSchema = z.object({
-	seconds_elapsed: z.number().default(0),
-	percent_done: z.number().default(0),
-	total_files: z.number().default(0),
-	files_restored: z.number().default(0),
-	total_bytes: z.number().default(0),
-	bytes_restored: z.number().default(0),
-});
-
 const backupStartedEventSchema = backupEventBaseSchema;
 
 export const backupProgressEventSchema = backupEventBaseSchema.extend(resticBackupProgressMetricsSchema.shape);
@@ -45,36 +28,16 @@ const backupCompletedEventSchema = backupEventBaseSchema.extend({
 	summary: resticBackupRunSummarySchema.optional(),
 });
 
-const restoreStartedEventSchema = restoreEventBaseSchema;
-
-const restoreProgressEventSchema = restoreEventBaseSchema.extend(restoreProgressMetricsSchema.shape);
-
-const restoreCompletedEventSchema = restoreEventBaseSchema.extend({
-	status: restoreEventStatusSchema,
-	error: z.string().optional(),
-	filesRestored: z.number().optional(),
-	filesSkipped: z.number().optional(),
-});
-
 const serverBackupStartedEventSchema = organizationScopedSchema.extend(backupStartedEventSchema.shape);
 
 const serverBackupProgressEventSchema = organizationScopedSchema.extend(backupProgressEventSchema.shape);
 
 const serverBackupCompletedEventSchema = organizationScopedSchema.extend(backupCompletedEventSchema.shape);
 
-const serverRestoreStartedEventSchema = organizationScopedSchema.extend(restoreStartedEventSchema.shape);
-
-const serverRestoreProgressEventSchema = organizationScopedSchema.extend(restoreProgressEventSchema.shape);
-
-const serverRestoreCompletedEventSchema = organizationScopedSchema.extend(restoreCompletedEventSchema.shape);
-
 const serverDumpStartedEventSchema = organizationScopedSchema.extend(dumpStartedEventSchema.shape);
 
 export type BackupProgressEventDto = z.infer<typeof backupProgressEventSchema>;
 export type ServerBackupStartedEventDto = z.infer<typeof serverBackupStartedEventSchema>;
 export type ServerBackupProgressEventDto = z.infer<typeof serverBackupProgressEventSchema>;
 export type ServerBackupCompletedEventDto = z.infer<typeof serverBackupCompletedEventSchema>;
-export type ServerRestoreStartedEventDto = z.infer<typeof serverRestoreStartedEventSchema>;
-export type ServerRestoreProgressEventDto = z.infer<typeof serverRestoreProgressEventSchema>;
-export type ServerRestoreCompletedEventDto = z.infer<typeof serverRestoreCompletedEventSchema>;
 export type ServerDumpStartedEventDto = z.infer<typeof serverDumpStartedEventSchema>;

app/schemas/server-events.ts

@@ -3,9 +3,6 @@ import type {
 	ServerBackupProgressEventDto,
 	ServerBackupStartedEventDto,
 	ServerDumpStartedEventDto,
-	ServerRestoreCompletedEventDto,
-	ServerRestoreProgressEventDto,
-	ServerRestoreStartedEventDto,
 } from "~/schemas/events-dto";
 import type { TaskKind, TaskResourceType, TaskStatus } from "~/schemas/tasks";
 
@@ -24,9 +21,6 @@ export const serverEventPayloads = {
 	"backup:started": payload<ServerBackupStartedEventDto>(),
 	"backup:progress": payload<ServerBackupProgressEventDto>(),
 	"backup:completed": payload<ServerBackupCompletedEventDto>(),
-	"restore:started": payload<ServerRestoreStartedEventDto>(),
-	"restore:progress": payload<ServerRestoreProgressEventDto>(),
-	"restore:completed": payload<ServerRestoreCompletedEventDto>(),
 	"dump:started": payload<ServerDumpStartedEventDto>(),
 	"mirror:started": payload<{
 		organizationId: string;

app/server/modules/repositories/repositories.service.ts

@@ -72,21 +72,11 @@ const assertAllowedControllerLocalRestoreTarget = (target: string) => {
 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) => {
+const updateActiveRestoreTask = (restoreId: string, eventName: string, update: () => void) => {
 	try {
-		return asRestoreTask(update(), restoreId, eventName);
+		update();
 	} catch (error) {
 		logger.warn(`Received ${eventName} for inactive restore ${restoreId}: ${toMessage(error)}`);
-		return null;
 	}
 };
 
@@ -125,89 +115,31 @@ const findActiveDoctorTask = (organizationId: string, repositoryShortId: string)
 	});
 };
 
-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);
+	updateActiveRestoreTask(restoreId, "restore.started", () => taskStore.markRunning(restoreId));
 };
 
 const updateRestoreProgress = (restoreId: string, progress: RestoreExecutionProgress) => {
-	const task = updateActiveRestoreTask(restoreId, "restore.progress", () =>
+	updateActiveRestoreTask(restoreId, "restore.progress", () =>
 		taskStore.updateProgress(restoreId, { kind: "restore", progress }),
 	);
-	if (!task) return;
-
-	emitRestoreProgress(task, progress);
 };
 
 const completeRestoreTask = (
 	restoreId: string,
 	result: Extract<RestoreExecutionResult, { status: "completed" }>["result"],
 ) => {
-	const task = updateActiveRestoreTask(restoreId, "restore.completed", () =>
+	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 });
+	updateActiveRestoreTask(restoreId, "restore.failed", () => taskStore.fail(restoreId, 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 });
+	updateActiveRestoreTask(restoreId, "restore.cancelled", () => taskStore.cancel(restoreId, message ?? null));
 };
 
 const finishRestoreExecution = async (restoreId: string, resultPromise: Promise<RestoreExecutionResult>) => {