From 23a2a168becf582169cd2d50081c15b12bc11b77 Mon Sep 17 00:00:00 2001 From: Nicolas Meienberger Date: Wed, 1 Apr 2026 22:06:08 +0200 Subject: [PATCH] fix: cancel started backups when agent disconnects --- .../controller-agent-session.test.ts | 86 +++++++++++++++++++ .../agents/controller-agent-session.ts | 35 ++++++++ 2 files changed, 121 insertions(+) create mode 100644 app/server/modules/agents/__tests__/controller-agent-session.test.ts diff --git a/app/server/modules/agents/__tests__/controller-agent-session.test.ts b/app/server/modules/agents/__tests__/controller-agent-session.test.ts new file mode 100644 index 00000000..9f526f2e --- /dev/null +++ b/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[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); + } +}); diff --git a/app/server/modules/agents/controller-agent-session.ts b/app/server/modules/agents/controller-agent-session.ts index 4a2fe63c..7c8f0479 100644 --- a/app/server/modules/agents/controller-agent-session.ts +++ b/app/server/modules/agents/controller-agent-session.ts @@ -52,6 +52,7 @@ export const createControllerAgentSession = ( handlers: ControllerAgentSessionHandlers = {}, ): ControllerAgentSession => { const outboundQueue = Effect.runSync(Queue.bounded(64)); + const activeBackupJobs = Effect.runSync(Ref.make>(new Map())); const state = Effect.runSync( Ref.make({ 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(() => {});