From 37039d3b67ba6fddd75f871eea24f11aad6ab13c 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 | 121 +++++++++ app/server/modules/agents/agent-tokens.ts | 12 + .../agents/controller-agent-session.ts | 247 ++++++++++++++++++ 3 files changed, 380 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/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..c2a35dd5 --- /dev/null +++ b/app/server/modules/agents/__tests__/controller-agent-session.test.ts @@ -0,0 +1,121 @@ +import { expect, mock, test } from "bun:test"; +import { createAgentMessage } from "@zerobyte/contracts/agent-protocol"; +import { createControllerAgentSession } from "../controller-agent-session"; + +const createSocket = () => { + return { + data: { id: "connection-1", agentId: "local", organizationId: null, agentName: "Local Agent" }, + send: mock(() => undefined), + } as unknown as Parameters[0]; +}; + +test("close emits a synthetic backup.cancelled for a started backup", () => { + const onBackupCancelled = mock(() => undefined); + 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("close does not emit a synthetic backup.cancelled after a terminal event", () => { + for (const testCase of [ + { + jobId: "job-1", + scheduleId: "schedule-1", + terminalMessage: createAgentMessage("backup.completed", { + jobId: "job-1", + scheduleId: "schedule-1", + exitCode: 0, + result: null, + }), + expectedCancelledCalls: 0, + }, + { + jobId: "job-2", + scheduleId: "schedule-2", + terminalMessage: createAgentMessage("backup.failed", { + jobId: "job-2", + scheduleId: "schedule-2", + error: "backup failed", + }), + expectedCancelledCalls: 0, + }, + { + jobId: "job-3", + scheduleId: "schedule-3", + terminalMessage: createAgentMessage("backup.cancelled", { + jobId: "job-3", + scheduleId: "schedule-3", + message: "Backup was cancelled", + }), + expectedCancelledCalls: 1, + }, + ]) { + const onBackupCancelled = mock(() => undefined); + 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 = mock(() => undefined); + 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/controller-agent-session.ts b/app/server/modules/agents/controller-agent-session.ts new file mode 100644 index 00000000..900aee72 --- /dev/null +++ b/app/server/modules/agents/controller-agent-session.ts @@ -0,0 +1,247 @@ +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 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) => string; + sendBackupCancel: (payload: BackupCancelPayload) => void; + isReady: () => boolean; + close: () => void; +}; + +export const createControllerAgentSession = ( + socket: AgentSocket, + handlers: ControllerAgentSessionHandlers = {}, +): ControllerAgentSession => { + const outboundQueue = Effect.runSync(Queue.bounded(64)); + const activeBackupJobs = Effect.runSync(Ref.make>(new Map())); + const pendingBackupJobs = 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 setActiveBackupJob = (jobId: string, scheduleId: string) => { + Effect.runSync( + Ref.update(activeBackupJobs, (current) => { + const next = new Map(current); + next.set(jobId, scheduleId); + return next; + }), + ); + }; + + const setPendingBackupJob = (jobId: string, scheduleId: string) => { + Effect.runSync( + Ref.update(pendingBackupJobs, (current) => { + const next = new Map(current); + next.set(jobId, scheduleId); + return next; + }), + ); + }; + + const deleteActiveBackupJob = (jobId: string) => { + Effect.runSync( + Ref.update(activeBackupJobs, (current) => { + const next = new Map(current); + next.delete(jobId); + return next; + }), + ); + }; + + const deletePendingBackupJob = (jobId: string) => { + Effect.runSync( + Ref.update(pendingBackupJobs, (current) => { + const next = new Map(current); + next.delete(jobId); + return next; + }), + ); + }; + + 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": { + deletePendingBackupJob(message.payload.jobId); + setActiveBackupJob(message.payload.jobId, message.payload.scheduleId); + 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": { + deletePendingBackupJob(message.payload.jobId); + deleteActiveBackupJob(message.payload.jobId); + handlers.onBackupCompleted?.(message.payload); + break; + } + case "backup.failed": { + deletePendingBackupJob(message.payload.jobId); + deleteActiveBackupJob(message.payload.jobId); + handlers.onBackupFailed?.(message.payload); + break; + } + case "backup.cancelled": { + deletePendingBackupJob(message.payload.jobId); + deleteActiveBackupJob(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) => { + setPendingBackupJob(payload.jobId, payload.scheduleId); + offerOutbound(createControllerMessage("backup.run", payload)); + return payload.jobId; + }, + sendBackupCancel: (payload) => { + offerOutbound(createControllerMessage("backup.cancel", payload)); + }, + isReady: () => Effect.runSync(Ref.get(state)).isReady, + close: () => { + updateState((current) => ({ ...current, isReady: false })); + const pendingJobs = Effect.runSync(Ref.get(pendingBackupJobs)); + Effect.runSync(Ref.set(pendingBackupJobs, new Map())); + for (const [jobId, scheduleId] of pendingJobs) { + handlers.onBackupCancelled?.({ + jobId, + scheduleId, + message: + "The connection to the backup agent was lost before this backup started. Restart the backup to ensure it completes.", + }); + } + + const activeJobs = Effect.runSync(Ref.get(activeBackupJobs)); + Effect.runSync(Ref.set(activeBackupJobs, new Map())); + for (const [jobId, scheduleId] of activeJobs) { + handlers.onBackupCancelled?.({ + jobId, + scheduleId, + message: + "The connection to the backup agent was lost while this backup was running. Restart the backup to ensure it completes.", + }); + } + void Effect.runPromise(Fiber.interrupt(writerFiber)).catch(() => {}); + void Effect.runPromise(Fiber.interrupt(heartbeatFiber)).catch(() => {}); + void Effect.runPromise(Queue.shutdown(outboundQueue)).catch(() => {}); + }, + }; +};