diff --git a/app/server/modules/agents/agent-protocol.ts b/app/server/modules/agents/agent-protocol.ts new file mode 100644 index 00000000..ab2e5f40 --- /dev/null +++ b/app/server/modules/agents/agent-protocol.ts @@ -0,0 +1,98 @@ +import { z } from "zod"; +import { safeJsonParse } from "~/server/utils/json"; + +const backupCommandSchema = z + .object({ + type: z.literal("backup"), + payload: z.object({ scheduleId: z.string() }), + }) + .strict(); + +const agentReadySchema = z + .object({ + type: z.literal("agent.ready"), + payload: z.object({ agentId: z.string() }), + }) + .strict(); + +const backupStartedSchema = z + .object({ + type: z.literal("backup.started"), + payload: z.object({ scheduleId: z.string() }), + }) + .strict(); + +const controllerMessageSchema = z.discriminatedUnion("type", [backupCommandSchema]); +const agentMessageSchema = z.discriminatedUnion("type", [agentReadySchema, backupStartedSchema]); + +export type BackupCommandPayload = z.infer["payload"]; +export type ControllerMessage = z.infer; +export type AgentMessage = z.infer; + +type Brand = TValue & { + readonly __brand: TBrand; +}; + +type MessageSender = { + send(message: string): unknown; +}; + +export type ControllerWireMessage = Brand; +export type AgentWireMessage = Brand; + +type PayloadForMessage = Extract< + TMessage, + { type: TType } +>["payload"]; + +const parseJsonMessage = (data: string) => safeJsonParse(data); + +export const parseControllerMessage = (data: ControllerWireMessage) => { + const parsed = parseJsonMessage(data); + if (parsed === null) { + return null; + } + + return controllerMessageSchema.safeParse(parsed); +}; + +export const parseAgentMessage = (data: string) => { + const parsed = parseJsonMessage(data); + if (parsed === null) { + return null; + } + + return agentMessageSchema.safeParse(parsed); +}; + +export const createControllerMessage = ( + type: TType, + payload: PayloadForMessage, +) => + JSON.stringify( + controllerMessageSchema.parse({ + type, + payload, + }), + ) as ControllerWireMessage; + +export const createAgentMessage = ( + type: TType, + payload: PayloadForMessage, +) => + JSON.stringify( + agentMessageSchema.parse({ + type, + payload, + }), + ) as AgentWireMessage; + +export const sendControllerMessage = (target: MessageSender, message: ControllerWireMessage) => { + target.send(message); +}; + +export const sendAgentMessage = (target: MessageSender, message: AgentWireMessage) => { + target.send(message); +}; + +export type ControllerData = MessageEvent; diff --git a/app/server/modules/agents/agents-manager.ts b/app/server/modules/agents/agents-manager.ts new file mode 100644 index 00000000..29b21b3e --- /dev/null +++ b/app/server/modules/agents/agents-manager.ts @@ -0,0 +1,158 @@ +import { spawn } from "node:child_process"; +import path from "node:path"; +import { + createControllerMessage, + parseAgentMessage, + sendControllerMessage, + type BackupCommandPayload, +} from "./agent-protocol"; +import { logger } from "~/server/utils/logger"; + +type AgentConnectionData = { + id: string; + agentId?: string; +}; + +type AgentSocket = Bun.ServerWebSocket; +type AgentServer = ReturnType>; +const AGENT_SERVER_KEY = Symbol.for("zerobyte.agent-manager.server"); +const AGENT_SOCKETS_KEY = Symbol.for("zerobyte.agent-manager.sockets"); + +const globalState = globalThis as typeof globalThis & { + [AGENT_SERVER_KEY]?: AgentServer; + [AGENT_SOCKETS_KEY]?: Map; +}; + +const getServer = () => globalState[AGENT_SERVER_KEY] ?? null; +const getAgentSockets = () => { + globalState[AGENT_SOCKETS_KEY] ??= new Map(); + return globalState[AGENT_SOCKETS_KEY]; +}; +const clearAgentSockets = () => { + getAgentSockets().clear(); +}; + +const setServer = (server: AgentServer | null) => { + if (server) { + globalState[AGENT_SERVER_KEY] = server; + return; + } + + delete globalState[AGENT_SERVER_KEY]; +}; + +export const spawnLocalAgent = () => { + const wsUrl = `ws://localhost:3001`; + + const agentEntryPoint = path.join(process.cwd(), "app", "server", "modules", "agents", "local-agent.ts"); + + const localAgent = spawn("bun", ["run", agentEntryPoint], { + env: { + ...process.env, + ZEROBYTE_CONTROLLER_URL: wsUrl, + }, + stdio: ["ignore", "pipe", "pipe"], + }); + + localAgent.stdout?.on("data", (data: Buffer) => { + const line = data.toString().trim(); + if (line) logger.info(`[agent] ${line}`); + }); + + localAgent.stderr?.on("data", (data: Buffer) => { + const line = data.toString().trim(); + if (line) logger.error(`[agent] ${line}`); + }); + + localAgent.on("exit", (code, signal) => { + logger.info(`Agent process exited with code ${code} and signal ${signal}`); + }); +}; + +export const agentManager = { + start: () => { + const existingServer = getServer(); + if (existingServer) { + existingServer.stop(true); + setServer(null); + clearAgentSockets(); + } + + logger.info("Starting Agent Manager..."); + const server = Bun.serve({ + port: 3001, + fetch(req, srv) { + const upgraded = srv.upgrade(req, { data: { id: Bun.randomUUIDv7() } }); + if (upgraded) return undefined; + return new Response("Agent WebSocket endpoint", { status: 200 }); + }, + websocket: { + open: (ws) => logger.info(`WebSocket opened with id: ${ws.data.id}`), + message: (ws, data) => { + if (typeof data !== "string") { + logger.warn(`Ignoring non-text message from agent connection ${ws.data.id}`); + return; + } + + const parsed = parseAgentMessage(data); + + if (parsed === null) { + logger.warn(`Invalid JSON from agent connection ${ws.data.id}`); + return; + } + + if (!parsed.success) { + logger.warn(`Invalid agent message on connection ${ws.data.id}: ${parsed.error.message}`); + return; + } + + switch (parsed.data.type) { + case "agent.ready": { + ws.data.agentId = parsed.data.payload.agentId; + getAgentSockets().set(parsed.data.payload.agentId, ws); + logger.info(`Backup agent ${parsed.data.payload.agentId} is ready on connection ${ws.data.id}`); + break; + } + case "backup.started": { + logger.info( + `Backup started on agent ${ws.data.agentId ?? ws.data.id} for schedule ${parsed.data.payload.scheduleId}`, + ); + break; + } + } + }, + close: (ws) => { + if (ws.data.agentId && getAgentSockets().get(ws.data.agentId) === ws) { + getAgentSockets().delete(ws.data.agentId); + } + + logger.info(`WebSocket closed for agent ${ws.data.agentId ?? ws.data.id}`); + }, + }, + }); + setServer(server); + + logger.info(`Agent Manager listening on port ${server.port}`); + }, + sendBackup: (agentId: string, payload: BackupCommandPayload) => { + const agentSocket = getAgentSockets().get(agentId); + + if (!agentSocket) { + logger.warn(`Cannot send backup command. Agent ${agentId} is not connected.`); + return false; + } + + sendControllerMessage(agentSocket, createControllerMessage("backup", payload)); + logger.info(`Sent backup command to agent ${agentId} for schedule ${payload.scheduleId}`); + return true; + }, + stop: () => { + const server = getServer(); + if (!server) return; + + logger.info("Stopping Agent Manager..."); + server.stop(true); + setServer(null); + clearAgentSockets(); + }, +}; diff --git a/app/server/modules/agents/local-agent.ts b/app/server/modules/agents/local-agent.ts new file mode 100644 index 00000000..2e33f5f6 --- /dev/null +++ b/app/server/modules/agents/local-agent.ts @@ -0,0 +1,61 @@ +import { logger } from "~/server/utils/logger"; +import { createAgentMessage, parseControllerMessage, sendAgentMessage } from "./agent-protocol"; + +const controllerUrl = process.env.ZEROBYTE_CONTROLLER_URL; + +class Agent { + private ws: WebSocket | null = null; + + constructor(public id: string) { + this.connect(); + } + + private connect() { + if (!controllerUrl) { + throw new Error("Env variable ZEROBYTE_CONTROLLER_URL is not set"); + } + + this.ws = new WebSocket(controllerUrl); + this.ws.onopen = () => { + logger.info(`Agent ${this.id} connected to controller`); + + if (this.ws) { + sendAgentMessage(this.ws, createAgentMessage("agent.ready", { agentId: this.id })); + } + }; + + this.ws.onmessage = (event) => { + const parsed = parseControllerMessage(event.data); + + if (parsed === null) { + console.error(`Agent ${this.id} received invalid JSON`); + return; + } + + if (!parsed.success) { + console.error(`Agent ${this.id} received an invalid message: ${parsed.error.message}`); + return; + } + + switch (parsed.data.type) { + case "backup": + logger.info(`Agent ${this.id} starting backup for schedule ${parsed.data.payload.scheduleId}`); + if (this.ws) { + sendAgentMessage( + this.ws, + createAgentMessage("backup.started", { scheduleId: parsed.data.payload.scheduleId }), + ); + } + break; + } + }; + this.ws.onclose = () => { + logger.info(`Agent ${this.id} disconnected from controller`); + }; + this.ws.onerror = (error) => { + logger.error(`Agent ${this.id} encountered an error:`, error); + }; + } +} + +new Agent(Bun.randomUUIDv7()); diff --git a/app/server/modules/lifecycle/bootstrap.ts b/app/server/modules/lifecycle/bootstrap.ts index bb8f44ea..9d1e4ab9 100644 --- a/app/server/modules/lifecycle/bootstrap.ts +++ b/app/server/modules/lifecycle/bootstrap.ts @@ -1,4 +1,5 @@ import { runDbMigrations } from "../../db/db"; +import { agentManager, spawnLocalAgent } from "../agents/agents-manager"; import { runMigrations } from "./migrations"; import { startup } from "./startup"; @@ -8,6 +9,8 @@ const runBootstrap = async () => { await runDbMigrations(); await runMigrations(); await startup(); + agentManager.start(); + spawnLocalAgent(); }; export const bootstrapApplication = async () => {