From 1eda48916763f7b5a7f229b32280b08b40cc5d96 Mon Sep 17 00:00:00 2001 From: Nicolas Meienberger Date: Tue, 7 Apr 2026 22:00:30 +0200 Subject: [PATCH] feat(controller): add agent manager and session handling --- .../controller-agent-session.test.ts | 123 +++++++ app/server/modules/agents/agent-tokens.ts | 12 + app/server/modules/agents/agents-manager.ts | 328 ++++++++++++++++++ .../agents/controller-agent-session.ts | 234 +++++++++++++ 4 files changed, 697 insertions(+) create mode 100644 app/server/modules/agents/__tests__/controller-agent-session.test.ts create mode 100644 app/server/modules/agents/agent-tokens.ts create mode 100644 app/server/modules/agents/agents-manager.ts create mode 100644 app/server/modules/agents/controller-agent-session.ts diff --git a/app/server/modules/agents/__tests__/controller-agent-session.test.ts b/app/server/modules/agents/__tests__/controller-agent-session.test.ts new file mode 100644 index 00000000..8c37c270 --- /dev/null +++ b/app/server/modules/agents/__tests__/controller-agent-session.test.ts @@ -0,0 +1,123 @@ +import { expect, test, vi } from "vitest"; +import { fromPartial } from "@total-typescript/shoehorn"; +import { createAgentMessage } from "@zerobyte/contracts/agent-protocol"; +import { createControllerAgentSession } from "../controller-agent-session"; + +const createSocket = () => { + return fromPartial[0]>({ + data: { id: "connection-1", agentId: "local", organizationId: null, agentName: "Local Agent" }, + send: vi.fn(), + }); +}; + +test("close emits a synthetic backup.cancelled for a started backup", () => { + const onBackupCancelled = vi.fn(); + const session = createControllerAgentSession(createSocket(), { + onBackupCancelled, + }); + + session.handleMessage( + createAgentMessage("backup.started", { + jobId: "job-1", + scheduleId: "schedule-1", + }), + ); + + session.close(); + + expect(onBackupCancelled).toHaveBeenCalledTimes(1); + expect(onBackupCancelled).toHaveBeenCalledWith({ + jobId: "job-1", + scheduleId: "schedule-1", + message: + "The connection to the backup agent was lost while this backup was running. Restart the backup to ensure it completes.", + }); +}); + +test.each([ + { + name: "backup.completed", + jobId: "job-1", + scheduleId: "schedule-1", + terminalMessage: createAgentMessage("backup.completed", { + jobId: "job-1", + scheduleId: "schedule-1", + exitCode: 0, + result: null, + }), + expectedCancelledCalls: 0, + }, + { + name: "backup.failed", + jobId: "job-2", + scheduleId: "schedule-2", + terminalMessage: createAgentMessage("backup.failed", { + jobId: "job-2", + scheduleId: "schedule-2", + error: "backup failed", + }), + expectedCancelledCalls: 0, + }, + { + name: "backup.cancelled", + jobId: "job-3", + scheduleId: "schedule-3", + terminalMessage: createAgentMessage("backup.cancelled", { + jobId: "job-3", + scheduleId: "schedule-3", + message: "Backup was cancelled", + }), + expectedCancelledCalls: 1, + }, +])("close does not emit an extra synthetic backup.cancelled after $name", (testCase) => { + const onBackupCancelled = vi.fn(); + const session = createControllerAgentSession(createSocket(), { + onBackupCancelled, + }); + + session.handleMessage( + createAgentMessage("backup.started", { + jobId: testCase.jobId, + scheduleId: testCase.scheduleId, + }), + ); + session.handleMessage(testCase.terminalMessage); + session.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(), { + 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: [], + }, + }); + + session.close(); + + expect(onBackupCancelled).toHaveBeenCalledTimes(1); + expect(onBackupCancelled).toHaveBeenCalledWith({ + jobId: "job-queued", + scheduleId: "schedule-queued", + message: + "The connection to the backup agent was lost before this backup started. Restart the backup to ensure it completes.", + }); +}); diff --git a/app/server/modules/agents/agent-tokens.ts b/app/server/modules/agents/agent-tokens.ts new file mode 100644 index 00000000..febb9d94 --- /dev/null +++ b/app/server/modules/agents/agent-tokens.ts @@ -0,0 +1,12 @@ +import { cryptoUtils } from "~/server/utils/crypto"; + +export const deriveLocalAgentToken = async () => { + return cryptoUtils.deriveSecret("zerobyte:local-agent-token"); +}; + +export const validateAgentToken = async (token: string) => { + const localToken = await deriveLocalAgentToken(); + if (token === localToken) { + return { agentId: "local", organizationId: null, agentName: "local" }; + } +}; diff --git a/app/server/modules/agents/agents-manager.ts b/app/server/modules/agents/agents-manager.ts new file mode 100644 index 00000000..1964daf9 --- /dev/null +++ b/app/server/modules/agents/agents-manager.ts @@ -0,0 +1,328 @@ +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 { logger } from "@zerobyte/core/node"; +import type { + BackupCancelPayload, + BackupCancelledPayload, + BackupCompletedPayload, + BackupFailedPayload, + BackupProgressPayload, + BackupRunPayload, + BackupStartedPayload, +} from "@zerobyte/contracts/agent-protocol"; +import { config } from "../../core/config"; +import { validateAgentToken, deriveLocalAgentToken } from "./agent-tokens"; +import { + createControllerAgentSession, + type AgentConnectionData, + type ControllerAgentSession, +} from "./controller-agent-session"; + +type AgentBackupEventContext = { + agentId: string; + agentName: string; + payload: + | BackupStartedPayload + | BackupProgressPayload + | BackupCompletedPayload + | BackupFailedPayload + | BackupCancelledPayload; +}; + +export type AgentBackupEventHandlers = { + onBackupStarted?: (context: AgentBackupEventContext & { payload: BackupStartedPayload }) => void; + onBackupProgress?: (context: AgentBackupEventContext & { payload: BackupProgressPayload }) => void; + onBackupCompleted?: (context: AgentBackupEventContext & { payload: BackupCompletedPayload }) => void; + onBackupFailed?: (context: AgentBackupEventContext & { payload: BackupFailedPayload }) => void; + onBackupCancelled?: (context: AgentBackupEventContext & { payload: BackupCancelledPayload }) => void; +}; + +type AgentManagerRuntime = ReturnType; +type AgentRuntimeState = { + agentManager: AgentManagerRuntime; + localAgent: ChildProcess | null; +}; + +type ProcessWithAgentRuntime = NodeJS.Process & { + __zerobyteAgentRuntime?: AgentRuntimeState; +}; + +const getAgentRuntimeState = () => { + const runtimeProcess = process as ProcessWithAgentRuntime; + const existingRuntime = runtimeProcess.__zerobyteAgentRuntime; + + if (existingRuntime) { + return existingRuntime; + } + + const runtime = { + agentManager: createAgentManagerRuntime(), + localAgent: null, + }; + + runtimeProcess.__zerobyteAgentRuntime = runtime; + return runtime; +}; + +const getAgentManagerRuntime = () => getAgentRuntimeState().agentManager; + +export const spawnLocalAgent = async () => { + await stopLocalAgent(); + + const sourceEntryPoint = path.join(process.cwd(), "apps", "agent", "src", "index.ts"); + const productionEntryPoint = path.join(process.cwd(), ".output", "agent", "index.mjs"); + + if (config.__prod__ && !existsSync(productionEntryPoint)) { + throw new Error(`Local agent entrypoint not found at ${productionEntryPoint}`); + } + + const agentEntryPoint = config.__prod__ ? productionEntryPoint : sourceEntryPoint; + const agentToken = await deriveLocalAgentToken(); + const args = config.__prod__ ? ["run", agentEntryPoint] : ["run", "--watch", agentEntryPoint]; + + const runtime = getAgentRuntimeState(); + const agentProcess = spawn("bun", args, { + env: { + PATH: process.env.PATH, + ZEROBYTE_CONTROLLER_URL: "ws://localhost:3001", + ZEROBYTE_AGENT_TOKEN: agentToken, + }, + stdio: ["ignore", "pipe", "pipe"], + }); + + runtime.localAgent = agentProcess; + + agentProcess.stdout?.on("data", (data: Buffer) => { + const line = data.toString().trim(); + if (line) logger.info(`[agent] ${line}`); + }); + + agentProcess.stderr?.on("data", (data: Buffer) => { + const line = data.toString().trim(); + if (line) logger.error(`[agent] ${line}`); + }); + + agentProcess.on("exit", (code, signal) => { + if (runtime.localAgent === agentProcess) { + runtime.localAgent = null; + } + logger.info(`Agent process exited with code ${code} and signal ${signal}`); + }); +}; + +export const stopLocalAgent = async () => { + const runtime = getAgentRuntimeState(); + if (!runtime.localAgent) { + return; + } + + const agentProcess = runtime.localAgent; + runtime.localAgent = null; + + if (agentProcess.exitCode !== null || agentProcess.signalCode !== null) { + return; + } + + const exited = new Promise((resolve) => { + agentProcess.once("exit", () => { + resolve(); + }); + }); + + agentProcess.kill(); + await exited; +}; + +const createAgentManagerRuntime = () => { + 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 getSession = (agentId: string) => sessions.get(agentId); + + const setSession = (agentId: string, session: ControllerAgentSession) => { + const existingSession = getSession(agentId); + if (existingSession) { + existingSession.close(); + } + + sessions.set(agentId, session); + }; + + const removeSession = (agentId: string, connectionId: string) => { + const session = getSession(agentId); + if (!session || session.connectionId !== connectionId) { + return; + } + + session.close(); + sessions.delete(agentId); + }; + + const createSessionHandlers = (ws: Bun.ServerWebSocket) => { + const agentId = ws.data.agentId; + const agentName = ws.data.agentName; + + return { + onBackupStarted: (payload: BackupStartedPayload) => { + backupHandlers.onBackupStarted?.({ agentId, agentName, payload }); + }, + onBackupProgress: (payload: BackupProgressPayload) => { + backupHandlers.onBackupProgress?.({ agentId, agentName, payload }); + }, + onBackupCompleted: (payload: BackupCompletedPayload) => { + backupHandlers.onBackupCompleted?.({ agentId, agentName, payload }); + }, + onBackupFailed: (payload: BackupFailedPayload) => { + backupHandlers.onBackupFailed?.({ agentId, agentName, payload }); + }, + onBackupCancelled: (payload: BackupCancelledPayload) => { + backupHandlers.onBackupCancelled?.({ agentId, agentName, payload }); + }, + }; + }; + + const acquireServer = Effect.acquireRelease( + Effect.sync(() => + Bun.serve({ + port: 3001, + async fetch(req, srv) { + const url = new URL(req.url); + const token = url.searchParams.get("token"); + + if (!token) { + return new Response("Missing token", { status: 401 }); + } + + const result = await validateAgentToken(token); + if (!result) { + return new Response("Invalid or revoked token", { status: 401 }); + } + + const upgraded = srv.upgrade(req, { + data: { + id: Bun.randomUUIDv7(), + agentId: result.agentId, + organizationId: result.organizationId, + agentName: result.agentName, + }, + }); + if (upgraded) return undefined; + return new Response("WebSocket upgrade failed", { status: 400 }); + }, + websocket: { + open: (ws) => { + setSession(ws.data.agentId, createControllerAgentSession(ws, createSessionHandlers(ws))); + 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; + } + + session.handleMessage(data); + }, + close: (ws) => { + removeSession(ws.data.agentId, ws.data.id); + logger.info(`Agent "${ws.data.agentName}" (${ws.data.agentId}) disconnected`); + }, + }, + }), + ), + (server) => + Effect.sync(() => { + closeAllSessions(); + void server.stop(true); + }), + ); + + const stop = () => { + if (!runtimeScope) { + return; + } + + logger.info("Stopping Agent Manager..."); + const scope = runtimeScope; + runtimeScope = null; + Effect.runSync(Scope.close(scope, Exit.succeed(undefined))); + }; + + const start = () => { + if (runtimeScope) { + stop(); + } + + logger.info("Starting Agent Manager..."); + const scope = Effect.runSync(Scope.make()); + + try { + const server = Effect.runSync(Scope.extend(acquireServer, scope)); + runtimeScope = scope; + logger.info(`Agent Manager listening on port ${server.port}`); + } catch (error) { + Effect.runSync(Scope.close(scope, Exit.fail(error))); + throw error; + } + }; + + return { + start, + sendBackup: (agentId: string, payload: BackupRunPayload) => { + const session = getSession(agentId); + + if (!session) { + logger.warn(`Cannot send backup command. Agent ${agentId} is not connected.`); + return false; + } + + if (!session.isReady()) { + logger.warn(`Cannot send backup command. Agent ${agentId} is not ready.`); + return false; + } + + session.sendBackup(payload); + logger.info(`Sent backup command ${payload.jobId} to agent ${agentId} for schedule ${payload.scheduleId}`); + return true; + }, + cancelBackup: (agentId: string, payload: BackupCancelPayload) => { + const session = getSession(agentId); + + if (!session) { + logger.warn(`Cannot cancel backup command. Agent ${agentId} is not connected.`); + return false; + } + + session.sendBackupCancel(payload); + logger.info(`Sent backup cancel for command ${payload.jobId} to agent ${agentId}`); + return true; + }, + setBackupEventHandlers: (handlers: AgentBackupEventHandlers) => { + backupHandlers = handlers; + }, + getBackupEventHandlers: () => backupHandlers, + stop, + }; +}; + +export const agentManager = getAgentManagerRuntime(); + +export const stopAgentRuntime = async () => { + getAgentManagerRuntime().stop(); + await stopLocalAgent(); +}; diff --git a/app/server/modules/agents/controller-agent-session.ts b/app/server/modules/agents/controller-agent-session.ts new file mode 100644 index 00000000..78e03ba4 --- /dev/null +++ b/app/server/modules/agents/controller-agent-session.ts @@ -0,0 +1,234 @@ +import { Effect, Fiber, Queue, Ref } from "effect"; +import { + createControllerMessage, + parseAgentMessage, + type AgentMessage, + type BackupCancelledPayload, + type BackupCompletedPayload, + type BackupFailedPayload, + type BackupProgressPayload, + type BackupRunPayload, + type BackupCancelPayload, + type BackupStartedPayload, + type ControllerWireMessage, +} from "@zerobyte/contracts/agent-protocol"; +import { logger } from "@zerobyte/core/node"; +import { toMessage } from "@zerobyte/core/utils"; + +export type AgentConnectionData = { + id: string; + agentId: string; + organizationId: string | null; + agentName: string; +}; + +type AgentSocket = Bun.ServerWebSocket; + +type SessionState = { + isReady: boolean; + lastSeenAt: number | null; + lastPongAt: number | null; +}; + +type TrackedBackupJob = { + scheduleId: string; + state: "pending" | "active"; +}; + +type ControllerAgentSessionHandlers = { + onBackupStarted?: (payload: BackupStartedPayload) => void; + onBackupProgress?: (payload: BackupProgressPayload) => void; + onBackupCompleted?: (payload: BackupCompletedPayload) => void; + onBackupFailed?: (payload: BackupFailedPayload) => void; + onBackupCancelled?: (payload: BackupCancelledPayload) => void; +}; + +export type ControllerAgentSession = { + readonly connectionId: string; + handleMessage: (data: string) => void; + sendBackup: (payload: BackupRunPayload) => void; + sendBackupCancel: (payload: BackupCancelPayload) => void; + isReady: () => boolean; + close: () => void; +}; + +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({ + 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 setTrackedBackupJob = (jobId: string, trackedBackupJob: TrackedBackupJob) => { + Effect.runSync( + 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 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, + }); + } + void Effect.runPromise(Fiber.interrupt(writerFiber)).catch(() => {}); + void Effect.runPromise(Fiber.interrupt(heartbeatFiber)).catch(() => {}); + void Effect.runPromise(Queue.shutdown(outboundQueue)).catch(() => {}); + }, + }; +};