refactor: improve error handling

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

@@ -24,12 +24,12 @@ vi.mock("../helpers/tokens", () => ({
 	validateAgentToken: tokenMocks.validateAgentToken,
 }));
 
-const createSocket = (id: string) => ({
+const createSocket = (id: string, agentId = LOCAL_AGENT_ID) => ({
 	data: {
 		id,
-		agentId: LOCAL_AGENT_ID,
+		agentId,
 		organizationId: null,
-		agentName: LOCAL_AGENT_NAME,
+		agentName: agentId === LOCAL_AGENT_ID ? LOCAL_AGENT_NAME : `${LOCAL_AGENT_NAME} ${agentId}`,
 		agentKind: LOCAL_AGENT_KIND,
 	},
 	send: vi.fn(() => 1),
@@ -162,6 +162,26 @@ test("websocket lifecycle updates agent connection status", async () => {
 	expect(stop).toHaveBeenCalledWith(true);
 });
 
+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());
+	const serve = vi.spyOn(Bun, "serve").mockReturnValue(fromPartial({ port: 3001, stop }));
+	const { runtime, onEvent } = await startRuntime(vi.fn());
+	const websocket = serve.mock.calls[0]?.[0].websocket;
+	const firstSocket = createSocket("connection-1", "agent-1");
+	const secondSocket = createSocket("connection-2", "agent-2");
+
+	await websocket?.open?.(fromPartial(firstSocket));
+	await websocket?.open?.(fromPartial(secondSocket));
+	await Effect.runPromise(runtime.stop);
+
+	expect(agentsServiceMocks.markAgentOffline).toHaveBeenCalledWith("agent-1");
+	expect(agentsServiceMocks.markAgentOffline).toHaveBeenCalledWith("agent-2");
+	expect(onEvent).toHaveBeenCalledWith(expect.objectContaining({ type: "agent.disconnected", agentId: "agent-1" }));
+	expect(onEvent).toHaveBeenCalledWith(expect.objectContaining({ type: "agent.disconnected", agentId: "agent-2" }));
+	expect(stop).toHaveBeenCalledWith(true);
+});
+
 test("closing a replaced connection does not report the active agent as disconnected", async () => {
 	const serve = vi
 		.spyOn(Bun, "serve")

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

@@ -51,10 +51,20 @@ export function createAgentManagerRuntime(onEvent: (event: AgentManagerEvent) =>
 			});
 		});
 
+	const markAgentOfflineForShutdown = (agentId: string) =>
+		Effect.tryPromise({
+			try: () => agentsService.markAgentOffline(agentId),
+			catch: (error) => new StopAgentManagerServerError({ cause: error }),
+		}).pipe(
+			Effect.catchAll((error) =>
+				logger.effect.error(`Failed to mark agent ${agentId} offline during shutdown: ${toMessage(error)}`),
+			),
+		);
+
 	const closeAllSessions = Effect.gen(function* () {
 		const currentSessions = [...sessions.entries()];
 		for (const [agentId, sessionHandle] of currentSessions) {
-			yield* Effect.promise(() => agentsService.markAgentOffline(agentId));
+			yield* markAgentOfflineForShutdown(agentId);
 			yield* closeSession(sessionHandle);
 		}
 		sessions = new Map();
@@ -174,6 +184,21 @@ export function createAgentManagerRuntime(onEvent: (event: AgentManagerEvent) =>
 			yield* logger.effect.info(`Agent "${ws.data.agentName}" (${ws.data.agentId}) disconnected`);
 		});
 
+	const runWebSocketHandler = (
+		ws: Bun.ServerWebSocket<AgentConnectionData>,
+		event: string,
+		effect: Effect.Effect<void>,
+	) =>
+		Effect.runPromise(
+			effect.pipe(
+				Effect.catchAllCause((cause) =>
+					logger.effect.error(
+						`Agent websocket ${event} failed for ${ws.data.agentId} on ${ws.data.id}: ${toMessage(cause)}`,
+					),
+				),
+			),
+		);
+
 	const acquireServer = Effect.acquireRelease(
 		Effect.sync(() =>
 			Bun.serve<AgentConnectionData>({
@@ -205,13 +230,13 @@ export function createAgentManagerRuntime(onEvent: (event: AgentManagerEvent) =>
 				},
 				websocket: {
 					open: async (ws) => {
-						await Effect.runPromise(handleOpen(ws));
+						await runWebSocketHandler(ws, "open", handleOpen(ws));
 					},
 					message: async (ws, data) => {
-						await Effect.runPromise(handleMessage(ws, data));
+						await runWebSocketHandler(ws, "message", handleMessage(ws, data));
 					},
 					close: async (ws) => {
-						await Effect.runPromise(handleClose(ws));
+						await runWebSocketHandler(ws, "close", handleClose(ws));
 					},
 				},
 			}),