diff --git a/app/server/modules/agents/__tests__/controller-runtime.test.ts b/app/server/modules/agents/__tests__/controller-runtime.test.ts index 1841f36d..a039bba3 100644 --- a/app/server/modules/agents/__tests__/controller-runtime.test.ts +++ b/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") diff --git a/app/server/modules/agents/controller/server.ts b/app/server/modules/agents/controller/server.ts index 2e306941..8e647798 100644 --- a/app/server/modules/agents/controller/server.ts +++ b/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, + event: string, + effect: Effect.Effect, + ) => + 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({ @@ -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)); }, }, }),