feat: init agent/controller minimal
This commit is contained in:
parent
95aadf6e73
commit
4991d3e2ba
4 changed files with 320 additions and 0 deletions
98
app/server/modules/agents/agent-protocol.ts
Normal file
98
app/server/modules/agents/agent-protocol.ts
Normal 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>;
|
||||||
158
app/server/modules/agents/agents-manager.ts
Normal file
158
app/server/modules/agents/agents-manager.ts
Normal 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();
|
||||||
|
},
|
||||||
|
};
|
||||||
61
app/server/modules/agents/local-agent.ts
Normal file
61
app/server/modules/agents/local-agent.ts
Normal 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());
|
||||||
|
|
@ -1,4 +1,5 @@
|
||||||
import { runDbMigrations } from "../../db/db";
|
import { runDbMigrations } from "../../db/db";
|
||||||
|
import { agentManager, spawnLocalAgent } from "../agents/agents-manager";
|
||||||
import { runMigrations } from "./migrations";
|
import { runMigrations } from "./migrations";
|
||||||
import { startup } from "./startup";
|
import { startup } from "./startup";
|
||||||
|
|
||||||
|
|
@ -8,6 +9,8 @@ const runBootstrap = async () => {
|
||||||
await runDbMigrations();
|
await runDbMigrations();
|
||||||
await runMigrations();
|
await runMigrations();
|
||||||
await startup();
|
await startup();
|
||||||
|
agentManager.start();
|
||||||
|
spawnLocalAgent();
|
||||||
};
|
};
|
||||||
|
|
||||||
export const bootstrapApplication = async () => {
|
export const bootstrapApplication = async () => {
|
||||||
|
|
|
||||||
Loading…
Reference in a new issue