fix(agent): harden disconnect and send-failure handling

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

@@ -0,0 +1,87 @@
+import { EventEmitter } from "node:events";
+import type * as childProcess from "node:child_process";
+import { PassThrough } from "node:stream";
+import { afterEach, expect, test, vi } from "vitest";
+
+const spawnMock = vi.fn();
+
+vi.mock("node:child_process", async () => {
+	const actual = await vi.importActual<typeof import("node:child_process")>("node:child_process");
+	return {
+		...actual,
+		spawn: spawnMock,
+	};
+});
+
+const { spawnLocalAgent, stopLocalAgent } = await import("../agents-manager");
+
+const flushMicrotasks = async () => {
+	await Promise.resolve();
+};
+
+type FakeChildProcess = EventEmitter & {
+	stdout: PassThrough;
+	stderr: PassThrough;
+	exitCode: number | null;
+	signalCode: NodeJS.Signals | null;
+	kill: ReturnType<typeof vi.fn>;
+};
+
+const createFakeChild = () => {
+	const child = new EventEmitter() as FakeChildProcess;
+
+	child.stdout = new PassThrough();
+	child.stderr = new PassThrough();
+	child.exitCode = null;
+	child.signalCode = null;
+	child.kill = vi.fn(() => {
+		child.exitCode = 0;
+		child.emit("exit", 0, null);
+		return true;
+	});
+
+	return child;
+};
+
+afterEach(async () => {
+	await stopLocalAgent();
+	spawnMock.mockReset();
+	vi.restoreAllMocks();
+	vi.useRealTimers();
+});
+
+test("respawns the local agent after an unexpected exit", async () => {
+	vi.useFakeTimers();
+
+	const firstChild = createFakeChild();
+	const secondChild = createFakeChild();
+	spawnMock
+		.mockReturnValueOnce(firstChild as unknown as ReturnType<typeof childProcess.spawn>)
+		.mockReturnValueOnce(secondChild as unknown as ReturnType<typeof childProcess.spawn>);
+
+	await spawnLocalAgent();
+
+	firstChild.exitCode = 1;
+	firstChild.emit("exit", 1, null);
+
+	vi.advanceTimersByTime(1_000);
+	await flushMicrotasks();
+
+	expect(spawnMock).toHaveBeenCalledTimes(2);
+});
+
+test("does not respawn the local agent after an intentional stop", async () => {
+	vi.useFakeTimers();
+
+	const child = createFakeChild();
+	spawnMock.mockReturnValue(child as unknown as ReturnType<typeof childProcess.spawn>);
+
+	await spawnLocalAgent();
+	await stopLocalAgent();
+
+	vi.advanceTimersByTime(1_000);
+	await flushMicrotasks();
+
+	expect(spawnMock).toHaveBeenCalledTimes(1);
+	expect(child.kill).toHaveBeenCalledTimes(1);
+});

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

@@ -1,121 +0,0 @@
-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);
-	}
-});
-
-test("close emits a synthetic backup.cancelled for a queued backup", () => {
-	const onBackupCancelled = mock(() => undefined);
-	const session = createControllerAgentSession(createSocket(), {
-		onBackupCancelled,
-	});
-
-	session.sendBackup({
-		jobId: "job-queued",
-		scheduleId: "schedule-queued",
-		organizationId: "org-1",
-		sourcePath: "/tmp/source",
-		repositoryConfig: {
-			backend: "local",
-			path: "/tmp/repository",
-		},
-		options: {},
-		runtime: {
-			password: "password",
-			cacheDir: "/tmp/cache",
-			passFile: "/tmp/pass",
-			defaultExcludes: [],
-		},
-	});
-
-	session.close();
-
-	expect(onBackupCancelled).toHaveBeenCalledTimes(1);
-	expect(onBackupCancelled).toHaveBeenCalledWith({
-		jobId: "job-queued",
-		scheduleId: "schedule-queued",
-		message:
-			"The connection to the backup agent was lost before this backup started. Restart the backup to ensure it completes.",
-	});
-});

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

@@ -1,24 +1,36 @@
 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 { createControllerAgentSession } from "../controller/session";
 
