diff --git a/app/server/modules/agents/agents-manager.ts b/app/server/modules/agents/agents-manager.ts index 260ac0e5..c5f87fef 100644 --- a/app/server/modules/agents/agents-manager.ts +++ b/app/server/modules/agents/agents-manager.ts @@ -1,13 +1,14 @@ -import { spawn } from "node:child_process"; +import { type ChildProcess, spawn } from "node:child_process"; import path from "node:path"; -import { - createControllerMessage, - parseAgentMessage, - sendControllerMessage, - type BackupCommandPayload, -} from "@zerobyte/contracts/agent-protocol"; -import { logger } from "@zerobyte/core/utils"; +import { Effect, Exit, Ref, Scope } from "effect"; +import { logger } from "@zerobyte/core/node"; +import { config } from "../../core/config"; import { validateAgentToken, deriveLocalAgentToken } from "./agent-tokens"; +import { + createControllerAgentSession, + type BackupDispatchPayload, + type ControllerAgentSession, +} from "./controller-agent-session"; type AgentConnectionData = { id: string; @@ -16,39 +17,17 @@ type AgentConnectionData = { agentName: 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; +export const spawnLocalAgent = async () => { + const previousAgent = (globalThis as Record).__localAgent as ChildProcess | undefined; + if (previousAgent) { + previousAgent.kill(); } - delete globalState[AGENT_SERVER_KEY]; -}; - -export const spawnLocalAgent = async () => { const agentEntryPoint = path.join(process.cwd(), "apps", "agent", "src", "index.ts"); const agentToken = await deriveLocalAgentToken(); + const args = config.__prod__ ? ["run", agentEntryPoint] : ["run", "--watch", agentEntryPoint]; - const localAgent = spawn("bun", ["run", agentEntryPoint], { + const localAgent = spawn("bun", args, { env: { PATH: process.env.PATH, ZEROBYTE_CONTROLLER_URL: "ws://localhost:3001", @@ -57,6 +36,8 @@ export const spawnLocalAgent = async () => { stdio: ["ignore", "pipe", "pipe"], }); + (globalThis as Record).__localAgent = localAgent; + localAgent.stdout?.on("data", (data: Buffer) => { const line = data.toString().trim(); if (line) logger.info(`[agent] ${line}`); @@ -72,108 +53,161 @@ export const spawnLocalAgent = async () => { }); }; -export const agentManager = { - start: () => { - const existingServer = getServer(); - if (existingServer) { - existingServer.stop(true); - setServer(null); - clearAgentSockets(); +const createAgentManagerRuntime = () => { + const sessionsRef = Effect.runSync(Ref.make>(new Map())); + let runtimeScope: Scope.CloseableScope | null = null; + + const getSessions = () => Effect.runSync(Ref.get(sessionsRef)); + const setSessions = (sessions: Map) => { + Effect.runSync(Ref.set(sessionsRef, sessions)); + }; + + const closeAllSessions = () => { + const sessions = getSessions(); + for (const session of sessions.values()) { + session.close(); + } + setSessions(new Map()); + }; + + const getSession = (agentId: string) => getSessions().get(agentId); + + const setSession = (agentId: string, session: ControllerAgentSession) => { + const existingSession = getSession(agentId); + if (existingSession) { + existingSession.close(); + } + + const nextSessions = new Map(getSessions()); + nextSessions.set(agentId, session); + setSessions(nextSessions); + }; + + const removeSession = (agentId: string, connectionId: string) => { + const session = getSession(agentId); + if (!session || session.connectionId !== connectionId) { + return; + } + + session.close(); + const nextSessions = new Map(getSessions()); + nextSessions.delete(agentId); + setSessions(nextSessions); + }; + + 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)); + 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(); + 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 server = Bun.serve({ - port: 3001, - async fetch(req, srv) { - const url = new URL(req.url); - const token = url.searchParams.get("token"); + const scope = Effect.runSync(Scope.make()); - 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) => { - getAgentSockets().set(ws.data.agentId, 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 parsed = parseAgentMessage(data); - - if (parsed === null) { - logger.warn(`Invalid JSON from agent ${ws.data.agentId}`); - return; - } - - if (!parsed.success) { - logger.warn(`Invalid agent message from ${ws.data.agentId}: ${parsed.error.message}`); - return; - } - - switch (parsed.data.type) { - case "agent.ready": { - logger.info(`Agent "${ws.data.agentName}" (${ws.data.agentId}) is ready`); - break; - } - case "backup.started": { - logger.info(`Backup started on agent ${ws.data.agentId} for schedule ${parsed.data.payload.scheduleId}`); - break; - } - } - }, - close: (ws) => { - if (getAgentSockets().get(ws.data.agentId) === ws) { - getAgentSockets().delete(ws.data.agentId); - } - - logger.info(`Agent "${ws.data.agentName}" (${ws.data.agentId}) disconnected`); - }, - }, - }); - - 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; + 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; } + }; - 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; + return { + start, + sendBackup: (agentId: string, payload: BackupDispatchPayload) => { + const session = getSession(agentId); - logger.info("Stopping Agent Manager..."); - server.stop(true); - setServer(null); - clearAgentSockets(); - }, + if (!session) { + logger.warn(`Cannot send backup command. Agent ${agentId} is not connected.`); + return false; + } + + const jobId = session.sendBackup(payload); + logger.info(`Sent backup command ${jobId} to agent ${agentId} for schedule ${payload.scheduleId}`); + return true; + }, + stop, + }; }; + +const previous = (globalThis as Record).__agentManager as + | ReturnType + | undefined; + +previous?.stop(); + +export const agentManager = createAgentManagerRuntime(); +(globalThis as Record).__agentManager = agentManager; 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..5e7cc723 --- /dev/null +++ b/app/server/modules/agents/controller-agent-session.ts @@ -0,0 +1,157 @@ +import { Effect, Fiber, Queue, Ref } from "effect"; +import { + createControllerMessage, + parseAgentMessage, + type AgentMessage, + type ControllerWireMessage, +} from "@zerobyte/contracts/agent-protocol"; +import { logger } from "@zerobyte/core/node"; + +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; +}; + +const toMessage = (error: unknown) => { + if (error instanceof Error) { + return error.message; + } + + return String(error); +}; + +export type BackupDispatchPayload = { + scheduleId: string; + jobId?: string; +}; + +export type ControllerAgentSession = { + readonly connectionId: string; + handleMessage: (data: string) => void; + sendBackup: (payload: BackupDispatchPayload) => string; + isReady: () => boolean; + close: () => void; +}; + +export const createControllerAgentSession = (socket: AgentSocket): ControllerAgentSession => { + const outboundQueue = Effect.runSync(Queue.bounded(64)); + 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 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) => { + switch (message.type) { + case "agent.ready": { + updateState((current) => ({ ...current, isReady: true, lastSeenAt: Date.now() })); + logger.info(`Agent "${socket.data.agentName}" (${socket.data.agentId}) is ready`); + break; + } + case "backup.started": { + updateState((current) => ({ ...current, lastSeenAt: Date.now() })); + logger.info( + `Backup ${message.payload.jobId} started on agent ${socket.data.agentId} for schedule ${message.payload.scheduleId}`, + ); + break; + } + case "heartbeat.pong": { + updateState((current) => ({ + ...current, + lastSeenAt: Date.now(), + 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) => { + const jobId = payload.jobId ?? Bun.randomUUIDv7(); + offerOutbound( + createControllerMessage("backup.run", { + jobId, + scheduleId: payload.scheduleId, + }), + ); + return jobId; + }, + isReady: () => Effect.runSync(Ref.get(state)).isReady, + close: () => { + updateState((current) => ({ ...current, isReady: false })); + void Effect.runPromise(Fiber.interrupt(writerFiber)).catch(() => {}); + void Effect.runPromise(Fiber.interrupt(heartbeatFiber)).catch(() => {}); + void Effect.runPromise(Queue.shutdown(outboundQueue)).catch(() => {}); + }, + }; +}; diff --git a/apps/agent/package.json b/apps/agent/package.json index a92c4c2f..3751fd6a 100644 --- a/apps/agent/package.json +++ b/apps/agent/package.json @@ -5,7 +5,8 @@ "module": "index.ts", "dependencies": { "@zerobyte/contracts": "workspace:*", - "@zerobyte/core": "workspace:*" + "@zerobyte/core": "workspace:*", + "effect": "^3.18.4" }, "devDependencies": { "@types/bun": "latest" diff --git a/apps/agent/src/controller-session.ts b/apps/agent/src/controller-session.ts new file mode 100644 index 00000000..50c8e1a3 --- /dev/null +++ b/apps/agent/src/controller-session.ts @@ -0,0 +1,116 @@ +import { Effect, Fiber, Queue } from "effect"; +import { + createAgentMessage, + parseControllerMessage, + type AgentWireMessage, + type ControllerWireMessage, +} from "@zerobyte/contracts/agent-protocol"; +import { logger } from "@zerobyte/core/node"; + +const toMessage = (error: unknown) => { + if (error instanceof Error) { + return error.message; + } + + return String(error); +}; + +export type ControllerSession = { + onOpen: () => void; + onMessage: (data: unknown) => void; + close: () => void; +}; + +export const createControllerSession = (ws: WebSocket): ControllerSession => { + const outboundQueue = Effect.runSync(Queue.bounded(64)); + const inboundQueue = Effect.runSync(Queue.bounded(64)); + + const offerOutbound = (message: AgentWireMessage) => { + void Effect.runPromise(Queue.offer(outboundQueue, message)).catch((error) => { + logger.error(`Failed to queue outbound controller message: ${toMessage(error)}`); + }); + }; + + const offerInbound = (message: ControllerWireMessage) => { + void Effect.runPromise(Queue.offer(inboundQueue, message)).catch((error) => { + logger.error(`Failed to queue inbound controller message: ${toMessage(error)}`); + }); + }; + + const writerFiber = Effect.runFork( + Effect.forever( + Effect.gen(function* () { + const message = yield* Queue.take(outboundQueue); + yield* Effect.sync(() => { + try { + ws.send(message); + } catch (error) { + logger.error(`Failed to send controller message: ${toMessage(error)}`); + } + }); + }), + ), + ); + + const processorFiber = Effect.runFork( + Effect.forever( + Effect.gen(function* () { + const data = yield* Queue.take(inboundQueue); + const parsed = parseControllerMessage(data); + + if (parsed === null) { + logger.warn("Agent received invalid JSON"); + return; + } + + if (!parsed.success) { + logger.warn(`Agent received an invalid message: ${parsed.error.message}`); + return; + } + + switch (parsed.data.type) { + case "backup.run": { + logger.info(`Starting backup ${parsed.data.payload.jobId} for schedule ${parsed.data.payload.scheduleId}`); + yield* Queue.offer( + outboundQueue, + createAgentMessage("backup.started", { + jobId: parsed.data.payload.jobId, + scheduleId: parsed.data.payload.scheduleId, + }), + ); + break; + } + case "heartbeat.ping": { + yield* Queue.offer( + outboundQueue, + createAgentMessage("heartbeat.pong", { + sentAt: parsed.data.payload.sentAt, + }), + ); + break; + } + } + }), + ), + ); + + return { + onOpen: () => { + offerOutbound(createAgentMessage("agent.ready", { agentId: "" })); + }, + onMessage: (data) => { + if (typeof data !== "string") { + logger.warn("Agent received a non-text message"); + return; + } + + offerInbound(data as ControllerWireMessage); + }, + close: () => { + void Effect.runPromise(Fiber.interrupt(writerFiber)).catch(() => {}); + void Effect.runPromise(Fiber.interrupt(processorFiber)).catch(() => {}); + void Effect.runPromise(Queue.shutdown(outboundQueue)).catch(() => {}); + void Effect.runPromise(Queue.shutdown(inboundQueue)).catch(() => {}); + }, + }; +}; diff --git a/apps/agent/src/index.ts b/apps/agent/src/index.ts index 1bae4a99..1e5b004d 100644 --- a/apps/agent/src/index.ts +++ b/apps/agent/src/index.ts @@ -1,11 +1,12 @@ -import { createAgentMessage, parseControllerMessage, sendAgentMessage } from "@zerobyte/contracts/agent-protocol"; -import { logger } from "@zerobyte/core/utils"; +import { logger } from "@zerobyte/core/node"; +import { createControllerSession, type ControllerSession } from "./controller-session"; const controllerUrl = process.env.ZEROBYTE_CONTROLLER_URL; const agentToken = process.env.ZEROBYTE_AGENT_TOKEN; class Agent { private ws: WebSocket | null = null; + private controllerSession: ControllerSession | null = null; connect() { if (!controllerUrl) { @@ -20,40 +21,20 @@ class Agent { url.searchParams.set("token", agentToken); this.ws = new WebSocket(url.toString()); + this.controllerSession = createControllerSession(this.ws); + this.ws.onopen = () => { logger.info("Agent connected to controller"); - - if (this.ws) { - sendAgentMessage(this.ws, createAgentMessage("agent.ready", { agentId: "" })); - } + this.controllerSession?.onOpen(); }; this.ws.onmessage = (event) => { - const parsed = parseControllerMessage(event.data); - - if (parsed === null) { - console.error("Agent received invalid JSON"); - return; - } - - if (!parsed.success) { - console.error(`Agent received an invalid message: ${parsed.error.message}`); - return; - } - - switch (parsed.data.type) { - case "backup": - logger.info(`Starting backup for schedule ${parsed.data.payload.scheduleId}`); - if (this.ws) { - sendAgentMessage( - this.ws, - createAgentMessage("backup.started", { scheduleId: parsed.data.payload.scheduleId }), - ); - } - break; - } + this.controllerSession?.onMessage(event.data); }; this.ws.onclose = () => { + this.controllerSession?.close(); + this.controllerSession = null; + this.ws = null; logger.info("Agent disconnected from controller"); }; this.ws.onerror = (error) => { diff --git a/bun.lock b/bun.lock index a9f2df8f..900f8584 100644 --- a/bun.lock +++ b/bun.lock @@ -47,6 +47,7 @@ "dither-plugin": "^1.1.1", "dotenv": "^17.3.1", "drizzle-orm": "^1.0.0-beta.16-ea816b6", + "effect": "^3.18.4", "es-toolkit": "^1.45.1", "hono": "^4.12.7", "hono-openapi": "^1.3.0", @@ -116,6 +117,7 @@ "dependencies": { "@zerobyte/contracts": "workspace:*", "@zerobyte/core": "workspace:*", + "effect": "^3.18.4", }, "devDependencies": { "@types/bun": "latest", diff --git a/package.json b/package.json index ae94ddaa..5fdc1081 100644 --- a/package.json +++ b/package.json @@ -72,6 +72,7 @@ "dither-plugin": "^1.1.1", "dotenv": "^17.3.1", "drizzle-orm": "^1.0.0-beta.16-ea816b6", + "effect": "^3.18.4", "es-toolkit": "^1.45.1", "hono": "^4.12.7", "hono-openapi": "^1.3.0", diff --git a/packages/contracts/src/agent-protocol.ts b/packages/contracts/src/agent-protocol.ts index ec01483a..231b9e99 100644 --- a/packages/contracts/src/agent-protocol.ts +++ b/packages/contracts/src/agent-protocol.ts @@ -3,8 +3,15 @@ import { safeJsonParse } from "@zerobyte/core/utils"; const backupCommandSchema = z .object({ - type: z.literal("backup"), - payload: z.object({ scheduleId: z.string() }), + type: z.literal("backup.run"), + payload: z.object({ jobId: z.string(), scheduleId: z.string() }), + }) + .strict(); + +const heartbeatPingSchema = z + .object({ + type: z.literal("heartbeat.ping"), + payload: z.object({ sentAt: z.number() }), }) .strict(); @@ -18,12 +25,19 @@ const agentReadySchema = z const backupStartedSchema = z .object({ type: z.literal("backup.started"), - payload: z.object({ scheduleId: z.string() }), + payload: z.object({ jobId: z.string(), scheduleId: z.string() }), }) .strict(); -const controllerMessageSchema = z.discriminatedUnion("type", [backupCommandSchema]); -const agentMessageSchema = z.discriminatedUnion("type", [agentReadySchema, backupStartedSchema]); +const heartbeatPongSchema = z + .object({ + type: z.literal("heartbeat.pong"), + payload: z.object({ sentAt: z.number() }), + }) + .strict(); + +const controllerMessageSchema = z.discriminatedUnion("type", [backupCommandSchema, heartbeatPingSchema]); +const agentMessageSchema = z.discriminatedUnion("type", [agentReadySchema, backupStartedSchema, heartbeatPongSchema]); export type BackupCommandPayload = z.infer["payload"]; export type ControllerMessage = z.infer;