From 91286eb84f472209b19c97e410c45ae7bb1a6a38 Mon Sep 17 00:00:00 2001 From: Nicolas Meienberger Date: Fri, 1 May 2026 18:20:22 +0200 Subject: [PATCH] refactor: centralize agent backup lifecycle state --- .../modules/agents/__tests__/session.test.ts | 112 +++--------------- app/server/modules/agents/agents-manager.ts | 18 +++ .../modules/agents/controller/server.ts | 53 +++++---- .../modules/agents/controller/session.ts | 63 ++-------- .../modules/agents/helpers/runtime-state.ts | 1 + app/test/setup-shared.ts | 7 ++ packages/core/test/setup.ts | 7 ++ 7 files changed, 93 insertions(+), 168 deletions(-) diff --git a/app/server/modules/agents/__tests__/session.test.ts b/app/server/modules/agents/__tests__/session.test.ts index b6b88a84..139f5ea3 100644 --- a/app/server/modules/agents/__tests__/session.test.ts +++ b/app/server/modules/agents/__tests__/session.test.ts @@ -2,7 +2,6 @@ import { Effect, Exit, Scope } from "effect"; import { 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"; import { createControllerAgentSession } from "../controller/session"; @@ -29,6 +28,7 @@ const createSession = ( const sessionHandlers: Parameters[1] = { onReady: () => Effect.void, onHeartbeatPong: () => Effect.void, + onDisconnect: () => Effect.void, onBackupStarted: () => Effect.void, onBackupProgress: () => Effect.void, onBackupCompleted: () => Effect.void, @@ -61,87 +61,22 @@ const createSession = ( } }; -test("close emits a synthetic backup.cancelled for a started backup", () => { - const onBackupCancelled = vi.fn(() => Effect.void); - const { session, close } = createSession({ - onBackupCancelled, - onBackupStarted: vi.fn(() => Effect.void), - }); - - Effect.runSync( - session.handleMessage( - createAgentMessage("backup.started", { - jobId: "job-1", - scheduleId: "schedule-1", - }), - ), - ); +test("close reports a transport disconnect", () => { + const onDisconnect = vi.fn(() => Effect.void); + const { close } = createSession({ onDisconnect }); close(); - expect(onBackupCancelled).toHaveBeenCalledTimes(1); - expect(onBackupCancelled).toHaveBeenCalledWith( + expect(onDisconnect).toHaveBeenCalledTimes(1); + expect(onDisconnect).toHaveBeenCalledWith( expect.objectContaining({ - jobId: "job-1", - scheduleId: "schedule-1", + agentId: LOCAL_AGENT_ID, + agentName: LOCAL_AGENT_NAME, }), ); }); -test.each([ - { - name: "backup.completed", - jobId: "job-1", - scheduleId: "schedule-1", - terminalMessage: createAgentMessage("backup.completed", { - jobId: "job-1", - scheduleId: "schedule-1", - exitCode: 0, - result: null, - }), - expectedCancelledCalls: 0, - }, - { - name: "backup.failed", - jobId: "job-2", - scheduleId: "schedule-2", - terminalMessage: createAgentMessage("backup.failed", { - jobId: "job-2", - scheduleId: "schedule-2", - error: "backup failed", - }), - expectedCancelledCalls: 0, - }, - { - name: "backup.cancelled", - jobId: "job-3", - scheduleId: "schedule-3", - terminalMessage: createAgentMessage("backup.cancelled", { - jobId: "job-3", - scheduleId: "schedule-3", - message: "Backup was cancelled", - }), - expectedCancelledCalls: 1, - }, -])("close does not emit an extra synthetic backup.cancelled after $name", (testCase) => { - const onBackupCancelled = vi.fn(() => Effect.void); - const { session, close } = createSession({ onBackupCancelled }); - - Effect.runSync( - session.handleMessage( - createAgentMessage("backup.started", { - jobId: testCase.jobId, - scheduleId: testCase.scheduleId, - }), - ), - ); - Effect.runSync(session.handleMessage(testCase.terminalMessage)); - close(); - - expect(onBackupCancelled).toHaveBeenCalledTimes(testCase.expectedCancelledCalls); -}); - -test("close emits a synthetic backup.cancelled for a queued backup", () => { +test("sendBackup only queues the transport message", () => { const onBackupCancelled = vi.fn(() => Effect.void); const { session, close } = createSession({ onBackupCancelled }); @@ -171,31 +106,17 @@ test("close emits a synthetic backup.cancelled for a queued backup", () => { close(); - expect(onBackupCancelled).toHaveBeenCalledTimes(1); - expect(onBackupCancelled).toHaveBeenLastCalledWith( - expect.objectContaining({ - jobId: "job-queued", - scheduleId: "schedule-queued", - }), - ); + expect(onBackupCancelled).not.toHaveBeenCalled(); }); -test("a dropped backup.cancel closes the session and emits a synthetic backup.cancelled", async () => { +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 onBackupCancelled = vi.fn(() => Effect.void); - const { session, run, closeAsync } = createSession({ onBackupCancelled }, socket); + const onDisconnect = vi.fn(() => Effect.void); + const { session, run, closeAsync } = createSession({ onDisconnect }, socket); try { run(); - Effect.runSync( - session.handleMessage( - createAgentMessage("backup.started", { - jobId: "job-1", - scheduleId: "schedule-1", - }), - ), - ); Effect.runSync( session.sendBackupCancel({ jobId: "job-1", @@ -206,11 +127,10 @@ test("a dropped backup.cancel closes the session and emits a synthetic backup.ca await waitForExpect(() => { expect(send).toHaveBeenCalledTimes(1); expect(socket.close).toHaveBeenCalledTimes(1); - expect(onBackupCancelled).toHaveBeenCalledTimes(1); - expect(onBackupCancelled).toHaveBeenCalledWith( + expect(onDisconnect).toHaveBeenCalledTimes(1); + expect(onDisconnect).toHaveBeenCalledWith( expect.objectContaining({ - jobId: "job-1", - scheduleId: "schedule-1", + agentId: LOCAL_AGENT_ID, }), ); }); diff --git a/app/server/modules/agents/agents-manager.ts b/app/server/modules/agents/agents-manager.ts index 519ab9f9..32dc616f 100644 --- a/app/server/modules/agents/agents-manager.ts +++ b/app/server/modules/agents/agents-manager.ts @@ -46,6 +46,17 @@ const resolveActiveBackupRun = (scheduleId: number, result: BackupExecutionResul return true; }; +const cancelActiveBackupRunsForAgent = (agentId: string, message: string) => { + const activeBackupsByScheduleId = getActiveBackupsByScheduleId(); + const matchingScheduleIds = [...activeBackupsByScheduleId.values()] + .filter((activeBackupRun) => activeBackupRun.agentId === agentId) + .map((activeBackupRun) => activeBackupRun.scheduleId); + + for (const scheduleId of matchingScheduleIds) { + resolveActiveBackupRun(scheduleId, { status: "cancelled", message }); + } +}; + const getActiveBackupRun = (jobId: string, scheduleId: string, eventName: string, agentId: string) => { const trackedScheduleId = getActiveBackupScheduleIdsByJobId().get(jobId); if (trackedScheduleId === undefined) { @@ -99,6 +110,12 @@ const requestBackupCancellation = async (agentId: string, scheduleId: number) => }; 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); }, @@ -185,6 +202,7 @@ export const agentManager = { const completion = new Promise((resolve) => { getActiveBackupsByScheduleId().set(request.scheduleId, { + agentId, scheduleId: request.scheduleId, jobId: request.payload.jobId, scheduleShortId: request.payload.scheduleId, diff --git a/app/server/modules/agents/controller/server.ts b/app/server/modules/agents/controller/server.ts index 3aec9a7f..46fec4a3 100644 --- a/app/server/modules/agents/controller/server.ts +++ b/app/server/modules/agents/controller/server.ts @@ -26,6 +26,7 @@ type AgentBackupEventContext = { }; 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; @@ -55,28 +56,39 @@ export function createAgentManagerRuntime() { }); const closeAllSessions = Effect.gen(function* () { - const currentSessions = sessions; - sessions = new Map(); - for (const sessionHandle of currentSessions.values()) { + const currentSessions = [...sessions.entries()]; + for (const [agentId, sessionHandle] of currentSessions) { + yield* Effect.promise(() => agentsService.markAgentOffline(agentId)); yield* closeSession(sessionHandle); } + sessions = new Map(); }); const getSessionHandle = (agentId: string) => sessions.get(agentId); const getSession = (agentId: string) => getSessionHandle(agentId)?.session; - const createSessionHandlers = (ws: Bun.ServerWebSocket) => { + const createSessionHandlers = (ws: Bun.ServerWebSocket, markConnecting: Promise) => { const agentId = ws.data.agentId; const agentName = ws.data.agentName; return { onReady: ({ at }: { at: number }) => { - return Effect.promise(() => agentsService.markAgentOnline(agentId, at)); + return Effect.promise(async () => { + await markConnecting; + await agentsService.markAgentOnline(agentId, at); + }); }, onHeartbeatPong: ({ at }: { at: number }) => { return Effect.promise(() => agentsService.markAgentSeen(agentId, at)); }, + onDisconnect: () => { + if (getSession(ws.data.agentId)?.connectionId !== ws.data.id) { + return Effect.void; + } + + return Effect.sync(() => backupHandlers.onAgentDisconnected?.({ agentId, agentName })); + }, onBackupStarted: (payload: BackupStartedPayload) => { return Effect.sync(() => backupHandlers.onBackupStarted?.({ agentId, agentName, payload })); }, @@ -95,12 +107,14 @@ export function createAgentManagerRuntime() { }; }; - const createSession = (ws: Bun.ServerWebSocket) => { + const createSession = (ws: Bun.ServerWebSocket, markConnecting: Promise) => { // Manual scope management because we are out of Effect const scope = Effect.runSync(Scope.make()); try { - const session = Effect.runSync(Scope.extend(createControllerAgentSession(ws, createSessionHandlers(ws)), scope)); + const session = Effect.runSync( + Scope.extend(createControllerAgentSession(ws, createSessionHandlers(ws, markConnecting)), scope), + ); const runFiber = Effect.runFork(Scope.extend(session.run, scope)); return { session, runFiber, scope }; @@ -112,13 +126,13 @@ export function createAgentManagerRuntime() { const setSession = (agentId: string, sessionHandle: ControllerAgentSessionHandle) => { const existingSession = getSessionHandle(agentId); + sessions.set(agentId, sessionHandle); + if (existingSession) { void Effect.runPromise(closeSession(existingSession)).catch((error) => { logger.error(`Failed to close existing agent session for ${agentId}: ${toMessage(error)}`); }); } - - sessions.set(agentId, sessionHandle); }; const removeSession = (agentId: string, connectionId: string) => { @@ -165,17 +179,16 @@ export function createAgentManagerRuntime() { }, websocket: { open: (ws) => { - setSession(ws.data.agentId, createSession(ws)); - void agentsService - .markAgentConnecting({ - agentId: ws.data.agentId, - organizationId: ws.data.organizationId, - agentName: ws.data.agentName, - agentKind: ws.data.agentKind, - }) - .catch((error) => { - logger.error(`Failed to mark agent ${ws.data.agentId} as connecting: ${toMessage(error)}`); - }); + const markConnecting = agentsService.markAgentConnecting({ + agentId: ws.data.agentId, + organizationId: ws.data.organizationId, + agentName: ws.data.agentName, + agentKind: ws.data.agentKind, + }); + void markConnecting.catch((error) => { + logger.error(`Failed to mark agent ${ws.data.agentId} as connecting: ${toMessage(error)}`); + }); + setSession(ws.data.agentId, createSession(ws, markConnecting)); logger.info(`Agent "${ws.data.agentName}" (${ws.data.agentId}) connected on ${ws.data.id}`); }, message: (ws, data) => { diff --git a/app/server/modules/agents/controller/session.ts b/app/server/modules/agents/controller/session.ts index b8f0bdc0..cb355e6f 100644 --- a/app/server/modules/agents/controller/session.ts +++ b/app/server/modules/agents/controller/session.ts @@ -32,11 +32,6 @@ type SessionState = { lastPongAt: number | null; }; -type TrackedBackupJob = { - scheduleId: string; - state: "pending" | "active"; -}; - type AgentRuntimeEventPayload = { agentId: string; agentName: string; @@ -48,6 +43,7 @@ type AgentRuntimeEventPayload = { 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; @@ -71,7 +67,6 @@ export const createControllerAgentSession = ( Effect.gen(function* () { let isClosed = false; const outboundQueue = yield* Queue.bounded(64); - const trackedBackupJobs = yield* Ref.make>(new Map()); const state = yield* Ref.make({ isReady: false, lastSeenAt: null, @@ -90,35 +85,16 @@ export const createControllerAgentSession = ( const updateState = (update: (current: SessionState) => SessionState) => Ref.update(state, update); - const setTrackedBackupJob = (jobId: string, trackedBackupJob: TrackedBackupJob) => { - return Ref.update(trackedBackupJobs, (current) => { - const next = new Map(current); - next.set(jobId, trackedBackupJob); - return next; - }); - }; - - const deleteTrackedBackupJob = (jobId: string) => { - return Ref.update(trackedBackupJobs, (current) => { - const next = new Map(current); - next.delete(jobId); - return next; - }); - }; - - const takeTrackedBackupJobs = Ref.modify( - trackedBackupJobs, - (current) => [current, new Map()] as const, - ); - const releaseSession = Effect.gen(function* () { - yield* updateState((current) => ({ ...current, isReady: false })); - const trackedJobs = yield* takeTrackedBackupJobs; - for (const [jobId, trackedJob] of trackedJobs) { - const message = "The connection to the backup agent was lost. Restart the backup to ensure it completes."; - - yield* handlers.onBackupCancelled({ jobId, scheduleId: trackedJob.scheduleId, message }); - } + 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* Queue.shutdown(outboundQueue); }); @@ -200,10 +176,6 @@ export const createControllerAgentSession = ( break; } case "backup.started": { - yield* setTrackedBackupJob(message.payload.jobId, { - scheduleId: message.payload.scheduleId, - state: "active", - }); yield* logger.effect.info( `Backup ${message.payload.jobId} started on agent ${socket.data.agentId} for schedule ${message.payload.scheduleId}`, ); @@ -215,17 +187,14 @@ export const createControllerAgentSession = ( break; } case "backup.completed": { - yield* deleteTrackedBackupJob(message.payload.jobId); yield* handlers.onBackupCompleted(message.payload); break; } case "backup.failed": { - yield* deleteTrackedBackupJob(message.payload.jobId); yield* handlers.onBackupFailed(message.payload); break; } case "backup.cancelled": { - yield* deleteTrackedBackupJob(message.payload.jobId); yield* handlers.onBackupCancelled(message.payload); break; } @@ -264,17 +233,7 @@ export const createControllerAgentSession = ( yield* handleAgentMessage(parsed.data); }); }, - sendBackup: (payload) => { - return Effect.gen(function* () { - const queued = yield* offerOutbound(createControllerMessage("backup.run", payload)); - - if (queued) { - yield* setTrackedBackupJob(payload.jobId, { scheduleId: payload.scheduleId, state: "pending" }); - } - - return queued; - }); - }, + sendBackup: (payload) => offerOutbound(createControllerMessage("backup.run", payload)), sendBackupCancel: (payload) => offerOutbound(createControllerMessage("backup.cancel", payload)), isReady: () => Ref.get(state).pipe(Effect.map((current) => current.isReady)), run, diff --git a/app/server/modules/agents/helpers/runtime-state.ts b/app/server/modules/agents/helpers/runtime-state.ts index 51de3fb4..13fc2a15 100644 --- a/app/server/modules/agents/helpers/runtime-state.ts +++ b/app/server/modules/agents/helpers/runtime-state.ts @@ -11,6 +11,7 @@ export type BackupExecutionResult = | { status: "cancelled"; message?: string }; type ActiveBackupRun = { + agentId: string; scheduleId: number; jobId: string; scheduleShortId: string; diff --git a/app/test/setup-shared.ts b/app/test/setup-shared.ts index b9e12c55..81320690 100644 --- a/app/test/setup-shared.ts +++ b/app/test/setup-shared.ts @@ -1,4 +1,5 @@ import { vi } from "vitest"; +import { Effect } from "effect"; process.env.BASE_URL = "http://localhost:3000"; process.env.TRUSTED_ORIGINS = "http://localhost:3000"; @@ -13,6 +14,12 @@ vi.mock(import("@zerobyte/core/node"), async () => { info: () => {}, warn: () => {}, error: () => {}, + effect: { + debug: () => Effect.void, + info: () => Effect.void, + warn: () => Effect.void, + error: () => Effect.void, + }, }, }; }); diff --git a/packages/core/test/setup.ts b/packages/core/test/setup.ts index 37612ad0..5ccc1b2d 100644 --- a/packages/core/test/setup.ts +++ b/packages/core/test/setup.ts @@ -1,4 +1,5 @@ import { vi } from "vitest"; +import { Effect } from "effect"; vi.mock(import("../src/node/logger.ts"), () => ({ logger: { @@ -6,5 +7,11 @@ vi.mock(import("../src/node/logger.ts"), () => ({ info: () => {}, warn: () => {}, error: () => {}, + effect: { + debug: () => Effect.void, + info: () => Effect.void, + warn: () => Effect.void, + error: () => Effect.void, + }, }, }));