danvics/zerobyte
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions 1

refactor(agents): make the controller-agent effectful

Browse source
This commit is contained in:
Nicolas Meienberger 2026-04-08 23:23:01 +02:00
parent 1eda489167
commit 6cda41df7c
No known key found for this signature in database
3 changed files with 283 additions and 220 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

93
app/server/modules/agents/__tests__/controller-agent-session.test.ts
View file

@ -1,3 +1,4 @@
import { Effect, Exit, Scope } from "effect";
import { expect, test, vi } from "vitest"; import { expect, test, vi } from "vitest";
import { fromPartial } from "@total-typescript/shoehorn"; import { fromPartial } from "@total-typescript/shoehorn";
import { createAgentMessage } from "@zerobyte/contracts/agent-protocol"; import { createAgentMessage } from "@zerobyte/contracts/agent-protocol";
@ -10,20 +11,40 @@ const createSocket = () => {
}); });
}; };
const createSession = (handlers: Parameters<typeof createControllerAgentSession>[1] = {}) => {
const scope = Effect.runSync(Scope.make());
try {
const session = Effect.runSync(Scope.extend(createControllerAgentSession(createSocket(), handlers), scope));
return {
session,
close: () => {
Effect.runSync(Scope.close(scope, Exit.succeed(undefined)));
},
};
} catch (error) {
Effect.runSync(Scope.close(scope, Exit.fail(error)));
throw error;
}
};
test("close emits a synthetic backup.cancelled for a started backup", () => { test("close emits a synthetic backup.cancelled for a started backup", () => {
const onBackupCancelled = vi.fn(); const onBackupCancelled = vi.fn();
const session = createControllerAgentSession(createSocket(), { const { session, close } = createSession({
onBackupCancelled, onBackupCancelled,
}); });
session.handleMessage( Effect.runSync(
createAgentMessage("backup.started", { session.handleMessage(
jobId: "job-1", createAgentMessage("backup.started", {
scheduleId: "schedule-1", jobId: "job-1",
}), scheduleId: "schedule-1",
}),
),
); );
session.close(); close();
expect(onBackupCancelled).toHaveBeenCalledTimes(1); expect(onBackupCancelled).toHaveBeenCalledTimes(1);
expect(onBackupCancelled).toHaveBeenCalledWith({ expect(onBackupCancelled).toHaveBeenCalledWith({
@ -71,47 +92,51 @@ test.each([
}, },
])("close does not emit an extra synthetic backup.cancelled after $name", (testCase) => { ])("close does not emit an extra synthetic backup.cancelled after $name", (testCase) => {
const onBackupCancelled = vi.fn(); const onBackupCancelled = vi.fn();
const session = createControllerAgentSession(createSocket(), { const { session, close } = createSession({
onBackupCancelled, onBackupCancelled,
}); });
session.handleMessage( Effect.runSync(
createAgentMessage("backup.started", { session.handleMessage(
jobId: testCase.jobId, createAgentMessage("backup.started", {
scheduleId: testCase.scheduleId, jobId: testCase.jobId,
}), scheduleId: testCase.scheduleId,
}),
),
); );
session.handleMessage(testCase.terminalMessage); Effect.runSync(session.handleMessage(testCase.terminalMessage));
session.close(); close();
expect(onBackupCancelled).toHaveBeenCalledTimes(testCase.expectedCancelledCalls); expect(onBackupCancelled).toHaveBeenCalledTimes(testCase.expectedCancelledCalls);
}); });
test("close emits a synthetic backup.cancelled for a queued backup", () => { test("close emits a synthetic backup.cancelled for a queued backup", () => {
const onBackupCancelled = vi.fn(); const onBackupCancelled = vi.fn();
const session = createControllerAgentSession(createSocket(), { const { session, close } = createSession({
onBackupCancelled, onBackupCancelled,
}); });
session.sendBackup({ Effect.runSync(
jobId: "job-queued", session.sendBackup({
scheduleId: "schedule-queued", jobId: "job-queued",
organizationId: "org-1", scheduleId: "schedule-queued",
sourcePath: "/tmp/source", organizationId: "org-1",
repositoryConfig: { sourcePath: "/tmp/source",
backend: "local", repositoryConfig: {
path: "/tmp/repository", backend: "local",
}, path: "/tmp/repository",
options: {}, },
runtime: { options: {},
password: "password", runtime: {
cacheDir: "/tmp/cache", password: "password",
passFile: "/tmp/pass", cacheDir: "/tmp/cache",
defaultExcludes: [], passFile: "/tmp/pass",
}, defaultExcludes: [],
}); },
}),
);
session.close(); close();
expect(onBackupCancelled).toHaveBeenCalledTimes(1); expect(onBackupCancelled).toHaveBeenCalledTimes(1);
expect(onBackupCancelled).toHaveBeenCalledWith({ expect(onBackupCancelled).toHaveBeenCalledWith({

69
app/server/modules/agents/agents-manager.ts
View file

@ -1,7 +1,7 @@
import { type ChildProcess, spawn } from "node:child_process"; import { type ChildProcess, spawn } from "node:child_process";
import { existsSync } from "node:fs"; import { existsSync } from "node:fs";
import path from "node:path"; import path from "node:path";
import { Effect, Exit, Scope } from "effect"; import { Effect, Exit, Fiber, Scope } from "effect";
import { logger } from "@zerobyte/core/node"; import { logger } from "@zerobyte/core/node";
import type { import type {
BackupCancelPayload, BackupCancelPayload,
@ -45,6 +45,12 @@ type AgentRuntimeState = {
localAgent: ChildProcess | null; localAgent: ChildProcess | null;
}; };
type ControllerAgentSessionHandle = {
session: ControllerAgentSession;
runFiber: Fiber.RuntimeFiber<void, never>;
scope: Scope.CloseableScope;
};
type ProcessWithAgentRuntime = NodeJS.Process & { type ProcessWithAgentRuntime = NodeJS.Process & {
__zerobyteAgentRuntime?: AgentRuntimeState; __zerobyteAgentRuntime?: AgentRuntimeState;
}; };
@ -136,36 +142,59 @@ export const stopLocalAgent = async () => {
}; };
const createAgentManagerRuntime = () => { const createAgentManagerRuntime = () => {
let sessions = new Map<string, ControllerAgentSession>(); let sessions = new Map<string, ControllerAgentSessionHandle>();
let backupHandlers: AgentBackupEventHandlers = {}; let backupHandlers: AgentBackupEventHandlers = {};
let runtimeScope: Scope.CloseableScope | null = null; let runtimeScope: Scope.CloseableScope | null = null;
const closeAllSessions = () => { const closeSession = (sessionHandle: ControllerAgentSessionHandle) => {
for (const session of sessions.values()) { Effect.runSync(Fiber.interrupt(sessionHandle.runFiber));
session.close(); Effect.runSync(Scope.close(sessionHandle.scope, Exit.succeed(undefined)));
}
sessions = new Map();
}; };
const getSession = (agentId: string) => sessions.get(agentId); const closeAllSessions = () => {
const currentSessions = sessions;
sessions = new Map();
for (const sessionHandle of currentSessions.values()) {
closeSession(sessionHandle);
}
};
const setSession = (agentId: string, session: ControllerAgentSession) => { const getSessionHandle = (agentId: string) => sessions.get(agentId);
const existingSession = getSession(agentId);
const getSession = (agentId: string) => getSessionHandle(agentId)?.session;
const createSession = (ws: Bun.ServerWebSocket<AgentConnectionData>) => {
// Manual scope management because we are out of Effect
const scope = Effect.runSync(Scope.make());
try {
const session = Effect.runSync(Scope.extend(createControllerAgentSession(ws, createSessionHandlers(ws)), scope));
const runFiber = Effect.runFork(Scope.extend(session.run, scope));
return { session, runFiber, scope };
} catch (error) {
Effect.runSync(Scope.close(scope, Exit.fail(error)));
throw error;
}
};
const setSession = (agentId: string, sessionHandle: ControllerAgentSessionHandle) => {
const existingSession = getSessionHandle(agentId);
if (existingSession) { if (existingSession) {
existingSession.close(); closeSession(existingSession);
} }
sessions.set(agentId, session); sessions.set(agentId, sessionHandle);
}; };
const removeSession = (agentId: string, connectionId: string) => { const removeSession = (agentId: string, connectionId: string) => {
const session = getSession(agentId); const sessionHandle = getSessionHandle(agentId);
if (!session || session.connectionId !== connectionId) { if (!sessionHandle || sessionHandle.session.connectionId !== connectionId) {
return; return;
} }
session.close();
sessions.delete(agentId); sessions.delete(agentId);
closeSession(sessionHandle);
}; };
const createSessionHandlers = (ws: Bun.ServerWebSocket<AgentConnectionData>) => { const createSessionHandlers = (ws: Bun.ServerWebSocket<AgentConnectionData>) => {
@ -221,7 +250,7 @@ const createAgentManagerRuntime = () => {
}, },
websocket: { websocket: {
open: (ws) => { open: (ws) => {
setSession(ws.data.agentId, createControllerAgentSession(ws, createSessionHandlers(ws))); setSession(ws.data.agentId, createSession(ws));
logger.info(`Agent "${ws.data.agentName}" (${ws.data.agentId}) connected on ${ws.data.id}`); logger.info(`Agent "${ws.data.agentName}" (${ws.data.agentId}) connected on ${ws.data.id}`);
}, },
message: (ws, data) => { message: (ws, data) => {
@ -236,7 +265,7 @@ const createAgentManagerRuntime = () => {
return; return;
} }
session.handleMessage(data); Effect.runSync(session.handleMessage(data));
}, },
close: (ws) => { close: (ws) => {
removeSession(ws.data.agentId, ws.data.id); removeSession(ws.data.agentId, ws.data.id);
@ -291,12 +320,12 @@ const createAgentManagerRuntime = () => {
return false; return false;
} }
if (!session.isReady()) { if (!Effect.runSync(session.isReady())) {
logger.warn(`Cannot send backup command. Agent ${agentId} is not ready.`); logger.warn(`Cannot send backup command. Agent ${agentId} is not ready.`);
return false; return false;
} }
session.sendBackup(payload); Effect.runSync(session.sendBackup(payload));
logger.info(`Sent backup command ${payload.jobId} to agent ${agentId} for schedule ${payload.scheduleId}`); logger.info(`Sent backup command ${payload.jobId} to agent ${agentId} for schedule ${payload.scheduleId}`);
return true; return true;
}, },
@ -308,7 +337,7 @@ const createAgentManagerRuntime = () => {
return false; return false;
} }
session.sendBackupCancel(payload); Effect.runSync(session.sendBackupCancel(payload));
logger.info(`Sent backup cancel for command ${payload.jobId} to agent ${agentId}`); logger.info(`Sent backup cancel for command ${payload.jobId} to agent ${agentId}`);
return true; return true;
}, },

341
app/server/modules/agents/controller-agent-session.ts
View file

@ -1,4 +1,4 @@
import { Effect, Fiber, Queue, Ref } from "effect"; import { Effect, Queue, Ref, type Scope } from "effect";
import { import {
createControllerMessage, createControllerMessage,
parseAgentMessage, parseAgentMessage,
@ -45,190 +45,199 @@ type ControllerAgentSessionHandlers = {
export type ControllerAgentSession = { export type ControllerAgentSession = {
readonly connectionId: string; readonly connectionId: string;
handleMessage: (data: string) => void; handleMessage: (data: string) => Effect.Effect<void>;
sendBackup: (payload: BackupRunPayload) => void; sendBackup: (payload: BackupRunPayload) => Effect.Effect<void>;
sendBackupCancel: (payload: BackupCancelPayload) => void; sendBackupCancel: (payload: BackupCancelPayload) => Effect.Effect<void>;
isReady: () => boolean; isReady: () => Effect.Effect<boolean>;
close: () => void; run: Effect.Effect<void, never, Scope.Scope>;
}; };
export const createControllerAgentSession = ( export const createControllerAgentSession = (
socket: AgentSocket, socket: AgentSocket,
handlers: ControllerAgentSessionHandlers = {}, handlers: ControllerAgentSessionHandlers = {},
): ControllerAgentSession => { ): Effect.Effect<ControllerAgentSession, never, Scope.Scope> =>
const outboundQueue = Effect.runSync(Queue.bounded<ControllerWireMessage>(64)); Effect.gen(function* () {
const trackedBackupJobs = Effect.runSync(Ref.make<Map<string, TrackedBackupJob>>(new Map())); const outboundQueue = yield* Queue.bounded<ControllerWireMessage>(64);
const state = Effect.runSync( const trackedBackupJobs = yield* Ref.make<Map<string, TrackedBackupJob>>(new Map());
Ref.make<SessionState>({ const state = yield* Ref.make<SessionState>({
isReady: false, isReady: false,
lastSeenAt: null, lastSeenAt: null,
lastPongAt: 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) => { const offerOutbound = (message: ControllerWireMessage) =>
Effect.runSync(Ref.update(state, update)); Queue.offer(outboundQueue, message).pipe(
}; Effect.catchAllCause((cause) =>
Effect.sync(() => {
logger.error(`Failed to queue outbound message for agent ${socket.data.agentId}: ${toMessage(cause)}`);
}),
),
);
const setTrackedBackupJob = (jobId: string, trackedBackupJob: TrackedBackupJob) => { const updateState = (update: (current: SessionState) => SessionState) => Ref.update(state, update);
Effect.runSync(
Ref.update(trackedBackupJobs, (current) => { const setTrackedBackupJob = (jobId: string, trackedBackupJob: TrackedBackupJob) => {
return Ref.update(trackedBackupJobs, (current) => {
const next = new Map(current); const next = new Map(current);
next.set(jobId, trackedBackupJob); next.set(jobId, trackedBackupJob);
return next; return next;
}), });
); };
};
const deleteTrackedBackupJob = (jobId: string) => { const deleteTrackedBackupJob = (jobId: string) => {
Effect.runSync( return Ref.update(trackedBackupJobs, (current) => {
Ref.update(trackedBackupJobs, (current) => {
const next = new Map(current); const next = new Map(current);
next.delete(jobId); next.delete(jobId);
return next; return next;
}),
);
};
const takeTrackedBackupJobs = () => {
return Effect.runSync(
Ref.modify(trackedBackupJobs, (current) => [current, new Map<string, TrackedBackupJob>()] as const),
);
};
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) => {
updateState((current) => ({ ...current, lastSeenAt: Date.now() }));
switch (message.type) {
case "agent.ready": {
updateState((current) => ({ ...current, isReady: true }));
logger.info(`Agent "${socket.data.agentName}" (${socket.data.agentId}) is ready`);
break;
}
case "backup.started": {
setTrackedBackupJob(message.payload.jobId, {
scheduleId: message.payload.scheduleId,
state: "active",
});
logger.info(
`Backup ${message.payload.jobId} started on agent ${socket.data.agentId} for schedule ${message.payload.scheduleId}`,
);
handlers.onBackupStarted?.(message.payload);
break;
}
case "backup.progress": {
handlers.onBackupProgress?.(message.payload);
break;
}
case "backup.completed": {
deleteTrackedBackupJob(message.payload.jobId);
handlers.onBackupCompleted?.(message.payload);
break;
}
case "backup.failed": {
deleteTrackedBackupJob(message.payload.jobId);
handlers.onBackupFailed?.(message.payload);
break;
}
case "backup.cancelled": {
deleteTrackedBackupJob(message.payload.jobId);
handlers.onBackupCancelled?.(message.payload);
break;
}
case "heartbeat.pong": {
updateState((current) => ({ ...current, 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) => {
setTrackedBackupJob(payload.jobId, {
scheduleId: payload.scheduleId,
state: "pending",
}); });
offerOutbound(createControllerMessage("backup.run", payload)); };
},
sendBackupCancel: (payload) => {
offerOutbound(createControllerMessage("backup.cancel", payload));
},
isReady: () => Effect.runSync(Ref.get(state)).isReady,
close: () => {
updateState((current) => ({ ...current, isReady: false }));
const trackedJobs = takeTrackedBackupJobs();
for (const [jobId, trackedJob] of trackedJobs) {
let message: string;
if (trackedJob.state === "pending") {
message =
"The connection to the backup agent was lost before this backup started. Restart the backup to ensure it completes.";
} else {
message =
"The connection to the backup agent was lost while this backup was running. Restart the backup to ensure it completes.";
}
handlers.onBackupCancelled?.({ const takeTrackedBackupJobs = Ref.modify(
jobId, trackedBackupJobs,
scheduleId: trackedJob.scheduleId, (current) => [current, new Map<string, TrackedBackupJob>()] as const,
message, );
const releaseSession = Effect.gen(function* () {
yield* updateState((current) => ({ ...current, isReady: false }));
const trackedJobs = yield* takeTrackedBackupJobs;
for (const [jobId, trackedJob] of trackedJobs) {
let message = "The connection to the backup agent was lost. Restart the backup to ensure it completes.";
yield* Effect.sync(() => {
handlers.onBackupCancelled?.({ jobId, scheduleId: trackedJob.scheduleId, message });
}); });
} }
void Effect.runPromise(Fiber.interrupt(writerFiber)).catch(() => {});
void Effect.runPromise(Fiber.interrupt(heartbeatFiber)).catch(() => {}); yield* Queue.shutdown(outboundQueue);
void Effect.runPromise(Queue.shutdown(outboundQueue)).catch(() => {}); });
},
}; yield* Effect.addFinalizer(() => releaseSession);
};
const run = Effect.gen(function* () {
yield* Effect.forkScoped(
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)}`,
);
}
});
}),
),
);
yield* Effect.forkScoped(
Effect.forever(
Effect.gen(function* () {
yield* Effect.sleep("15 seconds");
yield* Queue.offer(
outboundQueue,
createControllerMessage("heartbeat.ping", {
sentAt: Date.now(),
}),
);
}),
),
);
return yield* Effect.never;
});
const handleAgentMessage = (message: AgentMessage) =>
Effect.gen(function* () {
yield* updateState((current) => ({ ...current, lastSeenAt: Date.now() }));
switch (message.type) {
case "agent.ready": {
yield* updateState((current) => ({ ...current, isReady: true }));
yield* Effect.sync(() => {
logger.info(`Agent "${socket.data.agentName}" (${socket.data.agentId}) is ready`);
});
break;
}
case "backup.started": {
yield* setTrackedBackupJob(message.payload.jobId, {
scheduleId: message.payload.scheduleId,
state: "active",
});
yield* Effect.sync(() => {
logger.info(
`Backup ${message.payload.jobId} started on agent ${socket.data.agentId} for schedule ${message.payload.scheduleId}`,
);
handlers.onBackupStarted?.(message.payload);
});
break;
}
case "backup.progress": {
yield* Effect.sync(() => {
handlers.onBackupProgress?.(message.payload);
});
break;
}
case "backup.completed": {
yield* deleteTrackedBackupJob(message.payload.jobId);
yield* Effect.sync(() => {
handlers.onBackupCompleted?.(message.payload);
});
break;
}
case "backup.failed": {
yield* deleteTrackedBackupJob(message.payload.jobId);
yield* Effect.sync(() => {
handlers.onBackupFailed?.(message.payload);
});
break;
}
case "backup.cancelled": {
yield* deleteTrackedBackupJob(message.payload.jobId);
yield* Effect.sync(() => {
handlers.onBackupCancelled?.(message.payload);
});
break;
}
case "heartbeat.pong": {
yield* updateState((current) => ({ ...current, lastPongAt: message.payload.sentAt }));
break;
}
}
});
return {
connectionId: socket.data.id,
handleMessage: (data: string) => {
return Effect.gen(function* () {
const parsed = parseAgentMessage(data);
if (parsed === null) {
yield* Effect.sync(() => {
logger.warn(`Invalid JSON from agent ${socket.data.agentId}`);
});
return;
}
if (!parsed.success) {
yield* Effect.sync(() => {
logger.warn(`Invalid agent message from ${socket.data.agentId}: ${parsed.error.message}`);
});
return;
}
yield* handleAgentMessage(parsed.data);
});
},
sendBackup: (payload) => {
return Effect.gen(function* () {
yield* setTrackedBackupJob(payload.jobId, { scheduleId: payload.scheduleId, state: "pending" });
yield* offerOutbound(createControllerMessage("backup.run", payload));
});
},
sendBackupCancel: (payload) => offerOutbound(createControllerMessage("backup.cancel", payload)),
isReady: () => Ref.get(state).pipe(Effect.map((current) => current.isReady)),
run,
};
});