fix: cancel started backups when agent disconnects

app/server/modules/agents/__tests__/controller-agent-session.test.ts

@@ -0,0 +1,86 @@
+import { expect, mock, test } from "bun:test";
+import { createAgentMessage } from "@zerobyte/contracts/agent-protocol";
+import { createControllerAgentSession } from "../controller-agent-session";
+
+const createSocket = () => {
+	return {
+		data: { id: "connection-1", agentId: "local", organizationId: null, agentName: "Local Agent" },
+		send: mock(() => undefined),
+	} as unknown as Parameters<typeof createControllerAgentSession>[0];
+};
+
+test("close emits a synthetic backup.cancelled for a started backup", () => {
+	const onBackupCancelled = mock(() => undefined);
+	const session = createControllerAgentSession(createSocket(), {
+		onBackupCancelled,
+	});
+
+	session.handleMessage(
+		createAgentMessage("backup.started", {
+			jobId: "job-1",
+			scheduleId: "schedule-1",
+		}),
+	);
+
+	session.close();
+
+	expect(onBackupCancelled).toHaveBeenCalledTimes(1);
+	expect(onBackupCancelled).toHaveBeenCalledWith({
+		jobId: "job-1",
+		scheduleId: "schedule-1",
+		message:
+			"The connection to the backup agent was lost while this backup was running. Restart the backup to ensure it completes.",
+	});
+});
+
+test("close does not emit a synthetic backup.cancelled after a terminal event", () => {
+	for (const testCase of [
+		{
+			jobId: "job-1",
+			scheduleId: "schedule-1",
+			terminalMessage: createAgentMessage("backup.completed", {
+				jobId: "job-1",
+				scheduleId: "schedule-1",
+				exitCode: 0,
+				result: null,
+			}),
+			expectedCancelledCalls: 0,
+		},
+		{
+			jobId: "job-2",
+			scheduleId: "schedule-2",
+			terminalMessage: createAgentMessage("backup.failed", {
+				jobId: "job-2",
+				scheduleId: "schedule-2",
+				error: "backup failed",
+			}),
+			expectedCancelledCalls: 0,
+		},
+		{
+			jobId: "job-3",
+			scheduleId: "schedule-3",
+			terminalMessage: createAgentMessage("backup.cancelled", {
+				jobId: "job-3",
+				scheduleId: "schedule-3",
+				message: "Backup was cancelled",
+			}),
+			expectedCancelledCalls: 1,
+		},
+	]) {
+		const onBackupCancelled = mock(() => undefined);
+		const session = createControllerAgentSession(createSocket(), {
+			onBackupCancelled,
+		});
+
+		session.handleMessage(
+			createAgentMessage("backup.started", {
+				jobId: testCase.jobId,
+				scheduleId: testCase.scheduleId,
+			}),
+		);
+		session.handleMessage(testCase.terminalMessage);
+		session.close();
+
+		expect(onBackupCancelled).toHaveBeenCalledTimes(testCase.expectedCancelledCalls);
+	}
+});

app/server/modules/agents/controller-agent-session.ts

@@ -52,6 +52,7 @@ export const createControllerAgentSession = (
 	handlers: ControllerAgentSessionHandlers = {},
 ): ControllerAgentSession => {
 	const outboundQueue = Effect.runSync(Queue.bounded<ControllerWireMessage>(64));
+	const activeBackupJobs = Effect.runSync(Ref.make<Map<string, string>>(new Map()));
 	const state = Effect.runSync(
 		Ref.make<SessionState>({
 			isReady: false,
@@ -70,6 +71,26 @@ export const createControllerAgentSession = (
 		Effect.runSync(Ref.update(state, update));
 	};
 
+	const setActiveBackupJob = (jobId: string, scheduleId: string) => {
+		Effect.runSync(
+			Ref.update(activeBackupJobs, (current) => {
+				const next = new Map(current);
+				next.set(jobId, scheduleId);
+				return next;
+			}),
+		);
+	};
+
+	const deleteActiveBackupJob = (jobId: string) => {
+		Effect.runSync(
+			Ref.update(activeBackupJobs, (current) => {
+				const next = new Map(current);
+				next.delete(jobId);
+				return next;
+			}),
+		);
+	};
+
 	const writerFiber = Effect.runFork(
 		Effect.forever(
 			Effect.gen(function* () {
@@ -110,6 +131,7 @@ export const createControllerAgentSession = (
 			}
 			case "backup.started": {
 				updateState((current) => ({ ...current, lastSeenAt: Date.now() }));
+				setActiveBackupJob(message.payload.jobId, message.payload.scheduleId);
 				logger.info(
 					`Backup ${message.payload.jobId} started on agent ${socket.data.agentId} for schedule ${message.payload.scheduleId}`,
 				);
@@ -123,16 +145,19 @@ export const createControllerAgentSession = (
 			}
 			case "backup.completed": {
 				updateState((current) => ({ ...current, lastSeenAt: Date.now() }));
+				deleteActiveBackupJob(message.payload.jobId);
 				handlers.onBackupCompleted?.(message.payload);
 				break;
 			}
 			case "backup.failed": {
 				updateState((current) => ({ ...current, lastSeenAt: Date.now() }));
+				deleteActiveBackupJob(message.payload.jobId);
 				handlers.onBackupFailed?.(message.payload);
 				break;
 			}
 			case "backup.cancelled": {
 				updateState((current) => ({ ...current, lastSeenAt: Date.now() }));
+				deleteActiveBackupJob(message.payload.jobId);
 				handlers.onBackupCancelled?.(message.payload);
 				break;
 			}
@@ -182,6 +207,16 @@ export const createControllerAgentSession = (
 		isReady: () => Effect.runSync(Ref.get(state)).isReady,
 		close: () => {
 			updateState((current) => ({ ...current, isReady: false }));
+			const pendingJobs = Effect.runSync(Ref.get(activeBackupJobs));
+			Effect.runSync(Ref.set(activeBackupJobs, new Map()));
+			for (const [jobId, scheduleId] of pendingJobs) {
+				handlers.onBackupCancelled?.({
+					jobId,
+					scheduleId,
+					message:
+						"The connection to the backup agent was lost while this backup was running. Restart the backup to ensure it completes.",
+				});
+			}
 			void Effect.runPromise(Fiber.interrupt(writerFiber)).catch(() => {});
 			void Effect.runPromise(Fiber.interrupt(heartbeatFiber)).catch(() => {});
 			void Effect.runPromise(Queue.shutdown(outboundQueue)).catch(() => {});