refactor: move effect / async boundary in one place
This commit is contained in:
Nicolas Meienberger
parent
88fe3abc11
commit
87e479aeea
8 changed files with 809 additions and 288 deletions
|
|
@ -1,6 +1,7 @@
|
|||
import { afterEach, expect, test, vi } from "vitest";
|
||||
import waitForExpect from "wait-for-expect";
|
||||
import { fromAny, fromPartial } from "@total-typescript/shoehorn";
|
||||
import { Effect } from "effect";
|
||||
import { agentManager, type ProcessWithAgentRuntime } from "../agents-manager";
|
||||
import type { AgentManagerRuntime } from "../controller/server";
|
||||
import type { BackupRunPayload } from "@zerobyte/contracts/agent-protocol";
|
||||
|
|
@ -20,8 +21,8 @@ afterEach(() => {
|
|||
});
|
||||
|
||||
test("cancelBackup resolves a running backup when the cancel command cannot be delivered", async () => {
|
||||
const sendBackup = vi.fn().mockResolvedValue(true);
|
||||
const cancelBackup = vi.fn().mockResolvedValue(false);
|
||||
const sendBackup = vi.fn(() => Effect.succeed(true));
|
||||
const cancelBackup = vi.fn(() => Effect.succeed(false));
|
||||
setAgentRuntime({ sendBackup, cancelBackup });
|
||||
|
||||
const resultPromise = agentManager.runBackup("local", {
|
||||
|
|
|
|||
|
|
@ -0,0 +1,235 @@
|
|||
import { afterEach, expect, test, vi } from "vitest";
|
||||
import { Effect } from "effect";
|
||||
import { fromAny, fromPartial } from "@total-typescript/shoehorn";
|
||||
import type { BackupRunPayload } from "@zerobyte/contracts/agent-protocol";
|
||||
import type { AgentManagerEvent } from "../controller/server";
|
||||
import type { ProcessWithAgentRuntime } from "../agents-manager";
|
||||
|
||||
const controllerMock = vi.hoisted(() => ({
|
||||
onEvent: null as null | ((event: AgentManagerEvent) => void),
|
||||
sendBackup: vi.fn(),
|
||||
cancelBackup: vi.fn(),
|
||||
stop: vi.fn(),
|
||||
}));
|
||||
|
||||
vi.mock("../controller/server", async () => {
|
||||
const { Effect } = await import("effect");
|
||||
return {
|
||||
createAgentManagerRuntime: vi.fn((onEvent: (event: AgentManagerEvent) => void) => {
|
||||
controllerMock.onEvent = onEvent;
|
||||
return {
|
||||
start: Effect.void,
|
||||
stop: Effect.sync(controllerMock.stop),
|
||||
sendBackup: controllerMock.sendBackup,
|
||||
cancelBackup: controllerMock.cancelBackup,
|
||||
};
|
||||
}),
|
||||
};
|
||||
});
|
||||
|
||||
const processWithAgentRuntime = process as ProcessWithAgentRuntime;
|
||||
|
||||
const resetAgentRuntime = () => {
|
||||
processWithAgentRuntime.__zerobyteAgentRuntime = {
|
||||
agentManager: null,
|
||||
localAgent: null,
|
||||
isStoppingLocalAgent: false,
|
||||
localAgentRestartTimeout: null,
|
||||
activeBackupsByScheduleId: new Map(),
|
||||
activeBackupScheduleIdsByJobId: new Map(),
|
||||
};
|
||||
};
|
||||
|
||||
const backupPayload = fromPartial<BackupRunPayload>({
|
||||
jobId: "job-1",
|
||||
scheduleId: "schedule-1",
|
||||
});
|
||||
|
||||
afterEach(() => {
|
||||
delete processWithAgentRuntime.__zerobyteAgentRuntime;
|
||||
controllerMock.onEvent = null;
|
||||
controllerMock.sendBackup.mockReset();
|
||||
controllerMock.cancelBackup.mockReset();
|
||||
controllerMock.stop.mockReset();
|
||||
vi.resetModules();
|
||||
vi.restoreAllMocks();
|
||||
});
|
||||
|
||||
test("backup progress is delivered to the running backup callback", async () => {
|
||||
resetAgentRuntime();
|
||||
controllerMock.sendBackup.mockImplementation(() => Effect.succeed(true));
|
||||
const { agentManager, startAgentController, stopAgentController } = await import("../agents-manager");
|
||||
const onProgress = vi.fn();
|
||||
|
||||
await startAgentController();
|
||||
const resultPromise = agentManager.runBackup("local", {
|
||||
scheduleId: 42,
|
||||
payload: backupPayload,
|
||||
signal: new AbortController().signal,
|
||||
onProgress,
|
||||
});
|
||||
|
||||
controllerMock.onEvent?.({
|
||||
type: "backup.progress",
|
||||
agentId: "local",
|
||||
agentName: "Local Agent",
|
||||
payload: fromAny({ jobId: "job-1", scheduleId: "schedule-1", progress: { percentDone: 0.5 } }),
|
||||
});
|
||||
controllerMock.onEvent?.({
|
||||
type: "backup.completed",
|
||||
agentId: "local",
|
||||
agentName: "Local Agent",
|
||||
payload: { jobId: "job-1", scheduleId: "schedule-1", exitCode: 0, result: null },
|
||||
});
|
||||
|
||||
await expect(resultPromise).resolves.toEqual({
|
||||
status: "completed",
|
||||
exitCode: 0,
|
||||
result: null,
|
||||
warningDetails: null,
|
||||
});
|
||||
expect(onProgress).toHaveBeenCalledWith({ percentDone: 0.5 });
|
||||
await stopAgentController();
|
||||
});
|
||||
|
||||
test("backup failed and cancelled events resolve the matching running backup", async () => {
|
||||
resetAgentRuntime();
|
||||
controllerMock.sendBackup.mockImplementation(() => Effect.succeed(true));
|
||||
const { agentManager, startAgentController, stopAgentController } = await import("../agents-manager");
|
||||
|
||||
await startAgentController();
|
||||
const failedPromise = agentManager.runBackup("local", {
|
||||
scheduleId: 42,
|
||||
payload: backupPayload,
|
||||
signal: new AbortController().signal,
|
||||
onProgress: vi.fn(),
|
||||
});
|
||||
controllerMock.onEvent?.({
|
||||
type: "backup.failed",
|
||||
agentId: "local",
|
||||
agentName: "Local Agent",
|
||||
payload: { jobId: "job-1", scheduleId: "schedule-1", error: "failed", errorDetails: "restic failed" },
|
||||
});
|
||||
await expect(failedPromise).resolves.toEqual({ status: "failed", error: "restic failed" });
|
||||
|
||||
const cancelledPromise = agentManager.runBackup("local", {
|
||||
scheduleId: 43,
|
||||
payload: fromPartial<BackupRunPayload>({ jobId: "job-2", scheduleId: "schedule-2" }),
|
||||
signal: new AbortController().signal,
|
||||
onProgress: vi.fn(),
|
||||
});
|
||||
controllerMock.onEvent?.({
|
||||
type: "backup.cancelled",
|
||||
agentId: "local",
|
||||
agentName: "Local Agent",
|
||||
payload: { jobId: "job-2", scheduleId: "schedule-2", message: "cancelled remotely" },
|
||||
});
|
||||
await expect(cancelledPromise).resolves.toEqual({ status: "cancelled", message: "cancelled remotely" });
|
||||
await stopAgentController();
|
||||
});
|
||||
|
||||
test("agent disconnect cancels only backups owned by that agent", async () => {
|
||||
resetAgentRuntime();
|
||||
controllerMock.sendBackup.mockImplementation(() => Effect.succeed(true));
|
||||
const { agentManager, startAgentController, stopAgentController } = await import("../agents-manager");
|
||||
|
||||
await startAgentController();
|
||||
const localPromise = agentManager.runBackup("local", {
|
||||
scheduleId: 42,
|
||||
payload: backupPayload,
|
||||
signal: new AbortController().signal,
|
||||
onProgress: vi.fn(),
|
||||
});
|
||||
const remotePromise = agentManager.runBackup("remote", {
|
||||
scheduleId: 43,
|
||||
payload: fromPartial<BackupRunPayload>({ jobId: "job-2", scheduleId: "schedule-2" }),
|
||||
signal: new AbortController().signal,
|
||||
onProgress: vi.fn(),
|
||||
});
|
||||
|
||||
controllerMock.onEvent?.({ type: "agent.disconnected", agentId: "local", agentName: "Local Agent" });
|
||||
controllerMock.onEvent?.({
|
||||
type: "backup.completed",
|
||||
agentId: "remote",
|
||||
agentName: "Remote Agent",
|
||||
payload: { jobId: "job-2", scheduleId: "schedule-2", exitCode: 0, result: null },
|
||||
});
|
||||
|
||||
await expect(localPromise).resolves.toEqual({
|
||||
status: "cancelled",
|
||||
message: "The connection to the backup agent was lost. Restart the backup to ensure it completes.",
|
||||
});
|
||||
await expect(remotePromise).resolves.toEqual({
|
||||
status: "completed",
|
||||
exitCode: 0,
|
||||
result: null,
|
||||
warningDetails: null,
|
||||
});
|
||||
await stopAgentController();
|
||||
});
|
||||
|
||||
test("runBackup returns unavailable and clears the active run when the command cannot be sent", async () => {
|
||||
resetAgentRuntime();
|
||||
controllerMock.sendBackup.mockImplementation(() => Effect.succeed(false));
|
||||
const { agentManager, startAgentController, stopAgentController } = await import("../agents-manager");
|
||||
|
||||
await startAgentController();
|
||||
const result = await agentManager.runBackup("local", {
|
||||
scheduleId: 42,
|
||||
payload: backupPayload,
|
||||
signal: new AbortController().signal,
|
||||
onProgress: vi.fn(),
|
||||
});
|
||||
|
||||
expect(result).toEqual({
|
||||
status: "unavailable",
|
||||
error: new Error("Failed to send backup command to agent local"),
|
||||
});
|
||||
await expect(agentManager.cancelBackup("local", 42)).resolves.toBe(false);
|
||||
await stopAgentController();
|
||||
});
|
||||
|
||||
test("runBackup rejects before sending when the abort signal is already aborted", async () => {
|
||||
resetAgentRuntime();
|
||||
controllerMock.sendBackup.mockImplementation(() => Effect.succeed(true));
|
||||
const { agentManager, startAgentController, stopAgentController } = await import("../agents-manager");
|
||||
const abortController = new AbortController();
|
||||
abortController.abort(new Error("cancelled before send"));
|
||||
|
||||
await startAgentController();
|
||||
await expect(
|
||||
agentManager.runBackup("local", {
|
||||
scheduleId: 42,
|
||||
payload: backupPayload,
|
||||
signal: abortController.signal,
|
||||
onProgress: vi.fn(),
|
||||
}),
|
||||
).rejects.toThrow("cancelled before send");
|
||||
expect(controllerMock.sendBackup).not.toHaveBeenCalled();
|
||||
await stopAgentController();
|
||||
});
|
||||
|
||||
test("runBackup requests cancellation when the abort signal fires while sending", async () => {
|
||||
resetAgentRuntime();
|
||||
const abortController = new AbortController();
|
||||
controllerMock.sendBackup.mockImplementation(() =>
|
||||
Effect.sync(() => {
|
||||
abortController.abort();
|
||||
return true;
|
||||
}),
|
||||
);
|
||||
controllerMock.cancelBackup.mockImplementation(() => Effect.succeed(false));
|
||||
const { agentManager, startAgentController, stopAgentController } = await import("../agents-manager");
|
||||
|
||||
await startAgentController();
|
||||
const result = await agentManager.runBackup("local", {
|
||||
scheduleId: 42,
|
||||
payload: backupPayload,
|
||||
signal: abortController.signal,
|
||||
onProgress: vi.fn(),
|
||||
});
|
||||
|
||||
expect(result).toEqual({ status: "cancelled" });
|
||||
expect(controllerMock.cancelBackup).toHaveBeenCalledWith("local", { jobId: "job-1", scheduleId: "schedule-1" });
|
||||
await stopAgentController();
|
||||
});
|
||||
|
|
@ -1,6 +1,7 @@
|
|||
import { beforeEach, expect, test } from "vitest";
|
||||
import { db } from "~/server/db/db";
|
||||
import { agentsTable } from "~/server/db/schema";
|
||||
import { LOCAL_AGENT_ID, LOCAL_AGENT_KIND, LOCAL_AGENT_NAME } from "../constants";
|
||||
import { agentsService } from "../agents.service";
|
||||
|
||||
beforeEach(async () => {
|
||||
|
|
@ -15,3 +16,57 @@ test("ensureLocalAgent seeds the built-in local agent once", async () => {
|
|||
|
||||
expect(agents).toHaveLength(1);
|
||||
});
|
||||
|
||||
test("markAgentConnecting creates and updates connection metadata", async () => {
|
||||
await agentsService.markAgentConnecting({
|
||||
agentId: "remote-agent",
|
||||
organizationId: null,
|
||||
agentName: "Remote Agent",
|
||||
agentKind: "remote",
|
||||
capabilities: { restic: true },
|
||||
connectedAt: 1_000,
|
||||
});
|
||||
await agentsService.markAgentConnecting({
|
||||
agentId: "remote-agent",
|
||||
organizationId: null,
|
||||
agentName: "Renamed Agent",
|
||||
agentKind: "remote",
|
||||
capabilities: { restic: true, webdav: true },
|
||||
connectedAt: 2_000,
|
||||
});
|
||||
|
||||
const agent = await agentsService.getAgent("remote-agent");
|
||||
|
||||
expect(agent).toMatchObject({
|
||||
id: "remote-agent",
|
||||
name: "Renamed Agent",
|
||||
kind: "remote",
|
||||
status: "connecting",
|
||||
capabilities: { restic: true, webdav: true },
|
||||
lastSeenAt: 2_000,
|
||||
updatedAt: 2_000,
|
||||
});
|
||||
});
|
||||
|
||||
test("agent runtime status moves from connecting to online, seen, and offline", async () => {
|
||||
await agentsService.markAgentConnecting({
|
||||
agentId: LOCAL_AGENT_ID,
|
||||
organizationId: null,
|
||||
agentName: LOCAL_AGENT_NAME,
|
||||
agentKind: LOCAL_AGENT_KIND,
|
||||
connectedAt: 1_000,
|
||||
});
|
||||
await agentsService.markAgentOnline(LOCAL_AGENT_ID, 2_000);
|
||||
await agentsService.markAgentSeen(LOCAL_AGENT_ID, 3_000);
|
||||
await agentsService.markAgentOffline(LOCAL_AGENT_ID, 4_000);
|
||||
|
||||
const agent = await agentsService.getAgent(LOCAL_AGENT_ID);
|
||||
|
||||
expect(agent).toMatchObject({
|
||||
id: LOCAL_AGENT_ID,
|
||||
status: "offline",
|
||||
lastSeenAt: 3_000,
|
||||
lastReadyAt: 2_000,
|
||||
updatedAt: 4_000,
|
||||
});
|
||||
});
|
||||
|
|
|
|||
219
app/server/modules/agents/__tests__/controller-runtime.test.ts
Normal file
219
app/server/modules/agents/__tests__/controller-runtime.test.ts
Normal file
|
|
@ -0,0 +1,219 @@
|
|||
import { Effect } from "effect";
|
||||
import { afterEach, expect, test, vi } from "vitest";
|
||||
import waitForExpect from "wait-for-expect";
|
||||
import { fromPartial } from "@total-typescript/shoehorn";
|
||||
import { createAgentMessage } from "@zerobyte/contracts/agent-protocol";
|
||||
import { LOCAL_AGENT_ID, LOCAL_AGENT_KIND, LOCAL_AGENT_NAME } from "../constants";
|
||||
|
||||
const agentsServiceMocks = vi.hoisted(() => ({
|
||||
markAgentConnecting: vi.fn(() => Promise.resolve()),
|
||||
markAgentOnline: vi.fn(() => Promise.resolve()),
|
||||
markAgentSeen: vi.fn(() => Promise.resolve()),
|
||||
markAgentOffline: vi.fn(() => Promise.resolve()),
|
||||
}));
|
||||
|
||||
const tokenMocks = vi.hoisted(() => ({
|
||||
validateAgentToken: vi.fn(),
|
||||
}));
|
||||
|
||||
vi.mock("../agents.service", () => ({
|
||||
agentsService: agentsServiceMocks,
|
||||
}));
|
||||
|
||||
vi.mock("../helpers/tokens", () => ({
|
||||
validateAgentToken: tokenMocks.validateAgentToken,
|
||||
}));
|
||||
|
||||
const createSocket = (id: string) => ({
|
||||
data: {
|
||||
id,
|
||||
agentId: LOCAL_AGENT_ID,
|
||||
organizationId: null,
|
||||
agentName: LOCAL_AGENT_NAME,
|
||||
agentKind: LOCAL_AGENT_KIND,
|
||||
},
|
||||
send: vi.fn(() => 1),
|
||||
close: vi.fn(),
|
||||
});
|
||||
|
||||
type CapturedFetch = NonNullable<Parameters<typeof Bun.serve>[0]["fetch"]>;
|
||||
|
||||
const invokeFetch = (fetch: CapturedFetch | undefined, request: Request, srv: Parameters<CapturedFetch>[1]) => {
|
||||
if (!fetch) {
|
||||
throw new Error("Bun.serve was not called with a fetch handler");
|
||||
}
|
||||
|
||||
return Reflect.apply(fetch, fromPartial<ThisParameterType<CapturedFetch>>({}), [
|
||||
request,
|
||||
srv,
|
||||
]) as ReturnType<CapturedFetch>;
|
||||
};
|
||||
|
||||
const startRuntime = async (onEvent = vi.fn()) => {
|
||||
const { createAgentManagerRuntime } = await import("../controller/server");
|
||||
const runtime = createAgentManagerRuntime(onEvent);
|
||||
await Effect.runPromise(runtime.start);
|
||||
return { runtime, onEvent };
|
||||
};
|
||||
|
||||
afterEach(() => {
|
||||
vi.restoreAllMocks();
|
||||
tokenMocks.validateAgentToken.mockReset();
|
||||
vi.resetModules();
|
||||
});
|
||||
|
||||
test("websocket fetch rejects requests without a bearer token", async () => {
|
||||
const serve = vi
|
||||
.spyOn(Bun, "serve")
|
||||
.mockReturnValue(fromPartial({ port: 3001, stop: vi.fn(() => Promise.resolve()) }));
|
||||
const { runtime } = await startRuntime();
|
||||
const fetch = serve.mock.calls[0]?.[0].fetch;
|
||||
const upgrade = vi.fn();
|
||||
const srv = fromPartial<Parameters<NonNullable<typeof fetch>>[1]>({ upgrade });
|
||||
|
||||
const response = await invokeFetch(fetch, new Request("http://localhost:3001/agent"), srv);
|
||||
await Effect.runPromise(runtime.stop);
|
||||
|
||||
expect(response?.status).toBe(401);
|
||||
expect(await response?.text()).toBe("Missing token");
|
||||
expect(upgrade).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
test("websocket fetch rejects invalid bearer tokens", async () => {
|
||||
tokenMocks.validateAgentToken.mockResolvedValue(undefined);
|
||||
const serve = vi
|
||||
.spyOn(Bun, "serve")
|
||||
.mockReturnValue(fromPartial({ port: 3001, stop: vi.fn(() => Promise.resolve()) }));
|
||||
const { runtime } = await startRuntime();
|
||||
const fetch = serve.mock.calls[0]?.[0].fetch;
|
||||
const upgrade = vi.fn();
|
||||
const srv = fromPartial<Parameters<NonNullable<typeof fetch>>[1]>({ upgrade });
|
||||
|
||||
const response = await invokeFetch(
|
||||
fetch,
|
||||
new Request("http://localhost:3001/agent", { headers: { authorization: "Bearer bad-token" } }),
|
||||
srv,
|
||||
);
|
||||
await Effect.runPromise(runtime.stop);
|
||||
|
||||
expect(response?.status).toBe(401);
|
||||
expect(await response?.text()).toBe("Invalid or revoked token");
|
||||
expect(tokenMocks.validateAgentToken).toHaveBeenCalledWith("bad-token");
|
||||
expect(upgrade).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
test("websocket fetch upgrades valid agent tokens with connection metadata", async () => {
|
||||
tokenMocks.validateAgentToken.mockResolvedValue({
|
||||
agentId: LOCAL_AGENT_ID,
|
||||
organizationId: null,
|
||||
agentName: LOCAL_AGENT_NAME,
|
||||
agentKind: LOCAL_AGENT_KIND,
|
||||
});
|
||||
const serve = vi
|
||||
.spyOn(Bun, "serve")
|
||||
.mockReturnValue(fromPartial({ port: 3001, stop: vi.fn(() => Promise.resolve()) }));
|
||||
const { runtime } = await startRuntime();
|
||||
const fetch = serve.mock.calls[0]?.[0].fetch;
|
||||
const upgrade = vi.fn(() => true);
|
||||
const srv = fromPartial<Parameters<NonNullable<typeof fetch>>[1]>({ upgrade });
|
||||
|
||||
const response = await invokeFetch(
|
||||
fetch,
|
||||
new Request("http://localhost:3001/agent", { headers: { authorization: "Bearer valid-token" } }),
|
||||
srv,
|
||||
);
|
||||
await Effect.runPromise(runtime.stop);
|
||||
|
||||
expect(response).toBeUndefined();
|
||||
expect(tokenMocks.validateAgentToken).toHaveBeenCalledWith("valid-token");
|
||||
expect(upgrade).toHaveBeenCalledWith(expect.any(Request), {
|
||||
data: expect.objectContaining({
|
||||
agentId: LOCAL_AGENT_ID,
|
||||
organizationId: null,
|
||||
agentName: LOCAL_AGENT_NAME,
|
||||
agentKind: LOCAL_AGENT_KIND,
|
||||
id: expect.any(String),
|
||||
}),
|
||||
});
|
||||
});
|
||||
|
||||
test("websocket lifecycle updates agent connection status", async () => {
|
||||
const stop = vi.fn(() => Promise.resolve());
|
||||
const serve = vi.spyOn(Bun, "serve").mockReturnValue(fromPartial({ port: 3001, stop }));
|
||||
const { runtime } = await startRuntime();
|
||||
const websocket = serve.mock.calls[0]?.[0].websocket;
|
||||
const socket = createSocket("connection-1");
|
||||
|
||||
await websocket?.open?.(fromPartial(socket));
|
||||
await websocket?.message?.(fromPartial(socket), createAgentMessage("agent.ready", { agentId: LOCAL_AGENT_ID }));
|
||||
await websocket?.message?.(fromPartial(socket), createAgentMessage("heartbeat.pong", { sentAt: 123 }));
|
||||
await websocket?.close?.(fromPartial(socket), 1000, "done");
|
||||
await Effect.runPromise(runtime.stop);
|
||||
|
||||
expect(agentsServiceMocks.markAgentConnecting).toHaveBeenCalledWith({
|
||||
agentId: LOCAL_AGENT_ID,
|
||||
organizationId: null,
|
||||
agentName: LOCAL_AGENT_NAME,
|
||||
agentKind: LOCAL_AGENT_KIND,
|
||||
});
|
||||
expect(agentsServiceMocks.markAgentOnline).toHaveBeenCalledWith(LOCAL_AGENT_ID, expect.any(Number));
|
||||
expect(agentsServiceMocks.markAgentSeen).toHaveBeenCalledWith(LOCAL_AGENT_ID, expect.any(Number));
|
||||
expect(agentsServiceMocks.markAgentOffline).toHaveBeenCalledWith(LOCAL_AGENT_ID);
|
||||
expect(stop).toHaveBeenCalledWith(true);
|
||||
});
|
||||
|
||||
test("closing a replaced connection does not report the active agent as disconnected", async () => {
|
||||
const serve = vi
|
||||
.spyOn(Bun, "serve")
|
||||
.mockReturnValue(fromPartial({ port: 3001, stop: vi.fn(() => Promise.resolve()) }));
|
||||
const { runtime, onEvent } = await startRuntime(vi.fn());
|
||||
const websocket = serve.mock.calls[0]?.[0].websocket;
|
||||
const oldSocket = createSocket("connection-1");
|
||||
const newSocket = createSocket("connection-2");
|
||||
|
||||
await websocket?.open?.(fromPartial(oldSocket));
|
||||
await websocket?.open?.(fromPartial(newSocket));
|
||||
await websocket?.close?.(fromPartial(oldSocket), 1000, "replaced");
|
||||
await Effect.runPromise(runtime.stop);
|
||||
|
||||
expect(onEvent).not.toHaveBeenCalledWith(
|
||||
expect.objectContaining({ type: "agent.disconnected", agentId: LOCAL_AGENT_ID }),
|
||||
);
|
||||
});
|
||||
|
||||
test("sendBackup is only delivered after the agent is ready", async () => {
|
||||
const serve = vi
|
||||
.spyOn(Bun, "serve")
|
||||
.mockReturnValue(fromPartial({ port: 3001, stop: vi.fn(() => Promise.resolve()) }));
|
||||
const { runtime } = await startRuntime();
|
||||
const websocket = serve.mock.calls[0]?.[0].websocket;
|
||||
const socket = createSocket("connection-1");
|
||||
const payload = {
|
||||
jobId: "job-1",
|
||||
scheduleId: "schedule-1",
|
||||
organizationId: "org-1",
|
||||
sourcePath: "/tmp/source",
|
||||
repositoryConfig: { backend: "local" as const, path: "/tmp/repository" },
|
||||
options: {},
|
||||
runtime: {
|
||||
password: "password",
|
||||
cacheDir: "/tmp/cache",
|
||||
passFile: "/tmp/pass",
|
||||
defaultExcludes: [],
|
||||
rcloneConfigFile: "/tmp/rclone.conf",
|
||||
},
|
||||
webhooks: { pre: null, post: null },
|
||||
webhookAllowedOrigins: [],
|
||||
};
|
||||
|
||||
await websocket?.open?.(fromPartial(socket));
|
||||
await expect(Effect.runPromise(runtime.sendBackup(LOCAL_AGENT_ID, payload))).resolves.toBe(false);
|
||||
|
||||
await websocket?.message?.(fromPartial(socket), createAgentMessage("agent.ready", { agentId: LOCAL_AGENT_ID }));
|
||||
await expect(Effect.runPromise(runtime.sendBackup(LOCAL_AGENT_ID, payload))).resolves.toBe(true);
|
||||
|
||||
await waitForExpect(() => {
|
||||
expect(socket.send).toHaveBeenCalledWith(expect.stringContaining('"type":"backup.run"'));
|
||||
});
|
||||
await Effect.runPromise(runtime.stop);
|
||||
});
|
||||
|
|
@ -2,6 +2,7 @@ import { Effect, Exit, Fiber, Scope } from "effect";
|
|||
import { expect, test, vi } from "vitest";
|
||||
import waitForExpect from "wait-for-expect";
|
||||
import { fromPartial } from "@total-typescript/shoehorn";
|
||||
import { createAgentMessage, type AgentMessage } from "@zerobyte/contracts/agent-protocol";
|
||||
import { LOCAL_AGENT_ID, LOCAL_AGENT_KIND, LOCAL_AGENT_NAME } from "../constants";
|
||||
import { createControllerAgentSession } from "../controller/session";
|
||||
|
||||
|
|
@ -21,26 +22,13 @@ const createSocket = (overrides: Partial<Parameters<typeof createControllerAgent
|
|||
};
|
||||
|
||||
const createSession = (
|
||||
handlers: Partial<Parameters<typeof createControllerAgentSession>[1]> = {},
|
||||
onEvent: Parameters<typeof createControllerAgentSession>[1] = () => Effect.void,
|
||||
socket = createSocket(),
|
||||
) => {
|
||||
const scope = Effect.runSync(Scope.make());
|
||||
const sessionHandlers: Parameters<typeof createControllerAgentSession>[1] = {
|
||||
onReady: () => Effect.void,
|
||||
onHeartbeatPong: () => Effect.void,
|
||||
onDisconnect: () => Effect.void,
|
||||
onBackupStarted: () => Effect.void,
|
||||
onBackupProgress: () => Effect.void,
|
||||
onBackupCompleted: () => Effect.void,
|
||||
onBackupFailed: () => Effect.void,
|
||||
onBackupCancelled: () => Effect.void,
|
||||
...handlers,
|
||||
};
|
||||
|
||||
try {
|
||||
const session = Effect.runSync(
|
||||
Scope.extend(createControllerAgentSession(fromPartial(socket), sessionHandlers), scope),
|
||||
);
|
||||
const session = Effect.runSync(Scope.extend(createControllerAgentSession(fromPartial(socket), onEvent), scope));
|
||||
|
||||
return {
|
||||
session,
|
||||
|
|
@ -74,23 +62,22 @@ test("closing the session scope interrupts the session runner", async () => {
|
|||
});
|
||||
|
||||
test("close reports a transport disconnect", () => {
|
||||
const onDisconnect = vi.fn(() => Effect.void);
|
||||
const { close } = createSession({ onDisconnect });
|
||||
const onEvent = vi.fn(() => Effect.void);
|
||||
const { close } = createSession(onEvent);
|
||||
|
||||
close();
|
||||
|
||||
expect(onDisconnect).toHaveBeenCalledTimes(1);
|
||||
expect(onDisconnect).toHaveBeenCalledWith(
|
||||
expect(onEvent).toHaveBeenCalledTimes(1);
|
||||
expect(onEvent).toHaveBeenCalledWith(
|
||||
expect.objectContaining({
|
||||
agentId: LOCAL_AGENT_ID,
|
||||
agentName: LOCAL_AGENT_NAME,
|
||||
type: "agent.disconnected",
|
||||
}),
|
||||
);
|
||||
});
|
||||
|
||||
test("sendBackup only queues the transport message", () => {
|
||||
const onBackupCancelled = vi.fn(() => Effect.void);
|
||||
const { session, close } = createSession({ onBackupCancelled });
|
||||
const onEvent = vi.fn(() => Effect.void);
|
||||
const { session, close } = createSession(onEvent);
|
||||
|
||||
Effect.runSync(
|
||||
session.sendBackup({
|
||||
|
|
@ -118,14 +105,74 @@ test("sendBackup only queues the transport message", () => {
|
|||
|
||||
close();
|
||||
|
||||
expect(onBackupCancelled).not.toHaveBeenCalled();
|
||||
expect(onEvent).not.toHaveBeenCalledWith(expect.objectContaining({ type: "backup.cancelled" }));
|
||||
});
|
||||
|
||||
test("invalid inbound messages are ignored", () => {
|
||||
const onEvent = vi.fn(() => Effect.void);
|
||||
const { session, close } = createSession(onEvent);
|
||||
|
||||
Effect.runSync(session.handleMessage("not json"));
|
||||
Effect.runSync(session.handleMessage(JSON.stringify({ type: "backup.progress", payload: {} })));
|
||||
|
||||
expect(onEvent).not.toHaveBeenCalled();
|
||||
close();
|
||||
});
|
||||
|
||||
test("agent.ready marks the session ready and forwards the event", () => {
|
||||
const onEvent = vi.fn(() => Effect.void);
|
||||
const { session, close } = createSession(onEvent);
|
||||
|
||||
expect(Effect.runSync(session.isReady())).toBe(false);
|
||||
Effect.runSync(session.handleMessage(createAgentMessage("agent.ready", { agentId: LOCAL_AGENT_ID })));
|
||||
|
||||
expect(Effect.runSync(session.isReady())).toBe(true);
|
||||
expect(onEvent).toHaveBeenCalledWith({ type: "agent.ready", payload: { agentId: LOCAL_AGENT_ID } });
|
||||
close();
|
||||
});
|
||||
|
||||
test("backup agent messages are forwarded unchanged", () => {
|
||||
const onEvent = vi.fn(() => Effect.void);
|
||||
const { session, close } = createSession(onEvent);
|
||||
const message = {
|
||||
type: "backup.progress" as const,
|
||||
payload: {
|
||||
jobId: "job-1",
|
||||
scheduleId: "schedule-1",
|
||||
progress: {
|
||||
message_type: "status" as const,
|
||||
seconds_elapsed: 0,
|
||||
seconds_remaining: 0,
|
||||
percent_done: 0.5,
|
||||
total_files: 0,
|
||||
files_done: 0,
|
||||
total_bytes: 0,
|
||||
bytes_done: 0,
|
||||
current_files: [],
|
||||
},
|
||||
},
|
||||
} satisfies Extract<AgentMessage, { type: "backup.progress" }>;
|
||||
|
||||
Effect.runSync(session.handleMessage(createAgentMessage(message.type, message.payload)));
|
||||
|
||||
expect(onEvent).toHaveBeenCalledWith(
|
||||
expect.objectContaining({
|
||||
type: message.type,
|
||||
payload: expect.objectContaining({
|
||||
jobId: message.payload.jobId,
|
||||
scheduleId: message.payload.scheduleId,
|
||||
progress: expect.objectContaining(message.payload.progress),
|
||||
}),
|
||||
}),
|
||||
);
|
||||
close();
|
||||
});
|
||||
|
||||
test("a dropped backup.cancel closes the session and reports a transport disconnect", async () => {
|
||||
const send = vi.fn(() => 0);
|
||||
const socket = createSocket({ send, close: vi.fn() });
|
||||
const onDisconnect = vi.fn(() => Effect.void);
|
||||
const { session, run, closeAsync } = createSession({ onDisconnect }, socket);
|
||||
const onEvent = vi.fn(() => Effect.void);
|
||||
const { session, run, closeAsync } = createSession(onEvent, socket);
|
||||
|
||||
try {
|
||||
run();
|
||||
|
|
@ -139,10 +186,10 @@ test("a dropped backup.cancel closes the session and reports a transport disconn
|
|||
await waitForExpect(() => {
|
||||
expect(send).toHaveBeenCalledTimes(1);
|
||||
expect(socket.close).toHaveBeenCalledTimes(1);
|
||||
expect(onDisconnect).toHaveBeenCalledTimes(1);
|
||||
expect(onDisconnect).toHaveBeenCalledWith(
|
||||
expect(onEvent).toHaveBeenCalledTimes(1);
|
||||
expect(onEvent).toHaveBeenCalledWith(
|
||||
expect.objectContaining({
|
||||
agentId: LOCAL_AGENT_ID,
|
||||
type: "agent.disconnected",
|
||||
}),
|
||||
);
|
||||
});
|
||||
|
|
|
|||
|
|
@ -1,7 +1,8 @@
|
|||
import { logger } from "@zerobyte/core/node";
|
||||
import type { BackupRunPayload } from "@zerobyte/contracts/agent-protocol";
|
||||
import { Effect } from "effect";
|
||||
import { config } from "../../core/config";
|
||||
import { createAgentManagerRuntime, type AgentBackupEventHandlers } from "./controller/server";
|
||||
import { createAgentManagerRuntime, type AgentManagerEvent } from "./controller/server";
|
||||
import { spawnLocalAgentProcess, stopLocalAgentProcess } from "./local/process";
|
||||
import type { BackupExecutionProgress, BackupExecutionResult } from "./helpers/runtime-state";
|
||||
import { createAgentRuntimeState } from "./helpers/runtime-state";
|
||||
|
|
@ -97,10 +98,12 @@ const requestBackupCancellation = async (agentId: string, scheduleId: number) =>
|
|||
}
|
||||
|
||||
if (
|
||||
await runtime.cancelBackup(agentId, {
|
||||
jobId: activeBackupRun.jobId,
|
||||
scheduleId: activeBackupRun.scheduleShortId,
|
||||
})
|
||||
await Effect.runPromise(
|
||||
runtime.cancelBackup(agentId, {
|
||||
jobId: activeBackupRun.jobId,
|
||||
scheduleId: activeBackupRun.scheduleShortId,
|
||||
}),
|
||||
)
|
||||
) {
|
||||
return true;
|
||||
}
|
||||
|
|
@ -109,71 +112,99 @@ const requestBackupCancellation = async (agentId: string, scheduleId: number) =>
|
|||
return true;
|
||||
};
|
||||
|
||||
const backupEventHandlers: AgentBackupEventHandlers = {
|
||||
onAgentDisconnected: ({ agentId }) => {
|
||||
cancelActiveBackupRunsForAgent(
|
||||
agentId,
|
||||
"The connection to the backup agent was lost. Restart the backup to ensure it completes.",
|
||||
);
|
||||
},
|
||||
onBackupStarted: ({ agentId, payload }) => {
|
||||
getActiveBackupRun(payload.jobId, payload.scheduleId, "backup.started", agentId);
|
||||
},
|
||||
onBackupProgress: ({ agentId, payload }) => {
|
||||
const activeBackupRun = getActiveBackupRun(payload.jobId, payload.scheduleId, "backup.progress", agentId);
|
||||
if (!activeBackupRun) {
|
||||
return;
|
||||
const handleAgentManagerEvent = (event: AgentManagerEvent) => {
|
||||
switch (event.type) {
|
||||
case "agent.disconnected": {
|
||||
cancelActiveBackupRunsForAgent(
|
||||
event.agentId,
|
||||
"The connection to the backup agent was lost. Restart the backup to ensure it completes.",
|
||||
);
|
||||
break;
|
||||
}
|
||||
|
||||
activeBackupRun.onProgress(payload.progress);
|
||||
},
|
||||
onBackupCompleted: ({ agentId, payload }) => {
|
||||
const activeBackupRun = getActiveBackupRun(payload.jobId, payload.scheduleId, "backup.completed", agentId);
|
||||
if (!activeBackupRun) {
|
||||
return;
|
||||
case "backup.started": {
|
||||
getActiveBackupRun(event.payload.jobId, event.payload.scheduleId, event.type, event.agentId);
|
||||
break;
|
||||
}
|
||||
case "backup.progress": {
|
||||
const activeBackupRun = getActiveBackupRun(
|
||||
event.payload.jobId,
|
||||
event.payload.scheduleId,
|
||||
event.type,
|
||||
event.agentId,
|
||||
);
|
||||
if (!activeBackupRun) {
|
||||
break;
|
||||
}
|
||||
|
||||
resolveActiveBackupRun(activeBackupRun.scheduleId, {
|
||||
status: "completed",
|
||||
exitCode: payload.exitCode,
|
||||
result: payload.result,
|
||||
warningDetails: payload.warningDetails ?? null,
|
||||
});
|
||||
},
|
||||
onBackupFailed: ({ agentId, payload }) => {
|
||||
const activeBackupRun = getActiveBackupRun(payload.jobId, payload.scheduleId, "backup.failed", agentId);
|
||||
if (!activeBackupRun) {
|
||||
return;
|
||||
activeBackupRun.onProgress(event.payload.progress);
|
||||
break;
|
||||
}
|
||||
case "backup.completed": {
|
||||
const activeBackupRun = getActiveBackupRun(
|
||||
event.payload.jobId,
|
||||
event.payload.scheduleId,
|
||||
event.type,
|
||||
event.agentId,
|
||||
);
|
||||
if (!activeBackupRun) {
|
||||
break;
|
||||
}
|
||||
|
||||
resolveActiveBackupRun(activeBackupRun.scheduleId, {
|
||||
status: "failed",
|
||||
error: payload.errorDetails ?? payload.error,
|
||||
});
|
||||
},
|
||||
onBackupCancelled: ({ agentId, payload }) => {
|
||||
const activeBackupRun = getActiveBackupRun(payload.jobId, payload.scheduleId, "backup.cancelled", agentId);
|
||||
if (!activeBackupRun) {
|
||||
return;
|
||||
resolveActiveBackupRun(activeBackupRun.scheduleId, {
|
||||
status: "completed",
|
||||
exitCode: event.payload.exitCode,
|
||||
result: event.payload.result,
|
||||
warningDetails: event.payload.warningDetails ?? null,
|
||||
});
|
||||
break;
|
||||
}
|
||||
case "backup.failed": {
|
||||
const activeBackupRun = getActiveBackupRun(
|
||||
event.payload.jobId,
|
||||
event.payload.scheduleId,
|
||||
event.type,
|
||||
event.agentId,
|
||||
);
|
||||
if (!activeBackupRun) {
|
||||
break;
|
||||
}
|
||||
|
||||
resolveActiveBackupRun(activeBackupRun.scheduleId, {
|
||||
status: "cancelled",
|
||||
message: activeBackupRun.cancellationRequested ? undefined : payload.message,
|
||||
});
|
||||
},
|
||||
resolveActiveBackupRun(activeBackupRun.scheduleId, {
|
||||
status: "failed",
|
||||
error: event.payload.errorDetails ?? event.payload.error,
|
||||
});
|
||||
break;
|
||||
}
|
||||
case "backup.cancelled": {
|
||||
const activeBackupRun = getActiveBackupRun(
|
||||
event.payload.jobId,
|
||||
event.payload.scheduleId,
|
||||
event.type,
|
||||
event.agentId,
|
||||
);
|
||||
if (!activeBackupRun) {
|
||||
break;
|
||||
}
|
||||
|
||||
resolveActiveBackupRun(activeBackupRun.scheduleId, {
|
||||
status: "cancelled",
|
||||
message: activeBackupRun.cancellationRequested ? undefined : event.payload.message,
|
||||
});
|
||||
break;
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
export const startAgentController = async () => {
|
||||
const runtime = getAgentRuntimeState();
|
||||
|
||||
if (runtime.agentManager) {
|
||||
await runtime.agentManager.stop();
|
||||
await Effect.runPromise(runtime.agentManager.stop);
|
||||
runtime.agentManager = null;
|
||||
}
|
||||
|
||||
const nextAgentManager = createAgentManagerRuntime(backupEventHandlers);
|
||||
await nextAgentManager.start();
|
||||
const nextAgentManager = createAgentManagerRuntime(handleAgentManagerEvent);
|
||||
await Effect.runPromise(nextAgentManager.start);
|
||||
runtime.agentManager = nextAgentManager;
|
||||
};
|
||||
|
||||
|
|
@ -181,7 +212,9 @@ export const stopAgentController = async () => {
|
|||
const runtime = getAgentRuntimeState();
|
||||
const agentManagerRuntime = runtime.agentManager;
|
||||
runtime.agentManager = null;
|
||||
await agentManagerRuntime?.stop();
|
||||
if (agentManagerRuntime) {
|
||||
await Effect.runPromise(agentManagerRuntime.stop);
|
||||
}
|
||||
};
|
||||
|
||||
export const agentManager = {
|
||||
|
|
@ -212,7 +245,7 @@ export const agentManager = {
|
|||
});
|
||||
|
||||
try {
|
||||
if (!(await runtime.sendBackup(agentId, request.payload))) {
|
||||
if (!(await Effect.runPromise(runtime.sendBackup(agentId, request.payload)))) {
|
||||
clearActiveBackupRun(request.scheduleId);
|
||||
return {
|
||||
status: "unavailable",
|
||||
|
|
|
|||
|
|
@ -1,38 +1,31 @@
|
|||
import { Data, Effect, Exit, Fiber, Scope } from "effect";
|
||||
import { logger } from "@zerobyte/core/node";
|
||||
import { toMessage } from "@zerobyte/core/utils";
|
||||
import type {
|
||||
BackupCancelPayload,
|
||||
BackupCancelledPayload,
|
||||
BackupCompletedPayload,
|
||||
BackupFailedPayload,
|
||||
BackupProgressPayload,
|
||||
BackupRunPayload,
|
||||
BackupStartedPayload,
|
||||
} from "@zerobyte/contracts/agent-protocol";
|
||||
import { createControllerAgentSession, type AgentConnectionData, type ControllerAgentSession } from "./session";
|
||||
import type { AgentMessage, BackupCancelPayload, BackupRunPayload } from "@zerobyte/contracts/agent-protocol";
|
||||
import {
|
||||
createControllerAgentSession,
|
||||
type AgentConnectionData,
|
||||
type ControllerAgentSession,
|
||||
type ControllerAgentSessionEvent,
|
||||
} from "./session";
|
||||
import { agentsService } from "../agents.service";
|
||||
import { validateAgentToken } from "../helpers/tokens";
|
||||
|
||||
type AgentBackupEventContext = {
|
||||
type AgentEventContext = {
|
||||
agentId: string;
|
||||
agentName: string;
|
||||
payload:
|
||||
| BackupStartedPayload
|
||||
| BackupProgressPayload
|
||||
| BackupCompletedPayload
|
||||
| BackupFailedPayload
|
||||
| BackupCancelledPayload;
|
||||
};
|
||||
|
||||
export type AgentBackupEventHandlers = {
|
||||
onAgentDisconnected?: (context: { agentId: string; agentName: string }) => void;
|
||||
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 AgentBackupMessage = Extract<
|
||||
AgentMessage,
|
||||
{
|
||||
type: "backup.started" | "backup.progress" | "backup.completed" | "backup.failed" | "backup.cancelled";
|
||||
}
|
||||
>;
|
||||
|
||||
export type AgentManagerEvent =
|
||||
| (AgentEventContext & { type: "agent.disconnected" })
|
||||
| (AgentEventContext & AgentBackupMessage);
|
||||
|
||||
type ControllerAgentSessionHandle = {
|
||||
agentId: string;
|
||||
|
|
@ -44,14 +37,18 @@ class StopAgentManagerServerError extends Data.TaggedError("StopAgentManagerServ
|
|||
cause: unknown;
|
||||
}> {}
|
||||
|
||||
export function createAgentManagerRuntime(handlers: AgentBackupEventHandlers) {
|
||||
export function createAgentManagerRuntime(onEvent: (event: AgentManagerEvent) => void) {
|
||||
let sessions = new Map<string, ControllerAgentSessionHandle>();
|
||||
let runtimeScope: Scope.CloseableScope | null = null;
|
||||
|
||||
const closeSession = (sessionHandle: ControllerAgentSessionHandle) =>
|
||||
Effect.gen(function* () {
|
||||
yield* Scope.close(sessionHandle.scope, Exit.succeed(undefined));
|
||||
yield* Effect.sync(() => sessions.delete(sessionHandle.agentId));
|
||||
yield* Effect.sync(() => {
|
||||
if (sessions.get(sessionHandle.agentId) === sessionHandle) {
|
||||
sessions.delete(sessionHandle.agentId);
|
||||
}
|
||||
});
|
||||
});
|
||||
|
||||
const closeAllSessions = Effect.gen(function* () {
|
||||
|
|
@ -66,40 +63,32 @@ export function createAgentManagerRuntime(handlers: AgentBackupEventHandlers) {
|
|||
const getSessionHandle = (agentId: string) => sessions.get(agentId);
|
||||
const getSession = (agentId: string) => getSessionHandle(agentId)?.session;
|
||||
|
||||
const createSessionHandlers = (params: { agentId: string; agentName: string; sessionId: string }) => {
|
||||
const handleSessionEvent = (params: { agentId: string; agentName: string; sessionId: string }) => {
|
||||
const { agentId, agentName, sessionId } = params;
|
||||
|
||||
return {
|
||||
onReady: ({ at }: { at: number }) => {
|
||||
return Effect.promise(async () => {
|
||||
await agentsService.markAgentOnline(agentId, at);
|
||||
});
|
||||
},
|
||||
onHeartbeatPong: ({ at }: { at: number }) => {
|
||||
return Effect.promise(() => agentsService.markAgentSeen(agentId, at));
|
||||
},
|
||||
onDisconnect: () => {
|
||||
if (getSession(agentId)?.connectionId !== sessionId) {
|
||||
return Effect.void;
|
||||
return (event: ControllerAgentSessionEvent) => {
|
||||
switch (event.type) {
|
||||
case "agent.ready": {
|
||||
const at = Date.now();
|
||||
return Effect.promise(async () => {
|
||||
await agentsService.markAgentOnline(agentId, at);
|
||||
});
|
||||
}
|
||||
case "heartbeat.pong": {
|
||||
const at = Date.now();
|
||||
return Effect.promise(() => agentsService.markAgentSeen(agentId, at));
|
||||
}
|
||||
case "agent.disconnected": {
|
||||
if (getSession(agentId)?.connectionId !== sessionId) {
|
||||
return Effect.void;
|
||||
}
|
||||
|
||||
return Effect.sync(() => handlers.onAgentDisconnected?.({ agentId, agentName }));
|
||||
},
|
||||
onBackupStarted: (payload: BackupStartedPayload) => {
|
||||
return Effect.sync(() => handlers.onBackupStarted?.({ agentId, agentName, payload }));
|
||||
},
|
||||
onBackupProgress: (payload: BackupProgressPayload) => {
|
||||
return Effect.sync(() => handlers.onBackupProgress?.({ agentId, agentName, payload }));
|
||||
},
|
||||
onBackupCompleted: (payload: BackupCompletedPayload) => {
|
||||
return Effect.sync(() => handlers.onBackupCompleted?.({ agentId, agentName, payload }));
|
||||
},
|
||||
onBackupFailed: (payload: BackupFailedPayload) => {
|
||||
return Effect.sync(() => handlers.onBackupFailed?.({ agentId, agentName, payload }));
|
||||
},
|
||||
onBackupCancelled: (payload: BackupCancelledPayload) => {
|
||||
return Effect.sync(() => handlers.onBackupCancelled?.({ agentId, agentName, payload }));
|
||||
},
|
||||
return Effect.sync(() => onEvent({ type: "agent.disconnected", agentId, agentName }));
|
||||
}
|
||||
default: {
|
||||
return Effect.sync(() => onEvent({ ...event, agentId, agentName }));
|
||||
}
|
||||
}
|
||||
};
|
||||
};
|
||||
|
||||
|
|
@ -110,7 +99,7 @@ export function createAgentManagerRuntime(handlers: AgentBackupEventHandlers) {
|
|||
const session = yield* Scope.extend(
|
||||
createControllerAgentSession(
|
||||
ws,
|
||||
createSessionHandlers({
|
||||
handleSessionEvent({
|
||||
agentId: ws.data.agentId,
|
||||
agentName: ws.data.agentName,
|
||||
sessionId: ws.data.id,
|
||||
|
|
@ -118,7 +107,7 @@ export function createAgentManagerRuntime(handlers: AgentBackupEventHandlers) {
|
|||
),
|
||||
scope,
|
||||
);
|
||||
const runFiber = yield* Effect.fork(Scope.extend(session.run, scope));
|
||||
const runFiber = yield* Effect.forkDaemon(Scope.extend(session.run, scope));
|
||||
yield* Scope.addFinalizer(scope, Fiber.interrupt(runFiber));
|
||||
|
||||
return { agentId: ws.data.agentId, session, scope };
|
||||
|
|
@ -163,6 +152,28 @@ export function createAgentManagerRuntime(handlers: AgentBackupEventHandlers) {
|
|||
yield* session.handleMessage(data);
|
||||
});
|
||||
|
||||
const handleOpen = (ws: Bun.ServerWebSocket<AgentConnectionData>) =>
|
||||
Effect.gen(function* () {
|
||||
yield* Effect.promise(() =>
|
||||
agentsService.markAgentConnecting({
|
||||
agentId: ws.data.agentId,
|
||||
organizationId: ws.data.organizationId,
|
||||
agentName: ws.data.agentName,
|
||||
agentKind: ws.data.agentKind,
|
||||
}),
|
||||
);
|
||||
|
||||
const sessionHandle = yield* createSession(ws);
|
||||
yield* setSession(sessionHandle);
|
||||
yield* logger.effect.info(`Agent "${ws.data.agentName}" (${ws.data.agentId}) connected on ${ws.data.id}`);
|
||||
});
|
||||
|
||||
const handleClose = (ws: Bun.ServerWebSocket<AgentConnectionData>) =>
|
||||
Effect.gen(function* () {
|
||||
yield* removeSession(ws.data.agentId, ws.data.id);
|
||||
yield* logger.effect.info(`Agent "${ws.data.agentName}" (${ws.data.agentId}) disconnected`);
|
||||
});
|
||||
|
||||
const acquireServer = Effect.acquireRelease(
|
||||
Effect.sync(() =>
|
||||
Bun.serve<AgentConnectionData>({
|
||||
|
|
@ -194,24 +205,13 @@ export function createAgentManagerRuntime(handlers: AgentBackupEventHandlers) {
|
|||
},
|
||||
websocket: {
|
||||
open: async (ws) => {
|
||||
await agentsService.markAgentConnecting({
|
||||
agentId: ws.data.agentId,
|
||||
organizationId: ws.data.organizationId,
|
||||
agentName: ws.data.agentName,
|
||||
agentKind: ws.data.agentKind,
|
||||
});
|
||||
|
||||
const sessionHandle = await Effect.runPromise(createSession(ws));
|
||||
await Effect.runPromise(setSession(sessionHandle));
|
||||
|
||||
logger.info(`Agent "${ws.data.agentName}" (${ws.data.agentId}) connected on ${ws.data.id}`);
|
||||
await Effect.runPromise(handleOpen(ws));
|
||||
},
|
||||
message: async (ws, data) => {
|
||||
await Effect.runPromise(handleMessage(ws, data));
|
||||
},
|
||||
close: async (ws) => {
|
||||
await Effect.runPromise(removeSession(ws.data.agentId, ws.data.id));
|
||||
logger.info(`Agent "${ws.data.agentName}" (${ws.data.agentId}) disconnected`);
|
||||
await Effect.runPromise(handleClose(ws));
|
||||
},
|
||||
},
|
||||
}),
|
||||
|
|
@ -230,7 +230,7 @@ export function createAgentManagerRuntime(handlers: AgentBackupEventHandlers) {
|
|||
),
|
||||
);
|
||||
|
||||
const stop = async () => {
|
||||
const stop = Effect.gen(function* () {
|
||||
if (!runtimeScope) {
|
||||
return;
|
||||
}
|
||||
|
|
@ -238,67 +238,67 @@ export function createAgentManagerRuntime(handlers: AgentBackupEventHandlers) {
|
|||
logger.info("Stopping Agent Manager...");
|
||||
const scope = runtimeScope;
|
||||
runtimeScope = null;
|
||||
await Effect.runPromise(Scope.close(scope, Exit.succeed(undefined)));
|
||||
};
|
||||
yield* Scope.close(scope, Exit.succeed(undefined));
|
||||
});
|
||||
|
||||
// TODO: Move the effect boundary up
|
||||
const start = async () => {
|
||||
const start = Effect.gen(function* () {
|
||||
if (runtimeScope) {
|
||||
await stop();
|
||||
yield* stop;
|
||||
}
|
||||
|
||||
logger.info("Starting Agent Manager...");
|
||||
const scope = Effect.runSync(Scope.make());
|
||||
const scope = yield* Scope.make();
|
||||
|
||||
try {
|
||||
const server = Effect.runSync(Scope.extend(acquireServer, scope));
|
||||
runtimeScope = scope;
|
||||
logger.info(`Agent Manager listening on port ${server.port}`);
|
||||
} catch (error) {
|
||||
await Effect.runPromise(Scope.close(scope, Exit.fail(error)));
|
||||
throw error;
|
||||
}
|
||||
};
|
||||
const server = yield* Scope.extend(acquireServer, scope).pipe(
|
||||
Effect.catchAllCause((cause) =>
|
||||
Scope.close(scope, Exit.failCause(cause)).pipe(Effect.andThen(Effect.failCause(cause))),
|
||||
),
|
||||
);
|
||||
runtimeScope = scope;
|
||||
logger.info(`Agent Manager listening on port ${server.port}`);
|
||||
});
|
||||
|
||||
return {
|
||||
start,
|
||||
sendBackup: async (agentId: string, payload: BackupRunPayload) => {
|
||||
const session = getSession(agentId);
|
||||
sendBackup: (agentId: string, payload: BackupRunPayload) =>
|
||||
Effect.gen(function* () {
|
||||
const session = getSession(agentId);
|
||||
|
||||
if (!session) {
|
||||
logger.warn(`Cannot send backup command. Agent ${agentId} is not connected.`);
|
||||
return false;
|
||||
}
|
||||
if (!session) {
|
||||
logger.warn(`Cannot send backup command. Agent ${agentId} is not connected.`);
|
||||
return false;
|
||||
}
|
||||
|
||||
if (!Effect.runSync(session.isReady())) {
|
||||
logger.warn(`Cannot send backup command. Agent ${agentId} is not ready.`);
|
||||
return false;
|
||||
}
|
||||
if (!(yield* session.isReady())) {
|
||||
logger.warn(`Cannot send backup command. Agent ${agentId} is not ready.`);
|
||||
return false;
|
||||
}
|
||||
|
||||
if (!(await Effect.runPromise(session.sendBackup(payload)))) {
|
||||
logger.warn(`Cannot send backup command. Agent ${agentId} is no longer accepting commands.`);
|
||||
return false;
|
||||
}
|
||||
if (!(yield* session.sendBackup(payload))) {
|
||||
logger.warn(`Cannot send backup command. Agent ${agentId} is no longer accepting commands.`);
|
||||
return false;
|
||||
}
|
||||
|
||||
logger.info(`Sent backup command ${payload.jobId} to agent ${agentId} for schedule ${payload.scheduleId}`);
|
||||
return true;
|
||||
},
|
||||
cancelBackup: async (agentId: string, payload: BackupCancelPayload) => {
|
||||
const session = getSession(agentId);
|
||||
logger.info(`Sent backup command ${payload.jobId} to agent ${agentId} for schedule ${payload.scheduleId}`);
|
||||
return true;
|
||||
}),
|
||||
cancelBackup: (agentId: string, payload: BackupCancelPayload) =>
|
||||
Effect.gen(function* () {
|
||||
const session = getSession(agentId);
|
||||
|
||||
if (!session) {
|
||||
logger.warn(`Cannot cancel backup command. Agent ${agentId} is not connected.`);
|
||||
return false;
|
||||
}
|
||||
if (!session) {
|
||||
logger.warn(`Cannot cancel backup command. Agent ${agentId} is not connected.`);
|
||||
return false;
|
||||
}
|
||||
|
||||
if (!(await Effect.runPromise(session.sendBackupCancel(payload)))) {
|
||||
logger.warn(`Cannot cancel backup command. Agent ${agentId} is no longer accepting commands.`);
|
||||
return false;
|
||||
}
|
||||
if (!(yield* session.sendBackupCancel(payload))) {
|
||||
logger.warn(`Cannot cancel backup command. Agent ${agentId} is no longer accepting commands.`);
|
||||
return false;
|
||||
}
|
||||
|
||||
logger.info(`Sent backup cancel for command ${payload.jobId} to agent ${agentId}`);
|
||||
return true;
|
||||
},
|
||||
logger.info(`Sent backup cancel for command ${payload.jobId} to agent ${agentId}`);
|
||||
return true;
|
||||
}),
|
||||
stop,
|
||||
};
|
||||
}
|
||||
|
|
|
|||
|
|
@ -5,12 +5,7 @@ import {
|
|||
parseAgentMessage,
|
||||
type AgentMessage,
|
||||
type BackupCancelPayload,
|
||||
type BackupCancelledPayload,
|
||||
type BackupCompletedPayload,
|
||||
type BackupFailedPayload,
|
||||
type BackupProgressPayload,
|
||||
type BackupRunPayload,
|
||||
type BackupStartedPayload,
|
||||
type ControllerWireMessage,
|
||||
} from "@zerobyte/contracts/agent-protocol";
|
||||
import { logger } from "@zerobyte/core/node";
|
||||
|
|
@ -32,24 +27,7 @@ type SessionState = {
|
|||
lastPongAt: number | null;
|
||||
};
|
||||
|
||||
type AgentRuntimeEventPayload = {
|
||||
agentId: string;
|
||||
agentName: string;
|
||||
organizationId: string | null;
|
||||
agentKind: AgentKind;
|
||||
at: number;
|
||||
};
|
||||
|
||||
type ControllerAgentSessionHandlers = {
|
||||
onReady: (payload: AgentRuntimeEventPayload) => Effect.Effect<void>;
|
||||
onHeartbeatPong: (payload: AgentRuntimeEventPayload) => Effect.Effect<void>;
|
||||
onDisconnect: (payload: AgentRuntimeEventPayload) => Effect.Effect<void>;
|
||||
onBackupStarted: (payload: BackupStartedPayload) => Effect.Effect<void>;
|
||||
onBackupProgress: (payload: BackupProgressPayload) => Effect.Effect<void>;
|
||||
onBackupCompleted: (payload: BackupCompletedPayload) => Effect.Effect<void>;
|
||||
onBackupFailed: (payload: BackupFailedPayload) => Effect.Effect<void>;
|
||||
onBackupCancelled: (payload: BackupCancelledPayload) => Effect.Effect<void>;
|
||||
};
|
||||
export type ControllerAgentSessionEvent = AgentMessage | { type: "agent.disconnected" };
|
||||
|
||||
export type ControllerAgentSession = {
|
||||
readonly connectionId: string;
|
||||
|
|
@ -62,7 +40,7 @@ export type ControllerAgentSession = {
|
|||
|
||||
export const createControllerAgentSession = (
|
||||
socket: AgentSocket,
|
||||
handlers: ControllerAgentSessionHandlers,
|
||||
onEvent: (event: ControllerAgentSessionEvent) => Effect.Effect<void>,
|
||||
): Effect.Effect<ControllerAgentSession, never, Scope.Scope> =>
|
||||
Effect.gen(function* () {
|
||||
let isClosed = false;
|
||||
|
|
@ -88,13 +66,7 @@ export const createControllerAgentSession = (
|
|||
const releaseSession = Effect.gen(function* () {
|
||||
const disconnectedAt = Date.now();
|
||||
yield* updateState((current) => ({ ...current, isReady: false, lastSeenAt: disconnectedAt }));
|
||||
yield* handlers.onDisconnect({
|
||||
agentId: socket.data.agentId,
|
||||
agentName: socket.data.agentName,
|
||||
organizationId: socket.data.organizationId,
|
||||
agentKind: socket.data.agentKind,
|
||||
at: disconnectedAt,
|
||||
});
|
||||
yield* onEvent({ type: "agent.disconnected" });
|
||||
|
||||
yield* Queue.shutdown(outboundQueue);
|
||||
});
|
||||
|
|
@ -159,59 +131,18 @@ export const createControllerAgentSession = (
|
|||
|
||||
const handleAgentMessage = (message: AgentMessage) =>
|
||||
Effect.gen(function* () {
|
||||
switch (message.type) {
|
||||
case "agent.ready": {
|
||||
const readyAt = Date.now();
|
||||
yield* updateState((current) => ({ ...current, isReady: true, lastSeenAt: readyAt }));
|
||||
|
||||
yield* handlers.onReady({
|
||||
agentId: socket.data.agentId,
|
||||
agentName: socket.data.agentName,
|
||||
organizationId: socket.data.organizationId,
|
||||
agentKind: socket.data.agentKind,
|
||||
at: readyAt,
|
||||
});
|
||||
|
||||
yield* logger.effect.info(`Agent "${socket.data.agentName}" (${socket.data.agentId}) is ready`);
|
||||
break;
|
||||
}
|
||||
case "backup.started": {
|
||||
yield* logger.effect.info(
|
||||
`Backup ${message.payload.jobId} started on agent ${socket.data.agentId} for schedule ${message.payload.scheduleId}`,
|
||||
);
|
||||
yield* handlers.onBackupStarted(message.payload);
|
||||
break;
|
||||
}
|
||||
case "backup.progress": {
|
||||
yield* handlers.onBackupProgress(message.payload);
|
||||
break;
|
||||
}
|
||||
case "backup.completed": {
|
||||
yield* handlers.onBackupCompleted(message.payload);
|
||||
break;
|
||||
}
|
||||
case "backup.failed": {
|
||||
yield* handlers.onBackupFailed(message.payload);
|
||||
break;
|
||||
}
|
||||
case "backup.cancelled": {
|
||||
yield* handlers.onBackupCancelled(message.payload);
|
||||
break;
|
||||
}
|
||||
case "heartbeat.pong": {
|
||||
const seenAt = Date.now();
|
||||
yield* updateState((current) => ({ ...current, lastSeenAt: seenAt, lastPongAt: message.payload.sentAt }));
|
||||
|
||||
yield* handlers.onHeartbeatPong({
|
||||
agentId: socket.data.agentId,
|
||||
agentName: socket.data.agentName,
|
||||
organizationId: socket.data.organizationId,
|
||||
agentKind: socket.data.agentKind,
|
||||
at: seenAt,
|
||||
});
|
||||
break;
|
||||
}
|
||||
if (message.type === "agent.ready") {
|
||||
const readyAt = Date.now();
|
||||
yield* updateState((current) => ({ ...current, isReady: true, lastSeenAt: readyAt }));
|
||||
yield* logger.effect.info(`Agent "${socket.data.agentName}" (${socket.data.agentId}) is ready`);
|
||||
}
|
||||
|
||||
if (message.type === "heartbeat.pong") {
|
||||
const seenAt = Date.now();
|
||||
yield* updateState((current) => ({ ...current, lastSeenAt: seenAt, lastPongAt: message.payload.sentAt }));
|
||||
}
|
||||
|
||||
yield* onEvent(message);
|
||||
});
|
||||
|
||||
return {
|
||||
|
|
|
|||
Loading…
Reference in a new issue