fix: pr feedbacks

apps/agent/src/commands/backup-run.test.ts

@@ -0,0 +1,117 @@
+import { afterEach, expect, mock, spyOn, test } from "bun:test";
+import { Effect } from "effect";
+import waitForExpect from "wait-for-expect";
+import { fromPartial } from "@total-typescript/shoehorn";
+import { parseAgentMessage, type BackupCancelPayload, type BackupRunPayload } from "@zerobyte/contracts/agent-protocol";
+import * as resticServer from "@zerobyte/core/restic/server";
+import { handleBackupCancelCommand } from "./backup-cancel";
+import { handleBackupRunCommand } from "./backup-run";
+import type { ControllerCommandContext, RunningJob } from "../context";
+
+const createDeferred = <T>() => {
+	let resolve!: (value: T) => void;
+	const promise = new Promise<T>((resolvePromise) => {
+		resolve = resolvePromise;
+	});
+
+	return { promise, resolve };
+};
+
+afterEach(() => {
+	mock.restore();
+});
+
+test("waits for running-job registration before returning to the processor loop", async () => {
+	const outboundMessages: string[] = [];
+	const runningJobs = new Map<string, RunningJob>();
+	const setRunningJobGate = createDeferred<void>();
+	const backupGate = createDeferred<{ exitCode: number; result: null; warningDetails: null }>();
+	let registeredAbortController: AbortController | undefined;
+
+	spyOn(resticServer, "createRestic").mockReturnValue(
+		fromPartial({
+			backup: () =>
+				Effect.async<{ exitCode: number; result: null; warningDetails: null }, never>((resume) => {
+					void backupGate.promise.then((result) => {
+						resume(Effect.succeed(result));
+					});
+				}),
+		}),
+	);
+
+	const context: ControllerCommandContext = {
+		getRunningJob: (jobId) => Effect.succeed(runningJobs.get(jobId)),
+		setRunningJob: (jobId, job) =>
+			Effect.async<void, never>((resume) => {
+				void setRunningJobGate.promise.then(() => {
+					runningJobs.set(jobId, job);
+					registeredAbortController = job.abortController;
+					resume(Effect.void);
+				});
+			}),
+		deleteRunningJob: (jobId) =>
+			Effect.sync(() => {
+				runningJobs.delete(jobId);
+			}),
+		offerOutbound: (message) =>
+			Effect.sync(() => {
+				outboundMessages.push(message);
+				return true;
+			}),
+	};
+
+	const runPayload = fromPartial<BackupRunPayload>({
+		jobId: "job-1",
+		scheduleId: "schedule-1",
+		organizationId: "org-1",
+		sourcePath: "/tmp/source",
+		repositoryConfig: {
+			backend: "local",
+			path: "/tmp/repository",
+		},
+		options: {},
+		runtime: {
+			password: "password",
+			cacheDir: "/tmp/restic-cache",
+			passFile: "/tmp/restic-pass",
+			defaultExcludes: [],
+		},
+	});
+	const cancelPayload = fromPartial<BackupCancelPayload>({
+		jobId: "job-1",
+		scheduleId: "schedule-1",
+	});
+
+	const runPromise = Effect.runPromise(handleBackupRunCommand(context, runPayload));
+
+	try {
+		const returnedBeforeRegistration = await Promise.race([
+			runPromise.then(() => true),
+			new Promise<false>((resolve) => {
+				setTimeout(() => resolve(false), 0);
+			}),
+		]);
+
+		expect(returnedBeforeRegistration).toBe(false);
+
+		setRunningJobGate.resolve(undefined);
+		await runPromise;
+
+		await Effect.runPromise(handleBackupCancelCommand(context, cancelPayload));
+		expect(registeredAbortController?.signal.aborted).toBe(true);
+
+		backupGate.resolve({ exitCode: 0, result: null, warningDetails: null });
+
+		await waitForExpect(() => {
+			const cancelledMessage = outboundMessages
+				.map((message) => parseAgentMessage(message))
+				.find((message) => message?.success && message.data.type === "backup.cancelled");
+
+			expect(cancelledMessage?.success).toBe(true);
+			expect(runningJobs.has("job-1")).toBe(false);
+		});
+	} finally {
+		setRunningJobGate.resolve(undefined);
+		backupGate.resolve({ exitCode: 0, result: null, warningDetails: null });
+	}
+});

apps/agent/src/commands/backup-run.ts

