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