refactor: effect
This commit is contained in:
Nicolas Meienberger
parent
904bbfc473
commit
52b58e2f93
8 changed files with 474 additions and 168 deletions
|
|
@ -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<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;
|
||||
export const spawnLocalAgent = async () => {
|
||||
const previousAgent = (globalThis as Record<string, unknown>).__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<string, unknown>).__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<Map<string, ControllerAgentSession>>(new Map()));
|
||||
let runtimeScope: Scope.CloseableScope | null = null;
|
||||
|
||||
const getSessions = () => Effect.runSync(Ref.get(sessionsRef));
|
||||
const setSessions = (sessions: Map<string, ControllerAgentSession>) => {
|
||||
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<AgentConnectionData>({
|
||||
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<AgentConnectionData>({
|
||||
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<string, unknown>).__agentManager as
|
||||
| ReturnType<typeof createAgentManagerRuntime>
|
||||
| undefined;
|
||||
|
||||
previous?.stop();
|
||||
|
||||
export const agentManager = createAgentManagerRuntime();
|
||||
(globalThis as Record<string, unknown>).__agentManager = agentManager;
|
||||
|
|
|
|||
157
app/server/modules/agents/controller-agent-session.ts
Normal file
157
app/server/modules/agents/controller-agent-session.ts
Normal file
|
|
@ -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<AgentConnectionData>;
|
||||
|
||||
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<ControllerWireMessage>(64));
|
||||
const state = Effect.runSync(
|
||||
Ref.make<SessionState>({
|
||||
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(() => {});
|
||||
},
|
||||
};
|
||||
};
|
||||
|
|
@ -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"
|
||||
|
|
|
|||
116
apps/agent/src/controller-session.ts
Normal file
116
apps/agent/src/controller-session.ts
Normal file
|
|
@ -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<AgentWireMessage>(64));
|
||||
const inboundQueue = Effect.runSync(Queue.bounded<ControllerWireMessage>(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(() => {});
|
||||
},
|
||||
};
|
||||
};
|
||||
|
|
@ -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) => {
|
||||
|
|
|
|||
2
bun.lock
2
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",
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
|
|
|
|||
|
|
@ -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<typeof backupCommandSchema>["payload"];
|
||||
export type ControllerMessage = z.infer<typeof controllerMessageSchema>;
|
||||
|
|
|
|||
Loading…
Reference in a new issue