feat(controller): add agent manager and session handling
This commit is contained in:
Nicolas Meienberger
parent
f9027033e5
commit
1eda489167
4 changed files with 697 additions and 0 deletions
|
|
@ -0,0 +1,123 @@
|
|||
import { expect, test, vi } from "vitest";
|
||||
import { fromPartial } from "@total-typescript/shoehorn";
|
||||
import { createAgentMessage } from "@zerobyte/contracts/agent-protocol";
|
||||
import { createControllerAgentSession } from "../controller-agent-session";
|
||||
|
||||
const createSocket = () => {
|
||||
return fromPartial<Parameters<typeof createControllerAgentSession>[0]>({
|
||||
data: { id: "connection-1", agentId: "local", organizationId: null, agentName: "Local Agent" },
|
||||
send: vi.fn(),
|
||||
});
|
||||
};
|
||||
|
||||
test("close emits a synthetic backup.cancelled for a started backup", () => {
|
||||
const onBackupCancelled = vi.fn();
|
||||
const session = createControllerAgentSession(createSocket(), {
|
||||
onBackupCancelled,
|
||||
});
|
||||
|
||||
session.handleMessage(
|
||||
createAgentMessage("backup.started", {
|
||||
jobId: "job-1",
|
||||
scheduleId: "schedule-1",
|
||||
}),
|
||||
);
|
||||
|
||||
session.close();
|
||||
|
||||
expect(onBackupCancelled).toHaveBeenCalledTimes(1);
|
||||
expect(onBackupCancelled).toHaveBeenCalledWith({
|
||||
jobId: "job-1",
|
||||
scheduleId: "schedule-1",
|
||||
message:
|
||||
"The connection to the backup agent was lost while this backup was running. Restart the backup to ensure it completes.",
|
||||
});
|
||||
});
|
||||
|
||||
test.each([
|
||||
{
|
||||
name: "backup.completed",
|
||||
jobId: "job-1",
|
||||
scheduleId: "schedule-1",
|
||||
terminalMessage: createAgentMessage("backup.completed", {
|
||||
jobId: "job-1",
|
||||
scheduleId: "schedule-1",
|
||||
exitCode: 0,
|
||||
result: null,
|
||||
}),
|
||||
expectedCancelledCalls: 0,
|
||||
},
|
||||
{
|
||||
name: "backup.failed",
|
||||
jobId: "job-2",
|
||||
scheduleId: "schedule-2",
|
||||
terminalMessage: createAgentMessage("backup.failed", {
|
||||
jobId: "job-2",
|
||||
scheduleId: "schedule-2",
|
||||
error: "backup failed",
|
||||
}),
|
||||
expectedCancelledCalls: 0,
|
||||
},
|
||||
{
|
||||
name: "backup.cancelled",
|
||||
jobId: "job-3",
|
||||
scheduleId: "schedule-3",
|
||||
terminalMessage: createAgentMessage("backup.cancelled", {
|
||||
jobId: "job-3",
|
||||
scheduleId: "schedule-3",
|
||||
message: "Backup was cancelled",
|
||||
}),
|
||||
expectedCancelledCalls: 1,
|
||||
},
|
||||
])("close does not emit an extra synthetic backup.cancelled after $name", (testCase) => {
|
||||
const onBackupCancelled = vi.fn();
|
||||
const session = createControllerAgentSession(createSocket(), {
|
||||
onBackupCancelled,
|
||||
});
|
||||
|
||||
session.handleMessage(
|
||||
createAgentMessage("backup.started", {
|
||||
jobId: testCase.jobId,
|
||||
scheduleId: testCase.scheduleId,
|
||||
}),
|
||||
);
|
||||
session.handleMessage(testCase.terminalMessage);
|
||||
session.close();
|
||||
|
||||
expect(onBackupCancelled).toHaveBeenCalledTimes(testCase.expectedCancelledCalls);
|
||||
});
|
||||
|
||||
test("close emits a synthetic backup.cancelled for a queued backup", () => {
|
||||
const onBackupCancelled = vi.fn();
|
||||
const session = createControllerAgentSession(createSocket(), {
|
||||
onBackupCancelled,
|
||||
});
|
||||
|
||||
session.sendBackup({
|
||||
jobId: "job-queued",
|
||||
scheduleId: "schedule-queued",
|
||||
organizationId: "org-1",
|
||||
sourcePath: "/tmp/source",
|
||||
repositoryConfig: {
|
||||
backend: "local",
|
||||
path: "/tmp/repository",
|
||||
},
|
||||
options: {},
|
||||
runtime: {
|
||||
password: "password",
|
||||
cacheDir: "/tmp/cache",
|
||||
passFile: "/tmp/pass",
|
||||
defaultExcludes: [],
|
||||
},
|
||||
});
|
||||
|
||||
session.close();
|
||||
|
||||
expect(onBackupCancelled).toHaveBeenCalledTimes(1);
|
||||
expect(onBackupCancelled).toHaveBeenCalledWith({
|
||||
jobId: "job-queued",
|
||||
scheduleId: "schedule-queued",
|
||||
message:
|
||||
"The connection to the backup agent was lost before this backup started. Restart the backup to ensure it completes.",
|
||||
});
|
||||
});
|
||||
12
app/server/modules/agents/agent-tokens.ts
Normal file
12
app/server/modules/agents/agent-tokens.ts
Normal file
|
|
@ -0,0 +1,12 @@
|
|||
import { cryptoUtils } from "~/server/utils/crypto";
|
||||
|
||||
export const deriveLocalAgentToken = async () => {
|
||||
return cryptoUtils.deriveSecret("zerobyte:local-agent-token");
|
||||
};
|
||||
|
||||
export const validateAgentToken = async (token: string) => {
|
||||
const localToken = await deriveLocalAgentToken();
|
||||
if (token === localToken) {
|
||||
return { agentId: "local", organizationId: null, agentName: "local" };
|
||||
}
|
||||
};
|
||||
328
app/server/modules/agents/agents-manager.ts
Normal file
328
app/server/modules/agents/agents-manager.ts
Normal file
|
|
@ -0,0 +1,328 @@
|
|||
import { type ChildProcess, spawn } from "node:child_process";
|
||||
import { existsSync } from "node:fs";
|
||||
import path from "node:path";
|
||||
import { Effect, Exit, Scope } from "effect";
|
||||
import { logger } from "@zerobyte/core/node";
|
||||
import type {
|
||||
BackupCancelPayload,
|
||||
BackupCancelledPayload,
|
||||
BackupCompletedPayload,
|
||||
BackupFailedPayload,
|
||||
BackupProgressPayload,
|
||||
BackupRunPayload,
|
||||
BackupStartedPayload,
|
||||
} from "@zerobyte/contracts/agent-protocol";
|
||||
import { config } from "../../core/config";
|
||||
import { validateAgentToken, deriveLocalAgentToken } from "./agent-tokens";
|
||||
import {
|
||||
createControllerAgentSession,
|
||||
type AgentConnectionData,
|
||||
type ControllerAgentSession,
|
||||
} from "./controller-agent-session";
|
||||
|
||||
type AgentBackupEventContext = {
|
||||
agentId: string;
|
||||
agentName: string;
|
||||
payload:
|
||||
| BackupStartedPayload
|
||||
| BackupProgressPayload
|
||||
| BackupCompletedPayload
|
||||
| BackupFailedPayload
|
||||
| BackupCancelledPayload;
|
||||
};
|
||||
|
||||
export type AgentBackupEventHandlers = {
|
||||
onBackupStarted?: (context: AgentBackupEventContext & { payload: BackupStartedPayload }) => void;
|
||||
onBackupProgress?: (context: AgentBackupEventContext & { payload: BackupProgressPayload }) => void;
|
||||
onBackupCompleted?: (context: AgentBackupEventContext & { payload: BackupCompletedPayload }) => void;
|
||||
onBackupFailed?: (context: AgentBackupEventContext & { payload: BackupFailedPayload }) => void;
|
||||
onBackupCancelled?: (context: AgentBackupEventContext & { payload: BackupCancelledPayload }) => void;
|
||||
};
|
||||
|
||||
type AgentManagerRuntime = ReturnType<typeof createAgentManagerRuntime>;
|
||||
type AgentRuntimeState = {
|
||||
agentManager: AgentManagerRuntime;
|
||||
localAgent: ChildProcess | null;
|
||||
};
|
||||
|
||||
type ProcessWithAgentRuntime = NodeJS.Process & {
|
||||
__zerobyteAgentRuntime?: AgentRuntimeState;
|
||||
};
|
||||
|
||||
const getAgentRuntimeState = () => {
|
||||
const runtimeProcess = process as ProcessWithAgentRuntime;
|
||||
const existingRuntime = runtimeProcess.__zerobyteAgentRuntime;
|
||||
|
||||
if (existingRuntime) {
|
||||
return existingRuntime;
|
||||
}
|
||||
|
||||
const runtime = {
|
||||
agentManager: createAgentManagerRuntime(),
|
||||
localAgent: null,
|
||||
};
|
||||
|
||||
runtimeProcess.__zerobyteAgentRuntime = runtime;
|
||||
return runtime;
|
||||
};
|
||||
|
||||
const getAgentManagerRuntime = () => getAgentRuntimeState().agentManager;
|
||||
|
||||
export const spawnLocalAgent = async () => {
|
||||
await stopLocalAgent();
|
||||
|
||||
const sourceEntryPoint = path.join(process.cwd(), "apps", "agent", "src", "index.ts");
|
||||
const productionEntryPoint = path.join(process.cwd(), ".output", "agent", "index.mjs");
|
||||
|
||||
if (config.__prod__ && !existsSync(productionEntryPoint)) {
|
||||
throw new Error(`Local agent entrypoint not found at ${productionEntryPoint}`);
|
||||
}
|
||||
|
||||
const agentEntryPoint = config.__prod__ ? productionEntryPoint : sourceEntryPoint;
|
||||
const agentToken = await deriveLocalAgentToken();
|
||||
const args = config.__prod__ ? ["run", agentEntryPoint] : ["run", "--watch", agentEntryPoint];
|
||||
|
||||
const runtime = getAgentRuntimeState();
|
||||
const agentProcess = spawn("bun", args, {
|
||||
env: {
|
||||
PATH: process.env.PATH,
|
||||
ZEROBYTE_CONTROLLER_URL: "ws://localhost:3001",
|
||||
ZEROBYTE_AGENT_TOKEN: agentToken,
|
||||
},
|
||||
stdio: ["ignore", "pipe", "pipe"],
|
||||
});
|
||||
|
||||
runtime.localAgent = agentProcess;
|
||||
|
||||
agentProcess.stdout?.on("data", (data: Buffer) => {
|
||||
const line = data.toString().trim();
|
||||
if (line) logger.info(`[agent] ${line}`);
|
||||
});
|
||||
|
||||
agentProcess.stderr?.on("data", (data: Buffer) => {
|
||||
const line = data.toString().trim();
|
||||
if (line) logger.error(`[agent] ${line}`);
|
||||
});
|
||||
|
||||
agentProcess.on("exit", (code, signal) => {
|
||||
if (runtime.localAgent === agentProcess) {
|
||||
runtime.localAgent = null;
|
||||
}
|
||||
logger.info(`Agent process exited with code ${code} and signal ${signal}`);
|
||||
});
|
||||
};
|
||||
|
||||
export const stopLocalAgent = async () => {
|
||||
const runtime = getAgentRuntimeState();
|
||||
if (!runtime.localAgent) {
|
||||
return;
|
||||
}
|
||||
|
||||
const agentProcess = runtime.localAgent;
|
||||
runtime.localAgent = null;
|
||||
|
||||
if (agentProcess.exitCode !== null || agentProcess.signalCode !== null) {
|
||||
return;
|
||||
}
|
||||
|
||||
const exited = new Promise<void>((resolve) => {
|
||||
agentProcess.once("exit", () => {
|
||||
resolve();
|
||||
});
|
||||
});
|
||||
|
||||
agentProcess.kill();
|
||||
await exited;
|
||||
};
|
||||
|
||||
const createAgentManagerRuntime = () => {
|
||||
let sessions = new Map<string, ControllerAgentSession>();
|
||||
let backupHandlers: AgentBackupEventHandlers = {};
|
||||
let runtimeScope: Scope.CloseableScope | null = null;
|
||||
|
||||
const closeAllSessions = () => {
|
||||
for (const session of sessions.values()) {
|
||||
session.close();
|
||||
}
|
||||
sessions = new Map();
|
||||
};
|
||||
|
||||
const getSession = (agentId: string) => sessions.get(agentId);
|
||||
|
||||
const setSession = (agentId: string, session: ControllerAgentSession) => {
|
||||
const existingSession = getSession(agentId);
|
||||
if (existingSession) {
|
||||
existingSession.close();
|
||||
}
|
||||
|
||||
sessions.set(agentId, session);
|
||||
};
|
||||
|
||||
const removeSession = (agentId: string, connectionId: string) => {
|
||||
const session = getSession(agentId);
|
||||
if (!session || session.connectionId !== connectionId) {
|
||||
return;
|
||||
}
|
||||
|
||||
session.close();
|
||||
sessions.delete(agentId);
|
||||
};
|
||||
|
||||
const createSessionHandlers = (ws: Bun.ServerWebSocket<AgentConnectionData>) => {
|
||||
const agentId = ws.data.agentId;
|
||||
const agentName = ws.data.agentName;
|
||||
|
||||
return {
|
||||
onBackupStarted: (payload: BackupStartedPayload) => {
|
||||
backupHandlers.onBackupStarted?.({ agentId, agentName, payload });
|
||||
},
|
||||
onBackupProgress: (payload: BackupProgressPayload) => {
|
||||
backupHandlers.onBackupProgress?.({ agentId, agentName, payload });
|
||||
},
|
||||
onBackupCompleted: (payload: BackupCompletedPayload) => {
|
||||
backupHandlers.onBackupCompleted?.({ agentId, agentName, payload });
|
||||
},
|
||||
onBackupFailed: (payload: BackupFailedPayload) => {
|
||||
backupHandlers.onBackupFailed?.({ agentId, agentName, payload });
|
||||
},
|
||||
onBackupCancelled: (payload: BackupCancelledPayload) => {
|
||||
backupHandlers.onBackupCancelled?.({ agentId, agentName, payload });
|
||||
},
|
||||
};
|
||||
};
|
||||
|
||||
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, createSessionHandlers(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();
|
||||
void 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 scope = Effect.runSync(Scope.make());
|
||||
|
||||
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;
|
||||
}
|
||||
};
|
||||
|
||||
return {
|
||||
start,
|
||||
sendBackup: (agentId: string, payload: BackupRunPayload) => {
|
||||
const session = getSession(agentId);
|
||||
|
||||
if (!session) {
|
||||
logger.warn(`Cannot send backup command. Agent ${agentId} is not connected.`);
|
||||
return false;
|
||||
}
|
||||
|
||||
if (!session.isReady()) {
|
||||
logger.warn(`Cannot send backup command. Agent ${agentId} is not ready.`);
|
||||
return false;
|
||||
}
|
||||
|
||||
session.sendBackup(payload);
|
||||
logger.info(`Sent backup command ${payload.jobId} to agent ${agentId} for schedule ${payload.scheduleId}`);
|
||||
return true;
|
||||
},
|
||||
cancelBackup: (agentId: string, payload: BackupCancelPayload) => {
|
||||
const session = getSession(agentId);
|
||||
|
||||
if (!session) {
|
||||
logger.warn(`Cannot cancel backup command. Agent ${agentId} is not connected.`);
|
||||
return false;
|
||||
}
|
||||
|
||||
session.sendBackupCancel(payload);
|
||||
logger.info(`Sent backup cancel for command ${payload.jobId} to agent ${agentId}`);
|
||||
return true;
|
||||
},
|
||||
setBackupEventHandlers: (handlers: AgentBackupEventHandlers) => {
|
||||
backupHandlers = handlers;
|
||||
},
|
||||
getBackupEventHandlers: () => backupHandlers,
|
||||
stop,
|
||||
};
|
||||
};
|
||||
|
||||
export const agentManager = getAgentManagerRuntime();
|
||||
|
||||
export const stopAgentRuntime = async () => {
|
||||
getAgentManagerRuntime().stop();
|
||||
await stopLocalAgent();
|
||||
};
|
||||
234
app/server/modules/agents/controller-agent-session.ts
Normal file
234
app/server/modules/agents/controller-agent-session.ts
Normal file
|
|
@ -0,0 +1,234 @@
|
|||
import { Effect, Fiber, Queue, Ref } from "effect";
|
||||
import {
|
||||
createControllerMessage,
|
||||
parseAgentMessage,
|
||||
type AgentMessage,
|
||||
type BackupCancelledPayload,
|
||||
type BackupCompletedPayload,
|
||||
type BackupFailedPayload,
|
||||
type BackupProgressPayload,
|
||||
type BackupRunPayload,
|
||||
type BackupCancelPayload,
|
||||
type BackupStartedPayload,
|
||||
type ControllerWireMessage,
|
||||
} from "@zerobyte/contracts/agent-protocol";
|
||||
import { logger } from "@zerobyte/core/node";
|
||||
import { toMessage } from "@zerobyte/core/utils";
|
||||
|
||||
export 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;
|
||||
};
|
||||
|
||||
type TrackedBackupJob = {
|
||||
scheduleId: string;
|
||||
state: "pending" | "active";
|
||||
};
|
||||
|
||||
type ControllerAgentSessionHandlers = {
|
||||
onBackupStarted?: (payload: BackupStartedPayload) => void;
|
||||
onBackupProgress?: (payload: BackupProgressPayload) => void;
|
||||
onBackupCompleted?: (payload: BackupCompletedPayload) => void;
|
||||
onBackupFailed?: (payload: BackupFailedPayload) => void;
|
||||
onBackupCancelled?: (payload: BackupCancelledPayload) => void;
|
||||
};
|
||||
|
||||
export type ControllerAgentSession = {
|
||||
readonly connectionId: string;
|
||||
handleMessage: (data: string) => void;
|
||||
sendBackup: (payload: BackupRunPayload) => void;
|
||||
sendBackupCancel: (payload: BackupCancelPayload) => void;
|
||||
isReady: () => boolean;
|
||||
close: () => void;
|
||||
};
|
||||
|
||||
export const createControllerAgentSession = (
|
||||
socket: AgentSocket,
|
||||
handlers: ControllerAgentSessionHandlers = {},
|
||||
): ControllerAgentSession => {
|
||||
const outboundQueue = Effect.runSync(Queue.bounded<ControllerWireMessage>(64));
|
||||
const trackedBackupJobs = Effect.runSync(Ref.make<Map<string, TrackedBackupJob>>(new Map()));
|
||||
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 setTrackedBackupJob = (jobId: string, trackedBackupJob: TrackedBackupJob) => {
|
||||
Effect.runSync(
|
||||
Ref.update(trackedBackupJobs, (current) => {
|
||||
const next = new Map(current);
|
||||
next.set(jobId, trackedBackupJob);
|
||||
return next;
|
||||
}),
|
||||
);
|
||||
};
|
||||
|
||||
const deleteTrackedBackupJob = (jobId: string) => {
|
||||
Effect.runSync(
|
||||
Ref.update(trackedBackupJobs, (current) => {
|
||||
const next = new Map(current);
|
||||
next.delete(jobId);
|
||||
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?.({
|
||||
jobId,
|
||||
scheduleId: trackedJob.scheduleId,
|
||||
message,
|
||||
});
|
||||
}
|
||||
void Effect.runPromise(Fiber.interrupt(writerFiber)).catch(() => {});
|
||||
void Effect.runPromise(Fiber.interrupt(heartbeatFiber)).catch(() => {});
|
||||
void Effect.runPromise(Queue.shutdown(outboundQueue)).catch(() => {});
|
||||
},
|
||||
};
|
||||
};
|
||||
Loading…
Reference in a new issue