feat: init agent/controller minimal

This commit is contained in:
Nicolas Meienberger 2026-03-10 19:42:40 +01:00
parent 95aadf6e73
commit 4991d3e2ba
4 changed files with 320 additions and 0 deletions

View file

@ -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<typeof backupCommandSchema>["payload"];
export type ControllerMessage = z.infer<typeof controllerMessageSchema>;
export type AgentMessage = z.infer<typeof agentMessageSchema>;
type Brand<TValue, TBrand extends string> = TValue & {
readonly __brand: TBrand;
};
type MessageSender = {
send(message: string): unknown;
};
export type ControllerWireMessage = Brand<string, "ControllerWireMessage">;
export type AgentWireMessage = Brand<string, "AgentWireMessage">;
type PayloadForMessage<TMessage extends { type: string; payload: unknown }, TType extends TMessage["type"]> = Extract<
TMessage,
{ type: TType }
>["payload"];
const parseJsonMessage = (data: string) => safeJsonParse<unknown>(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 = <TType extends ControllerMessage["type"]>(
type: TType,
payload: PayloadForMessage<ControllerMessage, TType>,
) =>
JSON.stringify(
controllerMessageSchema.parse({
type,
payload,
}),
) as ControllerWireMessage;
export const createAgentMessage = <TType extends AgentMessage["type"]>(
type: TType,
payload: PayloadForMessage<AgentMessage, TType>,
) =>
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<ControllerWireMessage>;

View file

@ -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<AgentConnectionData>;
type AgentServer = ReturnType<typeof Bun.serve<AgentConnectionData>>;
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<string, AgentSocket>;
};
const getServer = () => globalState[AGENT_SERVER_KEY] ?? null;
const getAgentSockets = () => {
globalState[AGENT_SOCKETS_KEY] ??= new Map<string, AgentSocket>();
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<AgentConnectionData>({
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();
},
};

View file

@ -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());

View file

@ -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 () => {