feat(restore): allow cancellation from UI (#1052)

* feat(restore): allow cancellation from UI

* fix: agent cancellation failure report

app/client/api-client/@tanstack/react-query.gen.ts

@@ -4,8 +4,8 @@
 import { type DefaultError, type InfiniteData, infiniteQueryOptions, queryOptions, type UseMutationOptions } from '@tanstack/react-query';
 
 import { client } from '../client.gen';
-import { browseFilesystem, cancelDoctor, createApiKey, createBackupSchedule, createDesktopSession, createNotificationDestination, createRepository, createVolume, deleteApiKey, deleteBackupSchedule, deleteNotificationDestination, deleteRepository, deleteSnapshot, deleteSnapshots, deleteSsoInvitation, deleteSsoProvider, deleteUserAccount, deleteVolume, downloadResticPassword, dumpSnapshot, getAdminUsers, getApiKeys, getBackupProgress, getBackupSchedule, getBackupScheduleForVolume, getDevPanel, getMirrorCompatibility, getMirrorSyncStatus, getNotificationDestination, getOrgMembers, getPasswordLoginStatus, getPublicSsoProviders, getRegistrationStatus, getRepository, getRepositoryStats, getScheduleMirrors, getScheduleNotifications, getSnapshotDetails, getSsoSettings, getStatus, getSystemInfo, getTask, getUpdates, getUserDeletionImpact, getUserSsoInvitations, getVolume, healthCheckVolume, listBackupSchedules, listFiles, listNotificationDestinations, listRcloneRemotes, listRepositories, listSnapshotFiles, listSnapshots, listTasks, listVolumes, mountVolume, type Options, refreshRepositoryStats, refreshSnapshots, removeOrgMember, reorderBackupSchedules, restoreSnapshot, runBackupNow, runForget, setPasswordLoginStatus, setRegistrationStatus, startDoctor, startInvitationSsoVerification, stopBackup, syncMirror, tagSnapshots, testConnection, testNotificationDestination, unlockRepository, unmountVolume, updateBackupSchedule, updateMemberRole, updateNotificationDestination, updateRepository, updateScheduleMirrors, updateScheduleNotifications, updateSsoProviderAutoLinking, updateVolume } from '../sdk.gen';
-import type { BrowseFilesystemData, BrowseFilesystemResponse, CancelDoctorData, CancelDoctorResponse, CreateApiKeyData, CreateApiKeyResponse, CreateBackupScheduleData, CreateBackupScheduleResponse, CreateDesktopSessionData, CreateNotificationDestinationData, CreateNotificationDestinationResponse, CreateRepositoryData, CreateRepositoryResponse, CreateVolumeData, CreateVolumeResponse, DeleteApiKeyData, DeleteBackupScheduleData, DeleteBackupScheduleResponse, DeleteNotificationDestinationData, DeleteNotificationDestinationResponse, DeleteRepositoryData, DeleteRepositoryResponse, DeleteSnapshotData, DeleteSnapshotResponse, DeleteSnapshotsData, DeleteSnapshotsResponse, DeleteSsoInvitationData, DeleteSsoProviderData, DeleteUserAccountData, DeleteVolumeData, DeleteVolumeResponse, DownloadResticPasswordData, DownloadResticPasswordResponse, DumpSnapshotData, DumpSnapshotResponse, GetAdminUsersData, GetAdminUsersResponse, GetApiKeysData, GetApiKeysResponse, GetBackupProgressData, GetBackupProgressResponse, GetBackupScheduleData, GetBackupScheduleForVolumeData, GetBackupScheduleForVolumeResponse, GetBackupScheduleResponse, GetDevPanelData, GetDevPanelResponse, GetMirrorCompatibilityData, GetMirrorCompatibilityResponse, GetMirrorSyncStatusData, GetMirrorSyncStatusResponse, GetNotificationDestinationData, GetNotificationDestinationResponse, GetOrgMembersData, GetOrgMembersResponse, GetPasswordLoginStatusData, GetPasswordLoginStatusResponse, GetPublicSsoProvidersData, GetPublicSsoProvidersResponse, GetRegistrationStatusData, GetRegistrationStatusResponse, GetRepositoryData, GetRepositoryResponse, GetRepositoryStatsData, GetRepositoryStatsResponse, GetScheduleMirrorsData, GetScheduleMirrorsResponse, GetScheduleNotificationsData, GetScheduleNotificationsResponse, GetSnapshotDetailsData, GetSnapshotDetailsResponse, GetSsoSettingsData, GetSsoSettingsResponse, GetStatusData, GetStatusResponse, GetSystemInfoData, GetSystemInfoResponse, GetTaskData, GetTaskResponse, GetUpdatesData, GetUpdatesResponse, GetUserDeletionImpactData, GetUserDeletionImpactResponse, GetUserSsoInvitationsData, GetUserSsoInvitationsResponse, GetVolumeData, GetVolumeResponse, HealthCheckVolumeData, HealthCheckVolumeResponse, ListBackupSchedulesData, ListBackupSchedulesResponse, ListFilesData, ListFilesResponse, ListNotificationDestinationsData, ListNotificationDestinationsResponse, ListRcloneRemotesData, ListRcloneRemotesResponse, ListRepositoriesData, ListRepositoriesResponse, ListSnapshotFilesData, ListSnapshotFilesResponse, ListSnapshotsData, ListSnapshotsResponse, ListTasksData, ListTasksResponse, ListVolumesData, ListVolumesResponse, MountVolumeData, MountVolumeResponse, RefreshRepositoryStatsData, RefreshRepositoryStatsResponse, RefreshSnapshotsData, RefreshSnapshotsResponse, RemoveOrgMemberData, ReorderBackupSchedulesData, ReorderBackupSchedulesResponse, RestoreSnapshotData, RestoreSnapshotResponse, RunBackupNowData, RunBackupNowResponse, RunForgetData, RunForgetResponse, SetPasswordLoginStatusData, SetPasswordLoginStatusResponse, SetRegistrationStatusData, SetRegistrationStatusResponse, StartDoctorData, StartDoctorResponse, StartInvitationSsoVerificationData, StopBackupData, StopBackupResponse, SyncMirrorData, SyncMirrorResponse, TagSnapshotsData, TagSnapshotsResponse, TestConnectionData, TestConnectionResponse, TestNotificationDestinationData, TestNotificationDestinationResponse, UnlockRepositoryData, UnlockRepositoryResponse, UnmountVolumeData, UnmountVolumeResponse, UpdateBackupScheduleData, UpdateBackupScheduleResponse, UpdateMemberRoleData, UpdateNotificationDestinationData, UpdateNotificationDestinationResponse, UpdateRepositoryData, UpdateRepositoryResponse, UpdateScheduleMirrorsData, UpdateScheduleMirrorsResponse, UpdateScheduleNotificationsData, UpdateScheduleNotificationsResponse, UpdateSsoProviderAutoLinkingData, UpdateVolumeData, UpdateVolumeResponse } from '../types.gen';
+import { browseFilesystem, cancelTask, createApiKey, createBackupSchedule, createDesktopSession, createNotificationDestination, createRepository, createVolume, deleteApiKey, deleteBackupSchedule, deleteNotificationDestination, deleteRepository, deleteSnapshot, deleteSnapshots, deleteSsoInvitation, deleteSsoProvider, deleteUserAccount, deleteVolume, downloadResticPassword, dumpSnapshot, getAdminUsers, getApiKeys, getBackupProgress, getBackupSchedule, getBackupScheduleForVolume, getDevPanel, getMirrorCompatibility, getMirrorSyncStatus, getNotificationDestination, getOrgMembers, getPasswordLoginStatus, getPublicSsoProviders, getRegistrationStatus, getRepository, getRepositoryStats, getScheduleMirrors, getScheduleNotifications, getSnapshotDetails, getSsoSettings, getStatus, getSystemInfo, getTask, getUpdates, getUserDeletionImpact, getUserSsoInvitations, getVolume, healthCheckVolume, listBackupSchedules, listFiles, listNotificationDestinations, listRcloneRemotes, listRepositories, listSnapshotFiles, listSnapshots, listTasks, listVolumes, mountVolume, type Options, refreshRepositoryStats, refreshSnapshots, removeOrgMember, reorderBackupSchedules, restoreSnapshot, runBackupNow, runForget, setPasswordLoginStatus, setRegistrationStatus, startDoctor, startInvitationSsoVerification, stopBackup, syncMirror, tagSnapshots, testConnection, testNotificationDestination, unlockRepository, unmountVolume, updateBackupSchedule, updateMemberRole, updateNotificationDestination, updateRepository, updateScheduleMirrors, updateScheduleNotifications, updateSsoProviderAutoLinking, updateVolume } from '../sdk.gen';
+import type { BrowseFilesystemData, BrowseFilesystemResponse, CancelTaskData, CancelTaskResponse, CreateApiKeyData, CreateApiKeyResponse, CreateBackupScheduleData, CreateBackupScheduleResponse, CreateDesktopSessionData, CreateNotificationDestinationData, CreateNotificationDestinationResponse, CreateRepositoryData, CreateRepositoryResponse, CreateVolumeData, CreateVolumeResponse, DeleteApiKeyData, DeleteBackupScheduleData, DeleteBackupScheduleResponse, DeleteNotificationDestinationData, DeleteNotificationDestinationResponse, DeleteRepositoryData, DeleteRepositoryResponse, DeleteSnapshotData, DeleteSnapshotResponse, DeleteSnapshotsData, DeleteSnapshotsResponse, DeleteSsoInvitationData, DeleteSsoProviderData, DeleteUserAccountData, DeleteVolumeData, DeleteVolumeResponse, DownloadResticPasswordData, DownloadResticPasswordResponse, DumpSnapshotData, DumpSnapshotResponse, GetAdminUsersData, GetAdminUsersResponse, GetApiKeysData, GetApiKeysResponse, GetBackupProgressData, GetBackupProgressResponse, GetBackupScheduleData, GetBackupScheduleForVolumeData, GetBackupScheduleForVolumeResponse, GetBackupScheduleResponse, GetDevPanelData, GetDevPanelResponse, GetMirrorCompatibilityData, GetMirrorCompatibilityResponse, GetMirrorSyncStatusData, GetMirrorSyncStatusResponse, GetNotificationDestinationData, GetNotificationDestinationResponse, GetOrgMembersData, GetOrgMembersResponse, GetPasswordLoginStatusData, GetPasswordLoginStatusResponse, GetPublicSsoProvidersData, GetPublicSsoProvidersResponse, GetRegistrationStatusData, GetRegistrationStatusResponse, GetRepositoryData, GetRepositoryResponse, GetRepositoryStatsData, GetRepositoryStatsResponse, GetScheduleMirrorsData, GetScheduleMirrorsResponse, GetScheduleNotificationsData, GetScheduleNotificationsResponse, GetSnapshotDetailsData, GetSnapshotDetailsResponse, GetSsoSettingsData, GetSsoSettingsResponse, GetStatusData, GetStatusResponse, GetSystemInfoData, GetSystemInfoResponse, GetTaskData, GetTaskResponse, GetUpdatesData, GetUpdatesResponse, GetUserDeletionImpactData, GetUserDeletionImpactResponse, GetUserSsoInvitationsData, GetUserSsoInvitationsResponse, GetVolumeData, GetVolumeResponse, HealthCheckVolumeData, HealthCheckVolumeResponse, ListBackupSchedulesData, ListBackupSchedulesResponse, ListFilesData, ListFilesResponse, ListNotificationDestinationsData, ListNotificationDestinationsResponse, ListRcloneRemotesData, ListRcloneRemotesResponse, ListRepositoriesData, ListRepositoriesResponse, ListSnapshotFilesData, ListSnapshotFilesResponse, ListSnapshotsData, ListSnapshotsResponse, ListTasksData, ListTasksResponse, ListVolumesData, ListVolumesResponse, MountVolumeData, MountVolumeResponse, RefreshRepositoryStatsData, RefreshRepositoryStatsResponse, RefreshSnapshotsData, RefreshSnapshotsResponse, RemoveOrgMemberData, ReorderBackupSchedulesData, ReorderBackupSchedulesResponse, RestoreSnapshotData, RestoreSnapshotResponse, RunBackupNowData, RunBackupNowResponse, RunForgetData, RunForgetResponse, SetPasswordLoginStatusData, SetPasswordLoginStatusResponse, SetRegistrationStatusData, SetRegistrationStatusResponse, StartDoctorData, StartDoctorResponse, StartInvitationSsoVerificationData, StopBackupData, StopBackupResponse, SyncMirrorData, SyncMirrorResponse, TagSnapshotsData, TagSnapshotsResponse, TestConnectionData, TestConnectionResponse, TestNotificationDestinationData, TestNotificationDestinationResponse, UnlockRepositoryData, UnlockRepositoryResponse, UnmountVolumeData, UnmountVolumeResponse, UpdateBackupScheduleData, UpdateBackupScheduleResponse, UpdateMemberRoleData, UpdateNotificationDestinationData, UpdateNotificationDestinationResponse, UpdateRepositoryData, UpdateRepositoryResponse, UpdateScheduleMirrorsData, UpdateScheduleMirrorsResponse, UpdateScheduleNotificationsData, UpdateScheduleNotificationsResponse, UpdateSsoProviderAutoLinkingData, UpdateVolumeData, UpdateVolumeResponse } from '../types.gen';
 
 export type QueryKey<TOptions extends Options> = [
     Pick<TOptions, 'baseUrl' | 'body' | 'headers' | 'path' | 'query'> & {
@@ -891,23 +891,6 @@ export const restoreSnapshotMutation = (options?: Partial<Options<RestoreSnapsho
     return mutationOptions;
 };
 
-/**
- * Cancel a running doctor operation on a repository
- */
-export const cancelDoctorMutation = (options?: Partial<Options<CancelDoctorData>>): UseMutationOptions<CancelDoctorResponse, DefaultError, Options<CancelDoctorData>> => {
-    const mutationOptions: UseMutationOptions<CancelDoctorResponse, DefaultError, Options<CancelDoctorData>> = {
-        mutationFn: async (fnOptions) => {
-            const { data } = await cancelDoctor({
-                ...options,
-                ...fnOptions,
-                throwOnError: true
-            });
-            return data;
-        }
-    };
-    return mutationOptions;
-};
-
 /**
  * Start an asynchronous doctor operation on a repository to fix common issues (unlock, check, repair index).
  */
@@ -1554,6 +1537,23 @@ export const getTaskOptions = (options: Options<GetTaskData>) => queryOptions<Ge
     queryKey: getTaskQueryKey(options)
 });
 
+/**
+ * Request cancellation of a running task
+ */
+export const cancelTaskMutation = (options?: Partial<Options<CancelTaskData>>): UseMutationOptions<CancelTaskResponse, DefaultError, Options<CancelTaskData>> => {
+    const mutationOptions: UseMutationOptions<CancelTaskResponse, DefaultError, Options<CancelTaskData>> = {
+        mutationFn: async (fnOptions) => {
+            const { data } = await cancelTask({
+                ...options,
+                ...fnOptions,
+                throwOnError: true
+            });
+            return data;
+        }
+    };
+    return mutationOptions;
+};
+
 /**
  * Create an authenticated desktop session
  */

app/client/api-client/index.ts

@@ -3,7 +3,7 @@
 
 export {
 	browseFilesystem,
-	cancelDoctor,
+	cancelTask,
 	createApiKey,
 	createBackupSchedule,
 	createDesktopSession,
@@ -94,10 +94,10 @@ export type {
 	BrowseFilesystemData,
 	BrowseFilesystemResponse,
 	BrowseFilesystemResponses,
-	CancelDoctorData,
-	CancelDoctorErrors,
-	CancelDoctorResponse,
-	CancelDoctorResponses,
+	CancelTaskData,
+	CancelTaskErrors,
+	CancelTaskResponse,
+	CancelTaskResponses,
 	ClientOptions,
 	CreateApiKeyData,
 	CreateApiKeyErrors,

app/client/api-client/types.gen.ts

@@ -2445,35 +2445,6 @@ export type RestoreSnapshotResponses = {
 
 export type RestoreSnapshotResponse = RestoreSnapshotResponses[keyof RestoreSnapshotResponses];
 
-export type CancelDoctorData = {
-    body?: never;
-    path: {
-        shortId: string;
-    };
-    query?: never;
-    url: '/api/v1/repositories/{shortId}/doctor';
-};
-
-export type CancelDoctorErrors = {
-    /**
-     * No doctor operation is currently running
-     */
-    409: unknown;
-};
-
-export type CancelDoctorResponses = {
-    /**
-     * Doctor operation cancelled
-     */
-    200: {
-        status: 'cancelled';
-    } | {
-        status: 'reset';
-    };
-};
-
-export type CancelDoctorResponse = CancelDoctorResponses[keyof CancelDoctorResponses];
-
 export type StartDoctorData = {
     body?: never;
     path: {
@@ -5911,6 +5882,33 @@ export type GetTaskResponses = {
 
 export type GetTaskResponse = GetTaskResponses[keyof GetTaskResponses];
 
+export type CancelTaskData = {
+    body?: never;
+    path: {
+        taskId: string;
+    };
+    query?: never;
+    url: '/api/v1/tasks/{taskId}/cancel';
+};
+
+export type CancelTaskErrors = {
+    /**
+     * Task is not cancellable or is no longer running
+     */
+    409: unknown;
+};
+
+export type CancelTaskResponses = {
+    /**
+     * Task cancellation requested
+     */
+    202: {
+        status: 'cancelling';
+    };
+};
+
+export type CancelTaskResponse = CancelTaskResponses[keyof CancelTaskResponses];
+
 export type CreateDesktopSessionData = {
     body: {
         dateFormat: 'MM/DD/YYYY' | 'DD/MM/YYYY' | 'YYYY/MM/DD';

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

@@ -153,8 +153,8 @@ describe("RestoreForm", () => {
 	test("recovers the active restore from the prefetched exact filtered collection", async () => {
 		const taskStream = await renderRestoreFormWithPrefetchedTask(createRestoreTask());
 
-		const restoreButton = await screen.findByRole("button", { name: "Restoring..." });
-		expect(restoreButton.hasAttribute("disabled")).toBe(true);
+		const restoreButton = await screen.findByRole("button", { name: "Cancel restore" });
+		expect(restoreButton.hasAttribute("disabled")).toBe(false);
 		expect(screen.getByText("Restore in progress")).toBeTruthy();
 		expect(taskStream?.url).toBe(
 			"/api/v1/tasks/events?kind=restore&resourceType=repository&resourceId=repo-1&operationKey=snap-1",
@@ -162,6 +162,23 @@ describe("RestoreForm", () => {
 		expect(MockEventSource.instances).toHaveLength(1);
 	});
 
+	test("cancels the exact active restore task", async () => {
+		let cancelledRestoreId: string | undefined;
+		server.use(
+			http.post("/api/v1/tasks/:taskId/cancel", ({ params }) => {
+				cancelledRestoreId = String(params.taskId);
+				return HttpResponse.json({ status: "cancelling" }, { status: 202 });
+			}),
+		);
+		await renderRestoreFormWithPrefetchedTask(createRestoreTask({ id: "restore-task-123" }));
+
+		await userEvent.click(screen.getByRole("button", { name: "Cancel restore" }));
+
+		await waitFor(() => {
+			expect(cancelledRestoreId).toBe("restore-task-123");
+		});
+	});
+
 	test("renders canonical task progress from the exact filtered stream", async () => {
 		const activeRestore = createRestoreTask();
 		const taskStream = await renderRestoreFormWithPrefetchedTask(activeRestore);
@@ -211,7 +228,7 @@ describe("RestoreForm", () => {
 			expect(queryClient.getMutationCache().getAll()[0]?.state.status).toBe("success");
 		});
 
-		expect(screen.getByRole("button", { name: "Restoring..." }).hasAttribute("disabled")).toBe(true);
+		expect(screen.getByRole("button", { name: "Cancel restore" }).hasAttribute("disabled")).toBe(false);
 		expect(screen.queryByRole("button", { name: "Restore All" })).toBeNull();
 
 		await waitFor(() => {
@@ -225,7 +242,37 @@ describe("RestoreForm", () => {
 			.find((eventSource) => eventSource.url === "/api/v1/tasks/task-restore/events")
 			?.emit(taskChangedEventName, createRestoreTask());
 
-		expect(await screen.findByRole("button", { name: "Restoring..." })).toBeTruthy();
+		expect(await screen.findByRole("button", { name: "Cancel restore" })).toBeTruthy();
+	});
+
+	test("uses the newest task update across the exact and filtered streams", async () => {
+		server.use(
+			snapshotFilesHandler,
+			http.post("/api/v1/repositories/:shortId/restore", () =>
+				HttpResponse.json({ restoreId: "task-restore", status: "started" }, { status: 202 }),
+			),
+		);
+		renderRestoreForm();
+
+		await userEvent.click(screen.getByRole("button", { name: "Restore All" }));
+		await waitFor(() => {
+			expect(MockEventSource.instances).toHaveLength(2);
+		});
+
+		const exactStream = MockEventSource.instances.find(
+			(eventSource) => eventSource.url === "/api/v1/tasks/task-restore/events",
+		);
+		const filteredStream = MockEventSource.instances.find((eventSource) =>
+			eventSource.url.includes("/tasks/events?"),
+		);
+		exactStream?.emit(taskChangedEventName, createRestoreTask({ updatedAt: 1711411200001 }));
+		filteredStream?.emit(
+			taskChangedEventName,
+			createRestoreTask({ status: "succeeded", updatedAt: 1711411200002 }),
+		);
+
+		expect(await screen.findByText("Restore completed")).toBeTruthy();
+		expect(screen.queryByRole("button", { name: "Cancel restore" })).toBeNull();
 	});
 
 	test.each([
@@ -254,8 +301,8 @@ describe("RestoreForm", () => {
 			status: "cancelled" as const,
 			result: null,
 			error: null,
-			title: "Restore failed",
-			description: "Snapshot snap-1 could not be restored.",
+			title: "Restore cancelled",
+			description: "Restore of snapshot snap-1 was cancelled.",
 		},
 		{
 			status: "stale" as const,

app/client/components/restore-form.tsx

@@ -1,6 +1,6 @@
 import { useCallback, useEffect, useState } from "react";
 import { useMutation } from "@tanstack/react-query";
-import { AlertTriangle, ChevronDown, Download, FolderOpen, RotateCcw } from "lucide-react";
+import { AlertTriangle, ChevronDown, Download, FolderOpen, RotateCcw, Square } from "lucide-react";
 import { Button } from "~/client/components/ui/button";
 import { Tooltip, TooltipContent, TooltipTrigger } from "~/client/components/ui/tooltip";
 import { Alert, AlertDescription, AlertTitle } from "~/client/components/ui/alert";
@@ -20,7 +20,7 @@ import {
 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 { cancelTaskMutation, restoreSnapshotMutation } from "~/client/api-client/@tanstack/react-query.gen";
 import { useRestoreTask, type RestoreTask } from "~/client/modules/repositories/restore-tasks";
 import { OVERWRITE_MODES, type OverwriteMode } from "@zerobyte/core/restic";
 import { isPathWithin } from "@zerobyte/core/utils";
@@ -92,9 +92,17 @@ export function RestoreForm({
 		restoreProgress,
 		finishedRestoreTask,
 		clearFinishedRestoreTask,
+		activeRestoreTaskId,
 		isRestoreRunning: isRestoreTaskRunning,
 	} = useRestoreTask(repository.shortId, snapshotId, restoreStart?.restoreId, initialActiveTask);
 
+	const cancelRestore = useMutation({
+		...cancelTaskMutation(),
+		onError: (error) => {
+			handleRepositoryError("Failed to cancel restore", error, repository.shortId);
+		},
+	});
+
 	const handleRestore = useCallback(() => {
 		const excludeXattrValues = excludeXattr
 			.split(",")
@@ -211,10 +219,25 @@ export function RestoreForm({
 							</p>
 						</TooltipContent>
 					</Tooltip>
-					<Button variant="primary" onClick={handleRestore} disabled={isRestoreRunning || !canRestore}>
-						<RotateCcw className="h-4 w-4 mr-2" />
-						{getRestoreButtonText()}
-					</Button>
+					{activeRestoreTaskId ? (
+						<Button
+							variant="destructive"
+							loading={cancelRestore.isPending}
+							onClick={() =>
+								cancelRestore.mutate({
+									path: { taskId: activeRestoreTaskId },
+								})
+							}
+						>
+							<Square className="h-4 w-4 mr-2" />
+							Cancel restore
+						</Button>
+					) : (
+						<Button variant="primary" onClick={handleRestore} disabled={isRestoreRunning || !canRestore}>
+							<RotateCcw className="h-4 w-4 mr-2" />
+							{getRestoreButtonText()}
+						</Button>
+					)}
 				</div>
 			</div>
 
@@ -372,12 +395,18 @@ export function RestoreForm({
 				<AlertDialogContent>
 					<AlertDialogHeader>
 						<AlertDialogTitle>
-							{finishedRestoreTask?.status === "succeeded" ? "Restore completed" : "Restore failed"}
+							{finishedRestoreTask?.status === "succeeded"
+								? "Restore completed"
+								: finishedRestoreTask?.status === "cancelled"
+									? "Restore cancelled"
+									: "Restore failed"}
 						</AlertDialogTitle>
 						<AlertDialogDescription>
 							{finishedRestoreTask?.status === "succeeded"
 								? `Snapshot ${snapshotId} was restored successfully.`
-								: finishedRestoreTask?.error || `Snapshot ${snapshotId} could not be restored.`}
+								: finishedRestoreTask?.status === "cancelled"
+									? finishedRestoreTask.error || `Restore of snapshot ${snapshotId} was cancelled.`
+									: finishedRestoreTask?.error || `Snapshot ${snapshotId} could not be restored.`}
 						</AlertDialogDescription>
 					</AlertDialogHeader>
 					<AlertDialogFooter>

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

@@ -33,34 +33,45 @@ export const getActiveRestoreTask = (tasks: ListTasksResponse): RestoreTask | nu
 	return task as RestoreTask;
 };
 
+const latestTask = (tasks: Array<RestoreTask | null>) =>
+	tasks.reduce<RestoreTask | null>((latest, task) => {
+		if (!task) return latest;
+		if (!latest || task.updatedAt > latest.updatedAt) return task;
+		if (task.updatedAt === latest.updatedAt && task.id >= latest.id) return task;
+		return latest;
+	}, null);
+
 export const useRestoreTask = (
 	repositoryId: string,
 	snapshotId: string,
 	startedTaskId?: string,
 	initialActiveTask?: RestoreTask | null,
 ) => {
-	const [retainedFinishedTask, setRetainedFinishedTask] = useState<RestoreTask | null>(null);
+	const [lastFinishedTask, setLastFinishedTask] = useState<RestoreTask | null>(null);
 	const filter = restoreTasksFilter(repositoryId, snapshotId);
 	const { data: activeRestoreTasks } = useActiveTasks(filter, {
 		initialTasks: initialActiveTask ? [initialActiveTask] : undefined,
-		onTaskFinished: setRetainedFinishedTask,
+		onTaskFinished: setLastFinishedTask,
 	});
-	const { task: exactStartedTask } = useTask<RestoreTask>(startedTaskId);
-	const activeRestoreTask = activeRestoreTasks?.[0] ?? null;
-	const restoreTask = exactStartedTask ?? activeRestoreTask ?? retainedFinishedTask;
-	const taskIsActive = restoreTask ? isTaskActive(restoreTask) : false;
-	const finishedRestoreTask = restoreTask && !taskIsActive ? restoreTask : null;
+	const { task: startedTask } = useTask<RestoreTask>(startedTaskId);
+	const restoreTask = latestTask([startedTask, activeRestoreTasks?.[0] ?? null, lastFinishedTask]);
+	const activeRestoreTaskId =
+		restoreTask && isTaskActive(restoreTask)
+			? restoreTask.id
+			: restoreTask === null && startedTaskId !== undefined
+				? startedTaskId
+				: null;
+	const taskIsActive = restoreTask !== null && activeRestoreTaskId === restoreTask.id;
 
 	const clearFinishedRestoreTask = useCallback(() => {
-		setRetainedFinishedTask(null);
+		setLastFinishedTask(null);
 	}, []);
 
-	const restoreProgress = taskIsActive ? (restoreTask?.progress?.progress ?? null) : null;
-
 	return {
-		restoreProgress,
-		finishedRestoreTask,
+		restoreProgress: taskIsActive ? (restoreTask?.progress?.progress ?? null) : null,
+		finishedRestoreTask: restoreTask && !taskIsActive ? restoreTask : null,
 		clearFinishedRestoreTask,
-		isRestoreRunning: (startedTaskId !== undefined && exactStartedTask === null) || taskIsActive,
+		activeRestoreTaskId,
+		isRestoreRunning: activeRestoreTaskId !== null,
 	};
 };

app/client/modules/repositories/routes/repository-details.tsx

@@ -4,7 +4,7 @@ import { useNavigate, useSearch } from "@tanstack/react-router";
 import { toast } from "sonner";
 import { ChevronDown, Database, Pencil, Square, Stethoscope, Trash2, Unlock } from "lucide-react";
 import {
-	cancelDoctorMutation,
+	cancelTaskMutation,
 	deleteRepositoryMutation,
 	getRepositoryOptions,
 	startDoctorMutation,
@@ -59,7 +59,7 @@ export default function RepositoryDetailsPage({
 	const { data: repository } = useSuspenseQuery({
 		...getRepositoryOptions({ path: { shortId: repositoryId } }),
 	});
-	const { isDoctorRunning: hasActiveDoctorTask } = useRepositoryDoctorTask(repositoryId);
+	const { activeDoctorTask } = useRepositoryDoctorTask(repositoryId);
 
 	const [showDeleteConfirm, setShowDeleteConfirm] = useState(false);
 
@@ -89,12 +89,7 @@ export default function RepositoryDetailsPage({
 	});
 
 	const cancelDoctor = useMutation({
-		...cancelDoctorMutation(),
-		onSuccess: (result) => {
-			if (result.status === "reset") {
-				toast.info("Doctor state reset");
-			}
-		},
+		...cancelTaskMutation(),
 		onError: (error) => {
 			toast.error("Failed to cancel doctor", {
 				description: parseError(error)?.message,
@@ -111,7 +106,7 @@ export default function RepositoryDetailsPage({
 		deleteRepo.mutate({ path: { shortId: repository.shortId } });
 	};
 
-	const isDoctorRunning = hasActiveDoctorTask || startDoctor.isPending;
+	const isDoctorRunning = activeDoctorTask !== null || startDoctor.isPending;
 	const displayStatus = isDoctorRunning ? "doctor" : repository.status;
 
 	return (
@@ -149,12 +144,12 @@ export default function RepositoryDetailsPage({
 							</div>
 						</div>
 						<div className="flex items-center gap-2">
-							{isDoctorRunning ? (
+							{activeDoctorTask ? (
 								<Button
 									type="button"
 									variant="destructive"
 									loading={cancelDoctor.isPending}
-									onClick={() => cancelDoctor.mutate({ path: { shortId: repository.shortId } })}
+									onClick={() => cancelDoctor.mutate({ path: { taskId: activeDoctorTask.id } })}
 								>
 									<Square className="h-4 w-4 mr-2" />
 									Cancel doctor

app/server/modules/agents/__tests__/agents-manager.events.test.ts

@@ -298,18 +298,64 @@ test("runBackup requests cancellation when the abort signal fires while sending"
 	await stopAgentController();
 });
 
+test("startRestore waits for agent confirmation when restore cancellation cannot be dispatched", async () => {
+	resetAgentRuntime();
+	controllerMock.sendRestore.mockImplementation(() => Effect.succeed(true));
+	controllerMock.cancelRestore.mockImplementation(() => Effect.succeed(false));
+	const { agentManager, startAgentController, stopAgentController } = await import("../agents-manager");
+	const abortController = new AbortController();
+
+	await startAgentController();
+	const started = await agentManager.startRestore("local", {
+		payload: restorePayload,
+		signal: abortController.signal,
+		onProgress: vi.fn(),
+	});
+	if (started.status !== "started") {
+		throw new Error("Expected restore to start");
+	}
+
+	let settled = false;
+	void started.result.finally(() => {
+		settled = true;
+	});
+	abortController.abort();
+
+	await vi.waitFor(() => {
+		expect(controllerMock.cancelRestore).toHaveBeenCalledWith("local", { restoreId: "restore-1" });
+	});
+	await Promise.resolve();
+	await Promise.resolve();
+	await Promise.resolve();
+	expect(settled).toBe(false);
+
+	controllerMock.onEvent?.({
+		type: "restore.cancelled",
+		agentId: "local",
+		agentName: "Local Agent",
+		payload: {
+			restoreId: "restore-1",
+			organizationId: "org-1",
+			repositoryId: "repo-1",
+			snapshotId: "snapshot-1",
+			message: "Restore was cancelled",
+		},
+	});
+
+	await expect(started.result).resolves.toEqual({ status: "cancelled" });
+	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") {
@@ -350,7 +396,6 @@ test("restore events are delivered to the running restore callbacks", async () =
 		status: "completed",
 		result: { message_type: "summary", files_restored: 2, files_skipped: 1 },
 	});
-	expect(onStarted).toHaveBeenCalledOnce();
 	expect(onProgress).toHaveBeenCalledWith({ percent_done: 0.5 });
 	await stopAgentController();
 });
@@ -364,13 +409,11 @@ test("agent disconnect cancels only restores owned by that agent", async () => {
 	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<RestoreRunPayload>({ restoreId: "restore-2" }),
 		signal: new AbortController().signal,
-		onStarted: vi.fn(),
 		onProgress: vi.fn(),
 	});
 	if (localStarted.status !== "started" || remoteStarted.status !== "started") {

app/server/modules/agents/agents-manager.ts

@@ -7,6 +7,7 @@ import type {
 } from "@zerobyte/contracts/agent-protocol";
 import { Effect } from "effect";
 import { config } from "../../core/config";
+import { toMessage } from "../../utils/errors";
 import { createAgentManagerRuntime, type AgentManagerEvent } from "./controller/server";
 import { LOCAL_AGENT_ID } from "./constants";
 import { spawnLocalAgentProcess, stopLocalAgentProcess } from "./local/process";
@@ -41,7 +42,6 @@ type AgentRunBackupRequest = {
 type AgentStartRestoreRequest = {
 	payload: RestoreRunPayload;
 	signal: AbortSignal;
-	onStarted: () => void;
 	onProgress: (progress: RestoreExecutionProgress) => void;
 };
 
@@ -226,12 +226,19 @@ const requestRestoreCancellation = async (agentId: string, restoreId: string) =>
 		return true;
 	}
 
-	if (await Effect.runPromise(runtime.cancelRestore(agentId, { restoreId }))) {
-		return true;
+	try {
+		if (await Effect.runPromise(runtime.cancelRestore(agentId, { restoreId }))) {
+			return true;
+		}
+
+		logger.warn(
+			`Failed to send restore cancellation for ${restoreId}; waiting for the agent to report its outcome`,
+		);
+	} catch (error) {
+		logger.warn(`Failed to send restore cancellation for ${restoreId}: ${toMessage(error)}`);
 	}
 
-	resolveActiveRestoreRun(restoreId, { status: "cancelled" });
-	return true;
+	return false;
 };
 
 const handleAgentManagerEvent = (event: AgentManagerEvent) => {
@@ -323,12 +330,7 @@ const handleAgentManagerEvent = (event: AgentManagerEvent) => {
 			break;
 		}
 		case "restore.started": {
-			const activeRestoreRun = getActiveRestoreRun(event.payload.restoreId, event.type, event.agentId);
-			if (!activeRestoreRun) {
-				break;
-			}
-
-			activeRestoreRun.onStarted();
+			getActiveRestoreRun(event.payload.restoreId, event.type, event.agentId);
 			break;
 		}
 		case "restore.progress": {
@@ -488,7 +490,6 @@ export const agentManager = {
 			getActiveRestoresByRestoreId().set(request.payload.restoreId, {
 				agentId,
 				restoreId: request.payload.restoreId,
-				onStarted: request.onStarted,
 				onProgress: request.onProgress,
 				resolve,
 				cancellationRequested: false,
@@ -504,11 +505,20 @@ export const agentManager = {
 				};
 			}
 
+			const cancelOnAbort = () => {
+				void requestRestoreCancellation(agentId, request.payload.restoreId);
+			};
+			request.signal.addEventListener("abort", cancelOnAbort, { once: true });
 			if (request.signal.aborted) {
-				await requestRestoreCancellation(agentId, request.payload.restoreId);
+				cancelOnAbort();
 			}
 
-			return { status: "started", result: completion };
+			return {
+				status: "started",
+				result: completion.finally(() => {
+					request.signal.removeEventListener("abort", cancelOnAbort);
+				}),
+			};
 		} catch (error) {
 			clearActiveRestoreRun(request.payload.restoreId);
 			throw error;

app/server/modules/agents/helpers/runtime-state.ts

@@ -38,7 +38,6 @@ type ActiveBackupRun = {
 type ActiveRestoreRun = {
 	agentId: string;
 	restoreId: string;
-	onStarted: () => void;
 	onProgress: (progress: RestoreExecutionProgress) => void;
 	resolve: (result: RestoreExecutionResult) => void;
 	cancellationRequested: boolean;

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

@@ -619,7 +619,7 @@ describe("repositoriesService.restoreSnapshot", () => {
 			]),
 		);
 
-		const restoreMock = vi.fn(() => Promise.resolve(createPendingRestoreStart()));
+		const restoreMock = vi.fn<typeof agentManager.startRestore>(() => Promise.resolve(createPendingRestoreStart()));
 		vi.spyOn(agentManager, "startRestore").mockImplementation(restoreMock);
 
 		return {
@@ -737,7 +737,7 @@ describe("repositoriesService.restoreSnapshot", () => {
 				resourceId: repositoryShortId,
 			});
 			expect(task?.id).toBe(restoreId);
-			expect(task?.status).toBe("queued");
+			expect(task?.status).toBe("running");
 		} finally {
 			releaseExclusive();
 			await fs.rm(targetPath, { recursive: true, force: true });

app/server/modules/repositories/commands/doctor-command.ts

@@ -10,7 +10,7 @@ import type { Repository } from "../../../db/schema";
 import { repositoriesTable } from "../../../db/schema";
 import { runEffectPromise, toMessage } from "../../../utils/errors";
 import type { TaskResult } from "~/schemas/tasks";
-import { requestTaskCancel, runTaskLifecycle, TaskCancelledError } from "../../tasks/tasks.lifecycle";
+import { runTaskLifecycle, TaskCancelledError } from "../../tasks/tasks.lifecycle";
 import { taskStore } from "../../tasks/tasks.store";
 
 type DoctorCommandParams = {
@@ -181,10 +181,6 @@ const executeDoctor = async (repository: Repository, signal: AbortSignal): Promi
 	}
 };
 
-export const cancelDoctorTask = (taskId: string) => {
-	return requestTaskCancel(taskId);
-};
-
 export const createDoctorCommand = (params: DoctorCommandParams) => {
 	return {
 		start: () => {
@@ -201,6 +197,7 @@ export const createDoctorCommand = (params: DoctorCommandParams) => {
 			void runTaskLifecycle({
 				taskId: task.id,
 				label: "doctor task",
+				cancellable: true,
 				onStarted: async () => {
 					await db
 						.update(repositoriesTable)

app/server/modules/repositories/commands/index.ts

@@ -1,10 +1,11 @@
 import { createDeleteSnapshotsCommand } from "./delete-snapshots-command";
-import { createDoctorCommand, cancelDoctorTask } from "./doctor-command";
+import { createDoctorCommand } from "./doctor-command";
 import { createTagSnapshotsCommand } from "./tag-snapshots-command";
+import { createRestoreCommand } from "./restore-command";
 
 export const commands = {
 	createDeleteSnapshots: createDeleteSnapshotsCommand,
 	createDoctor: createDoctorCommand,
+	createRestore: createRestoreCommand,
 	createTagSnapshots: createTagSnapshotsCommand,
-	cancelDoctor: cancelDoctorTask,
 };

app/server/modules/repositories/commands/restore-command.ts

@@ -0,0 +1,149 @@
+import type { RepositoryConfig } from "@zerobyte/core/restic";
+import type { RestoreRunPayload } from "@zerobyte/contracts/agent-protocol";
+import type { TaskResult } from "~/schemas/tasks";
+import { repoMutex } from "../../../core/repository-mutex";
+import { restic, resticDeps } from "../../../core/restic";
+import { runEffectPromise } from "../../../utils/errors";
+import { agentManager, type RestoreExecutionProgress } from "../../agents/agents-manager";
+import { runTaskLifecycle, TaskCancelledError } from "../../tasks/tasks.lifecycle";
+import { taskStore } from "../../tasks/tasks.store";
+
+type RestoreExecutionOptions = Omit<Parameters<typeof restic.restore>[3], "organizationId" | "signal" | "onProgress">;
+
+type RestoreExecutionTarget = { kind: "controller" } | { kind: "agent"; agentId: string };
+
+type RestoreCommandParams = {
+	organizationId: string;
+	repositoryId: string;
+	repositoryShortId: string;
+	repositoryConfig: RepositoryConfig;
+	snapshotId: string;
+	target: string;
+	executionTarget: RestoreExecutionTarget;
+	options: RestoreExecutionOptions;
+};
+
+type RestoreExecutionRequest = RestoreCommandParams & {
+	restoreId: string;
+	signal: AbortSignal;
+	onProgress: (progress: RestoreExecutionProgress) => void;
+};
+
+type RestoreTaskResult = Extract<TaskResult, { kind: "restore" }>;
+
+const createRestoreRunPayload = async (request: RestoreExecutionRequest): Promise<RestoreRunPayload> => {
+	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) => {
+	return await runEffectPromise(
+		restic.restore(request.repositoryConfig, request.snapshotId, request.target, {
+			...request.options,
+			organizationId: request.organizationId,
+			signal: request.signal,
+			onProgress: request.onProgress,
+		}),
+	);
+};
+
+const executeAgentRestore = async (request: RestoreExecutionRequest, agentId: string) => {
+	const payload = await createRestoreRunPayload(request);
+	const started = await agentManager.startRestore(agentId, {
+		payload,
+		signal: request.signal,
+		onProgress: request.onProgress,
+	});
+
+	if (started.status === "unavailable") {
+		throw started.error;
+	}
+
+	const result = await started.result;
+	switch (result.status) {
+		case "completed":
+			return result.result;
+		case "cancelled":
+			throw new TaskCancelledError(result.message ?? "Restore was cancelled");
+		case "failed":
+			throw new Error(result.error);
+		case "unavailable":
+			throw result.error;
+	}
+};
+
+const executeRestore = async (request: RestoreExecutionRequest) => {
+	const releaseLock = await repoMutex.acquireShared(
+		request.repositoryId,
+		`restore:${request.restoreId}`,
+		request.signal,
+	);
+
+	try {
+		if (request.executionTarget.kind === "controller") {
+			return await executeControllerRestore(request);
+		}
+
+		return await executeAgentRestore(request, request.executionTarget.agentId);
+	} finally {
+		releaseLock();
+	}
+};
+
+const runRestoreTask = async (
+	params: RestoreCommandParams,
+	restoreId: string,
+	signal: AbortSignal,
+): Promise<RestoreTaskResult> => {
+	const result = await executeRestore({
+		...params,
+		restoreId,
+		signal,
+		onProgress: (progress) => taskStore.updateProgress(restoreId, { kind: "restore", progress }),
+	});
+
+	return { kind: "restore", result };
+};
+
+export const createRestoreCommand = (params: RestoreCommandParams) => {
+	return {
+		start: () => {
+			const task = taskStore.create({
+				organizationId: params.organizationId,
+				resourceType: "repository",
+				resourceId: params.repositoryShortId,
+				operationKey: params.snapshotId,
+				targetAgentId: params.executionTarget.kind === "agent" ? params.executionTarget.agentId : null,
+				input: {
+					kind: "restore",
+					repositoryId: params.repositoryShortId,
+					snapshotId: params.snapshotId,
+					target: params.target,
+				},
+			});
+
+			void runTaskLifecycle({
+				taskId: task.id,
+				label: "restore task",
+				cancellable: true,
+				run: (signal) => runRestoreTask(params, task.id, signal),
+			});
+
+			return { restoreId: task.id, status: "started" as const };
+		},
+	};
+};

app/server/modules/repositories/repositories.controller.ts

@@ -11,7 +11,6 @@ import {
 	deleteSnapshotsBody,
 	deleteSnapshotsDto,
 	startDoctorDto,
-	cancelDoctorDto,
 	getRepositoryDto,
 	getRepositoryStatsDto,
 	refreshRepositoryStatsDto,
@@ -38,7 +37,6 @@ import {
 	type DeleteSnapshotDto,
 	type DeleteSnapshotsResponseDto,
 	type StartDoctorDto,
-	type CancelDoctorDto,
 	type GetRepositoryDto,
 	type GetRepositoryStatsDto,
 	type RefreshRepositoryStatsDto,
@@ -247,13 +245,6 @@ export const repositoriesController = new Hono()
 
 		return c.json<StartDoctorDto>(result, 202);
 	})
-	.delete("/:shortId/doctor", cancelDoctorDto, async (c) => {
-		const shortId = asShortId(c.req.param("shortId"));
-
-		const result = await repositoriesService.cancelDoctor(shortId);
-
-		return c.json<CancelDoctorDto>(result, 200);
-	})
 	.post("/:shortId/unlock", unlockRepositoryDto, async (c) => {
 		const shortId = asShortId(c.req.param("shortId"));
 

app/server/modules/repositories/repositories.dto.ts

@@ -397,32 +397,6 @@ export const startDoctorDto = describeRoute({
 	},
 });
 
-const cancelDoctorResponse = z.discriminatedUnion("status", [
-	z.object({ status: z.literal("cancelled") }),
-	z.object({ status: z.literal("reset") }),
-]);
-
-export type CancelDoctorDto = z.infer<typeof cancelDoctorResponse>;
-
-export const cancelDoctorDto = describeRoute({
-	description: "Cancel a running doctor operation on a repository",
-	tags: ["Repositories"],
-	operationId: "cancelDoctor",
-	responses: {
-		200: {
-			description: "Doctor operation cancelled",
-			content: {
-				"application/json": {
-					schema: resolver(cancelDoctorResponse),
-				},
-			},
-		},
-		409: {
-			description: "No doctor operation is currently running",
-		},
-	},
-});
-
 const rcloneRemoteSchema = z.object({
 	name: z.string(),
 	type: z.string(),

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

@@ -29,12 +29,10 @@ import type { DumpPathKind, UpdateRepositoryBody } from "./repositories.dto";
 import { findCommonAncestor } from "@zerobyte/core/utils";
 import { prepareSnapshotDump } from "./helpers/dump";
 import { emptyRepositoryStats, refreshStoredRepositoryStats } from "./helpers/repository-stats";
-import { restoreExecutor } from "./restore-executor";
 import type { ShortId } from "~/server/utils/branded";
 import { decryptRepositoryConfig, encryptRepositoryConfig } from "./repository-config-secrets";
 import { commands } from "./commands";
 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";
@@ -65,14 +63,6 @@ const assertAllowedControllerLocalRestoreTarget = (target: string) => {
 	}
 };
 
-const updateActiveRestoreTask = (restoreId: string, eventName: string, update: () => void) => {
-	try {
-		update();
-	} catch (error) {
-		logger.warn(`Received ${eventName} for inactive restore ${restoreId}: ${toMessage(error)}`);
-	}
-};
-
 const getLsLimiter = (repositoryId: string) => {
 	let limiter = lsLimiters.get(repositoryId);
 	if (!limiter) {
@@ -101,56 +91,6 @@ const findActiveDoctorTask = (organizationId: string, repositoryShortId: string)
 	});
 };
 
-const markRestoreStarted = (restoreId: string) => {
-	updateActiveRestoreTask(restoreId, "restore.started", () => taskStore.markRunning(restoreId));
-};
-
-const updateRestoreProgress = (restoreId: string, progress: RestoreExecutionProgress) => {
-	updateActiveRestoreTask(restoreId, "restore.progress", () =>
-		taskStore.updateProgress(restoreId, { kind: "restore", progress }),
-	);
-};
-
-const completeRestoreTask = (
-	restoreId: string,
-	result: Extract<RestoreExecutionResult, { status: "completed" }>["result"],
-) => {
-	updateActiveRestoreTask(restoreId, "restore.completed", () =>
-		taskStore.complete(restoreId, { kind: "restore", result }),
-	);
-};
-
-const failRestoreTask = (restoreId: string, error: string) => {
-	updateActiveRestoreTask(restoreId, "restore.failed", () => taskStore.fail(restoreId, error));
-};
-
-const cancelRestoreTask = (restoreId: string, message?: string) => {
-	updateActiveRestoreTask(restoreId, "restore.cancelled", () => taskStore.cancel(restoreId, message ?? null));
-};
-
-const finishRestoreExecution = async (restoreId: string, resultPromise: Promise<RestoreExecutionResult>) => {
-	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;
@@ -466,43 +406,25 @@ const restoreSnapshot = async (
 		throw new ConflictError("A restore is already running for this snapshot");
 	}
 
-	const task = taskStore.create({
+	const repositoryConfig = await decryptRepositoryConfig(repository.config);
+	const command = commands.createRestore({
 		organizationId,
-		resourceType: RESTORE_TASK_RESOURCE_TYPE,
-		resourceId: repository.shortId,
-		operationKey: snapshotId,
-		targetAgentId: useControllerLocalRestoreFallback ? null : executionAgentId,
-		input: { kind: "restore", repositoryId: repository.shortId, snapshotId, target },
+		repositoryId: repository.id,
+		repositoryShortId: repository.shortId,
+		repositoryConfig,
+		snapshotId,
+		target,
+		executionTarget: useControllerLocalRestoreFallback
+			? { kind: "controller" }
+			: { kind: "agent", agentId: executionAgentId },
+		options: {
+			basePath,
+			...restoreExecutionOptions,
+		},
 	});
-	const restoreId = task.id;
-	try {
-		const repositoryConfig = await decryptRepositoryConfig(repository.config);
-		const execution = restoreExecutor.start({
-			restoreId,
-			organizationId,
-			repositoryId: repository.id,
-			repositoryShortId: repository.shortId,
-			repositoryConfig,
-			snapshotId,
-			target,
-			executionAgentId,
-			options: {
-				basePath,
-				...restoreExecutionOptions,
-			},
-			onStarted: () => markRestoreStarted(restoreId),
-			onProgress: (progress) => updateRestoreProgress(restoreId, progress),
-		});
 
-		void finishRestoreExecution(restoreId, execution.result);
-
-		return { restoreId, status: "started" as const };
-	} catch (error) {
-		failRestoreTask(restoreId, toMessage(error));
-		throw error;
-	}
+	return command.start();
 };
-
 const dumpSnapshot = async (shortId: ShortId, snapshotId: string, path?: string, kind?: DumpPathKind) => {
 	const organizationId = getOrganizationId();
 	const repository = await findRepository(shortId);
@@ -648,34 +570,6 @@ const startDoctor = async (shortId: ShortId) => {
 	return command.start();
 };
 
-const cancelDoctor = async (shortId: ShortId) => {
-	const organizationId = getOrganizationId();
-	const repository = await findRepository(shortId);
-
-	if (!repository) {
-		throw new NotFoundError("Repository not found");
-	}
-
-	const activeDoctorTask = findActiveDoctorTask(organizationId, repository.shortId);
-	if (!activeDoctorTask) {
-		if (repository.status === "doctor") {
-			await db
-				.update(repositoriesTable)
-				.set({ status: "unknown" })
-				.where(eq(repositoriesTable.id, repository.id));
-			return { status: "reset" as const };
-		}
-
-		throw new ConflictError("No doctor operation is currently running");
-	}
-
-	if (!commands.cancelDoctor(activeDoctorTask.id)) {
-		throw new ConflictError("No doctor operation is currently running");
-	}
-
-	return { status: "cancelled" as const };
-};
-
 const deleteSnapshot = async (shortId: ShortId, snapshotId: string) => {
 	return deleteSnapshots(shortId, [snapshotId]);
 };
@@ -932,7 +826,6 @@ export const repositoriesService = {
 	getSnapshotDetails,
 	checkHealth,
 	startDoctor,
-	cancelDoctor,
 	deleteSnapshot,
 	deleteSnapshots,
 	tagSnapshots,

app/server/modules/repositories/restore-executor.ts

@@ -1,145 +0,0 @@
-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<Parameters<typeof restic.restore>[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<RestoreExecutionResult>;
-};
-
-const shouldRunInController = (agentId: string) => agentId === LOCAL_AGENT_ID && !appConfig.flags.enableLocalAgent;
-
-const createRestoreRunPayload = async (request: RestoreExecutionRequest): Promise<RestoreRunPayload> => {
-	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<RestoreExecutionResult> => {
-	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) };
-	}
-};
-
-const executeAgentRestore = async (
-	request: RestoreExecutionRequest,
-	signal: AbortSignal,
-): Promise<RestoreExecutionResult> => {
-	if (signal.aborted) {
-		return { status: "cancelled", message: "Restore was cancelled" };
-	}
-
-	try {
-		const payload = await createRestoreRunPayload(request);
-		const started = await agentManager.startRestore(request.executionAgentId, {
-			payload,
-			signal,
-			onStarted: request.onStarted,
-			onProgress: request.onProgress,
-		});
-
-		if (started.status === "unavailable") {
-			return started;
-		}
-
-		return await started.result;
-	} catch (error) {
-		if (signal.aborted) {
-			return { status: "cancelled", message: "Restore was cancelled" };
-		}
-
-		return { status: "failed", error: toMessage(error) };
-	}
-};
-
-const executeRestoreWithRepositoryLock = async (
-	request: RestoreExecutionRequest,
-	signal: AbortSignal,
-): Promise<RestoreExecutionResult> => {
-	let releaseLock: (() => void) | null = null;
-
-	try {
-		releaseLock = await repoMutex.acquireShared(request.repositoryId, `restore:${request.restoreId}`, signal);
-
-		if (shouldRunInController(request.executionAgentId)) {
-			return await executeControllerRestore(request, signal);
-		}
-
-		return await executeAgentRestore(request, signal);
-	} catch (error) {
-		if (signal.aborted) {
-			return { status: "cancelled", message: "Restore was cancelled" };
-		}
-
-		return { status: "failed", error: toMessage(error) };
-	} finally {
-		releaseLock?.();
-	}
-};
-
-export const restoreExecutor = {
-	start: (request: RestoreExecutionRequest): RestoreExecutionHandle => {
-		const abortController = new AbortController();
-
-		return {
-			result: executeRestoreWithRepositoryLock(request, abortController.signal),
-		};
-	},
-};

app/server/modules/tasks/__tests__/tasks.controller.test.ts

@@ -6,6 +6,7 @@ import { taskChangedEventName, tasksSnapshotEventName } from "~/schemas/task-eve
 import type { TaskResourceType } from "~/schemas/tasks";
 import { createTestSession } from "~/test/helpers/auth";
 import { taskStore } from "../tasks.store";
+import { runTaskLifecycle } from "../tasks.lifecycle";
 
 const app = createApp();
 
@@ -84,6 +85,46 @@ beforeEach(async () => {
 });
 
 describe("tasksController", () => {
+	test("requests cancellation for a cancellable running task", async () => {
+		const session = await createTestSession();
+		const task = createRestoreTask(session.organizationId);
+		const lifecycle = runTaskLifecycle({
+			taskId: task.id,
+			label: "restore task",
+			cancellable: true,
+			run: (signal) =>
+				new Promise<never>((_, reject) => {
+					signal.addEventListener("abort", () => reject(new DOMException("Cancelled", "AbortError")), {
+						once: true,
+					});
+				}),
+		});
+
+		const res = await app.request(`/api/v1/tasks/${task.id}/cancel`, {
+			method: "POST",
+			headers: session.headers,
+		});
+
+		expect(res.status).toBe(202);
+		await expect(res.json()).resolves.toEqual({ status: "cancelling" });
+		await lifecycle;
+		expect(taskStore.findById({ organizationId: session.organizationId, taskId: task.id })?.status).toBe(
+			"cancelled",
+		);
+	});
+
+	test("rejects cancellation for a task without cancellation support", async () => {
+		const session = await createTestSession();
+		const task = createTask(session.organizationId);
+
+		const res = await app.request(`/api/v1/tasks/${task.id}/cancel`, {
+			method: "POST",
+			headers: session.headers,
+		});
+
+		expect(res.status).toBe(409);
+	});
+
 	test("returns a task by id", async () => {
 		const session = await createTestSession();
 		const task = createTask(session.organizationId);

app/server/modules/tasks/__tests__/tasks.lifecycle.test.ts

@@ -123,6 +123,7 @@ describe("runTaskLifecycle", () => {
 		const lifecycle = runTaskLifecycle({
 			taskId: task.id,
 			label: "test task",
+			cancellable: true,
 			run: (signal) =>
 				new Promise<never>((_, reject) => {
 					signal.addEventListener(

app/server/modules/tasks/tasks.controller.ts

@@ -1,18 +1,21 @@
 import { Hono } from "hono";
 import { validator } from "hono-openapi";
-import { NotFoundError } from "http-errors-enhanced";
+import { ConflictError, NotFoundError } from "http-errors-enhanced";
 import { taskChangedEventName, tasksSnapshotEventName, type TaskEventPayloadMap } from "~/schemas/task-events";
 import { streamEvents } from "../events/server-event-stream";
 import { requireAuth } from "../auth/auth.middleware";
 import {
 	getTaskDto,
+	cancelTaskDto,
 	listTasksDto,
 	listTasksQuery,
 	streamTaskEventsDto,
 	streamTasksEventsDto,
 	type GetTaskDto,
+	type CancelTaskDto,
 	type ListTasksDto,
 } from "./tasks.dto";
+import { requestTaskCancel } from "./tasks.lifecycle";
 import { toTaskDto } from "./tasks.presenter";
 import { taskStore } from "./tasks.store";
 
@@ -106,4 +109,19 @@ export const tasksController = new Hono()
 		}
 
 		return c.json<GetTaskDto>(toTaskDto(task), 200);
+	})
+	.post("/:taskId/cancel", cancelTaskDto, async (c) => {
+		const organizationId = c.get("organizationId");
+		const taskId = c.req.param("taskId");
+		const task = taskStore.findById({ organizationId, taskId });
+
+		if (!task) {
+			throw new NotFoundError("Task not found");
+		}
+
+		if (!requestTaskCancel(taskId)) {
+			throw new ConflictError("Task is not cancellable or is no longer running");
+		}
+
+		return c.json<CancelTaskDto>({ status: "cancelling" }, 202);
 	});

app/server/modules/tasks/tasks.dto.ts

@@ -33,6 +33,12 @@ export type ListTasksDto = z.infer<typeof listTasksResponse>;
 
 export type GetTaskDto = z.infer<typeof taskResponse>;
 
+const cancelTaskResponse = z.object({
+	status: z.literal("cancelling"),
+});
+
+export type CancelTaskDto = z.infer<typeof cancelTaskResponse>;
+
 export const listTasksDto = describeRoute({
 	description: "List active tasks",
 	tags: ["Tasks"],
@@ -65,6 +71,25 @@ export const getTaskDto = describeRoute({
 	},
 });
 
+export const cancelTaskDto = describeRoute({
+	description: "Request cancellation of a running task",
+	tags: ["Tasks"],
+	operationId: "cancelTask",
+	responses: {
+		202: {
+			description: "Task cancellation requested",
+			content: {
+				"application/json": {
+					schema: resolver(cancelTaskResponse),
+				},
+			},
+		},
+		409: {
+			description: "Task is not cancellable or is no longer running",
+		},
+	},
+});
+
 export const streamTaskEventsDto = describeRoute({
 	description: "Subscribe to lifecycle events for one task",
 	tags: ["Tasks"],

app/server/modules/tasks/tasks.lifecycle.ts

@@ -17,6 +17,7 @@ export class TaskCancelledError<TResult extends TaskResult = TaskResult> extends
 type TaskLifecycleOptions<TResult extends TaskResult> = {
 	taskId: string;
 	label: string;
+	cancellable?: boolean;
 	run: (signal: AbortSignal) => Promise<TResult>;
 	onStarted?: (task: ParsedTask) => void | Promise<void>;
 	onSucceeded?: (task: ParsedTask, result: TResult) => void;
@@ -97,7 +98,9 @@ export const requestTaskCancel = (taskId: string) => {
 
 export const runTaskLifecycle = async <TResult extends TaskResult>(options: TaskLifecycleOptions<TResult>) => {
 	const abortController = new AbortController();
-	abortControllers.set(options.taskId, abortController);
+	if (options.cancellable) {
+		abortControllers.set(options.taskId, abortController);
+	}
 
 	try {
 		const startedTask = taskStore.markRunning(options.taskId);
@@ -133,6 +136,8 @@ export const runTaskLifecycle = async <TResult extends TaskResult>(options: Task
 			emitTaskLifecycleEvent("task:finished", failedTask);
 		}
 	} finally {
-		abortControllers.delete(options.taskId);
+		if (options.cancellable) {
+			abortControllers.delete(options.taskId);
+		}
 	}
 };