-const createSocket = () => {
+const createSocket = (
+	overrides: Partial<Parameters<typeof createControllerAgentSession>[0]> = {},
+) => {
 	return fromPartial<Parameters<typeof createControllerAgentSession>[0]>({
 		data: { id: "connection-1", agentId: "local", organizationId: null, agentName: "Local Agent" },
-		send: vi.fn(),
+		send: vi.fn(() => 1),
+		close: vi.fn(),
+		...overrides,
 	});
 };
 
-const createSession = (handlers: Parameters<typeof createControllerAgentSession>[1] = {}) => {
+const createSession = (
+	handlers: Parameters<typeof createControllerAgentSession>[1] = {},
+	socket = createSocket(),
+) => {
 	const scope = Effect.runSync(Scope.make());
 
 	try {
-		const session = Effect.runSync(Scope.extend(createControllerAgentSession(createSocket(), handlers), scope));
+		const session = Effect.runSync(Scope.extend(createControllerAgentSession(socket, handlers), scope));
 
 		return {
 			session,
+			run: () => {
+				Effect.runSync(Scope.extend(Effect.forkScoped(session.run), scope));
+			},
+			socket,
 			close: () => {
 				Effect.runSync(Scope.close(scope, Exit.succeed(undefined)));
 			},
@@ -50,7 +62,7 @@ test("close emits a synthetic backup.cancelled for a started backup", () => {
 	expect(onBackupCancelled).toHaveBeenCalledWith({
 		jobId: "job-1",
 		scheduleId: "schedule-1",
-		message: "The connection to the backup agent was lost. Restart the backup to ensure it completes.",
+		message: "The connection to the backup agent was lost while this backup was running. Restart the backup to ensure it completes.",
 	});
 });
 
@@ -142,6 +154,45 @@ test("close emits a synthetic backup.cancelled for a queued backup", () => {
 	expect(onBackupCancelled).toHaveBeenCalledWith({
 		jobId: "job-queued",
 		scheduleId: "schedule-queued",
-		message: "The connection to the backup agent was lost. Restart the backup to ensure it completes.",
+		message: "The connection to the backup agent was lost before this backup started. Restart the backup to ensure it completes.",
 	});
 });
+
+test("a dropped backup.cancel closes the session and emits a synthetic backup.cancelled", async () => {
+	const send = vi.fn(() => 0);
+	const socket = createSocket({ send, close: vi.fn() });
+	const onBackupCancelled = vi.fn();
+	const { session, run, close: closeSession } = createSession({ onBackupCancelled }, socket);
+
+	try {
+		run();
+		Effect.runSync(
+			session.handleMessage(
+				createAgentMessage("backup.started", {
+					jobId: "job-1",
+					scheduleId: "schedule-1",
+				}),
+			),
+		);
+		Effect.runSync(
+			session.sendBackupCancel({
+				jobId: "job-1",
+				scheduleId: "schedule-1",
+			}),
+		);
+
+		await waitForExpect(() => {
+			expect(send).toHaveBeenCalledTimes(1);
+			expect(socket.close).toHaveBeenCalledTimes(1);
+			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.",
+			});
+		});
+	} finally {
+		closeSession();
+	}
+});

app/server/modules/agents/agent-tokens.ts

@@ -1,12 +0,0 @@
-import { cryptoUtils } from "~/server/utils/crypto";
-
-export const deriveLocalAgentToken = async () => {
-	return cryptoUtils.deriveSecret("zerobyte:local-agent-token");
-};
-
-export const validateAgentToken = async (token: string) => {
-	const localToken = await deriveLocalAgentToken();
-	if (token === localToken) {
-		return { agentId: "local", organizationId: null, agentName: "local" };
-	}
-};

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

@@ -8,6 +8,8 @@ export type { AgentBackupEventHandlers } from "./controller/server";
 type AgentRuntimeState = {
 	agentManager: AgentManagerRuntime | null;
 	localAgent: ChildProcess | null;
+	isStoppingLocalAgent: boolean;
+	localAgentRestartTimeout: ReturnType<typeof setTimeout> | null;
 };
 
 type ProcessWithAgentRuntime = NodeJS.Process & {
@@ -25,6 +27,8 @@ const getAgentRuntimeState = () => {
 	const runtime = {
 		agentManager: null,
 		localAgent: null,
+		isStoppingLocalAgent: false,
+		localAgentRestartTimeout: null,
 	};
 
 	runtimeProcess.__zerobyteAgentRuntime = runtime;
@@ -43,11 +47,11 @@ export const startAgentRuntime = async () => {
 	}
 
 	const { createAgentManagerRuntime } = await import("./controller/server");
-	const agentManager = createAgentManagerRuntime();
+	const nextAgentManager = createAgentManagerRuntime();
-	agentManager.setBackupEventHandlers(backupEventHandlers);
+	nextAgentManager.setBackupEventHandlers(backupEventHandlers);
 
-	await agentManager.start();
+	await nextAgentManager.start();
-	runtime.agentManager = agentManager;
+	runtime.agentManager = nextAgentManager;
 };
 
 export const agentManager = {

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

@@ -1,247 +0,0 @@
-import { Effect, Fiber, Queue, Ref } from "effect";
-import {
-	createControllerMessage,
-	parseAgentMessage,
-	type AgentMessage,
-	type BackupCancelledPayload,
-	type BackupCompletedPayload,
-	type BackupFailedPayload,
-	type BackupProgressPayload,
-	type BackupRunPayload,
-	type BackupCancelPayload,
-	type BackupStartedPayload,
-	type ControllerWireMessage,
-} from "@zerobyte/contracts/agent-protocol";
-import { logger } from "@zerobyte/core/node";
-import { toMessage } from "@zerobyte/core/utils";
-
-export type AgentConnectionData = {
-	id: string;
-	agentId: string;
-	organizationId: string | null;
-	agentName: string;
-};
-
-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;
-	onBackupCompleted?: (payload: BackupCompletedPayload) => void;
-	onBackupFailed?: (payload: BackupFailedPayload) => void;
-	onBackupCancelled?: (payload: BackupCancelledPayload) => void;
-};
-
-export type ControllerAgentSession = {
-	readonly connectionId: string;
-	handleMessage: (data: string) => void;
-	sendBackup: (payload: BackupRunPayload) => string;
-	sendBackupCancel: (payload: BackupCancelPayload) => void;
-	isReady: () => boolean;
-	close: () => void;
-};
-
-export const createControllerAgentSession = (
-	socket: AgentSocket,
-	handlers: ControllerAgentSessionHandlers = {},
-): ControllerAgentSession => {
-	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(
-		Ref.make<SessionState>({
-			isReady: false,
-			lastSeenAt: null,
-			lastPongAt: null,
-		}),
-	);
-
-	const offerOutbound = (message: ControllerWireMessage) => {
-		void Effect.runPromise(Queue.offer(outboundQueue, message)).catch((error) => {
-			logger.error(`Failed to queue outbound message for agent ${socket.data.agentId}: ${toMessage(error)}`);
-		});
-	};
-
-	const updateState = (update: (current: SessionState) => SessionState) => {
-		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 setPendingBackupJob = (jobId: string, scheduleId: string) => {
-		Effect.runSync(
-			Ref.update(pendingBackupJobs, (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 deletePendingBackupJob = (jobId: string) => {
-		Effect.runSync(
-			Ref.update(pendingBackupJobs, (current) => {
-				const next = new Map(current);
-				next.delete(jobId);
-				return next;
-			}),
-		);
-	};
-
-	const writerFiber = Effect.runFork(
-		Effect.forever(
-			Effect.gen(function* () {
-				const message = yield* Queue.take(outboundQueue);
-				yield* Effect.sync(() => {
-					try {
-						socket.send(message);
-					} catch (error) {
-						logger.error(
-							`Failed to send message to agent ${socket.data.agentId} on ${socket.data.id}: ${toMessage(error)}`,
-						);
-					}
-				});
-			}),
-		),
-	);
-
-	const heartbeatFiber = Effect.runFork(
-		Effect.forever(
-			Effect.gen(function* () {
-				yield* Effect.sleep("15 seconds");
-				yield* Queue.offer(
-					outboundQueue,
-					createControllerMessage("heartbeat.ping", {
-						sentAt: Date.now(),
-					}),
-				);
-			}),
-		),
-	);
-
-	const handleAgentMessage = (message: AgentMessage) => {
-		updateState((current) => ({ ...current, lastSeenAt: Date.now() }));
-
-		switch (message.type) {
-			case "agent.ready": {
-				updateState((current) => ({ ...current, isReady: true }));
-				logger.info(`Agent "${socket.data.agentName}" (${socket.data.agentId}) is ready`);
-				break;
-			}
-			case "backup.started": {
-				deletePendingBackupJob(message.payload.jobId);
-				setActiveBackupJob(message.payload.jobId, message.payload.scheduleId);
-				logger.info(
-					`Backup ${message.payload.jobId} started on agent ${socket.data.agentId} for schedule ${message.payload.scheduleId}`,
-				);
-				handlers.onBackupStarted?.(message.payload);
-				break;
-			}
-			case "backup.progress": {
-				handlers.onBackupProgress?.(message.payload);
-				break;
-			}
-			case "backup.completed": {
-				deletePendingBackupJob(message.payload.jobId);
-				deleteActiveBackupJob(message.payload.jobId);
-				handlers.onBackupCompleted?.(message.payload);
-				break;
-			}
-			case "backup.failed": {
-				deletePendingBackupJob(message.payload.jobId);
-				deleteActiveBackupJob(message.payload.jobId);
-				handlers.onBackupFailed?.(message.payload);
-				break;
-			}
-			case "backup.cancelled": {
-				deletePendingBackupJob(message.payload.jobId);
-				deleteActiveBackupJob(message.payload.jobId);
-				handlers.onBackupCancelled?.(message.payload);
-				break;
-			}
-			case "heartbeat.pong": {
-				updateState((current) => ({ ...current, lastPongAt: message.payload.sentAt }));
-				break;
-			}
-		}
-	};
-
-	return {
-		connectionId: socket.data.id,
-		handleMessage: (data: string) => {
-			const parsed = parseAgentMessage(data);
-
-			if (parsed === null) {
-				logger.warn(`Invalid JSON from agent ${socket.data.agentId}`);
-				return;
-			}
-
-			if (!parsed.success) {
-				logger.warn(`Invalid agent message from ${socket.data.agentId}: ${parsed.error.message}`);
-				return;
-			}
-
-			handleAgentMessage(parsed.data);
-		},
-		sendBackup: (payload) => {
-			setPendingBackupJob(payload.jobId, payload.scheduleId);
-			offerOutbound(createControllerMessage("backup.run", payload));
-			return payload.jobId;
-		},
-		sendBackupCancel: (payload) => {
-			offerOutbound(createControllerMessage("backup.cancel", payload));
-		},
-		isReady: () => Effect.runSync(Ref.get(state)).isReady,
-		close: () => {
-			updateState((current) => ({ ...current, isReady: false }));
-			const pendingJobs = Effect.runSync(Ref.get(pendingBackupJobs));
-			Effect.runSync(Ref.set(pendingBackupJobs, new Map()));
-			for (const [jobId, scheduleId] of 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) {
-				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(() => {});
-		},
-	};
-};

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

@@ -102,7 +102,10 @@ export const createControllerAgentSession = (
 			yield* updateState((current) => ({ ...current, isReady: false }));
 			const trackedJobs = yield* takeTrackedBackupJobs;
 			for (const [jobId, trackedJob] of trackedJobs) {
-				let message = "The connection to the backup agent was lost. Restart the backup to ensure it completes.";
+				const message =
+					trackedJob.state === "pending"
+						? "The connection to the backup agent was lost before this backup started. Restart the backup to ensure it completes."
+						: "The connection to the backup agent was lost while this backup was running. Restart the backup to ensure it completes.";
 
 				yield* Effect.sync(() => {
 					handlers.onBackupCancelled?.({ jobId, scheduleId: trackedJob.scheduleId, message });
@@ -114,21 +117,34 @@ export const createControllerAgentSession = (
 
 		yield* Effect.addFinalizer(() => releaseSession);
 
+		const handleSendFailure = (reason: string) => {
+			logger.error(
+				`Closing session for agent ${socket.data.agentId} on ${socket.data.id} after an outbound websocket send failed: ${reason}`,
+			);
+
+			try {
+				socket.close();
+			} catch (error) {
+				logger.error(`Failed to close socket for agent ${socket.data.agentId} on ${socket.data.id}: ${toMessage(error)}`);
+			}
+		};
+
 		const run = Effect.gen(function* () {
 			yield* Effect.forkScoped(
 				Effect.forever(
-					Effect.gen(function* () {
+				Effect.gen(function* () {
-						const message = yield* Queue.take(outboundQueue);
+					const message = yield* Queue.take(outboundQueue);
-						yield* Effect.sync(() => {
+					yield* Effect.sync(() => {
-							try {
+						try {
-								socket.send(message);
+							const sendResult = socket.send(message);
-							} catch (error) {
+							if (sendResult <= 0) {
-								logger.error(
+								handleSendFailure(sendResult === 0 ? "connection issue" : "backpressure");
-									`Failed to send message to agent ${socket.data.agentId} on ${socket.data.id}: ${toMessage(error)}`,
-								);
 							}
-						});
+						} catch (error) {
-					}),
+							handleSendFailure(toMessage(error));
+						}
+					});
+				}),
 				),
 			);
 

app/server/modules/agents/local/process.ts

@@ -7,6 +7,8 @@ import { deriveLocalAgentToken } from "../helpers/tokens";
 
 type LocalAgentState = {
 	localAgent: ChildProcess | null;
+	isStoppingLocalAgent: boolean;
+	localAgentRestartTimeout: ReturnType<typeof setTimeout> | null;
 };
 
 export async function spawnLocalAgentProcess(runtime: LocalAgentState) {
@@ -48,27 +50,47 @@ export async function spawnLocalAgentProcess(runtime: LocalAgentState) {
 	});
 
 	agentProcess.on("exit", (code, signal) => {
+		const shouldRestart = runtime.localAgent === agentProcess && !runtime.isStoppingLocalAgent;
 		if (runtime.localAgent === agentProcess) {
 			runtime.localAgent = null;
 		}
 		logger.info(`Agent process exited with code ${code} and signal ${signal}`);
+
+		if (!shouldRestart) {
+			return;
+		}
+
+		runtime.localAgentRestartTimeout = setTimeout(() => {
+			runtime.localAgentRestartTimeout = null;
+			void spawnLocalAgentProcess(runtime).catch((error) => {
+				logger.error(`Failed to restart local agent: ${error instanceof Error ? error.message : String(error)}`);
+			});
+		}, 1_000);
 	});
 }
 
 export async function stopLocalAgentProcess(runtime: LocalAgentState) {
+	if (runtime.localAgentRestartTimeout) {
+		clearTimeout(runtime.localAgentRestartTimeout);
+		runtime.localAgentRestartTimeout = null;
+	}
+
 	if (!runtime.localAgent) {
 		return;
 	}
 
 	const agentProcess = runtime.localAgent;
 	runtime.localAgent = null;
+	runtime.isStoppingLocalAgent = true;
 
 	if (agentProcess.exitCode !== null || agentProcess.signalCode !== null) {
+		runtime.isStoppingLocalAgent = false;
 		return;
 	}
 
 	const exited = new Promise<void>((resolve) => {
 		agentProcess.once("exit", () => {
+			runtime.isStoppingLocalAgent = false;
 			resolve();
 		});
 	});

apps/agent/src/__tests__/controller-session.test.ts

@@ -70,3 +70,25 @@ test("emits backup.failed when a backup command hits a restic error", async () =
 		session.close();
 	}
 });
+
+test("closes the websocket when an outbound send throws", async () => {
+	const close = vi.fn(() => undefined);
+	const session = createControllerSession(
+		fromPartial({
+			send: () => {
+				throw new Error("socket write failed");
+			},
+			close,
+		}),
+	);
+
+	try {
+		session.onOpen();
+
+		await waitForExpect(() => {
+			expect(close).toHaveBeenCalledTimes(1);
+		});
+	} finally {
+		session.close();
+	}
+});

apps/agent/src/controller-session.ts

@@ -17,6 +17,7 @@ export type ControllerSession = {
 };
 
 export const createControllerSession = (ws: WebSocket): ControllerSession => {
+	let isClosed = false;
 	const outboundQueue = Effect.runSync(Queue.bounded<AgentWireMessage>(64));
 	const inboundQueue = Effect.runSync(Queue.bounded<ControllerWireMessage>(64));
 	const runningJobsRef = Effect.runSync(Ref.make<Map<string, RunningJob>>(new Map()));
@@ -63,6 +64,31 @@ export const createControllerSession = (ws: WebSocket): ControllerSession => {
 		offerOutbound,
 	};
 
+	const closeSession = () => {
+		if (isClosed) {
+			return;
+		}
+
+		isClosed = true;
+		void Effect.runPromise(abortRunningJobs).catch(() => {});
+		void Effect.runPromise(Fiber.interrupt(writerFiber)).catch(() => {});
+		void Effect.runPromise(Fiber.interrupt(processorFiber)).catch(() => {});
+		void Effect.runPromise(Queue.shutdown(outboundQueue)).catch(() => {});
+		void Effect.runPromise(Queue.shutdown(inboundQueue)).catch(() => {});
+	};
+
+	const handleSendFailure = (reason: string) => {
+		logger.error(`Closing agent session after an outbound websocket send failed: ${reason}`);
+
+		try {
+			ws.close();
+		} catch (error) {
+			logger.error(`Failed to close controller websocket after send failure: ${toMessage(error)}`);
+		}
+
+		closeSession();
+	};
+
 	const writerFiber = Effect.runFork(
 		Effect.forever(
 			Effect.gen(function* () {
@@ -71,7 +97,7 @@ export const createControllerSession = (ws: WebSocket): ControllerSession => {
 					try {
 						ws.send(message);
 					} catch (error) {
-						logger.error(`Failed to send controller message: ${toMessage(error)}`);
+						handleSendFailure(toMessage(error));
 					}
 				});
 			}),
@@ -106,16 +132,15 @@ export const createControllerSession = (ws: WebSocket): ControllerSession => {
 			});
 		},
 		onMessage: (data) => {
+			if (typeof data !== "string") {
+				logger.warn("Agent received a non-text message");
+				return;
+			}
+
 			void Effect.runPromise(offerInbound(data as ControllerWireMessage)).catch((error) => {
 				logger.error(`Failed to queue inbound message: ${toMessage(error)}`);
 			});
 		},
-		close: () => {
+		close: closeSession,
-			void Effect.runPromise(abortRunningJobs).catch(() => {});
-			void Effect.runPromise(Fiber.interrupt(writerFiber)).catch(() => {});
-			void Effect.runPromise(Fiber.interrupt(processorFiber)).catch(() => {});
-			void Effect.runPromise(Queue.shutdown(outboundQueue)).catch(() => {});
-			void Effect.runPromise(Queue.shutdown(inboundQueue)).catch(() => {});
-		},
 	};
 };

apps/agent/src/index.ts

@@ -5,7 +5,7 @@ const controllerUrl = process.env.ZEROBYTE_CONTROLLER_URL;
 const agentToken = process.env.ZEROBYTE_AGENT_TOKEN;
 const RECONNECT_DELAY_MS = 1000;
 
-class Agent {
+export class Agent {
 	private ws: WebSocket | null = null;
 	private controllerSession: ControllerSession | null = null;
 	private reconnectTimeout: ReturnType<typeof setTimeout> | null = null;
@@ -22,6 +22,10 @@ class Agent {
 	}
 
 	connect() {
+		if (this.reconnectTimeout) {
+			clearTimeout(this.reconnectTimeout);
+			this.reconnectTimeout = null;
+		}
 		if (this.ws) {
 			return;
 		}