diff --git a/app/server/modules/agents/__tests__/session.test.ts b/app/server/modules/agents/__tests__/session.test.ts index 139f5ea3..b624c201 100644 --- a/app/server/modules/agents/__tests__/session.test.ts +++ b/app/server/modules/agents/__tests__/session.test.ts @@ -1,4 +1,4 @@ -import { Effect, Exit, Scope } from "effect"; +import { Effect, Exit, Fiber, Scope } from "effect"; import { expect, test, vi } from "vitest"; import waitForExpect from "wait-for-expect"; import { fromPartial } from "@total-typescript/shoehorn"; @@ -45,7 +45,9 @@ const createSession = ( return { session, run: () => { - Effect.runFork(Scope.extend(session.run, scope)); + const fiber = Effect.runFork(Scope.extend(session.run, scope)); + Effect.runSync(Scope.addFinalizer(scope, Fiber.interrupt(fiber))); + return fiber; }, socket, close: () => { @@ -61,6 +63,16 @@ const createSession = ( } }; +test("closing the session scope interrupts the session runner", async () => { + const { run, closeAsync } = createSession(); + const fiber = run(); + + await closeAsync(); + + const exit = await Effect.runPromise(Fiber.await(fiber).pipe(Effect.timeout("100 millis"))); + expect(Exit.isInterrupted(exit)).toBe(true); +}); + test("close reports a transport disconnect", () => { const onDisconnect = vi.fn(() => Effect.void); const { close } = createSession({ onDisconnect }); diff --git a/app/server/modules/agents/agents-manager.ts b/app/server/modules/agents/agents-manager.ts index 32dc616f..c0c0c126 100644 --- a/app/server/modules/agents/agents-manager.ts +++ b/app/server/modules/agents/agents-manager.ts @@ -172,9 +172,7 @@ export const startAgentController = async () => { runtime.agentManager = null; } - const nextAgentManager = createAgentManagerRuntime(); - nextAgentManager.setBackupEventHandlers(backupEventHandlers); - + const nextAgentManager = createAgentManagerRuntime(backupEventHandlers); await nextAgentManager.start(); runtime.agentManager = nextAgentManager; }; diff --git a/app/server/modules/agents/controller/server.ts b/app/server/modules/agents/controller/server.ts index 46fec4a3..961d29e7 100644 --- a/app/server/modules/agents/controller/server.ts +++ b/app/server/modules/agents/controller/server.ts @@ -35,8 +35,8 @@ export type AgentBackupEventHandlers = { }; type ControllerAgentSessionHandle = { + agentId: string; session: ControllerAgentSession; - runFiber: Fiber.RuntimeFiber; scope: Scope.CloseableScope; }; @@ -44,15 +44,14 @@ class StopAgentManagerServerError extends Data.TaggedError("StopAgentManagerServ cause: unknown; }> {} -export function createAgentManagerRuntime() { +export function createAgentManagerRuntime(handlers: AgentBackupEventHandlers) { let sessions = new Map(); - let backupHandlers: AgentBackupEventHandlers = {}; let runtimeScope: Scope.CloseableScope | null = null; const closeSession = (sessionHandle: ControllerAgentSessionHandle) => Effect.gen(function* () { - yield* Fiber.interrupt(sessionHandle.runFiber); yield* Scope.close(sessionHandle.scope, Exit.succeed(undefined)); + yield* Effect.sync(() => sessions.delete(sessionHandle.agentId)); }); const closeAllSessions = Effect.gen(function* () { @@ -65,17 +64,14 @@ export function createAgentManagerRuntime() { }); const getSessionHandle = (agentId: string) => sessions.get(agentId); - const getSession = (agentId: string) => getSessionHandle(agentId)?.session; - const createSessionHandlers = (ws: Bun.ServerWebSocket, markConnecting: Promise) => { - const agentId = ws.data.agentId; - const agentName = ws.data.agentName; + const createSessionHandlers = (params: { agentId: string; agentName: string; sessionId: string }) => { + const { agentId, agentName, sessionId } = params; return { onReady: ({ at }: { at: number }) => { return Effect.promise(async () => { - await markConnecting; await agentsService.markAgentOnline(agentId, at); }); }, @@ -83,70 +79,89 @@ export function createAgentManagerRuntime() { return Effect.promise(() => agentsService.markAgentSeen(agentId, at)); }, onDisconnect: () => { - if (getSession(ws.data.agentId)?.connectionId !== ws.data.id) { + if (getSession(agentId)?.connectionId !== sessionId) { return Effect.void; } - return Effect.sync(() => backupHandlers.onAgentDisconnected?.({ agentId, agentName })); + return Effect.sync(() => handlers.onAgentDisconnected?.({ agentId, agentName })); }, onBackupStarted: (payload: BackupStartedPayload) => { - return Effect.sync(() => backupHandlers.onBackupStarted?.({ agentId, agentName, payload })); + return Effect.sync(() => handlers.onBackupStarted?.({ agentId, agentName, payload })); }, onBackupProgress: (payload: BackupProgressPayload) => { - return Effect.sync(() => backupHandlers.onBackupProgress?.({ agentId, agentName, payload })); + return Effect.sync(() => handlers.onBackupProgress?.({ agentId, agentName, payload })); }, onBackupCompleted: (payload: BackupCompletedPayload) => { - return Effect.sync(() => backupHandlers.onBackupCompleted?.({ agentId, agentName, payload })); + return Effect.sync(() => handlers.onBackupCompleted?.({ agentId, agentName, payload })); }, onBackupFailed: (payload: BackupFailedPayload) => { - return Effect.sync(() => backupHandlers.onBackupFailed?.({ agentId, agentName, payload })); + return Effect.sync(() => handlers.onBackupFailed?.({ agentId, agentName, payload })); }, onBackupCancelled: (payload: BackupCancelledPayload) => { - return Effect.sync(() => backupHandlers.onBackupCancelled?.({ agentId, agentName, payload })); + return Effect.sync(() => handlers.onBackupCancelled?.({ agentId, agentName, payload })); }, }; }; - const createSession = (ws: Bun.ServerWebSocket, markConnecting: Promise) => { - // Manual scope management because we are out of Effect - const scope = Effect.runSync(Scope.make()); + const createSession = (ws: Bun.ServerWebSocket) => + Effect.gen(function* () { + const scope = yield* Scope.make(); - try { - const session = Effect.runSync( - Scope.extend(createControllerAgentSession(ws, createSessionHandlers(ws, markConnecting)), scope), + const session = yield* Scope.extend( + createControllerAgentSession( + ws, + createSessionHandlers({ + agentId: ws.data.agentId, + agentName: ws.data.agentName, + sessionId: ws.data.id, + }), + ), + scope, ); - const runFiber = Effect.runFork(Scope.extend(session.run, scope)); + const runFiber = yield* Effect.fork(Scope.extend(session.run, scope)); + yield* Scope.addFinalizer(scope, Fiber.interrupt(runFiber)); - return { session, runFiber, scope }; - } catch (error) { - Effect.runSync(Scope.close(scope, Exit.fail(error))); - throw error; - } - }; - - const setSession = (agentId: string, sessionHandle: ControllerAgentSessionHandle) => { - const existingSession = getSessionHandle(agentId); - sessions.set(agentId, sessionHandle); - - if (existingSession) { - void Effect.runPromise(closeSession(existingSession)).catch((error) => { - logger.error(`Failed to close existing agent session for ${agentId}: ${toMessage(error)}`); - }); - } - }; - - const removeSession = (agentId: string, connectionId: string) => { - const sessionHandle = getSessionHandle(agentId); - if (!sessionHandle || sessionHandle.session.connectionId !== connectionId) { - return false; - } - - sessions.delete(agentId); - void Effect.runPromise(closeSession(sessionHandle)).catch((error) => { - logger.error(`Failed to close agent session for ${agentId}: ${toMessage(error)}`); + return { agentId: ws.data.agentId, session, scope }; + }); + + const setSession = (sessionHandle: ControllerAgentSessionHandle) => + Effect.gen(function* () { + const existingSession = sessions.get(sessionHandle.agentId); + sessions.set(sessionHandle.agentId, sessionHandle); + + if (existingSession) { + yield* closeSession(existingSession); + } + }); + + const removeSession = (agentId: string, connectionId: string) => + Effect.gen(function* () { + const handle = sessions.get(agentId); + if (!handle || handle.session.connectionId !== connectionId) { + return false; + } + + yield* closeSession(handle); + + yield* Effect.promise(() => agentsService.markAgentOffline(agentId)); + return true; + }); + + const handleMessage = (ws: Bun.ServerWebSocket, data: unknown) => + Effect.gen(function* () { + if (typeof data !== "string") { + yield* logger.effect.warn(`Ignoring non-text message from agent ${ws.data.agentId}`); + return; + } + + const session = getSession(ws.data.agentId); + if (!session || session.connectionId !== ws.data.id) { + yield* logger.effect.warn(`No active session for agent ${ws.data.agentId} on ${ws.data.id}`); + return; + } + + yield* session.handleMessage(data); }); - return true; - }; const acquireServer = Effect.acquireRelease( Effect.sync(() => @@ -178,43 +193,24 @@ export function createAgentManagerRuntime() { return new Response("WebSocket upgrade failed", { status: 400 }); }, websocket: { - open: (ws) => { - const markConnecting = agentsService.markAgentConnecting({ + open: async (ws) => { + await agentsService.markAgentConnecting({ agentId: ws.data.agentId, organizationId: ws.data.organizationId, agentName: ws.data.agentName, agentKind: ws.data.agentKind, }); - void markConnecting.catch((error) => { - logger.error(`Failed to mark agent ${ws.data.agentId} as connecting: ${toMessage(error)}`); - }); - setSession(ws.data.agentId, createSession(ws, markConnecting)); + + const sessionHandle = await Effect.runPromise(createSession(ws)); + await Effect.runPromise(setSession(sessionHandle)); + logger.info(`Agent "${ws.data.agentName}" (${ws.data.agentId}) connected on ${ws.data.id}`); }, - message: (ws, data) => { - if (typeof data !== "string") { - logger.warn(`Ignoring non-text message from agent ${ws.data.agentId}`); - return; - } - - const session = getSession(ws.data.agentId); - if (!session || session.connectionId !== ws.data.id) { - logger.warn(`No active session for agent ${ws.data.agentId} on ${ws.data.id}`); - return; - } - - void Effect.runPromise(session.handleMessage(data)).catch((error) => { - logger.error( - `Failed to handle message from agent ${ws.data.agentId} on ${ws.data.id}: ${toMessage(error)}`, - ); - }); + message: async (ws, data) => { + await Effect.runPromise(handleMessage(ws, data)); }, - close: (ws) => { - if (removeSession(ws.data.agentId, ws.data.id)) { - void agentsService.markAgentOffline(ws.data.agentId).catch((error) => { - logger.error(`Failed to mark agent ${ws.data.agentId} as offline: ${toMessage(error)}`); - }); - } + close: async (ws) => { + await Effect.runPromise(removeSession(ws.data.agentId, ws.data.id)); logger.info(`Agent "${ws.data.agentName}" (${ws.data.agentId}) disconnected`); }, }, @@ -303,10 +299,6 @@ export function createAgentManagerRuntime() { logger.info(`Sent backup cancel for command ${payload.jobId} to agent ${agentId}`); return true; }, - setBackupEventHandlers: (handlers: AgentBackupEventHandlers) => { - backupHandlers = handlers; - }, - getBackupEventHandlers: () => backupHandlers, stop, }; }