diff --git a/app/server/modules/agents/__tests__/agents-manager.backups.test.ts b/app/server/modules/agents/__tests__/agents-manager.backups.test.ts index 279f9b2e..25830203 100644 --- a/app/server/modules/agents/__tests__/agents-manager.backups.test.ts +++ b/app/server/modules/agents/__tests__/agents-manager.backups.test.ts @@ -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", { diff --git a/app/server/modules/agents/__tests__/agents-manager.events.test.ts b/app/server/modules/agents/__tests__/agents-manager.events.test.ts new file mode 100644 index 00000000..2b12ef5f --- /dev/null +++ b/app/server/modules/agents/__tests__/agents-manager.events.test.ts @@ -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({ + 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({ 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({ 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(); +}); diff --git a/app/server/modules/agents/__tests__/agents.service.test.ts b/app/server/modules/agents/__tests__/agents.service.test.ts index 9fdc1036..f6f87c8a 100644 --- a/app/server/modules/agents/__tests__/agents.service.test.ts +++ b/app/server/modules/agents/__tests__/agents.service.test.ts @@ -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, + }); +}); diff --git a/app/server/modules/agents/__tests__/controller-runtime.test.ts b/app/server/modules/agents/__tests__/controller-runtime.test.ts new file mode 100644 index 00000000..b96eb6b3 --- /dev/null +++ b/app/server/modules/agents/__tests__/controller-runtime.test.ts @@ -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[0]["fetch"]>; + +const invokeFetch = (fetch: CapturedFetch | undefined, request: Request, srv: Parameters[1]) => { + if (!fetch) { + throw new Error("Bun.serve was not called with a fetch handler"); + } + + return Reflect.apply(fetch, fromPartial>({}), [ + request, + srv, + ]) as ReturnType; +}; + +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>[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>[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>[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); +}); diff --git a/app/server/modules/agents/__tests__/session.test.ts b/app/server/modules/agents/__tests__/session.test.ts index b624c201..5c1766da 100644 --- a/app/server/modules/agents/__tests__/session.test.ts +++ b/app/server/modules/agents/__tests__/session.test.ts @@ -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[1]> = {}, + onEvent: Parameters[1] = () => Effect.void, socket = createSocket(), ) => { const scope = Effect.runSync(Scope.make()); - const sessionHandlers: Parameters[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; + + 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", }), ); }); diff --git a/app/server/modules/agents/agents-manager.ts b/app/server/modules/agents/agents-manager.ts index c0c0c126..e82b8dde 100644 --- a/app/server/modules/agents/agents-manager.ts +++ b/app/server/modules/agents/agents-manager.ts @@ -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", diff --git a/app/server/modules/agents/controller/server.ts b/app/server/modules/agents/controller/server.ts index 961d29e7..2e306941 100644 --- a/app/server/modules/agents/controller/server.ts +++ b/app/server/modules/agents/controller/server.ts @@ -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(); 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) => + 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) => + 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({ @@ -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, }; } diff --git a/app/server/modules/agents/controller/session.ts b/app/server/modules/agents/controller/session.ts index cb355e6f..c2c3d9c7 100644 --- a/app/server/modules/agents/controller/session.ts +++ b/app/server/modules/agents/controller/session.ts @@ -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; - onHeartbeatPong: (payload: AgentRuntimeEventPayload) => Effect.Effect; - onDisconnect: (payload: AgentRuntimeEventPayload) => Effect.Effect; - onBackupStarted: (payload: BackupStartedPayload) => Effect.Effect; - onBackupProgress: (payload: BackupProgressPayload) => Effect.Effect; - onBackupCompleted: (payload: BackupCompletedPayload) => Effect.Effect; - onBackupFailed: (payload: BackupFailedPayload) => Effect.Effect; - onBackupCancelled: (payload: BackupCancelledPayload) => Effect.Effect; -}; +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, ): Effect.Effect => 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 {