refactor: split active tasks vs tracked in their own hook

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

@@ -5250,6 +5250,7 @@ export type ListTasksData = {
         kind?: 'backup' | 'restore' | 'deleteSnapshots' | 'tagSnapshots' | 'doctor';
         resourceType?: 'backup_schedule' | 'repository';
         resourceId?: string;
+        operationKey?: string;
     };
     url: '/api/v1/tasks';
 };
@@ -5264,6 +5265,7 @@ export type ListTasksResponses = {
         status: 'queued' | 'running' | 'cancelling' | 'cancelled' | 'succeeded' | 'failed' | 'stale';
         resourceType: 'backup_schedule' | 'repository';
         resourceId: string;
+        operationKey: string | null;
         targetAgentId: string | null;
         input: {
             kind: 'backup';
@@ -5387,6 +5389,7 @@ export type StreamTasksEventsData = {
         kind?: 'backup' | 'restore' | 'deleteSnapshots' | 'tagSnapshots' | 'doctor';
         resourceType?: 'backup_schedule' | 'repository';
         resourceId?: string;
+        operationKey?: string;
     };
     url: '/api/v1/tasks/events';
 };
@@ -5401,6 +5404,7 @@ export type StreamTasksEventsResponses = {
         status: 'queued' | 'running' | 'cancelling' | 'cancelled' | 'succeeded' | 'failed' | 'stale';
         resourceType: 'backup_schedule' | 'repository';
         resourceId: string;
+        operationKey: string | null;
         targetAgentId: string | null;
         input: {
             kind: 'backup';
@@ -5518,6 +5522,7 @@ export type StreamTasksEventsResponses = {
         status: 'queued' | 'running' | 'cancelling' | 'cancelled' | 'succeeded' | 'failed' | 'stale';
         resourceType: 'backup_schedule' | 'repository';
         resourceId: string;
+        operationKey: string | null;
         targetAgentId: string | null;
         input: {
             kind: 'backup';
@@ -5653,6 +5658,7 @@ export type StreamTaskEventsResponses = {
         status: 'queued' | 'running' | 'cancelling' | 'cancelled' | 'succeeded' | 'failed' | 'stale';
         resourceType: 'backup_schedule' | 'repository';
         resourceId: string;
+        operationKey: string | null;
         targetAgentId: string | null;
         input: {
             kind: 'backup';
@@ -5788,6 +5794,7 @@ export type GetTaskResponses = {
         status: 'queued' | 'running' | 'cancelling' | 'cancelled' | 'succeeded' | 'failed' | 'stale';
         resourceType: 'backup_schedule' | 'repository';
         resourceId: string;
+        operationKey: string | null;
         targetAgentId: string | null;
         input: {
             kind: 'backup';

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

@@ -1,10 +1,12 @@
 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, createTestQueryClient, render, screen, userEvent, waitFor, within } from "~/test/test-utils";
 import { taskChangedEventName } from "~/schemas/task-events";
+import type { TaskOfKind } from "~/client/hooks/use-active-tasks";
+import type { Repository } from "~/client/lib/types";
 import { fromAny } from "@total-typescript/shoehorn";
+import { RestoreSnapshotPage } from "~/client/modules/repositories/routes/restore-snapshot";
 
 vi.mock("@tanstack/react-router", async (importOriginal) => {
 	const actual = await importOriginal<typeof import("@tanstack/react-router")>();
@@ -54,9 +56,9 @@ class MockEventSource {
 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" }>;
+type TaskResponse = TaskOfKind<"restore">;
+type RestoreProgress = NonNullable<TaskResponse["progress"]>["progress"];
+type RestoreResult = NonNullable<TaskResponse["result"]>;
 
 const createRestoreTask = (
 	options: {
@@ -79,6 +81,7 @@ const createRestoreTask = (
 		status,
 		resourceType: "repository",
 		resourceId: repositoryId,
+		operationKey: options.snapshotId ?? snapshotId,
 		targetAgentId: null,
 		input: {
 			kind: "restore",
@@ -119,13 +122,13 @@ const renderRestoreForm = (queryClient = createTestQueryClient()) => {
 	);
 };
 
-const renderRestoreFormWithPrefetchedTasks = async (tasks: ListTasksResponse) => {
+const renderRestoreFormWithPrefetchedTask = async (task: TaskResponse) => {
 	server.use(
-		http.get("/api/v1/tasks", () => HttpResponse.json(tasks)),
+		http.get("/api/v1/tasks", () => HttpResponse.json([task])),
 		snapshotFilesHandler,
 	);
 	const queryClient = createTestQueryClient();
-	await queryClient.ensureQueryData(restoreTasksOptions(repositoryId));
+	await queryClient.ensureQueryData(restoreTasksOptions(repositoryId, snapshotId));
 	renderRestoreForm(queryClient);
 
 	await waitFor(() => {
@@ -147,30 +150,23 @@ afterEach(() => {
 });
 
 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,
-		]);
+	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);
 		expect(screen.getByText("Restore in progress")).toBeTruthy();
-		expect(eventSource?.url).toBe("/api/v1/tasks/events?kind=restore&resourceType=repository&resourceId=repo-1");
+		expect(taskStream?.url).toBe(
+			"/api/v1/tasks/events?kind=restore&resourceType=repository&resourceId=repo-1&operationKey=snap-1",
+		);
+		expect(MockEventSource.instances).toHaveLength(1);
 	});
 
-	test("renders canonical task progress without creating another event stream", async () => {
+	test("renders canonical task progress from the exact filtered stream", async () => {
 		const activeRestore = createRestoreTask();
-		const eventSource = await renderRestoreFormWithPrefetchedTasks([activeRestore]);
+		const taskStream = await renderRestoreFormWithPrefetchedTask(activeRestore);
 
-		eventSource?.emit(
+		taskStream?.emit(
 			taskChangedEventName,
 			createRestoreTask({
 				progress: {
@@ -218,7 +214,16 @@ describe("RestoreForm", () => {
 		expect(screen.getByRole("button", { name: "Restoring..." }).hasAttribute("disabled")).toBe(true);
 		expect(screen.queryByRole("button", { name: "Restore All" })).toBeNull();
 
-		MockEventSource.instances[0]?.emit(taskChangedEventName, createRestoreTask());
+		await waitFor(() => {
+			expect(
+				MockEventSource.instances.some(
+					(eventSource) => eventSource.url === "/api/v1/tasks/task-restore/events",
+				),
+			).toBe(true);
+		});
+		MockEventSource.instances
+			.find((eventSource) => eventSource.url === "/api/v1/tasks/task-restore/events")
+			?.emit(taskChangedEventName, createRestoreTask());
 
 		expect(await screen.findByRole("button", { name: "Restoring..." })).toBeTruthy();
 	});
@@ -263,9 +268,9 @@ describe("RestoreForm", () => {
 		"clears restoring state and shows $status terminal feedback",
 		async ({ status, result, error, title, description }) => {
 			const activeRestore = createRestoreTask();
-			const eventSource = await renderRestoreFormWithPrefetchedTasks([activeRestore]);
+			const taskStream = await renderRestoreFormWithPrefetchedTask(activeRestore);
 
-			eventSource?.emit(
+			taskStream?.emit(
 				taskChangedEventName,
 				createRestoreTask({
 					status,
@@ -285,6 +290,38 @@ describe("RestoreForm", () => {
 		},
 	);
 
+	test("resets restore lifecycle state when the repository snapshot scope changes", async () => {
+		server.use(snapshotFilesHandler);
+		const repository: Repository = fromAny({ shortId: repositoryId, name: "Repo 1" });
+		const { rerender } = render(
+			<RestoreSnapshotPage repository={repository} snapshotId="snap-1" returnPath="/repositories/repo-1" />,
+		);
+		const firstStreamUrl =
+			"/api/v1/tasks/events?kind=restore&resourceType=repository&resourceId=repo-1&operationKey=snap-1";
+
+		await waitFor(() => {
+			expect(MockEventSource.instances.some(({ url }) => url === firstStreamUrl)).toBe(true);
+		});
+		const firstStream = MockEventSource.instances.find(({ url }) => url === firstStreamUrl);
+		firstStream?.emit(taskChangedEventName, createRestoreTask());
+		firstStream?.emit(taskChangedEventName, createRestoreTask({ status: "succeeded", updatedAt: 1711411200001 }));
+
+		expect(await screen.findByText("Restore completed")).toBeTruthy();
+
+		rerender(<RestoreSnapshotPage repository={repository} snapshotId="snap-2" returnPath="/repositories/repo-1" />);
+
+		expect(screen.queryByText("Restore completed")).toBeNull();
+		await waitFor(() => {
+			expect(
+				MockEventSource.instances.some(
+					({ url }) =>
+						url ===
+						"/api/v1/tasks/events?kind=restore&resourceType=repository&resourceId=repo-1&operationKey=snap-2",
+				),
+			).toBe(true);
+		});
+	});
+
 	test("restores the selected ancestor folder path from a broader display root", async () => {
 		let restoreRequestBody: unknown;
 

app/client/hooks/__tests__/use-active-tasks.test.tsx

@@ -1,9 +1,9 @@
 import { afterEach, beforeEach, describe, expect, test, vi } from "vitest";
 import type { ListTasksResponse } from "~/client/api-client";
 import { cleanup, createTestQueryClient, render, waitFor } from "~/test/test-utils";
-import { taskChangedEventName } from "~/schemas/task-events";
+import { taskChangedEventName, tasksSnapshotEventName } from "~/schemas/task-events";
 import type { TaskDto } from "~/schemas/tasks";
-import { taskEventsOptions, useTaskEvents } from "../use-task-events";
+import { taskEventsOptions, useActiveTasks, type TaskOfKind } from "../use-active-tasks";
 
 class MockEventSource {
 	static instances: MockEventSource[] = [];
@@ -42,22 +42,25 @@ class MockEventSource {
 const originalEventSource = globalThis.EventSource;
 
 const filter = {
-	kind: "deleteSnapshots",
+	kind: "restore",
 	resourceType: "repository",
 	resourceId: "repo-1",
+	operationKey: "snap-1",
 } as const;
 
 const activeTask: TaskDto = {
-	id: "task-delete",
-	kind: "deleteSnapshots",
+	id: "task-restore",
+	kind: "restore",
 	status: "running",
 	resourceType: "repository",
 	resourceId: "repo-1",
+	operationKey: "snap-1",
 	targetAgentId: null,
 	input: {
-		kind: "deleteSnapshots",
+		kind: "restore",
 		repositoryId: "repo-1",
-		snapshotIds: ["snap-1"],
+		snapshotId: "snap-1",
+		target: "/restore",
 	},
 	progress: null,
 	result: null,
@@ -72,17 +75,20 @@ const activeTask: TaskDto = {
 const finishedTask: TaskDto = {
 	...activeTask,
 	status: "succeeded",
-	result: { kind: "deleteSnapshots", deletedSnapshotIds: ["snap-1"] },
+	result: {
+		kind: "restore",
+		result: { message_type: "summary", files_restored: 1, files_skipped: 0 },
+	},
 	updatedAt: 1711411201000,
 	finishedAt: 1711411201000,
 };
 
-const TaskEventsConsumer = ({ onTaskFinished }: { onTaskFinished: (task: TaskDto) => void }) => {
-	useTaskEvents(filter, { onTaskFinished });
+const ActiveTasksConsumer = ({ onTaskFinished }: { onTaskFinished: (task: TaskOfKind<"restore">) => void }) => {
+	useActiveTasks(filter, { onTaskFinished });
 	return null;
 };
 
-describe("useTaskEvents", () => {
+describe("useActiveTasks", () => {
 	beforeEach(() => {
 		MockEventSource.reset();
 		globalThis.EventSource = MockEventSource as unknown as typeof EventSource;
@@ -94,20 +100,31 @@ describe("useTaskEvents", () => {
 		MockEventSource.reset();
 	});
 
-	test("ignores stale active events after a task has finished", async () => {
+	test("uses the exact operation URL and cache while reporting the finished restore once", async () => {
 		const queryClient = createTestQueryClient();
 		const onTaskFinished = vi.fn();
 
-		render(<TaskEventsConsumer onTaskFinished={onTaskFinished} />, { queryClient });
+		render(<ActiveTasksConsumer onTaskFinished={onTaskFinished} />, { queryClient });
 
 		await waitFor(() => {
 			expect(MockEventSource.instances).toHaveLength(1);
 		});
+		expect(MockEventSource.instances[0]?.url).toBe(
+			"/api/v1/tasks/events?kind=restore&resourceType=repository&resourceId=repo-1&operationKey=snap-1",
+		);
+
+		MockEventSource.instances[0]?.emit(tasksSnapshotEventName, [activeTask]);
+		await waitFor(() => {
+			expect(queryClient.getQueryData<ListTasksResponse>(taskEventsOptions(filter).queryKey)).toEqual([
+				activeTask,
+			]);
+		});
 
 		MockEventSource.instances[0]?.emit(taskChangedEventName, finishedTask);
 		await waitFor(() => {
 			expect(onTaskFinished).toHaveBeenCalledTimes(1);
 		});
+		expect(onTaskFinished.mock.calls[0]?.[0].input.snapshotId).toBe("snap-1");
 
 		MockEventSource.instances[0]?.emit(taskChangedEventName, activeTask);
 

app/client/hooks/use-active-tasks.ts

@@ -1,14 +1,23 @@
 import { useEffect, useMemo, useRef } from "react";
-import { useQuery, useQueryClient } from "@tanstack/react-query";
+import { useQuery, useQueryClient, type UseQueryResult } from "@tanstack/react-query";
 import type { ListTasksData, ListTasksResponse } from "~/client/api-client";
 import { getTaskOptions, listTasksOptions } from "~/client/api-client/@tanstack/react-query.gen";
 import { logger } from "~/client/lib/logger";
 import { taskChangedEventName, tasksSnapshotEventName } from "~/schemas/task-events";
-import { activeTaskStatuses, type TaskDto } from "~/schemas/tasks";
+import { activeTaskStatuses, type TaskDto, type TaskKind } from "~/schemas/tasks";
 
 export type TaskEventsQuery = NonNullable<ListTasksData["query"]>;
-type UseTaskEventsOptions = {
-	onTaskFinished?: (task: TaskDto) => void;
+export type TaskOfKind<K extends TaskKind> = TaskDto & {
+	kind: K;
+	input: Extract<TaskDto["input"], { kind: K }>;
+	progress: Extract<NonNullable<TaskDto["progress"]>, { kind: K }> | null;
+	result: Extract<NonNullable<TaskDto["result"]>, { kind: K }> | null;
+};
+
+type TaskForQuery<Q extends TaskEventsQuery> = Q extends { kind: infer K extends TaskKind } ? TaskOfKind<K> : TaskDto;
+
+type UseActiveTasksOptions<Q extends TaskEventsQuery> = {
+	onTaskFinished?: (task: TaskForQuery<Q>) => void;
 };
 
 const parseTaskEvent = (event: Event): TaskDto => {
@@ -19,7 +28,7 @@ const parseTasksSnapshotEvent = (event: Event): TaskDto[] => {
 	return JSON.parse((event as MessageEvent<string>).data) as TaskDto[];
 };
 
-const isActiveTask = (task: TaskDto) => {
+export const isTaskActive = (task: Pick<TaskDto, "status">) => {
 	return activeTaskStatuses.some((status) => status === task.status);
 };
 
@@ -28,6 +37,7 @@ const getTasksEventUrl = (query: TaskEventsQuery) => {
 	if (query.kind) params.set("kind", query.kind);
 	if (query.resourceType) params.set("resourceType", query.resourceType);
 	if (query.resourceId) params.set("resourceId", query.resourceId);
+	if (query.operationKey) params.set("operationKey", query.operationKey);
 
 	const queryString = params.toString();
 	if (!queryString) {
@@ -60,13 +70,14 @@ export const taskEventsOptions = (query: TaskEventsQuery) => {
 	return listTasksOptions({ query });
 };
 
-export const useTaskEvents = (query: TaskEventsQuery, options: UseTaskEventsOptions = {}) => {
+export const useActiveTasks = <const Q extends TaskEventsQuery>(query: Q, options: UseActiveTasksOptions<Q> = {}) => {
 	const queryClient = useQueryClient();
 	const onTaskFinishedRef = useRef(options.onTaskFinished);
 	const finishedTasksRef = useRef(new Map<string, TaskDto>());
 	const queryKind = query.kind;
 	const queryResourceType = query.resourceType;
 	const queryResourceId = query.resourceId;
+	const queryOperationKey = query.operationKey;
 	onTaskFinishedRef.current = options.onTaskFinished;
 
 	const taskListOptions = useMemo(() => {
@@ -74,16 +85,18 @@ export const useTaskEvents = (query: TaskEventsQuery, options: UseTaskEventsOpti
 			kind: queryKind,
 			resourceType: queryResourceType,
 			resourceId: queryResourceId,
+			operationKey: queryOperationKey,
 		});
-	}, [queryKind, queryResourceId, queryResourceType]);
+	}, [queryKind, queryOperationKey, queryResourceId, queryResourceType]);
 
 	const taskEventsUrl = useMemo(() => {
 		return getTasksEventUrl({
 			kind: queryKind,
 			resourceType: queryResourceType,
 			resourceId: queryResourceId,
+			operationKey: queryOperationKey,
 		});
-	}, [queryKind, queryResourceId, queryResourceType]);
+	}, [queryKind, queryOperationKey, queryResourceId, queryResourceType]);
 
 	const taskQueryKeyRef = useRef(taskListOptions.queryKey);
 	taskQueryKeyRef.current = taskListOptions.queryKey;
@@ -91,9 +104,24 @@ export const useTaskEvents = (query: TaskEventsQuery, options: UseTaskEventsOpti
 	const tasks = useQuery({ ...taskListOptions, enabled: false });
 
 	useEffect(() => {
-		const eventSource = new EventSource(taskEventsUrl);
+		const finishTask = (task: TaskDto) => {
+			finishedTasksRef.current.set(task.id, task);
+			onTaskFinishedRef.current?.(task as TaskForQuery<Q>);
+		};
 
-		eventSource.addEventListener(tasksSnapshotEventName, (event) => {
+		const reconcileMissingTask = async (missingTask: Pick<TaskDto, "id" | "updatedAt">) => {
+			const fetchedTask = (await queryClient.fetchQuery(
+				getTaskOptions({ path: { taskId: missingTask.id } }),
+			)) as TaskDto;
+
+			if (isTaskActive(fetchedTask)) return;
+			if (fetchedTask.updatedAt < missingTask.updatedAt) return;
+			if (hasTaskFinished(finishedTasksRef.current, fetchedTask)) return;
+
+			finishTask(fetchedTask);
+		};
+
+		const handleTasksSnapshot = (event: Event) => {
 			const snapshot = parseTasksSnapshotEvent(event);
 			const currentTasks = queryClient.getQueryData<ListTasksResponse>(taskQueryKeyRef.current) ?? [];
 
@@ -106,32 +134,17 @@ export const useTaskEvents = (query: TaskEventsQuery, options: UseTaskEventsOpti
 			);
 
 			for (const missingTask of missingTasks) {
-				void queryClient
-					.fetchQuery(getTaskOptions({ path: { taskId: missingTask.id } }))
-					.then((task) => {
-						const fetchedTask = task as TaskDto;
-						if (
-							isActiveTask(fetchedTask) ||
-							fetchedTask.updatedAt < missingTask.updatedAt ||
-							hasTaskFinished(finishedTasksRef.current, fetchedTask)
-						) {
-							return;
-						}
-
-						finishedTasksRef.current.set(fetchedTask.id, fetchedTask);
-						onTaskFinishedRef.current?.(fetchedTask);
-					})
-					.catch((error: unknown) => {
-						logger.error("[SSE] Failed to reconcile missing task:", error);
-					});
+				void reconcileMissingTask(missingTask).catch((error: unknown) => {
+					logger.error("[SSE] Failed to reconcile missing task:", error);
+				});
 			}
-		});
+		};
 
-		eventSource.addEventListener(taskChangedEventName, (event) => {
+		const handleTaskChanged = (event: Event) => {
 			const task = parseTaskEvent(event);
 			const activeTasks = queryClient.getQueryData<ListTasksResponse>(taskQueryKeyRef.current) ?? [];
 
-			if (isActiveTask(task)) {
+			if (isTaskActive(task)) {
 				if (hasTaskFinished(finishedTasksRef.current, task)) {
 					return;
 				}
@@ -154,9 +167,12 @@ export const useTaskEvents = (query: TaskEventsQuery, options: UseTaskEventsOpti
 				return;
 			}
 
-			finishedTasksRef.current.set(task.id, task);
-			onTaskFinishedRef.current?.(task);
-		});
+			finishTask(task);
+		};
+
+		const eventSource = new EventSource(taskEventsUrl);
+		eventSource.addEventListener(tasksSnapshotEventName, handleTasksSnapshot);
+		eventSource.addEventListener(taskChangedEventName, handleTaskChanged);
 
 		eventSource.onerror = (error) => {
 			logger.error("[SSE] Task stream connection error:", error);
@@ -167,5 +183,5 @@ export const useTaskEvents = (query: TaskEventsQuery, options: UseTaskEventsOpti
 		};
 	}, [queryClient, taskEventsUrl]);
 
-	return tasks;
+	return tasks as UseQueryResult<TaskForQuery<Q>[]>;
 };

app/client/hooks/use-task.ts

@@ -0,0 +1,50 @@
+import { useEffect, useState } from "react";
+import type { GetTaskResponse } from "~/client/api-client";
+import { logger } from "~/client/lib/logger";
+import { taskChangedEventName } from "~/schemas/task-events";
+
+type StreamedTask<T extends GetTaskResponse> = {
+	taskId: string;
+	task: T;
+};
+
+const parseTaskEvent = <T extends GetTaskResponse>(event: Event): T => {
+	return JSON.parse((event as MessageEvent<string>).data) as T;
+};
+
+export const useTask = <T extends GetTaskResponse = GetTaskResponse>(taskId: string | null | undefined) => {
+	const [streamedTask, setStreamedTask] = useState<StreamedTask<T> | null>(null);
+
+	useEffect(() => {
+		if (!taskId) return;
+
+		let isCurrent = true;
+		const eventSource = new EventSource(`/api/v1/tasks/${taskId}/events`);
+		const updateTask = (event: Event) => {
+			const nextTask = parseTaskEvent<T>(event);
+			if (!isCurrent || nextTask.id !== taskId) return;
+
+			setStreamedTask((current) => {
+				if (current?.taskId === taskId && nextTask.updatedAt < current.task.updatedAt) {
+					return current;
+				}
+
+				return { taskId, task: nextTask };
+			});
+		};
+
+		eventSource.addEventListener(taskChangedEventName, updateTask);
+		eventSource.onerror = (error) => {
+			logger.error(`[SSE] Task ${taskId} connection error:`, error);
+		};
+
+		return () => {
+			isCurrent = false;
+			eventSource.close();
+		};
+	}, [taskId]);
+
+	return {
+		task: streamedTask && streamedTask.taskId === taskId ? streamedTask.task : null,
+	};
+};

app/client/modules/backups/routes/__tests__/backup-details.test.tsx

@@ -138,6 +138,7 @@ const deleteSnapshotsTask = {
 	status: "running",
 	resourceType: "repository",
 	resourceId: "repo-1",
+	operationKey: null,
 	targetAgentId: null,
 	input: {
 		kind: "deleteSnapshots",

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

@@ -1,6 +1,12 @@
 import { toast } from "sonner";
-import { taskEventsOptions, useTaskEvents, type TaskEventsQuery } from "~/client/hooks/use-task-events";
-import type { TaskDto } from "~/schemas/tasks";
+import {
+	taskEventsOptions,
+	useActiveTasks,
+	type TaskEventsQuery,
+	type TaskOfKind,
+} from "~/client/hooks/use-active-tasks";
+
+type DoctorTask = TaskOfKind<"doctor">;
 
 const doctorTasksFilter = (repositoryId: string) => {
 	return {
@@ -10,7 +16,7 @@ const doctorTasksFilter = (repositoryId: string) => {
 	} satisfies TaskEventsQuery;
 };
 
-const applyDoctorTaskFinished = (task: TaskDto) => {
+const applyDoctorTaskFinished = (task: DoctorTask) => {
 	if (task.status === "cancelled") {
 		toast.info("Doctor cancelled");
 		return;
@@ -23,7 +29,7 @@ const applyDoctorTaskFinished = (task: TaskDto) => {
 		return;
 	}
 
-	const result = task.result?.kind === "doctor" ? task.result : null;
+	const result = task.result;
 	if (result?.repositoryStatus === "healthy") {
 		toast.success("Doctor completed");
 		return;
@@ -39,7 +45,7 @@ export const doctorTasksOptions = (repositoryId: string) => {
 };
 
 export const useRepositoryDoctorTask = (repositoryId: string) => {
-	const doctorTasks = useTaskEvents(doctorTasksFilter(repositoryId), {
+	const doctorTasks = useActiveTasks(doctorTasksFilter(repositoryId), {
 		onTaskFinished: applyDoctorTaskFinished,
 	});
 	const activeDoctorTask = doctorTasks.data?.[0] ?? null;

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

@@ -1,71 +1,50 @@
-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";
+import { useCallback, useState } from "react";
+import {
+	isTaskActive,
+	taskEventsOptions,
+	useActiveTasks,
+	type TaskEventsQuery,
+	type TaskOfKind,
+} from "~/client/hooks/use-active-tasks";
+import { useTask } from "~/client/hooks/use-task";
 
-type TaskResponse = ListTasksResponse[number];
+type RestoreTask = TaskOfKind<"restore">;
 
-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 restoreTasksFilter = (repositoryId: string, snapshotId: string) => {
+	return {
+		kind: "restore",
+		resourceType: "repository",
+		resourceId: repositoryId,
+		operationKey: snapshotId,
+	} satisfies TaskEventsQuery;
 };
 
-const isActive = (task: TaskResponse) => {
-	return activeTaskStatuses.some((status) => status === task.status);
-};
-
-export const restoreTasksOptions = (repositoryId: string) => {
-	return taskEventsOptions(restoreTasksFilter(repositoryId));
+export const restoreTasksOptions = (repositoryId: string, snapshotId: string) => {
+	return taskEventsOptions(restoreTasksFilter(repositoryId, snapshotId));
 };
 
 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 [retainedFinishedTask, setRetainedFinishedTask] = useState<RestoreTask | null>(null);
+	const filter = restoreTasksFilter(repositoryId, snapshotId);
+	const { data: activeRestoreTasks } = useActiveTasks(filter, {
+		onTaskFinished: setRetainedFinishedTask,
 	});
-
-	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 { 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 clearFinishedRestoreTask = useCallback(() => {
-		setStreamedFinishedTask(null);
+		setRetainedFinishedTask(null);
 	}, []);
 
-	const restoreProgress =
-		restoreIsActive && restoreTask?.progress?.kind === "restore" ? restoreTask.progress.progress : null;
+	const restoreProgress = taskIsActive ? (restoreTask?.progress?.progress ?? null) : null;
 
 	return {
 		restoreProgress,
-		finishedRestoreTask: restoreIsActive ? null : restoreTask,
+		finishedRestoreTask,
 		clearFinishedRestoreTask,
-		isRestoreRunning: restoreTask ? restoreIsActive : startedTaskId !== undefined,
+		isRestoreRunning: (startedTaskId !== undefined && exactStartedTask === null) || taskIsActive,
 	};
 };

app/client/modules/repositories/routes/restore-snapshot.tsx

@@ -15,6 +15,7 @@ export function RestoreSnapshotPage(props: Props) {
 
 	return (
 		<RestoreForm
+			key={`${repository.shortId}:${snapshotId}`}
 			repository={repository}
 			snapshotId={snapshotId}
 			returnPath={returnPath}

app/client/modules/repositories/snapshots/delete-tasks.ts

@@ -2,27 +2,25 @@ import { useMemo } from "react";
 import type { QueryClient } from "@tanstack/react-query";
 import { useQueryClient } from "@tanstack/react-query";
 import { toast } from "sonner";
-import type { ListSnapshotsResponse, ListTasksData } from "~/client/api-client";
+import type { ListSnapshotsResponse } from "~/client/api-client";
 import { listSnapshotsQueryKey } from "~/client/api-client/@tanstack/react-query.gen";
-import { taskEventsOptions, useTaskEvents } from "~/client/hooks/use-task-events";
-import type { TaskDto } from "~/schemas/tasks";
+import {
+	taskEventsOptions,
+	useActiveTasks,
+	type TaskEventsQuery,
+	type TaskOfKind,
+} from "~/client/hooks/use-active-tasks";
 
-type DeleteSnapshotsTaskInput = Extract<TaskDto["input"], { kind: "deleteSnapshots" }>;
-type DeleteSnapshotsTaskResult = Extract<NonNullable<TaskDto["result"]>, { kind: "deleteSnapshots" }>;
-type DeleteSnapshotsTask = TaskDto & {
-	input: DeleteSnapshotsTaskInput;
-	result: DeleteSnapshotsTaskResult | null;
-};
-type DeleteSnapshotTasksFilter = NonNullable<ListTasksData["query"]>;
+type DeleteSnapshotsTask = TaskOfKind<"deleteSnapshots">;
 
 const emptyDeletingSnapshotIds = new Set<string>();
 
-const deleteSnapshotTasksFilter = (repositoryId: string): DeleteSnapshotTasksFilter => {
+const deleteSnapshotTasksFilter = (repositoryId: string) => {
 	return {
 		kind: "deleteSnapshots",
 		resourceType: "repository",
 		resourceId: repositoryId,
-	};
+	} satisfies TaskEventsQuery;
 };
 
 const removeSnapshotsFromCache = (queryClient: QueryClient, repositoryId: string, snapshotIds: string[]) => {
@@ -62,8 +60,8 @@ export const deleteSnapshotTasksOptions = (repositoryId: string) => {
 export const useDeletingSnapshots = (repositoryId: string) => {
 	const queryClient = useQueryClient();
 	const filter = deleteSnapshotTasksFilter(repositoryId);
-	const deleteTasks = useTaskEvents(filter, {
-		onTaskFinished: (task) => applyDeleteSnapshotsTaskFinished(queryClient, task as DeleteSnapshotsTask),
+	const deleteTasks = useActiveTasks(filter, {
+		onTaskFinished: (task) => applyDeleteSnapshotsTaskFinished(queryClient, task),
 	});
 
 	const deletingSnapshotIds = useMemo(() => {
@@ -72,9 +70,7 @@ export const useDeletingSnapshots = (repositoryId: string) => {
 		}
 
 		const snapshotIds = new Set<string>();
-		const deleteSnapshotTasks = deleteTasks.data as DeleteSnapshotsTask[];
-
-		for (const task of deleteSnapshotTasks) {
+		for (const task of deleteTasks.data) {
 			for (const snapshotId of task.input.snapshotIds) {
 				snapshotIds.add(snapshotId);
 			}

app/client/modules/repositories/snapshots/tag-tasks.ts

@@ -2,28 +2,26 @@ import { useMemo } from "react";
 import type { QueryClient } from "@tanstack/react-query";
 import { useQueryClient } from "@tanstack/react-query";
 import { toast } from "sonner";
-import type { ListSnapshotsResponse, ListTasksData } from "~/client/api-client";
+import type { ListSnapshotsResponse } from "~/client/api-client";
 import { listSnapshotsQueryKey } from "~/client/api-client/@tanstack/react-query.gen";
-import { taskEventsOptions, useTaskEvents } from "~/client/hooks/use-task-events";
+import {
+	taskEventsOptions,
+	useActiveTasks,
+	type TaskEventsQuery,
+	type TaskOfKind,
+} from "~/client/hooks/use-active-tasks";
 import type { BackupSchedule } from "~/client/lib/types";
-import type { TaskDto } from "~/schemas/tasks";
 
-type TagSnapshotsTaskInput = Extract<TaskDto["input"], { kind: "tagSnapshots" }>;
-type TagSnapshotsTaskResult = Extract<NonNullable<TaskDto["result"]>, { kind: "tagSnapshots" }>;
-type TagSnapshotsTask = TaskDto & {
-	input: TagSnapshotsTaskInput;
-	result: TagSnapshotsTaskResult | null;
-};
-type TagSnapshotTasksFilter = NonNullable<ListTasksData["query"]>;
+type TagSnapshotsTask = TaskOfKind<"tagSnapshots">;
 
 const emptyTaggingSnapshotIds = new Set<string>();
 
-const tagSnapshotTasksFilter = (repositoryId: string): TagSnapshotTasksFilter => {
+const tagSnapshotTasksFilter = (repositoryId: string) => {
 	return {
 		kind: "tagSnapshots",
 		resourceType: "repository",
 		resourceId: repositoryId,
-	};
+	} satisfies TaskEventsQuery;
 };
 
 const applyTaskTagsToSnapshot = (snapshot: ListSnapshotsResponse[number], task: TagSnapshotsTask) => {
@@ -131,8 +129,8 @@ export const tagSnapshotTasksOptions = (repositoryId: string) => {
 export const useTaggingSnapshots = (repositoryId: string, backups: BackupSchedule[]) => {
 	const queryClient = useQueryClient();
 	const filter = tagSnapshotTasksFilter(repositoryId);
-	const tagTasks = useTaskEvents(filter, {
-		onTaskFinished: (task) => applyTagSnapshotsTaskFinished(queryClient, task as TagSnapshotsTask, backups),
+	const tagTasks = useActiveTasks(filter, {
+		onTaskFinished: (task) => applyTagSnapshotsTaskFinished(queryClient, task, backups),
 	});
 
 	const taggingSnapshotIds = useMemo(() => {
@@ -141,9 +139,7 @@ export const useTaggingSnapshots = (repositoryId: string, backups: BackupSchedul
 		}
 
 		const snapshotIds = new Set<string>();
-		const tagSnapshotTasks = tagTasks.data as TagSnapshotsTask[];
-
-		for (const task of tagSnapshotTasks) {
+		for (const task of tagTasks.data) {
 			for (const snapshotId of task.input.snapshotIds) {
 				snapshotIds.add(snapshotId);
 			}

app/drizzle/20260710122522_brown_living_lightning/migration.sql

@@ -0,0 +1 @@
+ALTER TABLE `tasks` ADD `operation_key` text;

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

@@ -16,7 +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 restoreTaskOptions = restoreTasksOptions(schedule.data.repository.shortId, params.snapshotId);
 		const [snapshot, repository] = await Promise.all([
 			context.queryClient.ensureQueryData({
 				...getSnapshotDetailsOptions({

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

@@ -10,7 +10,7 @@ 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 restoreTaskOptions = restoreTasksOptions(params.repositoryId, params.snapshotId);
 		const [snapshot, repository] = await Promise.all([
 			context.queryClient.ensureQueryData({
 				...getSnapshotDetailsOptions({ path: { shortId: params.repositoryId, snapshotId: params.snapshotId } }),

app/schemas/tasks.ts

@@ -93,6 +93,7 @@ const taskShape = {
 	status: taskStatusSchema,
 	resourceType: taskResourceTypeSchema,
 	resourceId: z.string(),
+	operationKey: z.string().nullable(),
 	targetAgentId: z.string().nullable(),
 	input: taskInputSchema,
 	progress: taskProgressSchema.nullable(),

app/server/db/schema.ts

@@ -368,6 +368,7 @@ export const tasksTable = sqliteTable(
 		status: text("status").notNull(),
 		resourceType: text("resource_type").notNull(),
 		resourceId: text("resource_id").notNull(),
+		operationKey: text("operation_key"),
 		targetAgentId: text("target_agent_id"),
 		input: text("input", { mode: "json" }).$type<TaskJson>().notNull(),
 		progress: text("progress", { mode: "json" }).$type<TaskJson | null>(),

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

@@ -38,15 +38,11 @@ import type { RestoreExecutionProgress, RestoreExecutionResult } from "../agents
 import { agentsService } from "../agents/agents.service";
 import { LOCAL_AGENT_ID } from "../agents/constants";
 import { taskStore } from "../tasks/tasks.store";
-import type { ParsedTask, TaskInput } from "~/schemas/tasks";
 import { Effect } from "effect";
 
 const lsLimiters = new Map<string, Effect.Semaphore>();
 const RESTORE_TASK_RESOURCE_TYPE = "repository";
 
-type RestoreTaskInput = Extract<TaskInput, { kind: "restore" }>;
-type RestoreTask = ParsedTask & { kind: "restore"; input: RestoreTaskInput };
-
 const getBlockedRestoreTargets = () => {
 	return [
 		...RESTORE_BLOCKED_ROOTS,
@@ -69,9 +65,6 @@ const assertAllowedControllerLocalRestoreTarget = (target: string) => {
 	}
 };
 
-const isRestoreTask = (task: ParsedTask): task is RestoreTask =>
-	task.kind === "restore" && task.input.kind === "restore";
-
 const updateActiveRestoreTask = (restoreId: string, eventName: string, update: () => void) => {
 	try {
 		update();
@@ -89,21 +82,14 @@ const getLsLimiter = (repositoryId: string) => {
 	return limiter;
 };
 
-const findActiveRestoreTask = (
-	organizationId: string,
-	repositoryShortId: string,
-	snapshotId: string,
-): RestoreTask | null => {
-	return (
-		taskStore
-			.listActiveByResource({
-				organizationId,
-				kind: "restore",
-				resourceType: RESTORE_TASK_RESOURCE_TYPE,
-				resourceId: repositoryShortId,
-			})
-			.find((task): task is RestoreTask => isRestoreTask(task) && task.input.snapshotId === snapshotId) ?? null
-	);
+const findActiveRestoreTask = (organizationId: string, repositoryShortId: string, snapshotId: string) => {
+	return taskStore.findActiveByResource({
+		organizationId,
+		kind: "restore",
+		resourceType: RESTORE_TASK_RESOURCE_TYPE,
+		resourceId: repositoryShortId,
+		operationKey: snapshotId,
+	});
 };
 
 const findActiveDoctorTask = (organizationId: string, repositoryShortId: string) => {
@@ -484,6 +470,7 @@ const restoreSnapshot = async (
 		organizationId,
 		resourceType: RESTORE_TASK_RESOURCE_TYPE,
 		resourceId: repository.shortId,
+		operationKey: snapshotId,
 		targetAgentId: useControllerLocalRestoreFallback ? null : executionAgentId,
 		input: { kind: "restore", repositoryId: repository.shortId, snapshotId, target },
 	});

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

@@ -28,15 +28,16 @@ const createTask = (
 	});
 };
 
-const createRestoreTask = (organizationId: string, repositoryId = "repo-short") => {
+const createRestoreTask = (organizationId: string, repositoryId = "repo-short", snapshotId = "snapshot-1") => {
 	return taskStore.create({
 		organizationId,
 		resourceType: "repository",
 		resourceId: repositoryId,
+		operationKey: snapshotId,
 		input: {
 			kind: "restore",
 			repositoryId,
-			snapshotId: "snapshot-1",
+			snapshotId,
 			target: "/tmp/restore",
 		},
 	});
@@ -231,4 +232,55 @@ describe("tasksController", () => {
 		const byKindAndResourceBody = await byKindAndResource.json();
 		expect(byKindAndResourceBody.map((task: { id: string }) => task.id).sort()).toEqual([repoTask.id]);
 	});
+
+	test("lists only the task for the exact operation key", async () => {
+		const session = await createTestSession();
+		const matchingTask = createRestoreTask(session.organizationId, "repo-short", "snapshot-1");
+		createRestoreTask(session.organizationId, "repo-short", "snapshot-2");
+
+		const res = await app.request(
+			"/api/v1/tasks?kind=restore&resourceType=repository&resourceId=repo-short&operationKey=snapshot-1",
+			{ headers: session.headers },
+		);
+
+		expect(res.status).toBe(200);
+		const body = await res.json();
+		expect(body.map((task: { id: string }) => task.id)).toEqual([matchingTask.id]);
+	});
+
+	test("streams task snapshots and updates only for the exact operation key", async () => {
+		const session = await createTestSession();
+		const matchingTask = createRestoreTask(session.organizationId, "repo-short", "snapshot-1");
+		const otherSnapshotTask = createRestoreTask(session.organizationId, "repo-short", "snapshot-2");
+
+		const res = await app.request(
+			"/api/v1/tasks/events?kind=restore&resourceType=repository&resourceId=repo-short&operationKey=snapshot-1",
+			{ headers: session.headers },
+		);
+
+		expect(res.status).toBe(200);
+
+		const reader = res.body!.getReader();
+		const decoder = new TextDecoder();
+		let text = "";
+
+		try {
+			text = await readReaderUntil(reader, decoder, text, matchingTask.id);
+
+			expect(text).toContain(`event: ${tasksSnapshotEventName}`);
+			expect(text).toContain(matchingTask.id);
+			expect(text).not.toContain(otherSnapshotTask.id);
+
+			taskStore.fail(otherSnapshotTask.id, "other snapshot failed");
+			taskStore.fail(matchingTask.id, "matching snapshot failed");
+			text = await readReaderUntil(reader, decoder, text, "matching snapshot failed");
+
+			expect(text).toContain(`event: ${taskChangedEventName}`);
+			expect(text).toContain("matching snapshot failed");
+			expect(text).not.toContain("other snapshot failed");
+		} finally {
+			void reader.cancel();
+			reader.releaseLock();
+		}
+	});
 });

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

@@ -16,39 +16,19 @@ import {
 import { toTaskDto } from "./tasks.presenter";
 import { taskStore } from "./tasks.store";
 
-type TaskFilter = NonNullable<Parameters<typeof taskStore.listActive>[0]>;
-
-const taskMatchesFilter = (task: ReturnType<typeof taskStore.listActive>[number], filter: TaskFilter) => {
-	if (filter.organizationId && task.organizationId !== filter.organizationId) {
-		return false;
-	}
-
-	if (filter.kind && task.kind !== filter.kind) {
-		return false;
-	}
-
-	if (filter.resourceType && task.resourceType !== filter.resourceType) {
-		return false;
-	}
-
-	if (filter.resourceId && task.resourceId !== filter.resourceId) {
-		return false;
-	}
-
-	return true;
-};
-
 export const tasksController = new Hono()
 	.use(requireAuth)
 	.get("/", validator("query", listTasksQuery), listTasksDto, async (c) => {
 		const organizationId = c.get("organizationId");
 		const query = c.req.valid("query");
-		const tasks = taskStore.listActive({
+		const filter = {
 			organizationId,
 			kind: query.kind,
 			resourceType: query.resourceType,
 			resourceId: query.resourceId,
-		});
+			operationKey: query.operationKey,
+		};
+		const tasks = taskStore.listActive(filter);
 		const response = tasks.map(toTaskDto);
 
 		return c.json<ListTasksDto>(response, 200);
@@ -61,6 +41,7 @@ export const tasksController = new Hono()
 			kind: query.kind,
 			resourceType: query.resourceType,
 			resourceId: query.resourceId,
+			operationKey: query.operationKey,
 		};
 
 		return streamEvents<TaskEventPayloadMap, typeof taskChangedEventName>(c, {
@@ -76,11 +57,7 @@ export const tasksController = new Hono()
 				});
 			},
 			subscribe: (_eventName, handler) => {
-				return taskStore.subscribeToAllChanges((changedTask) => {
-					if (!taskMatchesFilter(changedTask, filter)) {
-						return;
-					}
-
+				return taskStore.subscribeToAllChanges(filter, (changedTask) => {
 					const taskData = toTaskDto(changedTask);
 					void handler(taskData);
 				});

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

@@ -7,6 +7,7 @@ export const listTasksQuery = z
 		kind: taskKindSchema.optional(),
 		resourceType: taskResourceTypeSchema.optional(),
 		resourceId: z.string().optional(),
+		operationKey: z.string().optional(),
 	})
 	.superRefine((query, ctx) => {
 		const hasResourceType = Boolean(query.resourceType);

app/server/modules/tasks/tasks.store.ts

@@ -20,6 +20,7 @@ type TaskResource = {
 	kind: TaskKind;
 	resourceType: TaskResourceType;
 	resourceId: string;
+	operationKey?: string;
 };
 
 type CreateTaskParams = {
@@ -27,6 +28,7 @@ type CreateTaskParams = {
 	organizationId: string;
 	resourceType: TaskResourceType;
 	resourceId: string;
+	operationKey?: string | null;
 	targetAgentId?: string | null;
 	input: TaskInput;
 };
@@ -89,6 +91,16 @@ const subscribeToTaskChanges = (taskId: string, listener: TaskChangeListener) =>
 	};
 };
 
+const taskMatchesFilter = (task: ParsedTask, filter: Partial<TaskResource>) => {
+	if (filter.organizationId && task.organizationId !== filter.organizationId) return false;
+	if (filter.kind && task.kind !== filter.kind) return false;
+	if (filter.resourceType && task.resourceType !== filter.resourceType) return false;
+	if (filter.resourceId && task.resourceId !== filter.resourceId) return false;
+	if (filter.operationKey && task.operationKey !== filter.operationKey) return false;
+
+	return true;
+};
+
 const activeStatusCondition = () => inArray(tasksTable.status, activeTaskStatuses);
 
 const byIdCondition = (id: string) => eq(tasksTable.id, id);
@@ -100,6 +112,7 @@ const buildActiveConditions = (params: Partial<TaskResource> = {}) => {
 	if (params.kind) conditions.push(eq(tasksTable.kind, params.kind));
 	if (params.resourceType) conditions.push(eq(tasksTable.resourceType, params.resourceType));
 	if (params.resourceId) conditions.push(eq(tasksTable.resourceId, params.resourceId));
+	if (params.operationKey) conditions.push(eq(tasksTable.operationKey, params.operationKey));
 
 	return conditions;
 };
@@ -137,6 +150,7 @@ export const taskStore = {
 				status: "queued",
 				resourceType: params.resourceType,
 				resourceId: params.resourceId,
+				operationKey: params.operationKey ?? null,
 				targetAgentId: params.targetAgentId ?? null,
 				input,
 				progress: null,
@@ -268,10 +282,6 @@ export const taskStore = {
 		return row ? parseTask(row) : null;
 	},
 
-	listActiveByResource: (params: TaskResource): ParsedTask[] => {
-		return listActiveTasks(params);
-	},
-
 	listActive: (params: ListActiveTasksParams = {}): ParsedTask[] => {
 		return listActiveTasks(params);
 	},
@@ -280,8 +290,12 @@ export const taskStore = {
 		return subscribeToTaskChanges(taskId, listener);
 	},
 
-	subscribeToAllChanges: (listener: TaskChangeListener) => {
-		return subscribeToAllTaskChanges(listener);
+	subscribeToAllChanges: (filter: ListActiveTasksParams, listener: TaskChangeListener) => {
+		return subscribeToAllTaskChanges((task) => {
+			if (taskMatchesFilter(task, filter)) {
+				listener(task);
+			}
+		});
 	},
 
 	findById: (params: FindTaskParams): ParsedTask | null => {