refactor: context as effectful callbacks

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

@@ -24,12 +24,6 @@ export type AgentConnectionData = {
 
 type AgentSocket = Bun.ServerWebSocket<AgentConnectionData>;
 
-type SessionState = {
-	isReady: boolean;
-	lastSeenAt: number | null;
-	lastPongAt: number | null;
-};
-
 type ControllerAgentSessionHandlers = {
 	onBackupStarted?: (payload: BackupStartedPayload) => void;
 	onBackupProgress?: (payload: BackupProgressPayload) => void;
@@ -55,13 +49,7 @@ export const createControllerAgentSession = (
 	const outboundQueue = Effect.runSync(Queue.bounded<ControllerWireMessage>(64));
 	const activeBackupJobs = Effect.runSync(Ref.make<Map<string, string>>(new Map()));
 	const pendingBackupJobs = Effect.runSync(Ref.make<Map<string, string>>(new Map()));
-	const state = Effect.runSync(
+	const isReadyRef = Effect.runSync(Ref.make(false));
-		Ref.make<SessionState>({
-			isReady: false,
-			lastSeenAt: null,
-			lastPongAt: null,
-		}),
-	);
 
 	const offerOutbound = (message: ControllerWireMessage) => {
 		void Effect.runPromise(Queue.offer(outboundQueue, message)).catch((error) => {
@@ -69,8 +57,8 @@ export const createControllerAgentSession = (
 		});
 	};
 
-	const updateState = (update: (current: SessionState) => SessionState) => {
+	const setIsReady = (isReady: boolean) => {
-		Effect.runSync(Ref.update(state, update));
+		Effect.runSync(Ref.set(isReadyRef, isReady));
 	};
 
 	const setActiveBackupJob = (jobId: string, scheduleId: string) => {
@@ -113,35 +101,41 @@ export const createControllerAgentSession = (
 		);
 	};
 
+	const clearTrackedBackupJob = (jobId: string) => {
+		deletePendingBackupJob(jobId);
+		deleteActiveBackupJob(jobId);
+	};
+
+	const cancelTrackedJobs = (jobs: Map<string, string>, message: string) => {
+		for (const [jobId, scheduleId] of jobs) {
+			handlers.onBackupCancelled?.({
+				jobId,
+				scheduleId,
+				message,
+			});
+		}
+	};
+
 	const closeSession = () => {
 		if (isClosed) {
 			return;
 		}
 
 		isClosed = true;
-		updateState((current) => ({ ...current, isReady: false }));
+		setIsReady(false);
 		const pendingJobs = Effect.runSync(Ref.get(pendingBackupJobs));
 		Effect.runSync(Ref.set(pendingBackupJobs, new Map()));
-
+		cancelTrackedJobs(
-		for (const [jobId, scheduleId] of pendingJobs) {
+			pendingJobs,
-			handlers.onBackupCancelled?.({
-				jobId,
-				scheduleId,
-				message:
 			"The connection to the backup agent was lost before this backup started. Restart the backup to ensure it completes.",
-			});
+		);
-		}
 
 		const activeJobs = Effect.runSync(Ref.get(activeBackupJobs));
 		Effect.runSync(Ref.set(activeBackupJobs, new Map()));
-		for (const [jobId, scheduleId] of activeJobs) {
+		cancelTrackedJobs(
-			handlers.onBackupCancelled?.({
+			activeJobs,
-				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(() => {});
@@ -190,11 +184,9 @@ export const createControllerAgentSession = (
 	);
 
 	const handleAgentMessage = (message: AgentMessage) => {
-		updateState((current) => ({ ...current, lastSeenAt: Date.now() }));
-
 		switch (message.type) {
 			case "agent.ready": {
-				updateState((current) => ({ ...current, isReady: true }));
+				setIsReady(true);
 				logger.info(`Agent "${socket.data.agentName}" (${socket.data.agentId}) is ready`);
 				break;
 			}
@@ -212,25 +204,21 @@ export const createControllerAgentSession = (
 				break;
 			}
 			case "backup.completed": {
-				deletePendingBackupJob(message.payload.jobId);
+				clearTrackedBackupJob(message.payload.jobId);
-				deleteActiveBackupJob(message.payload.jobId);
 				handlers.onBackupCompleted?.(message.payload);
 				break;
 			}
 			case "backup.failed": {
-				deletePendingBackupJob(message.payload.jobId);
+				clearTrackedBackupJob(message.payload.jobId);
-				deleteActiveBackupJob(message.payload.jobId);
 				handlers.onBackupFailed?.(message.payload);
 				break;
 			}
 			case "backup.cancelled": {
-				deletePendingBackupJob(message.payload.jobId);
+				clearTrackedBackupJob(message.payload.jobId);
-				deleteActiveBackupJob(message.payload.jobId);
 				handlers.onBackupCancelled?.(message.payload);
 				break;
 			}
 			case "heartbeat.pong": {
-				updateState((current) => ({ ...current, lastPongAt: message.payload.sentAt }));
 				break;
 			}
 		}
@@ -261,7 +249,7 @@ export const createControllerAgentSession = (
 		sendBackupCancel: (payload) => {
 			offerOutbound(createControllerMessage("backup.cancel", payload));
 		},
-		isReady: () => Effect.runSync(Ref.get(state)).isReady,
+		isReady: () => Effect.runSync(Ref.get(isReadyRef)),
 		close: closeSession,
 	};
 };

app/server/modules/backups/backup-executor.ts

@@ -99,17 +99,25 @@ const clearActiveExecution = (jobId: string) => {
 	return activeExecution;
 };
 
-const getActiveExecution = (jobId: string, scheduleId: string, eventName: string, executorId: string) => {
+const clearExecutionState = (jobId: string, scheduleId: number) => {
+	requestedCancellationsByScheduleId.delete(scheduleId);
+	clearActiveExecution(jobId);
+};
+
+const resolveExecution = (jobId: string, activeExecution: ActiveBackupExecution, result: BackupExecutionResult) => {
+	clearExecutionState(jobId, activeExecution.scheduleId);
+	activeExecution.resolve(result);
+};
+
+const getActiveExecution = (jobId: string, scheduleId: string, eventName: string, agentId: string) => {
 	const activeExecution = activeExecutionsByJobId.get(jobId);
 	if (!activeExecution) {
-		logger.warn(`Received ${eventName} for unknown job ${jobId} from executor ${executorId}`);
+		logger.warn(`Received ${eventName} for unknown job ${jobId} from agent ${agentId}`);
 		return null;
 	}
 
 	if (activeExecution.scheduleShortId !== scheduleId) {
-		logger.warn(
+		logger.warn(`Ignoring ${eventName} for job ${jobId} due to schedule mismatch ${scheduleId} from agent ${agentId}`);
-			`Ignoring ${eventName} for job ${jobId} due to schedule mismatch ${scheduleId} from executor ${executorId}`,
-		);
 		return null;
 	}
 
@@ -134,9 +142,7 @@ agentManager.setBackupEventHandlers({
 			return;
 		}
 
-		requestedCancellationsByScheduleId.delete(activeExecution.scheduleId);
+		resolveExecution(payload.jobId, activeExecution, {
-		clearActiveExecution(payload.jobId);
-		activeExecution.resolve({
 			status: "completed",
 			exitCode: payload.exitCode,
 			result: payload.result,
@@ -149,9 +155,7 @@ agentManager.setBackupEventHandlers({
 			return;
 		}
 
-		requestedCancellationsByScheduleId.delete(activeExecution.scheduleId);
+		resolveExecution(payload.jobId, activeExecution, {
-		clearActiveExecution(payload.jobId);
-		activeExecution.resolve({
 			status: "failed",
 			error: payload.errorDetails ?? payload.error,
 		});
@@ -163,9 +167,7 @@ agentManager.setBackupEventHandlers({
 		}
 
 		const wasRequested = requestedCancellationsByScheduleId.has(activeExecution.scheduleId);
-		requestedCancellationsByScheduleId.delete(activeExecution.scheduleId);
+		resolveExecution(payload.jobId, activeExecution, {
-		clearActiveExecution(payload.jobId);
-		activeExecution.resolve({
 			status: "cancelled",
 			message: wasRequested ? undefined : payload.message,
 		});
@@ -207,8 +209,7 @@ export const backupExecutor = {
 			}
 
 			if (!agentManager.sendBackup(LOCAL_AGENT_ID, payload)) {
-				requestedCancellationsByScheduleId.delete(request.scheduleId);
+				clearExecutionState(jobId, request.scheduleId);
-				clearActiveExecution(jobId);
 				return {
 					status: "unavailable",
 					error: new Error("Local backup agent is not connected"),
@@ -217,8 +218,7 @@ export const backupExecutor = {
 
 			return completion;
 		} catch (error) {
-			requestedCancellationsByScheduleId.delete(request.scheduleId);
+			clearExecutionState(jobId, request.scheduleId);
-			clearActiveExecution(jobId);
 			throw error;
 		}
 	},

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

@@ -4,8 +4,8 @@ import { logger } from "@zerobyte/core/node";
 import type { ControllerCommandContext } from "../context";
 
 export const handleBackupCancelCommand = (context: ControllerCommandContext, payload: BackupCancelPayload) =>
-	Effect.sync(() => {
+	Effect.gen(function* () {
-		const running = context.getRunningJob(payload.jobId);
+		const running = yield* context.getRunningJob(payload.jobId);
 		if (!running) {
 			logger.warn(`Backup ${payload.jobId} is not running`);
 			return;

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

@@ -9,9 +9,9 @@ import type { ControllerCommandContext } from "../context";
 export const handleBackupRunCommand = (context: ControllerCommandContext, payload: BackupRunPayload) =>
 	Effect.fork(
 		Effect.gen(function* () {
-			const existing = context.getRunningJob(payload.jobId);
+			const existing = yield* context.getRunningJob(payload.jobId);
 			if (existing) {
-				context.offerOutbound(
+				yield* context.offerOutbound(
 					createAgentMessage("backup.failed", {
 						jobId: payload.jobId,
 						scheduleId: payload.scheduleId,
@@ -23,9 +23,19 @@ export const handleBackupRunCommand = (context: ControllerCommandContext, payloa
 
 			logger.info(`Starting backup ${payload.jobId} for schedule ${payload.scheduleId}`);
 			const abortController = new AbortController();
-			context.setRunningJob(payload.jobId, { scheduleId: payload.scheduleId, abortController });
+			yield* context.setRunningJob(payload.jobId, { scheduleId: payload.scheduleId, abortController });
 
-			context.offerOutbound(
+			const sendCancelled = () => {
+				return context.offerOutbound(
+					createAgentMessage("backup.cancelled", {
+						jobId: payload.jobId,
+						scheduleId: payload.scheduleId,
+						message: "Backup was cancelled",
+					}),
+				);
+			};
+
+			yield* context.offerOutbound(
 				createAgentMessage("backup.started", {
 					jobId: payload.jobId,
 					scheduleId: payload.scheduleId,
@@ -49,31 +59,27 @@ export const handleBackupRunCommand = (context: ControllerCommandContext, payloa
 					...payload.options,
 					signal: abortController.signal,
 					onProgress: (progress) => {
+						void Effect.runPromise(
 							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) => {
-							return Effect.sync(() => {
 							if (abortController.signal.aborted) {
-									context.offerOutbound(
+								return sendCancelled();
-										createAgentMessage("backup.cancelled", {
-											jobId: payload.jobId,
-											scheduleId: payload.scheduleId,
-											message: "Backup was cancelled",
-										}),
-									);
-									return;
 							}
 
-								context.offerOutbound(
+							return context.offerOutbound(
 								createAgentMessage("backup.completed", {
 									jobId: payload.jobId,
 									scheduleId: payload.scheduleId,
@@ -82,22 +88,13 @@ export const handleBackupRunCommand = (context: ControllerCommandContext, payloa
 									warningDetails: result.warningDetails ?? undefined,
 								}),
 							);
-							});
 						},
 						onFailure: (error) => {
-							return Effect.sync(() => {
 							if (abortController.signal.aborted) {
-									context.offerOutbound(
+								return sendCancelled();
-										createAgentMessage("backup.cancelled", {
-											jobId: payload.jobId,
-											scheduleId: payload.scheduleId,
-											message: "Backup was cancelled",
-										}),
-									);
-									return;
 							}
 
-								context.offerOutbound(
+							return context.offerOutbound(
 								createAgentMessage("backup.failed", {
 									jobId: payload.jobId,
 									scheduleId: payload.scheduleId,
@@ -105,14 +102,9 @@ export const handleBackupRunCommand = (context: ControllerCommandContext, payloa
 									errorDetails: toErrorDetails(error),
 								}),
 							);
-							});
 						},
 					}),
-					Effect.ensuring(
+					Effect.ensuring(context.deleteRunningJob(payload.jobId)),
-						Effect.sync(() => {
-							context.deleteRunningJob(payload.jobId);
-						}),
-					),
 				);
 		}),
 	).pipe(Effect.asVoid);

apps/agent/src/commands/heartbeat-ping.ts

@@ -5,8 +5,8 @@ import type { ControllerCommandContext } from "../context";
 type HeartbeatPingPayload = Extract<ControllerMessage, { type: "heartbeat.ping" }>["payload"];
 
 export const handleHeartbeatPingCommand = (context: ControllerCommandContext, payload: HeartbeatPingPayload) =>
-	Effect.sync(() => {
+	Effect.gen(function* () {
-		context.offerOutbound(
+		yield* context.offerOutbound(
 			createAgentMessage("heartbeat.pong", {
 				sentAt: payload.sentAt,
 			}),

apps/agent/src/context.ts

@@ -1,4 +1,5 @@
 import type { AgentWireMessage } from "@zerobyte/contracts/agent-protocol";
+import type { Effect } from "effect";
 
 export type RunningJob = {
 	scheduleId: string;
@@ -6,8 +7,8 @@ export type RunningJob = {
 };
 
 export type ControllerCommandContext = {
-	getRunningJob: (jobId: string) => RunningJob | undefined;
+	getRunningJob: (jobId: string) => Effect.Effect<RunningJob | undefined, never, never>;
-	setRunningJob: (jobId: string, job: RunningJob) => void;
+	setRunningJob: (jobId: string, job: RunningJob) => Effect.Effect<void, never, never>;
-	deleteRunningJob: (jobId: string) => void;
+	deleteRunningJob: (jobId: string) => Effect.Effect<void, never, never>;
-	offerOutbound: (message: AgentWireMessage) => void;
+	offerOutbound: (message: AgentWireMessage) => Effect.Effect<boolean, never, never>;
 };

apps/agent/src/controller-session.ts

@@ -8,6 +8,7 @@ import {
 import { logger } from "@zerobyte/core/node";
 import { toMessage } from "@zerobyte/core/utils";
 import { handleControllerCommand } from "./commands";
+import type { ControllerCommandContext, RunningJob } from "./context";
 
 export type ControllerSession = {
 	onOpen: () => void;
@@ -18,45 +19,44 @@ export type ControllerSession = {
 export const createControllerSession = (ws: WebSocket): ControllerSession => {
 	const outboundQueue = Effect.runSync(Queue.bounded<AgentWireMessage>(64));
 	const inboundQueue = Effect.runSync(Queue.bounded<ControllerWireMessage>(64));
-	const runningJobsRef = Effect.runSync(
+	const runningJobsRef = Effect.runSync(Ref.make<Map<string, RunningJob>>(new Map()));
-		Ref.make<Map<string, { scheduleId: string; abortController: AbortController }>>(new Map()),
-	);
 
-	const getRunningJob = (jobId: string) => Effect.runSync(Ref.get(runningJobsRef)).get(jobId);
+	const getRunningJob = (jobId: string) => Ref.get(runningJobsRef).pipe(Effect.map((map) => map.get(jobId)));
 
-	const setRunningJob = (jobId: string, job: { scheduleId: string; abortController: AbortController }) => {
+	const setRunningJob = (jobId: string, job: RunningJob) => {
-		Effect.runSync(
+		return Ref.update(runningJobsRef, (current) => {
-			Ref.update(runningJobsRef, (current) => {
 			const next = new Map(current);
 			next.set(jobId, job);
 			return next;
-			}),
+		});
-		);
 	};
 
+	const abortRunningJobs = Effect.gen(function* () {
+		const runningJobs = yield* Ref.modify(runningJobsRef, (current) => [current, new Map()]);
+		yield* Effect.sync(() => {
+			for (const runningJob of runningJobs.values()) {
+				runningJob.abortController.abort();
+			}
+		});
+	});
+
 	const deleteRunningJob = (jobId: string) => {
-		Effect.runSync(
+		return Ref.update(runningJobsRef, (current) => {
-			Ref.update(runningJobsRef, (current) => {
 			const next = new Map(current);
 			next.delete(jobId);
 			return next;
-			}),
+		});
-		);
 	};
 
 	const offerOutbound = (message: AgentWireMessage) => {
-		void Effect.runPromise(Queue.offer(outboundQueue, message)).catch((error) => {
+		return Queue.offer(outboundQueue, message);
-			logger.error(`Failed to queue outbound controller message: ${toMessage(error)}`);
-		});
 	};
 
 	const offerInbound = (message: ControllerWireMessage) => {
-		void Effect.runPromise(Queue.offer(inboundQueue, message)).catch((error) => {
+		return Queue.offer(inboundQueue, message);
-			logger.error(`Failed to queue inbound controller message: ${toMessage(error)}`);
-		});
 	};
 
-	const commandContext = {
+	const commandContext: ControllerCommandContext = {
 		getRunningJob,
 		setRunningJob,
 		deleteRunningJob,
@@ -101,7 +101,9 @@ export const createControllerSession = (ws: WebSocket): ControllerSession => {
 
 	return {
 		onOpen: () => {
-			offerOutbound(createAgentMessage("agent.ready", { agentId: "" }));
+			void Effect.runPromise(offerOutbound(createAgentMessage("agent.ready", { agentId: "" }))).catch((error) => {
+				logger.error(`Failed to queue ready message: ${toMessage(error)}`);
+			});
 		},
 		onMessage: (data) => {
 			if (typeof data !== "string") {
@@ -109,14 +111,12 @@ export const createControllerSession = (ws: WebSocket): ControllerSession => {
 				return;
 			}
 
-			offerInbound(data as ControllerWireMessage);
+			void Effect.runPromise(offerInbound(data as ControllerWireMessage)).catch((error) => {
+				logger.error(`Failed to queue inbound message: ${toMessage(error)}`);
+			});
 		},
 		close: () => {
-			const runningJobs = Effect.runSync(Ref.get(runningJobsRef));
+			void Effect.runPromise(abortRunningJobs).catch(() => {});
-			for (const running of runningJobs.values()) {
-				running.abortController.abort();
-			}
-			Effect.runSync(Ref.set(runningJobsRef, new Map()));
 			void Effect.runPromise(Fiber.interrupt(writerFiber)).catch(() => {});
 			void Effect.runPromise(Fiber.interrupt(processorFiber)).catch(() => {});
 			void Effect.runPromise(Queue.shutdown(outboundQueue)).catch(() => {});