From 9fba2d083c7325550b8ce714b1a01ffa5f0c6be6 Mon Sep 17 00:00:00 2001
From: Nicolas Meienberger <github@thisprops.com>
Date: Tue, 7 Apr 2026 18:13:48 +0200
Subject: [PATCH] refactor: correctly close session on send failure

---
 apps/agent/package.json                       |  3 ++
 .../src/__tests__/controller-session.test.ts  | 22 ++++++++++++
 apps/agent/src/controller-session.ts          | 36 ++++++++++++++-----
 3 files changed, 53 insertions(+), 8 deletions(-)

diff --git a/apps/agent/package.json b/apps/agent/package.json
index 3751fd6a..40361ac1 100644
--- a/apps/agent/package.json
+++ b/apps/agent/package.json
@@ -3,6 +3,9 @@
 	"private": true,
 	"type": "module",
 	"module": "index.ts",
+	"scripts": {
+		"tsc": "tsc --noEmit"
+	},
 	"dependencies": {
 		"@zerobyte/contracts": "workspace:*",
 		"@zerobyte/core": "workspace:*",
diff --git a/apps/agent/src/__tests__/controller-session.test.ts b/apps/agent/src/__tests__/controller-session.test.ts
index 87d1fdbc..9c79d3b5 100644
--- a/apps/agent/src/__tests__/controller-session.test.ts
+++ b/apps/agent/src/__tests__/controller-session.test.ts
@@ -69,3 +69,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 = mock(() => 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();
+	}
+});
diff --git a/apps/agent/src/controller-session.ts b/apps/agent/src/controller-session.ts
index c1b45b9d..d8f82e24 100644
--- a/apps/agent/src/controller-session.ts
+++ b/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));
 					}
 				});
 			}),
@@ -115,12 +141,6 @@ export const createControllerSession = (ws: WebSocket): ControllerSession => {
 				logger.error(`Failed to queue inbound message: ${toMessage(error)}`);
 			});
 		},
-		close: () => {
-			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(() => {});
-		},
+		close: closeSession,
 	};
 };