From 6cda41df7c90c98aff62d355c12ef1868374b5d0 Mon Sep 17 00:00:00 2001 From: Nicolas Meienberger Date: Wed, 8 Apr 2026 23:23:01 +0200 Subject: [PATCH] refactor(agents): make the controller-agent effectful --- .../controller-agent-session.test.ts | 93 +++-- app/server/modules/agents/agents-manager.ts | 69 +++- .../agents/controller-agent-session.ts | 341 +++++++++--------- 3 files changed, 283 insertions(+), 220 deletions(-) diff --git a/app/server/modules/agents/__tests__/controller-agent-session.test.ts b/app/server/modules/agents/__tests__/controller-agent-session.test.ts index 8c37c270..a62067ae 100644 --- a/app/server/modules/agents/__tests__/controller-agent-session.test.ts +++ b/app/server/modules/agents/__tests__/controller-agent-session.test.ts @@ -1,3 +1,4 @@ +import { Effect, Exit, Scope } from "effect"; import { expect, test, vi } from "vitest"; import { fromPartial } from "@total-typescript/shoehorn"; import { createAgentMessage } from "@zerobyte/contracts/agent-protocol"; @@ -10,20 +11,40 @@ const createSocket = () => { }); }; +const createSession = (handlers: Parameters[1] = {}) => { + const scope = Effect.runSync(Scope.make()); + + try { + const session = Effect.runSync(Scope.extend(createControllerAgentSession(createSocket(), handlers), scope)); + + return { + session, + close: () => { + Effect.runSync(Scope.close(scope, Exit.succeed(undefined))); + }, + }; + } catch (error) { + Effect.runSync(Scope.close(scope, Exit.fail(error))); + throw error; + } +}; + test("close emits a synthetic backup.cancelled for a started backup", () => { const onBackupCancelled = vi.fn(); - const session = createControllerAgentSession(createSocket(), { + const { session, close } = createSession({ onBackupCancelled, }); - session.handleMessage( - createAgentMessage("backup.started", { - jobId: "job-1", - scheduleId: "schedule-1", - }), + Effect.runSync( + session.handleMessage( + createAgentMessage("backup.started", { + jobId: "job-1", + scheduleId: "schedule-1", + }), + ), ); - session.close(); + close(); expect(onBackupCancelled).toHaveBeenCalledTimes(1); expect(onBackupCancelled).toHaveBeenCalledWith({ @@ -71,47 +92,51 @@ test.each([ }, ])("close does not emit an extra synthetic backup.cancelled after $name", (testCase) => { const onBackupCancelled = vi.fn(); - const session = createControllerAgentSession(createSocket(), { + const { session, close } = createSession({ onBackupCancelled, }); - session.handleMessage( - createAgentMessage("backup.started", { - jobId: testCase.jobId, - scheduleId: testCase.scheduleId, - }), + Effect.runSync( + session.handleMessage( + createAgentMessage("backup.started", { + jobId: testCase.jobId, + scheduleId: testCase.scheduleId, + }), + ), ); - session.handleMessage(testCase.terminalMessage); - session.close(); + Effect.runSync(session.handleMessage(testCase.terminalMessage)); + close(); expect(onBackupCancelled).toHaveBeenCalledTimes(testCase.expectedCancelledCalls); }); test("close emits a synthetic backup.cancelled for a queued backup", () => { const onBackupCancelled = vi.fn(); - const session = createControllerAgentSession(createSocket(), { + const { session, close } = createSession({ onBackupCancelled, }); - session.sendBackup({ - jobId: "job-queued", - scheduleId: "schedule-queued", - organizationId: "org-1", - sourcePath: "/tmp/source", - repositoryConfig: { - backend: "local", - path: "/tmp/repository", - }, - options: {}, - runtime: { - password: "password", - cacheDir: "/tmp/cache", - passFile: "/tmp/pass", - defaultExcludes: [], - }, - }); + Effect.runSync( + session.sendBackup({ + jobId: "job-queued", + scheduleId: "schedule-queued", + organizationId: "org-1", + sourcePath: "/tmp/source", + repositoryConfig: { + backend: "local", + path: "/tmp/repository", + }, + options: {}, + runtime: { + password: "password", + cacheDir: "/tmp/cache", + passFile: "/tmp/pass", + defaultExcludes: [], + }, + }), + ); - session.close(); + close(); expect(onBackupCancelled).toHaveBeenCalledTimes(1); expect(onBackupCancelled).toHaveBeenCalledWith({ diff --git a/app/server/modules/agents/agents-manager.ts b/app/server/modules/agents/agents-manager.ts index 1964daf9..b5eb6d7a 100644 --- a/app/server/modules/agents/agents-manager.ts +++ b/app/server/modules/agents/agents-manager.ts @@ -1,7 +1,7 @@ import { type ChildProcess, spawn } from "node:child_process"; import { existsSync } from "node:fs"; import path from "node:path"; -import { Effect, Exit, Scope } from "effect"; +import { Effect, Exit, Fiber, Scope } from "effect"; import { logger } from "@zerobyte/core/node"; import type { BackupCancelPayload, @@ -45,6 +45,12 @@ type AgentRuntimeState = { localAgent: ChildProcess | null; }; +type ControllerAgentSessionHandle = { + session: ControllerAgentSession; + runFiber: Fiber.RuntimeFiber; + scope: Scope.CloseableScope; +}; + type ProcessWithAgentRuntime = NodeJS.Process & { __zerobyteAgentRuntime?: AgentRuntimeState; }; @@ -136,36 +142,59 @@ export const stopLocalAgent = async () => { }; const createAgentManagerRuntime = () => { - let sessions = new Map(); + let sessions = new Map(); let backupHandlers: AgentBackupEventHandlers = {}; let runtimeScope: Scope.CloseableScope | null = null; - const closeAllSessions = () => { - for (const session of sessions.values()) { - session.close(); - } - sessions = new Map(); + const closeSession = (sessionHandle: ControllerAgentSessionHandle) => { + Effect.runSync(Fiber.interrupt(sessionHandle.runFiber)); + Effect.runSync(Scope.close(sessionHandle.scope, Exit.succeed(undefined))); }; - const getSession = (agentId: string) => sessions.get(agentId); + const closeAllSessions = () => { + const currentSessions = sessions; + sessions = new Map(); + for (const sessionHandle of currentSessions.values()) { + closeSession(sessionHandle); + } + }; - const setSession = (agentId: string, session: ControllerAgentSession) => { - const existingSession = getSession(agentId); + const getSessionHandle = (agentId: string) => sessions.get(agentId); + + const getSession = (agentId: string) => getSessionHandle(agentId)?.session; + + const createSession = (ws: Bun.ServerWebSocket) => { + // Manual scope management because we are out of Effect + const scope = Effect.runSync(Scope.make()); + + try { + const session = Effect.runSync(Scope.extend(createControllerAgentSession(ws, createSessionHandlers(ws)), scope)); + const runFiber = Effect.runFork(Scope.extend(session.run, scope)); + + 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); if (existingSession) { - existingSession.close(); + closeSession(existingSession); } - sessions.set(agentId, session); + sessions.set(agentId, sessionHandle); }; const removeSession = (agentId: string, connectionId: string) => { - const session = getSession(agentId); - if (!session || session.connectionId !== connectionId) { + const sessionHandle = getSessionHandle(agentId); + if (!sessionHandle || sessionHandle.session.connectionId !== connectionId) { return; } - session.close(); sessions.delete(agentId); + closeSession(sessionHandle); }; const createSessionHandlers = (ws: Bun.ServerWebSocket) => { @@ -221,7 +250,7 @@ const createAgentManagerRuntime = () => { }, websocket: { open: (ws) => { - setSession(ws.data.agentId, createControllerAgentSession(ws, createSessionHandlers(ws))); + setSession(ws.data.agentId, createSession(ws)); logger.info(`Agent "${ws.data.agentName}" (${ws.data.agentId}) connected on ${ws.data.id}`); }, message: (ws, data) => { @@ -236,7 +265,7 @@ const createAgentManagerRuntime = () => { return; } - session.handleMessage(data); + Effect.runSync(session.handleMessage(data)); }, close: (ws) => { removeSession(ws.data.agentId, ws.data.id); @@ -291,12 +320,12 @@ const createAgentManagerRuntime = () => { return false; } - if (!session.isReady()) { + if (!Effect.runSync(session.isReady())) { logger.warn(`Cannot send backup command. Agent ${agentId} is not ready.`); return false; } - session.sendBackup(payload); + Effect.runSync(session.sendBackup(payload)); logger.info(`Sent backup command ${payload.jobId} to agent ${agentId} for schedule ${payload.scheduleId}`); return true; }, @@ -308,7 +337,7 @@ const createAgentManagerRuntime = () => { return false; } - session.sendBackupCancel(payload); + Effect.runSync(session.sendBackupCancel(payload)); logger.info(`Sent backup cancel for command ${payload.jobId} to agent ${agentId}`); return true; }, diff --git a/app/server/modules/agents/controller-agent-session.ts b/app/server/modules/agents/controller-agent-session.ts index 78e03ba4..0b2529ec 100644 --- a/app/server/modules/agents/controller-agent-session.ts +++ b/app/server/modules/agents/controller-agent-session.ts @@ -1,4 +1,4 @@ -import { Effect, Fiber, Queue, Ref } from "effect"; +import { Effect, Queue, Ref, type Scope } from "effect"; import { createControllerMessage, parseAgentMessage, @@ -45,190 +45,199 @@ type ControllerAgentSessionHandlers = { export type ControllerAgentSession = { readonly connectionId: string; - handleMessage: (data: string) => void; - sendBackup: (payload: BackupRunPayload) => void; - sendBackupCancel: (payload: BackupCancelPayload) => void; - isReady: () => boolean; - close: () => void; + handleMessage: (data: string) => Effect.Effect; + sendBackup: (payload: BackupRunPayload) => Effect.Effect; + sendBackupCancel: (payload: BackupCancelPayload) => Effect.Effect; + isReady: () => Effect.Effect; + run: Effect.Effect; }; export const createControllerAgentSession = ( socket: AgentSocket, handlers: ControllerAgentSessionHandlers = {}, -): ControllerAgentSession => { - const outboundQueue = Effect.runSync(Queue.bounded(64)); - const trackedBackupJobs = Effect.runSync(Ref.make>(new Map())); - const state = Effect.runSync( - Ref.make({ +): Effect.Effect => + Effect.gen(function* () { + const outboundQueue = yield* Queue.bounded(64); + const trackedBackupJobs = yield* Ref.make>(new Map()); + const state = yield* Ref.make({ isReady: false, lastSeenAt: null, lastPongAt: null, - }), - ); - - const offerOutbound = (message: ControllerWireMessage) => { - void Effect.runPromise(Queue.offer(outboundQueue, message)).catch((error) => { - logger.error(`Failed to queue outbound message for agent ${socket.data.agentId}: ${toMessage(error)}`); }); - }; - const updateState = (update: (current: SessionState) => SessionState) => { - Effect.runSync(Ref.update(state, update)); - }; + const offerOutbound = (message: ControllerWireMessage) => + Queue.offer(outboundQueue, message).pipe( + Effect.catchAllCause((cause) => + Effect.sync(() => { + logger.error(`Failed to queue outbound message for agent ${socket.data.agentId}: ${toMessage(cause)}`); + }), + ), + ); - const setTrackedBackupJob = (jobId: string, trackedBackupJob: TrackedBackupJob) => { - Effect.runSync( - Ref.update(trackedBackupJobs, (current) => { + const updateState = (update: (current: SessionState) => SessionState) => Ref.update(state, update); + + const setTrackedBackupJob = (jobId: string, trackedBackupJob: TrackedBackupJob) => { + return Ref.update(trackedBackupJobs, (current) => { const next = new Map(current); next.set(jobId, trackedBackupJob); return next; - }), - ); - }; + }); + }; - const deleteTrackedBackupJob = (jobId: string) => { - Effect.runSync( - Ref.update(trackedBackupJobs, (current) => { + const deleteTrackedBackupJob = (jobId: string) => { + return Ref.update(trackedBackupJobs, (current) => { const next = new Map(current); next.delete(jobId); return next; - }), - ); - }; - - const takeTrackedBackupJobs = () => { - return Effect.runSync( - Ref.modify(trackedBackupJobs, (current) => [current, new Map()] as const), - ); - }; - - const writerFiber = Effect.runFork( - Effect.forever( - Effect.gen(function* () { - const message = yield* Queue.take(outboundQueue); - yield* Effect.sync(() => { - try { - socket.send(message); - } catch (error) { - logger.error( - `Failed to send message to agent ${socket.data.agentId} on ${socket.data.id}: ${toMessage(error)}`, - ); - } - }); - }), - ), - ); - - const heartbeatFiber = Effect.runFork( - Effect.forever( - Effect.gen(function* () { - yield* Effect.sleep("15 seconds"); - yield* Queue.offer( - outboundQueue, - createControllerMessage("heartbeat.ping", { - sentAt: Date.now(), - }), - ); - }), - ), - ); - - const handleAgentMessage = (message: AgentMessage) => { - updateState((current) => ({ ...current, lastSeenAt: Date.now() })); - - switch (message.type) { - case "agent.ready": { - updateState((current) => ({ ...current, isReady: true })); - logger.info(`Agent "${socket.data.agentName}" (${socket.data.agentId}) is ready`); - break; - } - case "backup.started": { - setTrackedBackupJob(message.payload.jobId, { - scheduleId: message.payload.scheduleId, - state: "active", - }); - logger.info( - `Backup ${message.payload.jobId} started on agent ${socket.data.agentId} for schedule ${message.payload.scheduleId}`, - ); - handlers.onBackupStarted?.(message.payload); - break; - } - case "backup.progress": { - handlers.onBackupProgress?.(message.payload); - break; - } - case "backup.completed": { - deleteTrackedBackupJob(message.payload.jobId); - handlers.onBackupCompleted?.(message.payload); - break; - } - case "backup.failed": { - deleteTrackedBackupJob(message.payload.jobId); - handlers.onBackupFailed?.(message.payload); - break; - } - case "backup.cancelled": { - deleteTrackedBackupJob(message.payload.jobId); - handlers.onBackupCancelled?.(message.payload); - break; - } - case "heartbeat.pong": { - updateState((current) => ({ ...current, lastPongAt: message.payload.sentAt })); - break; - } - } - }; - - return { - connectionId: socket.data.id, - handleMessage: (data: string) => { - const parsed = parseAgentMessage(data); - - if (parsed === null) { - logger.warn(`Invalid JSON from agent ${socket.data.agentId}`); - return; - } - - if (!parsed.success) { - logger.warn(`Invalid agent message from ${socket.data.agentId}: ${parsed.error.message}`); - return; - } - - handleAgentMessage(parsed.data); - }, - sendBackup: (payload) => { - setTrackedBackupJob(payload.jobId, { - scheduleId: payload.scheduleId, - state: "pending", }); - offerOutbound(createControllerMessage("backup.run", payload)); - }, - sendBackupCancel: (payload) => { - offerOutbound(createControllerMessage("backup.cancel", payload)); - }, - isReady: () => Effect.runSync(Ref.get(state)).isReady, - close: () => { - updateState((current) => ({ ...current, isReady: false })); - const trackedJobs = takeTrackedBackupJobs(); - for (const [jobId, trackedJob] of trackedJobs) { - let message: string; - if (trackedJob.state === "pending") { - message = - "The connection to the backup agent was lost before this backup started. Restart the backup to ensure it completes."; - } else { - message = - "The connection to the backup agent was lost while this backup was running. Restart the backup to ensure it completes."; - } + }; - handlers.onBackupCancelled?.({ - jobId, - scheduleId: trackedJob.scheduleId, - message, + const takeTrackedBackupJobs = Ref.modify( + trackedBackupJobs, + (current) => [current, new Map()] as const, + ); + + const releaseSession = Effect.gen(function* () { + yield* updateState((current) => ({ ...current, isReady: false })); + const trackedJobs = yield* takeTrackedBackupJobs; + for (const [jobId, trackedJob] of trackedJobs) { + let message = "The connection to the backup agent was lost. Restart the backup to ensure it completes."; + + yield* Effect.sync(() => { + handlers.onBackupCancelled?.({ jobId, scheduleId: trackedJob.scheduleId, message }); }); } - void Effect.runPromise(Fiber.interrupt(writerFiber)).catch(() => {}); - void Effect.runPromise(Fiber.interrupt(heartbeatFiber)).catch(() => {}); - void Effect.runPromise(Queue.shutdown(outboundQueue)).catch(() => {}); - }, - }; -}; + + yield* Queue.shutdown(outboundQueue); + }); + + yield* Effect.addFinalizer(() => releaseSession); + + const run = Effect.gen(function* () { + yield* Effect.forkScoped( + Effect.forever( + Effect.gen(function* () { + const message = yield* Queue.take(outboundQueue); + yield* Effect.sync(() => { + try { + socket.send(message); + } catch (error) { + logger.error( + `Failed to send message to agent ${socket.data.agentId} on ${socket.data.id}: ${toMessage(error)}`, + ); + } + }); + }), + ), + ); + + yield* Effect.forkScoped( + Effect.forever( + Effect.gen(function* () { + yield* Effect.sleep("15 seconds"); + yield* Queue.offer( + outboundQueue, + createControllerMessage("heartbeat.ping", { + sentAt: Date.now(), + }), + ); + }), + ), + ); + + return yield* Effect.never; + }); + + const handleAgentMessage = (message: AgentMessage) => + Effect.gen(function* () { + yield* updateState((current) => ({ ...current, lastSeenAt: Date.now() })); + + switch (message.type) { + case "agent.ready": { + yield* updateState((current) => ({ ...current, isReady: true })); + yield* Effect.sync(() => { + logger.info(`Agent "${socket.data.agentName}" (${socket.data.agentId}) is ready`); + }); + break; + } + case "backup.started": { + yield* setTrackedBackupJob(message.payload.jobId, { + scheduleId: message.payload.scheduleId, + state: "active", + }); + yield* Effect.sync(() => { + logger.info( + `Backup ${message.payload.jobId} started on agent ${socket.data.agentId} for schedule ${message.payload.scheduleId}`, + ); + handlers.onBackupStarted?.(message.payload); + }); + break; + } + case "backup.progress": { + yield* Effect.sync(() => { + handlers.onBackupProgress?.(message.payload); + }); + break; + } + case "backup.completed": { + yield* deleteTrackedBackupJob(message.payload.jobId); + yield* Effect.sync(() => { + handlers.onBackupCompleted?.(message.payload); + }); + break; + } + case "backup.failed": { + yield* deleteTrackedBackupJob(message.payload.jobId); + yield* Effect.sync(() => { + handlers.onBackupFailed?.(message.payload); + }); + break; + } + case "backup.cancelled": { + yield* deleteTrackedBackupJob(message.payload.jobId); + yield* Effect.sync(() => { + handlers.onBackupCancelled?.(message.payload); + }); + break; + } + case "heartbeat.pong": { + yield* updateState((current) => ({ ...current, lastPongAt: message.payload.sentAt })); + break; + } + } + }); + + return { + connectionId: socket.data.id, + handleMessage: (data: string) => { + return Effect.gen(function* () { + const parsed = parseAgentMessage(data); + + if (parsed === null) { + yield* Effect.sync(() => { + logger.warn(`Invalid JSON from agent ${socket.data.agentId}`); + }); + return; + } + + if (!parsed.success) { + yield* Effect.sync(() => { + logger.warn(`Invalid agent message from ${socket.data.agentId}: ${parsed.error.message}`); + }); + return; + } + + yield* handleAgentMessage(parsed.data); + }); + }, + sendBackup: (payload) => { + return Effect.gen(function* () { + yield* setTrackedBackupJob(payload.jobId, { scheduleId: payload.scheduleId, state: "pending" }); + yield* offerOutbound(createControllerMessage("backup.run", payload)); + }); + }, + sendBackupCancel: (payload) => offerOutbound(createControllerMessage("backup.cancel", payload)), + isReady: () => Ref.get(state).pipe(Effect.map((current) => current.isReady)), + run, + }; + });