refactor: move effect / async boundary in one place

app/server/modules/agents/__tests__/agents-manager.backups.test.ts

@@ -1,6 +1,7 @@
 import { afterEach, expect, test, vi } from "vitest";
 import waitForExpect from "wait-for-expect";
 import { fromAny, fromPartial } from "@total-typescript/shoehorn";
+import { Effect } from "effect";
 import { agentManager, type ProcessWithAgentRuntime } from "../agents-manager";
 import type { AgentManagerRuntime } from "../controller/server";
 import type { BackupRunPayload } from "@zerobyte/contracts/agent-protocol";
@@ -20,8 +21,8 @@ afterEach(() => {
 });
 
 test("cancelBackup resolves a running backup when the cancel command cannot be delivered", async () => {
-	const sendBackup = vi.fn().mockResolvedValue(true);
+	const sendBackup = vi.fn(() => Effect.succeed(true));
-	const cancelBackup = vi.fn().mockResolvedValue(false);
+	const cancelBackup = vi.fn(() => Effect.succeed(false));
 	setAgentRuntime({ sendBackup, cancelBackup });
 
 	const resultPromise = agentManager.runBackup("local", {

app/server/modules/agents/__tests__/agents-manager.events.test.ts

@@ -0,0 +1,235 @@
+import { afterEach, expect, test, vi } from "vitest";
+import { Effect } from "effect";
+import { fromAny, fromPartial } from "@total-typescript/shoehorn";
+import type { BackupRunPayload } from "@zerobyte/contracts/agent-protocol";
+import type { AgentManagerEvent } from "../controller/server";
+import type { ProcessWithAgentRuntime } from "../agents-manager";
+
+const controllerMock = vi.hoisted(() => ({
+	onEvent: null as null | ((event: AgentManagerEvent) => void),
+	sendBackup: vi.fn(),
+	cancelBackup: vi.fn(),
+	stop: vi.fn(),
+}));
+
+vi.mock("../controller/server", async () => {
+	const { Effect } = await import("effect");
+	return {
+		createAgentManagerRuntime: vi.fn((onEvent: (event: AgentManagerEvent) => void) => {
+			controllerMock.onEvent = onEvent;
+			return {
+				start: Effect.void,
+				stop: Effect.sync(controllerMock.stop),
+				sendBackup: controllerMock.sendBackup,
+				cancelBackup: controllerMock.cancelBackup,
+			};
+		}),
+	};
+});
+
+const processWithAgentRuntime = process as ProcessWithAgentRuntime;
+
+const resetAgentRuntime = () => {
+	processWithAgentRuntime.__zerobyteAgentRuntime = {
+		agentManager: null,
+		localAgent: null,
+		isStoppingLocalAgent: false,
+		localAgentRestartTimeout: null,
+		activeBackupsByScheduleId: new Map(),
+		activeBackupScheduleIdsByJobId: new Map(),
+	};
+};
+
+const backupPayload = fromPartial<BackupRunPayload>({
+	jobId: "job-1",
+	scheduleId: "schedule-1",
+});
+
+afterEach(() => {
+	delete processWithAgentRuntime.__zerobyteAgentRuntime;
+	controllerMock.onEvent = null;
+	controllerMock.sendBackup.mockReset();
+	controllerMock.cancelBackup.mockReset();
+	controllerMock.stop.mockReset();
+	vi.resetModules();
+	vi.restoreAllMocks();
+});
+
+test("backup progress is delivered to the running backup callback", async () => {
+	resetAgentRuntime();
+	controllerMock.sendBackup.mockImplementation(() => Effect.succeed(true));
+	const { agentManager, startAgentController, stopAgentController } = await import("../agents-manager");
+	const onProgress = vi.fn();
+
+	await startAgentController();
+	const resultPromise = agentManager.runBackup("local", {
+		scheduleId: 42,
+		payload: backupPayload,
+		signal: new AbortController().signal,
+		onProgress,
+	});
+
+	controllerMock.onEvent?.({
+		type: "backup.progress",
+		agentId: "local",
+		agentName: "Local Agent",
+		payload: fromAny({ jobId: "job-1", scheduleId: "schedule-1", progress: { percentDone: 0.5 } }),
+	});
+	controllerMock.onEvent?.({
+		type: "backup.completed",
+		agentId: "local",
+		agentName: "Local Agent",
+		payload: { jobId: "job-1", scheduleId: "schedule-1", exitCode: 0, result: null },
+	});
+
+	await expect(resultPromise).resolves.toEqual({
+		status: "completed",
+		exitCode: 0,
+		result: null,
+		warningDetails: null,
+	});
+	expect(onProgress).toHaveBeenCalledWith({ percentDone: 0.5 });
+	await stopAgentController();
+});
+
+test("backup failed and cancelled events resolve the matching running backup", async () => {
+	resetAgentRuntime();
+	controllerMock.sendBackup.mockImplementation(() => Effect.succeed(true));
+	const { agentManager, startAgentController, stopAgentController } = await import("../agents-manager");
+
+	await startAgentController();
+	const failedPromise = agentManager.runBackup("local", {
+		scheduleId: 42,
+		payload: backupPayload,
+		signal: new AbortController().signal,
+		onProgress: vi.fn(),
+	});
+	controllerMock.onEvent?.({
+		type: "backup.failed",
+		agentId: "local",
+		agentName: "Local Agent",
+		payload: { jobId: "job-1", scheduleId: "schedule-1", error: "failed", errorDetails: "restic failed" },
+	});
+	await expect(failedPromise).resolves.toEqual({ status: "failed", error: "restic failed" });
+
+	const cancelledPromise = agentManager.runBackup("local", {
+		scheduleId: 43,
+		payload: fromPartial<BackupRunPayload>({ jobId: "job-2", scheduleId: "schedule-2" }),
+		signal: new AbortController().signal,
+		onProgress: vi.fn(),
+	});
+	controllerMock.onEvent?.({
+		type: "backup.cancelled",
+		agentId: "local",
+		agentName: "Local Agent",
+		payload: { jobId: "job-2", scheduleId: "schedule-2", message: "cancelled remotely" },
+	});
+	await expect(cancelledPromise).resolves.toEqual({ status: "cancelled", message: "cancelled remotely" });
+	await stopAgentController();
+});
+
+test("agent disconnect cancels only backups owned by that agent", async () => {
+	resetAgentRuntime();
+	controllerMock.sendBackup.mockImplementation(() => Effect.succeed(true));
+	const { agentManager, startAgentController, stopAgentController } = await import("../agents-manager");
+
+	await startAgentController();
+	const localPromise = agentManager.runBackup("local", {
+		scheduleId: 42,
+		payload: backupPayload,
+		signal: new AbortController().signal,
+		onProgress: vi.fn(),
+	});
+	const remotePromise = agentManager.runBackup("remote", {
+		scheduleId: 43,
+		payload: fromPartial<BackupRunPayload>({ jobId: "job-2", scheduleId: "schedule-2" }),
+		signal: new AbortController().signal,
+		onProgress: vi.fn(),
+	});
+
+	controllerMock.onEvent?.({ type: "agent.disconnected", agentId: "local", agentName: "Local Agent" });
+	controllerMock.onEvent?.({
+		type: "backup.completed",
+		agentId: "remote",
+		agentName: "Remote Agent",
+		payload: { jobId: "job-2", scheduleId: "schedule-2", exitCode: 0, result: null },
+	});
+
+	await expect(localPromise).resolves.toEqual({
+		status: "cancelled",
+		message: "The connection to the backup agent was lost. Restart the backup to ensure it completes.",
+	});
+	await expect(remotePromise).resolves.toEqual({
+		status: "completed",
+		exitCode: 0,
+		result: null,
+		warningDetails: null,
+	});
+	await stopAgentController();
+});
+
+test("runBackup returns unavailable and clears the active run when the command cannot be sent", async () => {
+	resetAgentRuntime();
+	controllerMock.sendBackup.mockImplementation(() => Effect.succeed(false));
+	const { agentManager, startAgentController, stopAgentController } = await import("../agents-manager");
+
+	await startAgentController();
+	const result = await agentManager.runBackup("local", {
+		scheduleId: 42,
+		payload: backupPayload,
+		signal: new AbortController().signal,
+		onProgress: vi.fn(),
+	});
+
+	expect(result).toEqual({
+		status: "unavailable",
+		error: new Error("Failed to send backup command to agent local"),
+	});
+	await expect(agentManager.cancelBackup("local", 42)).resolves.toBe(false);
+	await stopAgentController();
+});
+
+test("runBackup rejects before sending when the abort signal is already aborted", async () => {
+	resetAgentRuntime();
+	controllerMock.sendBackup.mockImplementation(() => Effect.succeed(true));
+	const { agentManager, startAgentController, stopAgentController } = await import("../agents-manager");
+	const abortController = new AbortController();
+	abortController.abort(new Error("cancelled before send"));
+
+	await startAgentController();
+	await expect(
+		agentManager.runBackup("local", {
+			scheduleId: 42,
+			payload: backupPayload,
+			signal: abortController.signal,
+			onProgress: vi.fn(),
+		}),
+	).rejects.toThrow("cancelled before send");
+	expect(controllerMock.sendBackup).not.toHaveBeenCalled();
+	await stopAgentController();
+});
+
+test("runBackup requests cancellation when the abort signal fires while sending", async () => {
+	resetAgentRuntime();
+	const abortController = new AbortController();
+	controllerMock.sendBackup.mockImplementation(() =>
+		Effect.sync(() => {
+			abortController.abort();
+			return true;
+		}),
+	);
+	controllerMock.cancelBackup.mockImplementation(() => Effect.succeed(false));
+	const { agentManager, startAgentController, stopAgentController } = await import("../agents-manager");
+
+	await startAgentController();
+	const result = await agentManager.runBackup("local", {
+		scheduleId: 42,
+		payload: backupPayload,
+		signal: abortController.signal,
+		onProgress: vi.fn(),
+	});
+
+	expect(result).toEqual({ status: "cancelled" });
+	expect(controllerMock.cancelBackup).toHaveBeenCalledWith("local", { jobId: "job-1", scheduleId: "schedule-1" });
+	await stopAgentController();
+});

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

@@ -1,6 +1,7 @@
 import { beforeEach, expect, test } from "vitest";
 import { db } from "~/server/db/db";
 import { agentsTable } from "~/server/db/schema";
+import { LOCAL_AGENT_ID, LOCAL_AGENT_KIND, LOCAL_AGENT_NAME } from "../constants";
 import { agentsService } from "../agents.service";
 
 beforeEach(async () => {
@@ -15,3 +16,57 @@ test("ensureLocalAgent seeds the built-in local agent once", async () => {
 
 	expect(agents).toHaveLength(1);
 });
+
+test("markAgentConnecting creates and updates connection metadata", async () => {
+	await agentsService.markAgentConnecting({
+		agentId: "remote-agent",
+		organizationId: null,
+		agentName: "Remote Agent",
+		agentKind: "remote",
+		capabilities: { restic: true },
+		connectedAt: 1_000,
+	});
+	await agentsService.markAgentConnecting({
+		agentId: "remote-agent",
+		organizationId: null,
+		agentName: "Renamed Agent",
+		agentKind: "remote",
+		capabilities: { restic: true, webdav: true },
+		connectedAt: 2_000,
+	});
+
+	const agent = await agentsService.getAgent("remote-agent");
+
+	expect(agent).toMatchObject({
+		id: "remote-agent",
+		name: "Renamed Agent",
+		kind: "remote",
+		status: "connecting",
+		capabilities: { restic: true, webdav: true },
+		lastSeenAt: 2_000,
+		updatedAt: 2_000,
+	});
+});
+
+test("agent runtime status moves from connecting to online, seen, and offline", async () => {
+	await agentsService.markAgentConnecting({
+		agentId: LOCAL_AGENT_ID,
+		organizationId: null,
+		agentName: LOCAL_AGENT_NAME,
+		agentKind: LOCAL_AGENT_KIND,
+		connectedAt: 1_000,
+	});
+	await agentsService.markAgentOnline(LOCAL_AGENT_ID, 2_000);
+	await agentsService.markAgentSeen(LOCAL_AGENT_ID, 3_000);
+	await agentsService.markAgentOffline(LOCAL_AGENT_ID, 4_000);
+
+	const agent = await agentsService.getAgent(LOCAL_AGENT_ID);
+
+	expect(agent).toMatchObject({
+		id: LOCAL_AGENT_ID,
+		status: "offline",
+		lastSeenAt: 3_000,
+		lastReadyAt: 2_000,
+		updatedAt: 4_000,
+	});
+});

app/server/modules/agents/__tests__/controller-runtime.test.ts

@@ -0,0 +1,219 @@
+import { Effect } from "effect";
+import { afterEach, 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";
+
+const agentsServiceMocks = vi.hoisted(() => ({
+	markAgentConnecting: vi.fn(() => Promise.resolve()),
+	markAgentOnline: vi.fn(() => Promise.resolve()),
+	markAgentSeen: vi.fn(() => Promise.resolve()),
+	markAgentOffline: vi.fn(() => Promise.resolve()),
+}));
+
+const tokenMocks = vi.hoisted(() => ({
+	validateAgentToken: vi.fn(),
+}));
+
+vi.mock("../agents.service", () => ({
+	agentsService: agentsServiceMocks,
+}));
+
+vi.mock("../helpers/tokens", () => ({
+	validateAgentToken: tokenMocks.validateAgentToken,
+}));
+
+const createSocket = (id: string) => ({
+	data: {
+		id,
+		agentId: LOCAL_AGENT_ID,
+		organizationId: null,
+		agentName: LOCAL_AGENT_NAME,
+		agentKind: LOCAL_AGENT_KIND,
+	},
+	send: vi.fn(() => 1),
+	close: vi.fn(),
+});
+
+type CapturedFetch = NonNullable<Parameters<typeof Bun.serve>[0]["fetch"]>;
+
+const invokeFetch = (fetch: CapturedFetch | undefined, request: Request, srv: Parameters<CapturedFetch>[1]) => {
+	if (!fetch) {
+		throw new Error("Bun.serve was not called with a fetch handler");
+	}
+
+	return Reflect.apply(fetch, fromPartial<ThisParameterType<CapturedFetch>>({}), [
+		request,
+		srv,
+	]) as ReturnType<CapturedFetch>;
+};
+
+const startRuntime = async (onEvent = vi.fn()) => {
+	const { createAgentManagerRuntime } = await import("../controller/server");
+	const runtime = createAgentManagerRuntime(onEvent);
+	await Effect.runPromise(runtime.start);
+	return { runtime, onEvent };
+};
+
+afterEach(() => {
+	vi.restoreAllMocks();
+	tokenMocks.validateAgentToken.mockReset();
+	vi.resetModules();
+});
+
+test("websocket fetch rejects requests without a bearer token", async () => {
+	const serve = vi
+		.spyOn(Bun, "serve")
+		.mockReturnValue(fromPartial({ port: 3001, stop: vi.fn(() => Promise.resolve()) }));
+	const { runtime } = await startRuntime();
+	const fetch = serve.mock.calls[0]?.[0].fetch;
+	const upgrade = vi.fn();
+	const srv = fromPartial<Parameters<NonNullable<typeof fetch>>[1]>({ upgrade });
+
+	const response = await invokeFetch(fetch, new Request("http://localhost:3001/agent"), srv);
+	await Effect.runPromise(runtime.stop);
+
+	expect(response?.status).toBe(401);
+	expect(await response?.text()).toBe("Missing token");
+	expect(upgrade).not.toHaveBeenCalled();
+});
+
+test("websocket fetch rejects invalid bearer tokens", async () => {
+	tokenMocks.validateAgentToken.mockResolvedValue(undefined);
+	const serve = vi
+		.spyOn(Bun, "serve")
+		.mockReturnValue(fromPartial({ port: 3001, stop: vi.fn(() => Promise.resolve()) }));
+	const { runtime } = await startRuntime();
+	const fetch = serve.mock.calls[0]?.[0].fetch;
+	const upgrade = vi.fn();
+	const srv = fromPartial<Parameters<NonNullable<typeof fetch>>[1]>({ upgrade });
+
+	const response = await invokeFetch(
+		fetch,
+		new Request("http://localhost:3001/agent", { headers: { authorization: "Bearer bad-token" } }),
+		srv,
+	);
+	await Effect.runPromise(runtime.stop);
+
+	expect(response?.status).toBe(401);
+	expect(await response?.text()).toBe("Invalid or revoked token");
+	expect(tokenMocks.validateAgentToken).toHaveBeenCalledWith("bad-token");
+	expect(upgrade).not.toHaveBeenCalled();
+});
+
+test("websocket fetch upgrades valid agent tokens with connection metadata", async () => {
+	tokenMocks.validateAgentToken.mockResolvedValue({
+		agentId: LOCAL_AGENT_ID,
+		organizationId: null,
+		agentName: LOCAL_AGENT_NAME,
+		agentKind: LOCAL_AGENT_KIND,
+	});
+	const serve = vi
+		.spyOn(Bun, "serve")
+		.mockReturnValue(fromPartial({ port: 3001, stop: vi.fn(() => Promise.resolve()) }));
+	const { runtime } = await startRuntime();
+	const fetch = serve.mock.calls[0]?.[0].fetch;
+	const upgrade = vi.fn(() => true);
+	const srv = fromPartial<Parameters<NonNullable<typeof fetch>>[1]>({ upgrade });
+
+	const response = await invokeFetch(
+		fetch,
+		new Request("http://localhost:3001/agent", { headers: { authorization: "Bearer valid-token" } }),
+		srv,
+	);
+	await Effect.runPromise(runtime.stop);
+
+	expect(response).toBeUndefined();
+	expect(tokenMocks.validateAgentToken).toHaveBeenCalledWith("valid-token");
+	expect(upgrade).toHaveBeenCalledWith(expect.any(Request), {
+		data: expect.objectContaining({
+			agentId: LOCAL_AGENT_ID,
+			organizationId: null,
+			agentName: LOCAL_AGENT_NAME,
+			agentKind: LOCAL_AGENT_KIND,
+			id: expect.any(String),
+		}),
+	});
+});
+
+test("websocket lifecycle updates agent connection status", async () => {
+	const stop = vi.fn(() => Promise.resolve());
+	const serve = vi.spyOn(Bun, "serve").mockReturnValue(fromPartial({ port: 3001, stop }));
+	const { runtime } = await startRuntime();
+	const websocket = serve.mock.calls[0]?.[0].websocket;
+	const socket = createSocket("connection-1");
+
+	await websocket?.open?.(fromPartial(socket));
+	await websocket?.message?.(fromPartial(socket), createAgentMessage("agent.ready", { agentId: LOCAL_AGENT_ID }));
+	await websocket?.message?.(fromPartial(socket), createAgentMessage("heartbeat.pong", { sentAt: 123 }));
+	await websocket?.close?.(fromPartial(socket), 1000, "done");
+	await Effect.runPromise(runtime.stop);
+
+	expect(agentsServiceMocks.markAgentConnecting).toHaveBeenCalledWith({
+		agentId: LOCAL_AGENT_ID,
+		organizationId: null,
+		agentName: LOCAL_AGENT_NAME,
+		agentKind: LOCAL_AGENT_KIND,
+	});
+	expect(agentsServiceMocks.markAgentOnline).toHaveBeenCalledWith(LOCAL_AGENT_ID, expect.any(Number));
+	expect(agentsServiceMocks.markAgentSeen).toHaveBeenCalledWith(LOCAL_AGENT_ID, expect.any(Number));
+	expect(agentsServiceMocks.markAgentOffline).toHaveBeenCalledWith(LOCAL_AGENT_ID);
+	expect(stop).toHaveBeenCalledWith(true);
+});
+
+test("closing a replaced connection does not report the active agent as disconnected", async () => {
+	const serve = vi
+		.spyOn(Bun, "serve")
+		.mockReturnValue(fromPartial({ port: 3001, stop: vi.fn(() => Promise.resolve()) }));
+	const { runtime, onEvent } = await startRuntime(vi.fn());
+	const websocket = serve.mock.calls[0]?.[0].websocket;
+	const oldSocket = createSocket("connection-1");
+	const newSocket = createSocket("connection-2");
+
+	await websocket?.open?.(fromPartial(oldSocket));
+	await websocket?.open?.(fromPartial(newSocket));
+	await websocket?.close?.(fromPartial(oldSocket), 1000, "replaced");
+	await Effect.runPromise(runtime.stop);
+
+	expect(onEvent).not.toHaveBeenCalledWith(
+		expect.objectContaining({ type: "agent.disconnected", agentId: LOCAL_AGENT_ID }),
+	);
+});
+
+test("sendBackup is only delivered after the agent is ready", async () => {
+	const serve = vi
+		.spyOn(Bun, "serve")
+		.mockReturnValue(fromPartial({ port: 3001, stop: vi.fn(() => Promise.resolve()) }));
+	const { runtime } = await startRuntime();
+	const websocket = serve.mock.calls[0]?.[0].websocket;
+	const socket = createSocket("connection-1");
+	const payload = {
+		jobId: "job-1",
+		scheduleId: "schedule-1",
+		organizationId: "org-1",
+		sourcePath: "/tmp/source",
+		repositoryConfig: { backend: "local" as const, path: "/tmp/repository" },
+		options: {},
+		runtime: {
+			password: "password",
+			cacheDir: "/tmp/cache",
+			passFile: "/tmp/pass",
+			defaultExcludes: [],
+			rcloneConfigFile: "/tmp/rclone.conf",
+		},
+		webhooks: { pre: null, post: null },
+		webhookAllowedOrigins: [],
+	};
+
+	await websocket?.open?.(fromPartial(socket));
+	await expect(Effect.runPromise(runtime.sendBackup(LOCAL_AGENT_ID, payload))).resolves.toBe(false);
+
+	await websocket?.message?.(fromPartial(socket), createAgentMessage("agent.ready", { agentId: LOCAL_AGENT_ID }));
+	await expect(Effect.runPromise(runtime.sendBackup(LOCAL_AGENT_ID, payload))).resolves.toBe(true);
+
+	await waitForExpect(() => {
+		expect(socket.send).toHaveBeenCalledWith(expect.stringContaining('"type":"backup.run"'));
+	});
+	await Effect.runPromise(runtime.stop);
+});

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

@@ -2,6 +2,7 @@ import { Effect, Exit, Fiber, Scope } from "effect";
 import { expect, test, vi } from "vitest";
 import waitForExpect from "wait-for-expect";
 import { fromPartial } from "@total-typescript/shoehorn";
+import { createAgentMessage, type AgentMessage } from "@zerobyte/contracts/agent-protocol";
 import { LOCAL_AGENT_ID, LOCAL_AGENT_KIND, LOCAL_AGENT_NAME } from "../constants";
 import { createControllerAgentSession } from "../controller/session";
 
@@ -21,26 +22,13 @@ const createSocket = (overrides: Partial<Parameters<typeof createControllerAgent
 };
 
 const createSession = (
-	handlers: Partial<Parameters<typeof createControllerAgentSession>[1]> = {},
+	onEvent: Parameters<typeof createControllerAgentSession>[1] = () => Effect.void,
 	socket = createSocket(),
 ) => {
 	const scope = Effect.runSync(Scope.make());
-	const sessionHandlers: Parameters<typeof createControllerAgentSession>[1] = {
-		onReady: () => Effect.void,
-		onHeartbeatPong: () => Effect.void,
-		onDisconnect: () => Effect.void,
-		onBackupStarted: () => Effect.void,
-		onBackupProgress: () => Effect.void,
-		onBackupCompleted: () => Effect.void,
-		onBackupFailed: () => Effect.void,
-		onBackupCancelled: () => Effect.void,
-		...handlers,
-	};
 
 	try {
-		const session = Effect.runSync(
+		const session = Effect.runSync(Scope.extend(createControllerAgentSession(fromPartial(socket), onEvent), scope));
-			Scope.extend(createControllerAgentSession(fromPartial(socket), sessionHandlers), scope),
-		);
 
 		return {
 			session,
@@ -74,23 +62,22 @@ test("closing the session scope interrupts the session runner", async () => {
 });
 
 test("close reports a transport disconnect", () => {
-	const onDisconnect = vi.fn(() => Effect.void);
+	const onEvent = vi.fn(() => Effect.void);
-	const { close } = createSession({ onDisconnect });
+	const { close } = createSession(onEvent);
 
 	close();
 
-	expect(onDisconnect).toHaveBeenCalledTimes(1);
+	expect(onEvent).toHaveBeenCalledTimes(1);
-	expect(onDisconnect).toHaveBeenCalledWith(
+	expect(onEvent).toHaveBeenCalledWith(
 		expect.objectContaining({
-			agentId: LOCAL_AGENT_ID,
+			type: "agent.disconnected",
-			agentName: LOCAL_AGENT_NAME,
 		}),
 	);
 });
 
 test("sendBackup only queues the transport message", () => {
-	const onBackupCancelled = vi.fn(() => Effect.void);
+	const onEvent = vi.fn(() => Effect.void);
-	const { session, close } = createSession({ onBackupCancelled });
+	const { session, close } = createSession(onEvent);
 
 	Effect.runSync(
 		session.sendBackup({
@@ -118,14 +105,74 @@ test("sendBackup only queues the transport message", () => {
 
 	close();
 
-	expect(onBackupCancelled).not.toHaveBeenCalled();
+	expect(onEvent).not.toHaveBeenCalledWith(expect.objectContaining({ type: "backup.cancelled" }));
+});
+
+test("invalid inbound messages are ignored", () => {
+	const onEvent = vi.fn(() => Effect.void);
+	const { session, close } = createSession(onEvent);
+
+	Effect.runSync(session.handleMessage("not json"));
+	Effect.runSync(session.handleMessage(JSON.stringify({ type: "backup.progress", payload: {} })));
+
+	expect(onEvent).not.toHaveBeenCalled();
+	close();
+});
+
+test("agent.ready marks the session ready and forwards the event", () => {
+	const onEvent = vi.fn(() => Effect.void);
+	const { session, close } = createSession(onEvent);
+
+	expect(Effect.runSync(session.isReady())).toBe(false);
+	Effect.runSync(session.handleMessage(createAgentMessage("agent.ready", { agentId: LOCAL_AGENT_ID })));
+
+	expect(Effect.runSync(session.isReady())).toBe(true);
+	expect(onEvent).toHaveBeenCalledWith({ type: "agent.ready", payload: { agentId: LOCAL_AGENT_ID } });
+	close();
+});
+
+test("backup agent messages are forwarded unchanged", () => {
+	const onEvent = vi.fn(() => Effect.void);
+	const { session, close } = createSession(onEvent);
+	const message = {
+		type: "backup.progress" as const,
+		payload: {
+			jobId: "job-1",
+			scheduleId: "schedule-1",
+			progress: {
+				message_type: "status" as const,
+				seconds_elapsed: 0,
+				seconds_remaining: 0,
+				percent_done: 0.5,
+				total_files: 0,
+				files_done: 0,
+				total_bytes: 0,
+				bytes_done: 0,
+				current_files: [],
+			},
+		},
+	} satisfies Extract<AgentMessage, { type: "backup.progress" }>;
+
+	Effect.runSync(session.handleMessage(createAgentMessage(message.type, message.payload)));
+
+	expect(onEvent).toHaveBeenCalledWith(
+		expect.objectContaining({
+			type: message.type,
+			payload: expect.objectContaining({
+				jobId: message.payload.jobId,
+				scheduleId: message.payload.scheduleId,
+				progress: expect.objectContaining(message.payload.progress),
+			}),
+		}),
+	);
+	close();
 });
 
 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 onDisconnect = vi.fn(() => Effect.void);
+	const onEvent = vi.fn(() => Effect.void);
-	const { session, run, closeAsync } = createSession({ onDisconnect }, socket);
+	const { session, run, closeAsync } = createSession(onEvent, socket);
 
 	try {
 		run();
@@ -139,10 +186,10 @@ test("a dropped backup.cancel closes the session and reports a transport disconn
 		await waitForExpect(() => {
 			expect(send).toHaveBeenCalledTimes(1);
 			expect(socket.close).toHaveBeenCalledTimes(1);
-			expect(onDisconnect).toHaveBeenCalledTimes(1);
+			expect(onEvent).toHaveBeenCalledTimes(1);
-			expect(onDisconnect).toHaveBeenCalledWith(
+			expect(onEvent).toHaveBeenCalledWith(
 				expect.objectContaining({
-					agentId: LOCAL_AGENT_ID,
+					type: "agent.disconnected",
 				}),
 			);
 		});

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

@@ -1,7 +1,8 @@
 import { logger } from "@zerobyte/core/node";
 import type { BackupRunPayload } from "@zerobyte/contracts/agent-protocol";
+import { Effect } from "effect";
 import { config } from "../../core/config";
-import { createAgentManagerRuntime, type AgentBackupEventHandlers } from "./controller/server";
+import { createAgentManagerRuntime, type AgentManagerEvent } from "./controller/server";
 import { spawnLocalAgentProcess, stopLocalAgentProcess } from "./local/process";
 import type { BackupExecutionProgress, BackupExecutionResult } from "./helpers/runtime-state";
 import { createAgentRuntimeState } from "./helpers/runtime-state";
@@ -97,10 +98,12 @@ const requestBackupCancellation = async (agentId: string, scheduleId: number) =>
 	}
 
 	if (
-		await runtime.cancelBackup(agentId, {
+		await Effect.runPromise(
+			runtime.cancelBackup(agentId, {
 				jobId: activeBackupRun.jobId,
 				scheduleId: activeBackupRun.scheduleShortId,
-		})
+			}),
+		)
 	) {
 		return true;
 	}
@@ -109,71 +112,99 @@ const requestBackupCancellation = async (agentId: string, scheduleId: number) =>
 	return true;
 };
 
-const backupEventHandlers: AgentBackupEventHandlers = {
+const handleAgentManagerEvent = (event: AgentManagerEvent) => {
-	onAgentDisconnected: ({ agentId }) => {
+	switch (event.type) {
+		case "agent.disconnected": {
 			cancelActiveBackupRunsForAgent(
-			agentId,
+				event.agentId,
 				"The connection to the backup agent was lost. Restart the backup to ensure it completes.",
 			);
-	},
+			break;
-	onBackupStarted: ({ agentId, payload }) => {
+		}
-		getActiveBackupRun(payload.jobId, payload.scheduleId, "backup.started", agentId);
+		case "backup.started": {
-	},
+			getActiveBackupRun(event.payload.jobId, event.payload.scheduleId, event.type, event.agentId);
-	onBackupProgress: ({ agentId, payload }) => {
+			break;
-		const activeBackupRun = getActiveBackupRun(payload.jobId, payload.scheduleId, "backup.progress", agentId);
+		}
+		case "backup.progress": {
+			const activeBackupRun = getActiveBackupRun(
+				event.payload.jobId,
+				event.payload.scheduleId,
+				event.type,
+				event.agentId,
+			);
 			if (!activeBackupRun) {
-			return;
+				break;
 			}
 
-		activeBackupRun.onProgress(payload.progress);
+			activeBackupRun.onProgress(event.payload.progress);
-	},
+			break;
-	onBackupCompleted: ({ agentId, payload }) => {
+		}
-		const activeBackupRun = getActiveBackupRun(payload.jobId, payload.scheduleId, "backup.completed", agentId);
+		case "backup.completed": {
+			const activeBackupRun = getActiveBackupRun(
+				event.payload.jobId,
+				event.payload.scheduleId,
+				event.type,
+				event.agentId,
+			);
 			if (!activeBackupRun) {
-			return;
+				break;
 			}
 
 			resolveActiveBackupRun(activeBackupRun.scheduleId, {
 				status: "completed",
-			exitCode: payload.exitCode,
+				exitCode: event.payload.exitCode,
-			result: payload.result,
+				result: event.payload.result,
-			warningDetails: payload.warningDetails ?? null,
+				warningDetails: event.payload.warningDetails ?? null,
 			});
-	},
+			break;
-	onBackupFailed: ({ agentId, payload }) => {
+		}
-		const activeBackupRun = getActiveBackupRun(payload.jobId, payload.scheduleId, "backup.failed", agentId);
+		case "backup.failed": {
+			const activeBackupRun = getActiveBackupRun(
+				event.payload.jobId,
+				event.payload.scheduleId,
+				event.type,
+				event.agentId,
+			);
 			if (!activeBackupRun) {
-			return;
+				break;
 			}
 
 			resolveActiveBackupRun(activeBackupRun.scheduleId, {
 				status: "failed",
-			error: payload.errorDetails ?? payload.error,
+				error: event.payload.errorDetails ?? event.payload.error,
 			});
-	},
+			break;
-	onBackupCancelled: ({ agentId, payload }) => {
+		}
-		const activeBackupRun = getActiveBackupRun(payload.jobId, payload.scheduleId, "backup.cancelled", agentId);
+		case "backup.cancelled": {
+			const activeBackupRun = getActiveBackupRun(
+				event.payload.jobId,
+				event.payload.scheduleId,
+				event.type,
+				event.agentId,
+			);
 			if (!activeBackupRun) {
-			return;
+				break;
 			}
 
 			resolveActiveBackupRun(activeBackupRun.scheduleId, {
 				status: "cancelled",
-			message: activeBackupRun.cancellationRequested ? undefined : payload.message,
+				message: activeBackupRun.cancellationRequested ? undefined : event.payload.message,
 			});
-	},
+			break;
+		}
+	}
 };
 
 export const startAgentController = async () => {
 	const runtime = getAgentRuntimeState();
 
 	if (runtime.agentManager) {
-		await runtime.agentManager.stop();
+		await Effect.runPromise(runtime.agentManager.stop);
 		runtime.agentManager = null;
 	}
 
-	const nextAgentManager = createAgentManagerRuntime(backupEventHandlers);
+	const nextAgentManager = createAgentManagerRuntime(handleAgentManagerEvent);
-	await nextAgentManager.start();
+	await Effect.runPromise(nextAgentManager.start);
 	runtime.agentManager = nextAgentManager;
 };
 
@@ -181,7 +212,9 @@ export const stopAgentController = async () => {
 	const runtime = getAgentRuntimeState();
 	const agentManagerRuntime = runtime.agentManager;
 	runtime.agentManager = null;
-	await agentManagerRuntime?.stop();
+	if (agentManagerRuntime) {
+		await Effect.runPromise(agentManagerRuntime.stop);
+	}
 };
 
 export const agentManager = {
@@ -212,7 +245,7 @@ export const agentManager = {
 		});
 
 		try {
-			if (!(await runtime.sendBackup(agentId, request.payload))) {
+			if (!(await Effect.runPromise(runtime.sendBackup(agentId, request.payload)))) {
 				clearActiveBackupRun(request.scheduleId);
 				return {
 					status: "unavailable",

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

@@ -1,38 +1,31 @@
 import { Data, Effect, Exit, Fiber, Scope } from "effect";
 import { logger } from "@zerobyte/core/node";
 import { toMessage } from "@zerobyte/core/utils";
-import type {
+import type { AgentMessage, BackupCancelPayload, BackupRunPayload } from "@zerobyte/contracts/agent-protocol";
-	BackupCancelPayload,
+import {
-	BackupCancelledPayload,
+	createControllerAgentSession,
-	BackupCompletedPayload,
+	type AgentConnectionData,
-	BackupFailedPayload,
+	type ControllerAgentSession,
-	BackupProgressPayload,
+	type ControllerAgentSessionEvent,
-	BackupRunPayload,
+} from "./session";
-	BackupStartedPayload,
-} from "@zerobyte/contracts/agent-protocol";
-import { createControllerAgentSession, type AgentConnectionData, type ControllerAgentSession } from "./session";
 import { agentsService } from "../agents.service";
 import { validateAgentToken } from "../helpers/tokens";
 
-type AgentBackupEventContext = {
+type AgentEventContext = {
 	agentId: string;
 	agentName: string;
-	payload:
-		| BackupStartedPayload
-		| BackupProgressPayload
-		| BackupCompletedPayload
-		| BackupFailedPayload
-		| BackupCancelledPayload;
 };
 
-export type AgentBackupEventHandlers = {
+type AgentBackupMessage = Extract<
-	onAgentDisconnected?: (context: { agentId: string; agentName: string }) => void;
+	AgentMessage,
-	onBackupStarted?: (context: AgentBackupEventContext & { payload: BackupStartedPayload }) => void;
+	{
-	onBackupProgress?: (context: AgentBackupEventContext & { payload: BackupProgressPayload }) => void;
+		type: "backup.started" | "backup.progress" | "backup.completed" | "backup.failed" | "backup.cancelled";
-	onBackupCompleted?: (context: AgentBackupEventContext & { payload: BackupCompletedPayload }) => void;
+	}
-	onBackupFailed?: (context: AgentBackupEventContext & { payload: BackupFailedPayload }) => void;
+>;
-	onBackupCancelled?: (context: AgentBackupEventContext & { payload: BackupCancelledPayload }) => void;
+
-};
+export type AgentManagerEvent =
+	| (AgentEventContext & { type: "agent.disconnected" })
+	| (AgentEventContext & AgentBackupMessage);
 
 type ControllerAgentSessionHandle = {
 	agentId: string;
@@ -44,14 +37,18 @@ class StopAgentManagerServerError extends Data.TaggedError("StopAgentManagerServ
 	cause: unknown;
 }> {}
 
-export function createAgentManagerRuntime(handlers: AgentBackupEventHandlers) {
+export function createAgentManagerRuntime(onEvent: (event: AgentManagerEvent) => void) {
 	let sessions = new Map<string, ControllerAgentSessionHandle>();
 	let runtimeScope: Scope.CloseableScope | null = null;
 
 	const closeSession = (sessionHandle: ControllerAgentSessionHandle) =>
 		Effect.gen(function* () {
 			yield* Scope.close(sessionHandle.scope, Exit.succeed(undefined));
-			yield* Effect.sync(() => sessions.delete(sessionHandle.agentId));
+			yield* Effect.sync(() => {
+				if (sessions.get(sessionHandle.agentId) === sessionHandle) {
+					sessions.delete(sessionHandle.agentId);
+				}
+			});
 		});
 
 	const closeAllSessions = Effect.gen(function* () {
@@ -66,40 +63,32 @@ export function createAgentManagerRuntime(handlers: AgentBackupEventHandlers) {
 	const getSessionHandle = (agentId: string) => sessions.get(agentId);
 	const getSession = (agentId: string) => getSessionHandle(agentId)?.session;
 
-	const createSessionHandlers = (params: { agentId: string; agentName: string; sessionId: string }) => {
+	const handleSessionEvent = (params: { agentId: string; agentName: string; sessionId: string }) => {
 		const { agentId, agentName, sessionId } = params;
 
-		return {
+		return (event: ControllerAgentSessionEvent) => {
-			onReady: ({ at }: { at: number }) => {
+			switch (event.type) {
+				case "agent.ready": {
+					const at = Date.now();
 					return Effect.promise(async () => {
 						await agentsService.markAgentOnline(agentId, at);
 					});
-			},
+				}
-			onHeartbeatPong: ({ at }: { at: number }) => {
+				case "heartbeat.pong": {
+					const at = Date.now();
 					return Effect.promise(() => agentsService.markAgentSeen(agentId, at));
-			},
+				}
-			onDisconnect: () => {
+				case "agent.disconnected": {
 					if (getSession(agentId)?.connectionId !== sessionId) {
 						return Effect.void;
 					}
 
-				return Effect.sync(() => handlers.onAgentDisconnected?.({ agentId, agentName }));
+					return Effect.sync(() => onEvent({ type: "agent.disconnected", agentId, agentName }));
-			},
+				}
-			onBackupStarted: (payload: BackupStartedPayload) => {
+				default: {
-				return Effect.sync(() => handlers.onBackupStarted?.({ agentId, agentName, payload }));
+					return Effect.sync(() => onEvent({ ...event, agentId, agentName }));
-			},
+				}
-			onBackupProgress: (payload: BackupProgressPayload) => {
+			}
-				return Effect.sync(() => handlers.onBackupProgress?.({ agentId, agentName, payload }));
-			},
-			onBackupCompleted: (payload: BackupCompletedPayload) => {
-				return Effect.sync(() => handlers.onBackupCompleted?.({ agentId, agentName, payload }));
-			},
-			onBackupFailed: (payload: BackupFailedPayload) => {
-				return Effect.sync(() => handlers.onBackupFailed?.({ agentId, agentName, payload }));
-			},
-			onBackupCancelled: (payload: BackupCancelledPayload) => {
-				return Effect.sync(() => handlers.onBackupCancelled?.({ agentId, agentName, payload }));
-			},
 		};
 	};
 
@@ -110,7 +99,7 @@ export function createAgentManagerRuntime(handlers: AgentBackupEventHandlers) {
 			const session = yield* Scope.extend(
 				createControllerAgentSession(
 					ws,
-					createSessionHandlers({
+					handleSessionEvent({
 						agentId: ws.data.agentId,
 						agentName: ws.data.agentName,
 						sessionId: ws.data.id,
@@ -118,7 +107,7 @@ export function createAgentManagerRuntime(handlers: AgentBackupEventHandlers) {
 				),
 				scope,
 			);
-			const runFiber = yield* Effect.fork(Scope.extend(session.run, scope));
+			const runFiber = yield* Effect.forkDaemon(Scope.extend(session.run, scope));
 			yield* Scope.addFinalizer(scope, Fiber.interrupt(runFiber));
 
 			return { agentId: ws.data.agentId, session, scope };
@@ -163,6 +152,28 @@ export function createAgentManagerRuntime(handlers: AgentBackupEventHandlers) {
 			yield* session.handleMessage(data);
 		});
 
+	const handleOpen = (ws: Bun.ServerWebSocket<AgentConnectionData>) =>
+		Effect.gen(function* () {
+			yield* Effect.promise(() =>
+				agentsService.markAgentConnecting({
+					agentId: ws.data.agentId,
+					organizationId: ws.data.organizationId,
+					agentName: ws.data.agentName,
+					agentKind: ws.data.agentKind,
+				}),
+			);
+
+			const sessionHandle = yield* createSession(ws);
+			yield* setSession(sessionHandle);
+			yield* logger.effect.info(`Agent "${ws.data.agentName}" (${ws.data.agentId}) connected on ${ws.data.id}`);
+		});
+
+	const handleClose = (ws: Bun.ServerWebSocket<AgentConnectionData>) =>
+		Effect.gen(function* () {
+			yield* removeSession(ws.data.agentId, ws.data.id);
+			yield* logger.effect.info(`Agent "${ws.data.agentName}" (${ws.data.agentId}) disconnected`);
+		});
+
 	const acquireServer = Effect.acquireRelease(
 		Effect.sync(() =>
 			Bun.serve<AgentConnectionData>({
@@ -194,24 +205,13 @@ export function createAgentManagerRuntime(handlers: AgentBackupEventHandlers) {
 				},
 				websocket: {
 					open: async (ws) => {
-						await agentsService.markAgentConnecting({
+						await Effect.runPromise(handleOpen(ws));
-							agentId: ws.data.agentId,
-							organizationId: ws.data.organizationId,
-							agentName: ws.data.agentName,
-							agentKind: ws.data.agentKind,
-						});
-
-						const sessionHandle = await Effect.runPromise(createSession(ws));
-						await Effect.runPromise(setSession(sessionHandle));
-
-						logger.info(`Agent "${ws.data.agentName}" (${ws.data.agentId}) connected on ${ws.data.id}`);
 					},
 					message: async (ws, data) => {
 						await Effect.runPromise(handleMessage(ws, data));
 					},
 					close: async (ws) => {
-						await Effect.runPromise(removeSession(ws.data.agentId, ws.data.id));
+						await Effect.runPromise(handleClose(ws));
-						logger.info(`Agent "${ws.data.agentName}" (${ws.data.agentId}) disconnected`);
 					},
 				},
 			}),
@@ -230,7 +230,7 @@ export function createAgentManagerRuntime(handlers: AgentBackupEventHandlers) {
 			),
 	);
 
-	const stop = async () => {
+	const stop = Effect.gen(function* () {
 		if (!runtimeScope) {
 			return;
 		}
@@ -238,31 +238,30 @@ export function createAgentManagerRuntime(handlers: AgentBackupEventHandlers) {
 		logger.info("Stopping Agent Manager...");
 		const scope = runtimeScope;
 		runtimeScope = null;
-		await Effect.runPromise(Scope.close(scope, Exit.succeed(undefined)));
+		yield* Scope.close(scope, Exit.succeed(undefined));
-	};
+	});
 
-	// TODO: Move the effect boundary up
+	const start = Effect.gen(function* () {
-	const start = async () => {
 		if (runtimeScope) {
-			await stop();
+			yield* stop;
 		}
 
 		logger.info("Starting Agent Manager...");
-		const scope = Effect.runSync(Scope.make());
+		const scope = yield* Scope.make();
 
-		try {
+		const server = yield* Scope.extend(acquireServer, scope).pipe(
-			const server = Effect.runSync(Scope.extend(acquireServer, scope));
+			Effect.catchAllCause((cause) =>
+				Scope.close(scope, Exit.failCause(cause)).pipe(Effect.andThen(Effect.failCause(cause))),
+			),
+		);
 		runtimeScope = scope;
 		logger.info(`Agent Manager listening on port ${server.port}`);
-		} catch (error) {
+	});
-			await Effect.runPromise(Scope.close(scope, Exit.fail(error)));
-			throw error;
-		}
-	};
 
 	return {
 		start,
-		sendBackup: async (agentId: string, payload: BackupRunPayload) => {
+		sendBackup: (agentId: string, payload: BackupRunPayload) =>
+			Effect.gen(function* () {
 				const session = getSession(agentId);
 
 				if (!session) {
@@ -270,20 +269,21 @@ export function createAgentManagerRuntime(handlers: AgentBackupEventHandlers) {
 					return false;
 				}
 
-			if (!Effect.runSync(session.isReady())) {
+				if (!(yield* session.isReady())) {
 					logger.warn(`Cannot send backup command. Agent ${agentId} is not ready.`);
 					return false;
 				}
 
-			if (!(await Effect.runPromise(session.sendBackup(payload)))) {
+				if (!(yield* session.sendBackup(payload))) {
 					logger.warn(`Cannot send backup command. Agent ${agentId} is no longer accepting commands.`);
 					return false;
 				}
 
 				logger.info(`Sent backup command ${payload.jobId} to agent ${agentId} for schedule ${payload.scheduleId}`);
 				return true;
-		},
+			}),
-		cancelBackup: async (agentId: string, payload: BackupCancelPayload) => {
+		cancelBackup: (agentId: string, payload: BackupCancelPayload) =>
+			Effect.gen(function* () {
 				const session = getSession(agentId);
 
 				if (!session) {
@@ -291,14 +291,14 @@ export function createAgentManagerRuntime(handlers: AgentBackupEventHandlers) {
 					return false;
 				}
 
-			if (!(await Effect.runPromise(session.sendBackupCancel(payload)))) {
+				if (!(yield* session.sendBackupCancel(payload))) {
 					logger.warn(`Cannot cancel backup command. Agent ${agentId} is no longer accepting commands.`);
 					return false;
 				}
 
 				logger.info(`Sent backup cancel for command ${payload.jobId} to agent ${agentId}`);
 				return true;
-		},
+			}),
 		stop,
 	};
 }

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

@@ -5,12 +5,7 @@ import {
 	parseAgentMessage,
 	type AgentMessage,
 	type BackupCancelPayload,
-	type BackupCancelledPayload,
-	type BackupCompletedPayload,
-	type BackupFailedPayload,
-	type BackupProgressPayload,
 	type BackupRunPayload,
-	type BackupStartedPayload,
 	type ControllerWireMessage,
 } from "@zerobyte/contracts/agent-protocol";
 import { logger } from "@zerobyte/core/node";
@@ -32,24 +27,7 @@ type SessionState = {
 	lastPongAt: number | null;
 };
 
-type AgentRuntimeEventPayload = {
+export type ControllerAgentSessionEvent = AgentMessage | { type: "agent.disconnected" };
-	agentId: string;
-	agentName: string;
-	organizationId: string | null;
-	agentKind: AgentKind;
-	at: number;
-};
-
-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>;
-	onBackupFailed: (payload: BackupFailedPayload) => Effect.Effect<void>;
-	onBackupCancelled: (payload: BackupCancelledPayload) => Effect.Effect<void>;
-};
 
 export type ControllerAgentSession = {
 	readonly connectionId: string;
@@ -62,7 +40,7 @@ export type ControllerAgentSession = {
 
 export const createControllerAgentSession = (
 	socket: AgentSocket,
-	handlers: ControllerAgentSessionHandlers,
+	onEvent: (event: ControllerAgentSessionEvent) => Effect.Effect<void>,
 ): Effect.Effect<ControllerAgentSession, never, Scope.Scope> =>
 	Effect.gen(function* () {
 		let isClosed = false;
@@ -88,13 +66,7 @@ export const createControllerAgentSession = (
 		const releaseSession = Effect.gen(function* () {
 			const disconnectedAt = Date.now();
 			yield* updateState((current) => ({ ...current, isReady: false, lastSeenAt: disconnectedAt }));
-			yield* handlers.onDisconnect({
+			yield* onEvent({ type: "agent.disconnected" });
-				agentId: socket.data.agentId,
-				agentName: socket.data.agentName,
-				organizationId: socket.data.organizationId,
-				agentKind: socket.data.agentKind,
-				at: disconnectedAt,
-			});
 
 			yield* Queue.shutdown(outboundQueue);
 		});
@@ -159,59 +131,18 @@ export const createControllerAgentSession = (
 
 		const handleAgentMessage = (message: AgentMessage) =>
 			Effect.gen(function* () {
-				switch (message.type) {
+				if (message.type === "agent.ready") {
-					case "agent.ready": {
 					const readyAt = Date.now();
 					yield* updateState((current) => ({ ...current, isReady: true, lastSeenAt: readyAt }));
-
-						yield* handlers.onReady({
-							agentId: socket.data.agentId,
-							agentName: socket.data.agentName,
-							organizationId: socket.data.organizationId,
-							agentKind: socket.data.agentKind,
-							at: readyAt,
-						});
-
 					yield* logger.effect.info(`Agent "${socket.data.agentName}" (${socket.data.agentId}) is ready`);
-						break;
 				}
-					case "backup.started": {
+
-						yield* logger.effect.info(
+				if (message.type === "heartbeat.pong") {
-							`Backup ${message.payload.jobId} started on agent ${socket.data.agentId} for schedule ${message.payload.scheduleId}`,
-						);
-						yield* handlers.onBackupStarted(message.payload);
-						break;
-					}
-					case "backup.progress": {
-						yield* handlers.onBackupProgress(message.payload);
-						break;
-					}
-					case "backup.completed": {
-						yield* handlers.onBackupCompleted(message.payload);
-						break;
-					}
-					case "backup.failed": {
-						yield* handlers.onBackupFailed(message.payload);
-						break;
-					}
-					case "backup.cancelled": {
-						yield* handlers.onBackupCancelled(message.payload);
-						break;
-					}
-					case "heartbeat.pong": {
 					const seenAt = Date.now();
 					yield* updateState((current) => ({ ...current, lastSeenAt: seenAt, lastPongAt: message.payload.sentAt }));
+				}
 
-						yield* handlers.onHeartbeatPong({
+				yield* onEvent(message);
-							agentId: socket.data.agentId,
-							agentName: socket.data.agentName,
-							organizationId: socket.data.organizationId,
-							agentKind: socket.data.agentKind,
-							at: seenAt,
-						});
-						break;
-					}
-				}
 			});
 
 		return {