@@ -7,107 +7,109 @@ import { toErrorDetails, toMessage } from "@zerobyte/core/utils";
 import type { ControllerCommandContext } from "../context";
 
 export const handleBackupRunCommand = (context: ControllerCommandContext, payload: BackupRunPayload) => {
-	return Effect.fork(
+	return Effect.gen(function* () {
-		Effect.gen(function* () {
+		const existing = yield* context.getRunningJob(payload.jobId);
-			const existing = yield* context.getRunningJob(payload.jobId);
+		if (existing) {
-			if (existing) {
-				yield* context.offerOutbound(
-					createAgentMessage("backup.failed", {
-						jobId: payload.jobId,
-						scheduleId: payload.scheduleId,
-						error: "Backup job is already running",
-					}),
-				);
-				return;
-			}
-
-			logger.info(`Starting backup ${payload.jobId} for schedule ${payload.scheduleId}`);
-			const abortController = new AbortController();
-			yield* context.setRunningJob(payload.jobId, { scheduleId: payload.scheduleId, abortController });
-
-			const sendCancelled = () => {
-				return context.offerOutbound(
-					createAgentMessage("backup.cancelled", {
-						jobId: payload.jobId,
-						scheduleId: payload.scheduleId,
-						message: "Backup was cancelled",
-					}),
-				);
-			};
-
 			yield* context.offerOutbound(
-				createAgentMessage("backup.started", {
+				createAgentMessage("backup.failed", {
 					jobId: payload.jobId,
 					scheduleId: payload.scheduleId,
+					error: "Backup job is already running",
 				}),
 			);
+			return;
+		}
 
-			const deps: ResticDeps = {
+		logger.info(`Starting backup ${payload.jobId} for schedule ${payload.scheduleId}`);
-				resolveSecret: async (encrypted) => encrypted,
+		const abortController = new AbortController();
-				getOrganizationResticPassword: async () => payload.runtime.password,
+		yield* context.setRunningJob(payload.jobId, { scheduleId: payload.scheduleId, abortController });
-				resticCacheDir: payload.runtime.cacheDir,
-				resticPassFile: payload.runtime.passFile,
-				defaultExcludes: payload.runtime.defaultExcludes,
-				hostname: payload.runtime.hostname,
-			};
 
-			const restic = createRestic(deps);
+		yield* Effect.fork(
-			const runtime = yield* Effect.runtime<never>();
+			Effect.gen(function* () {
+				const sendCancelled = () => {
+					return context.offerOutbound(
+						createAgentMessage("backup.cancelled", {
+							jobId: payload.jobId,
+							scheduleId: payload.scheduleId,
+							message: "Backup was cancelled",
+						}),
+					);
+				};
 
-			yield* restic
+				yield* context.offerOutbound(
-				.backup(payload.repositoryConfig, payload.sourcePath, {
+					createAgentMessage("backup.started", {
-					organizationId: payload.organizationId,
+						jobId: payload.jobId,
-					...payload.options,
+						scheduleId: payload.scheduleId,
-					signal: abortController.signal,
-					onProgress: (progress) => {
-						void Runtime.runPromise(
-							runtime,
-							context.offerOutbound(
-								createAgentMessage("backup.progress", {
-									jobId: payload.jobId,
-									scheduleId: payload.scheduleId,
-									progress,
-								}),
-							),
-						).catch((error) => {
-							logger.error(`Failed to send backup progress update: ${toMessage(error)}`);
-						});
-					},
-				})
-				.pipe(
-					Effect.matchEffect({
-						onSuccess: (result) => {
-							if (abortController.signal.aborted) {
-								return sendCancelled();
-							}
-
-							return context.offerOutbound(
-								createAgentMessage("backup.completed", {
-									jobId: payload.jobId,
-									scheduleId: payload.scheduleId,
-									exitCode: result.exitCode,
-									result: result.result,
-									warningDetails: result.warningDetails ?? undefined,
-								}),
-							);
-						},
-						onFailure: (error) => {
-							if (abortController.signal.aborted) {
-								return sendCancelled();
-							}
-
-							return context.offerOutbound(
-								createAgentMessage("backup.failed", {
-									jobId: payload.jobId,
-									scheduleId: payload.scheduleId,
-									error: toMessage(error),
-									errorDetails: toErrorDetails(error),
-								}),
-							);
-						},
 					}),
-					Effect.ensuring(context.deleteRunningJob(payload.jobId)),
 				);
-		}),
+
-	).pipe(Effect.asVoid);
+				const deps: ResticDeps = {
+					resolveSecret: async (encrypted) => encrypted,
+					getOrganizationResticPassword: async () => payload.runtime.password,
+					resticCacheDir: payload.runtime.cacheDir,
+					resticPassFile: payload.runtime.passFile,
+					defaultExcludes: payload.runtime.defaultExcludes,
+					hostname: payload.runtime.hostname,
+				};
+
+				const restic = createRestic(deps);
+				const runtime = yield* Effect.runtime<never>();
+
+				yield* restic
+					.backup(payload.repositoryConfig, payload.sourcePath, {
+						organizationId: payload.organizationId,
+						...payload.options,
+						signal: abortController.signal,
+						onProgress: (progress) => {
+							void Runtime.runPromise(
+								runtime,
+								context.offerOutbound(
+									createAgentMessage("backup.progress", {
+										jobId: payload.jobId,
+										scheduleId: payload.scheduleId,
+										progress,
+									}),
+								),
+							).catch((error) => {
+								logger.error(`Failed to send backup progress update: ${toMessage(error)}`);
+							});
+						},
+					})
+					.pipe(
+						Effect.matchEffect({
+							onSuccess: (result) => {
+								if (abortController.signal.aborted) {
+									return sendCancelled();
+								}
+
+								return context.offerOutbound(
+									createAgentMessage("backup.completed", {
+										jobId: payload.jobId,
+										scheduleId: payload.scheduleId,
+										exitCode: result.exitCode,
+										result: result.result,
+										warningDetails: result.warningDetails ?? undefined,
+									}),
+								);
+							},
+							onFailure: (error) => {
+								if (abortController.signal.aborted) {
+									return sendCancelled();
+								}
+
+								return context.offerOutbound(
+									createAgentMessage("backup.failed", {
+										jobId: payload.jobId,
+										scheduleId: payload.scheduleId,
+										error: toMessage(error),
+										errorDetails: toErrorDetails(error),
+									}),
+								);
+							},
+						}),
+						Effect.ensuring(context.deleteRunningJob(payload.jobId)),
+					);
+			}),
+		);
+	}).pipe(Effect.asVoid);
 };