refactor: centralize agent backup lifecycle state

app/server/modules/agents/__tests__/session.test.ts

@@ -2,7 +2,6 @@ import { Effect, Exit, Scope } from "effect";
 import { expect, test, vi } from "vitest";
 import waitForExpect from "wait-for-expect";
 import { fromPartial } from "@total-typescript/shoehorn";
-import { createAgentMessage } from "@zerobyte/contracts/agent-protocol";
 import { LOCAL_AGENT_ID, LOCAL_AGENT_KIND, LOCAL_AGENT_NAME } from "../constants";
 import { createControllerAgentSession } from "../controller/session";
 
@@ -29,6 +28,7 @@ const createSession = (
 	const sessionHandlers: Parameters<typeof createControllerAgentSession>[1] = {
 		onReady: () => Effect.void,
 		onHeartbeatPong: () => Effect.void,
+		onDisconnect: () => Effect.void,
 		onBackupStarted: () => Effect.void,
 		onBackupProgress: () => Effect.void,
 		onBackupCompleted: () => Effect.void,
@@ -61,87 +61,22 @@ const createSession = (
 	}
 };
 
-test("close emits a synthetic backup.cancelled for a started backup", () => {
-	const onBackupCancelled = vi.fn(() => Effect.void);
-	const { session, close } = createSession({
-		onBackupCancelled,
-		onBackupStarted: vi.fn(() => Effect.void),
-	});
-
-	Effect.runSync(
-		session.handleMessage(
-			createAgentMessage("backup.started", {
-				jobId: "job-1",
-				scheduleId: "schedule-1",
-			}),
-		),
-	);
+test("close reports a transport disconnect", () => {
+	const onDisconnect = vi.fn(() => Effect.void);
+	const { close } = createSession({ onDisconnect });
 
 	close();
 
-	expect(onBackupCancelled).toHaveBeenCalledTimes(1);
-	expect(onBackupCancelled).toHaveBeenCalledWith(
+	expect(onDisconnect).toHaveBeenCalledTimes(1);
+	expect(onDisconnect).toHaveBeenCalledWith(
 		expect.objectContaining({
-			jobId: "job-1",
-			scheduleId: "schedule-1",
+			agentId: LOCAL_AGENT_ID,
+			agentName: LOCAL_AGENT_NAME,
 		}),
 	);
 });
 
-test.each([
-	{
-		name: "backup.completed",
-		jobId: "job-1",
-		scheduleId: "schedule-1",
-		terminalMessage: createAgentMessage("backup.completed", {
-			jobId: "job-1",
-			scheduleId: "schedule-1",
-			exitCode: 0,
-			result: null,
-		}),
-		expectedCancelledCalls: 0,
-	},
-	{
-		name: "backup.failed",
-		jobId: "job-2",
-		scheduleId: "schedule-2",
-		terminalMessage: createAgentMessage("backup.failed", {
-			jobId: "job-2",
-			scheduleId: "schedule-2",
-			error: "backup failed",
-		}),
-		expectedCancelledCalls: 0,
-	},
-	{
-		name: "backup.cancelled",
-		jobId: "job-3",
-		scheduleId: "schedule-3",
-		terminalMessage: createAgentMessage("backup.cancelled", {
-			jobId: "job-3",
-			scheduleId: "schedule-3",
-			message: "Backup was cancelled",
-		}),
-		expectedCancelledCalls: 1,
-	},
-])("close does not emit an extra synthetic backup.cancelled after $name", (testCase) => {
-	const onBackupCancelled = vi.fn(() => Effect.void);
-	const { session, close } = createSession({ onBackupCancelled });
-
-	Effect.runSync(
-		session.handleMessage(
-			createAgentMessage("backup.started", {
-				jobId: testCase.jobId,
-				scheduleId: testCase.scheduleId,
-			}),
-		),
-	);
-	Effect.runSync(session.handleMessage(testCase.terminalMessage));
-	close();
-
-	expect(onBackupCancelled).toHaveBeenCalledTimes(testCase.expectedCancelledCalls);
-});
-
-test("close emits a synthetic backup.cancelled for a queued backup", () => {
+test("sendBackup only queues the transport message", () => {
 	const onBackupCancelled = vi.fn(() => Effect.void);
 	const { session, close } = createSession({ onBackupCancelled });
 
@@ -171,31 +106,17 @@ test("close emits a synthetic backup.cancelled for a queued backup", () => {
 
 	close();
 
-	expect(onBackupCancelled).toHaveBeenCalledTimes(1);
-	expect(onBackupCancelled).toHaveBeenLastCalledWith(
-		expect.objectContaining({
-			jobId: "job-queued",
-			scheduleId: "schedule-queued",
-		}),
-	);
+	expect(onBackupCancelled).not.toHaveBeenCalled();
 });
 
-test("a dropped backup.cancel closes the session and emits a synthetic backup.cancelled", async () => {
+test("a dropped backup.cancel closes the session and reports a transport disconnect", async () => {
 	const send = vi.fn(() => 0);
 	const socket = createSocket({ send, close: vi.fn() });
-	const onBackupCancelled = vi.fn(() => Effect.void);
-	const { session, run, closeAsync } = createSession({ onBackupCancelled }, socket);
+	const onDisconnect = vi.fn(() => Effect.void);
+	const { session, run, closeAsync } = createSession({ onDisconnect }, socket);
 
 	try {
 		run();
-		Effect.runSync(
-			session.handleMessage(
-				createAgentMessage("backup.started", {
-					jobId: "job-1",
-					scheduleId: "schedule-1",
-				}),
-			),
-		);
 		Effect.runSync(
 			session.sendBackupCancel({
 				jobId: "job-1",
@@ -206,11 +127,10 @@ test("a dropped backup.cancel closes the session and emits a synthetic backup.ca
 		await waitForExpect(() => {
 			expect(send).toHaveBeenCalledTimes(1);
 			expect(socket.close).toHaveBeenCalledTimes(1);
-			expect(onBackupCancelled).toHaveBeenCalledTimes(1);
-			expect(onBackupCancelled).toHaveBeenCalledWith(
+			expect(onDisconnect).toHaveBeenCalledTimes(1);
+			expect(onDisconnect).toHaveBeenCalledWith(
 				expect.objectContaining({
-					jobId: "job-1",
-					scheduleId: "schedule-1",
+					agentId: LOCAL_AGENT_ID,
 				}),
 			);
 		});

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

@@ -46,6 +46,17 @@ const resolveActiveBackupRun = (scheduleId: number, result: BackupExecutionResul
 	return true;
 };
 
+const cancelActiveBackupRunsForAgent = (agentId: string, message: string) => {
+	const activeBackupsByScheduleId = getActiveBackupsByScheduleId();
+	const matchingScheduleIds = [...activeBackupsByScheduleId.values()]
+		.filter((activeBackupRun) => activeBackupRun.agentId === agentId)
+		.map((activeBackupRun) => activeBackupRun.scheduleId);
+
+	for (const scheduleId of matchingScheduleIds) {
+		resolveActiveBackupRun(scheduleId, { status: "cancelled", message });
+	}
+};
+
 const getActiveBackupRun = (jobId: string, scheduleId: string, eventName: string, agentId: string) => {
 	const trackedScheduleId = getActiveBackupScheduleIdsByJobId().get(jobId);
 	if (trackedScheduleId === undefined) {
@@ -99,6 +110,12 @@ const requestBackupCancellation = async (agentId: string, scheduleId: number) =>
 };
 
 const backupEventHandlers: AgentBackupEventHandlers = {
+	onAgentDisconnected: ({ agentId }) => {
+		cancelActiveBackupRunsForAgent(
+			agentId,
+			"The connection to the backup agent was lost. Restart the backup to ensure it completes.",
+		);
+	},
 	onBackupStarted: ({ agentId, payload }) => {
 		getActiveBackupRun(payload.jobId, payload.scheduleId, "backup.started", agentId);
 	},
@@ -185,6 +202,7 @@ export const agentManager = {
 
 		const completion = new Promise<BackupExecutionResult>((resolve) => {
 			getActiveBackupsByScheduleId().set(request.scheduleId, {
+				agentId,
 				scheduleId: request.scheduleId,
 				jobId: request.payload.jobId,
 				scheduleShortId: request.payload.scheduleId,

app/server/modules/agents/controller/server.ts

@@ -26,6 +26,7 @@ type AgentBackupEventContext = {
 };
 
 export type AgentBackupEventHandlers = {
+	onAgentDisconnected?: (context: { agentId: string; agentName: string }) => void;
 	onBackupStarted?: (context: AgentBackupEventContext & { payload: BackupStartedPayload }) => void;
 	onBackupProgress?: (context: AgentBackupEventContext & { payload: BackupProgressPayload }) => void;
 	onBackupCompleted?: (context: AgentBackupEventContext & { payload: BackupCompletedPayload }) => void;
@@ -55,28 +56,39 @@ export function createAgentManagerRuntime() {
 		});
 
 	const closeAllSessions = Effect.gen(function* () {
-		const currentSessions = sessions;
-		sessions = new Map();
-		for (const sessionHandle of currentSessions.values()) {
+		const currentSessions = [...sessions.entries()];
+		for (const [agentId, sessionHandle] of currentSessions) {
+			yield* Effect.promise(() => agentsService.markAgentOffline(agentId));
 			yield* closeSession(sessionHandle);
 		}
+		sessions = new Map();
 	});
 
 	const getSessionHandle = (agentId: string) => sessions.get(agentId);
 
 	const getSession = (agentId: string) => getSessionHandle(agentId)?.session;
 
-	const createSessionHandlers = (ws: Bun.ServerWebSocket<AgentConnectionData>) => {
+	const createSessionHandlers = (ws: Bun.ServerWebSocket<AgentConnectionData>, markConnecting: Promise<unknown>) => {
 		const agentId = ws.data.agentId;
 		const agentName = ws.data.agentName;
 
 		return {
 			onReady: ({ at }: { at: number }) => {
-				return Effect.promise(() => agentsService.markAgentOnline(agentId, at));
+				return Effect.promise(async () => {
+					await markConnecting;
+					await agentsService.markAgentOnline(agentId, at);
+				});
 			},
 			onHeartbeatPong: ({ at }: { at: number }) => {
 				return Effect.promise(() => agentsService.markAgentSeen(agentId, at));
 			},
+			onDisconnect: () => {
+				if (getSession(ws.data.agentId)?.connectionId !== ws.data.id) {
+					return Effect.void;
+				}
+
+				return Effect.sync(() => backupHandlers.onAgentDisconnected?.({ agentId, agentName }));
+			},
 			onBackupStarted: (payload: BackupStartedPayload) => {
 				return Effect.sync(() => backupHandlers.onBackupStarted?.({ agentId, agentName, payload }));
 			},
@@ -95,12 +107,14 @@ export function createAgentManagerRuntime() {
 		};
 	};
 
-	const createSession = (ws: Bun.ServerWebSocket<AgentConnectionData>) => {
+	const createSession = (ws: Bun.ServerWebSocket<AgentConnectionData>, markConnecting: Promise<unknown>) => {
 		// Manual scope management because we are out of Effect
 		const scope = Effect.runSync(Scope.make());
 
 		try {
-			const session = Effect.runSync(Scope.extend(createControllerAgentSession(ws, createSessionHandlers(ws)), scope));
+			const session = Effect.runSync(
+				Scope.extend(createControllerAgentSession(ws, createSessionHandlers(ws, markConnecting)), scope),
+			);
 			const runFiber = Effect.runFork(Scope.extend(session.run, scope));
 
 			return { session, runFiber, scope };
@@ -112,13 +126,13 @@ export function createAgentManagerRuntime() {
 
 	const setSession = (agentId: string, sessionHandle: ControllerAgentSessionHandle) => {
 		const existingSession = getSessionHandle(agentId);
+		sessions.set(agentId, sessionHandle);
+
 		if (existingSession) {
 			void Effect.runPromise(closeSession(existingSession)).catch((error) => {
 				logger.error(`Failed to close existing agent session for ${agentId}: ${toMessage(error)}`);
 			});
 		}
-
-		sessions.set(agentId, sessionHandle);
 	};
 
 	const removeSession = (agentId: string, connectionId: string) => {
@@ -165,17 +179,16 @@ export function createAgentManagerRuntime() {
 				},
 				websocket: {
 					open: (ws) => {
-						setSession(ws.data.agentId, createSession(ws));
-						void agentsService
-							.markAgentConnecting({
-								agentId: ws.data.agentId,
-								organizationId: ws.data.organizationId,
-								agentName: ws.data.agentName,
-								agentKind: ws.data.agentKind,
-							})
-							.catch((error) => {
-								logger.error(`Failed to mark agent ${ws.data.agentId} as connecting: ${toMessage(error)}`);
-							});
+						const markConnecting = agentsService.markAgentConnecting({
+							agentId: ws.data.agentId,
+							organizationId: ws.data.organizationId,
+							agentName: ws.data.agentName,
+							agentKind: ws.data.agentKind,
+						});
+						void markConnecting.catch((error) => {
+							logger.error(`Failed to mark agent ${ws.data.agentId} as connecting: ${toMessage(error)}`);
+						});
+						setSession(ws.data.agentId, createSession(ws, markConnecting));
 						logger.info(`Agent "${ws.data.agentName}" (${ws.data.agentId}) connected on ${ws.data.id}`);
 					},
 					message: (ws, data) => {

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

@@ -32,11 +32,6 @@ type SessionState = {
 	lastPongAt: number | null;
 };
 
-type TrackedBackupJob = {
-	scheduleId: string;
-	state: "pending" | "active";
-};
-
 type AgentRuntimeEventPayload = {
 	agentId: string;
 	agentName: string;
@@ -48,6 +43,7 @@ type AgentRuntimeEventPayload = {
 type ControllerAgentSessionHandlers = {
 	onReady: (payload: AgentRuntimeEventPayload) => Effect.Effect<void>;
 	onHeartbeatPong: (payload: AgentRuntimeEventPayload) => Effect.Effect<void>;
+	onDisconnect: (payload: AgentRuntimeEventPayload) => Effect.Effect<void>;
 	onBackupStarted: (payload: BackupStartedPayload) => Effect.Effect<void>;
 	onBackupProgress: (payload: BackupProgressPayload) => Effect.Effect<void>;
 	onBackupCompleted: (payload: BackupCompletedPayload) => Effect.Effect<void>;
@@ -71,7 +67,6 @@ export const createControllerAgentSession = (
 	Effect.gen(function* () {
 		let isClosed = false;
 		const outboundQueue = yield* Queue.bounded<ControllerWireMessage>(64);
-		const trackedBackupJobs = yield* Ref.make<Map<string, TrackedBackupJob>>(new Map());
 		const state = yield* Ref.make<SessionState>({
 			isReady: false,
 			lastSeenAt: null,
@@ -90,35 +85,16 @@ export const createControllerAgentSession = (
 
 		const updateState = (update: (current: SessionState) => SessionState) => Ref.update(state, update);
 
-		const setTrackedBackupJob = (jobId: string, trackedBackupJob: TrackedBackupJob) => {
-			return Ref.update(trackedBackupJobs, (current) => {
-				const next = new Map(current);
-				next.set(jobId, trackedBackupJob);
-				return next;
-			});
-		};
-
-		const deleteTrackedBackupJob = (jobId: string) => {
-			return Ref.update(trackedBackupJobs, (current) => {
-				const next = new Map(current);
-				next.delete(jobId);
-				return next;
-			});
-		};
-
-		const takeTrackedBackupJobs = Ref.modify(
-			trackedBackupJobs,
-			(current) => [current, new Map<string, TrackedBackupJob>()] as const,
-		);
-
 		const releaseSession = Effect.gen(function* () {
-			yield* updateState((current) => ({ ...current, isReady: false }));
-			const trackedJobs = yield* takeTrackedBackupJobs;
-			for (const [jobId, trackedJob] of trackedJobs) {
-				const message = "The connection to the backup agent was lost. Restart the backup to ensure it completes.";
-
-				yield* handlers.onBackupCancelled({ jobId, scheduleId: trackedJob.scheduleId, message });
-			}
+			const disconnectedAt = Date.now();
+			yield* updateState((current) => ({ ...current, isReady: false, lastSeenAt: disconnectedAt }));
+			yield* handlers.onDisconnect({
+				agentId: socket.data.agentId,
+				agentName: socket.data.agentName,
+				organizationId: socket.data.organizationId,
+				agentKind: socket.data.agentKind,
+				at: disconnectedAt,
+			});
 
 			yield* Queue.shutdown(outboundQueue);
 		});
@@ -200,10 +176,6 @@ export const createControllerAgentSession = (
 						break;
 					}
 					case "backup.started": {
-						yield* setTrackedBackupJob(message.payload.jobId, {
-							scheduleId: message.payload.scheduleId,
-							state: "active",
-						});
 						yield* logger.effect.info(
 							`Backup ${message.payload.jobId} started on agent ${socket.data.agentId} for schedule ${message.payload.scheduleId}`,
 						);
@@ -215,17 +187,14 @@ export const createControllerAgentSession = (
 						break;
 					}
 					case "backup.completed": {
-						yield* deleteTrackedBackupJob(message.payload.jobId);
 						yield* handlers.onBackupCompleted(message.payload);
 						break;
 					}
 					case "backup.failed": {
-						yield* deleteTrackedBackupJob(message.payload.jobId);
 						yield* handlers.onBackupFailed(message.payload);
 						break;
 					}
 					case "backup.cancelled": {
-						yield* deleteTrackedBackupJob(message.payload.jobId);
 						yield* handlers.onBackupCancelled(message.payload);
 						break;
 					}
@@ -264,17 +233,7 @@ export const createControllerAgentSession = (
 					yield* handleAgentMessage(parsed.data);
 				});
 			},
-			sendBackup: (payload) => {
-				return Effect.gen(function* () {
-					const queued = yield* offerOutbound(createControllerMessage("backup.run", payload));
-
-					if (queued) {
-						yield* setTrackedBackupJob(payload.jobId, { scheduleId: payload.scheduleId, state: "pending" });
-					}
-
-					return queued;
-				});
-			},
+			sendBackup: (payload) => offerOutbound(createControllerMessage("backup.run", payload)),
 			sendBackupCancel: (payload) => offerOutbound(createControllerMessage("backup.cancel", payload)),
 			isReady: () => Ref.get(state).pipe(Effect.map((current) => current.isReady)),
 			run,

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

@@ -11,6 +11,7 @@ export type BackupExecutionResult =
 	| { status: "cancelled"; message?: string };
 
 type ActiveBackupRun = {
+	agentId: string;
 	scheduleId: number;
 	jobId: string;
 	scheduleShortId: string;

app/test/setup-shared.ts

@@ -1,4 +1,5 @@
 import { vi } from "vitest";
+import { Effect } from "effect";
 
 process.env.BASE_URL = "http://localhost:3000";
 process.env.TRUSTED_ORIGINS = "http://localhost:3000";
@@ -13,6 +14,12 @@ vi.mock(import("@zerobyte/core/node"), async () => {
 			info: () => {},
 			warn: () => {},
 			error: () => {},
+			effect: {
+				debug: () => Effect.void,
+				info: () => Effect.void,
+				warn: () => Effect.void,
+				error: () => Effect.void,
+			},
 		},
 	};
 });

packages/core/test/setup.ts

@@ -1,4 +1,5 @@
 import { vi } from "vitest";
+import { Effect } from "effect";
 
 vi.mock(import("../src/node/logger.ts"), () => ({
 	logger: {
@@ -6,5 +7,11 @@ vi.mock(import("../src/node/logger.ts"), () => ({
 		info: () => {},
 		warn: () => {},
 		error: () => {},
+		effect: {
+			debug: () => Effect.void,
+			info: () => Effect.void,
+			warn: () => Effect.void,
+			error: () => Effect.void,
+		},
 	},
 }));