From 1fc48a56860d60396c4864fe1af06b759f1e898c Mon Sep 17 00:00:00 2001
From: Nicolas Meienberger <github@thisprops.com>
Date: Tue, 5 May 2026 19:23:54 +0200
Subject: [PATCH] chore: pr feedback

---
 .../__tests__/controller-runtime.test.ts      | 44 ++++++++++++++++++-
 .../modules/agents/controller/server.ts       |  9 ++--
 2 files changed, 46 insertions(+), 7 deletions(-)

diff --git a/app/server/modules/agents/__tests__/controller-runtime.test.ts b/app/server/modules/agents/__tests__/controller-runtime.test.ts
index a039bba3..66523716 100644
--- a/app/server/modules/agents/__tests__/controller-runtime.test.ts
+++ b/app/server/modules/agents/__tests__/controller-runtime.test.ts
@@ -162,6 +162,21 @@ test("websocket lifecycle updates agent connection status", async () => {
 	expect(stop).toHaveBeenCalledWith(true);
 });
 
+test("websocket open failure closes the upgraded socket", async () => {
+	agentsServiceMocks.markAgentConnecting.mockRejectedValueOnce(new Error("db unavailable"));
+	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");
+
+	await websocket?.open?.(fromPartial(socket));
+
+	expect(socket.close).toHaveBeenCalled();
+	await Effect.runPromise(runtime.stop);
+});
+
 test("shutdown closes all sessions and stops the server when marking one agent offline fails", async () => {
 	agentsServiceMocks.markAgentOffline.mockRejectedValueOnce(new Error("db unavailable"));
 	const stop = vi.fn(() => Promise.resolve());
@@ -182,7 +197,7 @@ test("shutdown closes all sessions and stops the server when marking one agent o
 	expect(stop).toHaveBeenCalledWith(true);
 });
 
-test("closing a replaced connection does not report the active agent as disconnected", async () => {
+test("closing a replaced connection reports disconnect without marking the active agent offline", async () => {
 	const serve = vi
 		.spyOn(Bun, "serve")
 		.mockReturnValue(fromPartial({ port: 3001, stop: vi.fn(() => Promise.resolve()) }));
@@ -190,14 +205,39 @@ test("closing a replaced connection does not report the active agent as disconne
 	const websocket = serve.mock.calls[0]?.[0].websocket;
 	const oldSocket = createSocket("connection-1");
 	const newSocket = createSocket("connection-2");
+	const offlineCallsBeforeClose = agentsServiceMocks.markAgentOffline.mock.calls.length;
 
 	await websocket?.open?.(fromPartial(oldSocket));
 	await websocket?.open?.(fromPartial(newSocket));
+	await websocket?.message?.(fromPartial(newSocket), createAgentMessage("agent.ready", { agentId: LOCAL_AGENT_ID }));
 	await websocket?.close?.(fromPartial(oldSocket), 1000, "replaced");
 
-	expect(onEvent).not.toHaveBeenCalledWith(
+	expect(onEvent).toHaveBeenCalledWith(
 		expect.objectContaining({ type: "agent.disconnected", agentId: LOCAL_AGENT_ID }),
 	);
+	expect(agentsServiceMocks.markAgentOffline).toHaveBeenCalledTimes(offlineCallsBeforeClose);
+	expect(
+		await Effect.runPromise(
+			runtime.sendBackup(LOCAL_AGENT_ID, {
+				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: [],
+				webhookTimeoutMs: 60_000,
+			}),
+		),
+	).toBe(true);
 	await Effect.runPromise(runtime.stop);
 });
 
diff --git a/app/server/modules/agents/controller/server.ts b/app/server/modules/agents/controller/server.ts
index 8e647798..d5201fa4 100644
--- a/app/server/modules/agents/controller/server.ts
+++ b/app/server/modules/agents/controller/server.ts
@@ -74,7 +74,7 @@ export function createAgentManagerRuntime(onEvent: (event: AgentManagerEvent) =>
 	const getSession = (agentId: string) => getSessionHandle(agentId)?.session;
 
 	const handleSessionEvent = (params: { agentId: string; agentName: string; sessionId: string }) => {
-		const { agentId, agentName, sessionId } = params;
+		const { agentId, agentName } = params;
 
 		return (event: ControllerAgentSessionEvent) => {
 			switch (event.type) {
@@ -89,10 +89,6 @@ export function createAgentManagerRuntime(onEvent: (event: AgentManagerEvent) =>
 					return Effect.promise(() => agentsService.markAgentSeen(agentId, at));
 				}
 				case "agent.disconnected": {
-					if (getSession(agentId)?.connectionId !== sessionId) {
-						return Effect.void;
-					}
-
 					return Effect.sync(() => onEvent({ type: "agent.disconnected", agentId, agentName }));
 				}
 				default: {
@@ -231,6 +227,9 @@ export function createAgentManagerRuntime(onEvent: (event: AgentManagerEvent) =>
 				websocket: {
 					open: async (ws) => {
 						await runWebSocketHandler(ws, "open", handleOpen(ws));
+						if (getSession(ws.data.agentId)?.connectionId !== ws.data.id) {
+							ws.close();
+						}
 					},
 					message: async (ws, data) => {
 						await runWebSocketHandler(ws, "message", handleMessage(ws, data));