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 c29ccd2a..279f9b2e 100644 --- a/app/server/modules/agents/__tests__/agents-manager.backups.test.ts +++ b/app/server/modules/agents/__tests__/agents-manager.backups.test.ts @@ -1,22 +1,13 @@ import { afterEach, expect, test, vi } from "vitest"; import waitForExpect from "wait-for-expect"; -import { fromPartial } from "@total-typescript/shoehorn"; -import { agentManager } from "../agents-manager"; +import { fromAny, fromPartial } from "@total-typescript/shoehorn"; +import { agentManager, type ProcessWithAgentRuntime } from "../agents-manager"; import type { AgentManagerRuntime } from "../controller/server"; import type { BackupRunPayload } from "@zerobyte/contracts/agent-protocol"; -type ProcessWithAgentRuntime = NodeJS.Process & { - __zerobyteAgentRuntime?: { - agentManager: AgentManagerRuntime | null; - localAgent: null; - isStoppingLocalAgent: boolean; - localAgentRestartTimeout: null; - }; -}; - const setAgentRuntime = (agentManagerRuntime: Partial | null) => { (process as ProcessWithAgentRuntime).__zerobyteAgentRuntime = { - agentManager: agentManagerRuntime ? fromPartial(agentManagerRuntime) : null, + agentManager: fromAny(agentManagerRuntime), localAgent: null, isStoppingLocalAgent: false, localAgentRestartTimeout: null, diff --git a/app/server/modules/agents/__tests__/agents-manager.test.ts b/app/server/modules/agents/__tests__/agents-manager.test.ts index 1d882916..dd9310f2 100644 --- a/app/server/modules/agents/__tests__/agents-manager.test.ts +++ b/app/server/modules/agents/__tests__/agents-manager.test.ts @@ -1,16 +1,11 @@ import { EventEmitter } from "node:events"; -import type * as childProcess from "node:child_process"; import { PassThrough } from "node:stream"; import { afterEach, expect, test, vi } from "vitest"; const spawnMock = vi.fn(); vi.mock("node:child_process", async () => { - const actual = await vi.importActual("node:child_process"); - return { - ...actual, - spawn: spawnMock, - }; + return { spawn: spawnMock }; }); const { spawnLocalAgent, stopLocalAgent } = await import("../agents-manager"); @@ -51,9 +46,7 @@ test("respawns the local agent after an unexpected exit", async () => { const firstChild = createFakeChild(); const secondChild = createFakeChild(); - spawnMock - .mockReturnValueOnce(firstChild as unknown as ReturnType) - .mockReturnValueOnce(secondChild as unknown as ReturnType); + spawnMock.mockReturnValueOnce(firstChild).mockReturnValueOnce(secondChild); await spawnLocalAgent(); @@ -69,7 +62,7 @@ test("does not respawn the local agent after an intentional stop", async () => { vi.useFakeTimers(); const child = createFakeChild(); - spawnMock.mockReturnValue(child as unknown as ReturnType); + spawnMock.mockReturnValue(child); await spawnLocalAgent(); await stopLocalAgent(); diff --git a/app/server/modules/agents/__tests__/session.test.ts b/app/server/modules/agents/__tests__/session.test.ts index e8913fd1..de184c01 100644 --- a/app/server/modules/agents/__tests__/session.test.ts +++ b/app/server/modules/agents/__tests__/session.test.ts @@ -6,19 +6,19 @@ import { createAgentMessage } from "@zerobyte/contracts/agent-protocol"; import { createControllerAgentSession } from "../controller/session"; const createSocket = (overrides: Partial[0]> = {}) => { - return fromPartial[0]>({ + return { data: { id: "connection-1", agentId: "local", organizationId: null, agentName: "Local Agent" }, send: vi.fn(() => 1), close: vi.fn(), ...overrides, - }); + }; }; const createSession = (handlers: Parameters[1] = {}, socket = createSocket()) => { const scope = Effect.runSync(Scope.make()); try { - const session = Effect.runSync(Scope.extend(createControllerAgentSession(socket, handlers), scope)); + const session = Effect.runSync(Scope.extend(createControllerAgentSession(fromPartial(socket), handlers), scope)); return { session, @@ -57,12 +57,12 @@ test("close emits a synthetic backup.cancelled for a started backup", () => { close(); expect(onBackupCancelled).toHaveBeenCalledTimes(1); - expect(onBackupCancelled).toHaveBeenCalledWith({ - jobId: "job-1", - scheduleId: "schedule-1", - message: - "The connection to the backup agent was lost while this backup was running. Restart the backup to ensure it completes.", - }); + expect(onBackupCancelled).toHaveBeenCalledWith( + expect.objectContaining({ + jobId: "job-1", + scheduleId: "schedule-1", + }), + ); }); test.each([ @@ -150,12 +150,12 @@ test("close emits a synthetic backup.cancelled for a queued backup", () => { close(); expect(onBackupCancelled).toHaveBeenCalledTimes(1); - expect(onBackupCancelled).toHaveBeenCalledWith({ - jobId: "job-queued", - scheduleId: "schedule-queued", - message: - "The connection to the backup agent was lost before this backup started. Restart the backup to ensure it completes.", - }); + expect(onBackupCancelled).toHaveBeenLastCalledWith( + expect.objectContaining({ + jobId: "job-queued", + scheduleId: "schedule-queued", + }), + ); }); test("a dropped backup.cancel closes the session and emits a synthetic backup.cancelled", async () => { @@ -185,12 +185,12 @@ test("a dropped backup.cancel closes the session and emits a synthetic backup.ca expect(send).toHaveBeenCalledTimes(1); expect(socket.close).toHaveBeenCalledTimes(1); expect(onBackupCancelled).toHaveBeenCalledTimes(1); - expect(onBackupCancelled).toHaveBeenCalledWith({ - jobId: "job-1", - scheduleId: "schedule-1", - message: - "The connection to the backup agent was lost while this backup was running. Restart the backup to ensure it completes.", - }); + expect(onBackupCancelled).toHaveBeenCalledWith( + expect.objectContaining({ + jobId: "job-1", + scheduleId: "schedule-1", + }), + ); }); } finally { await closeAsync(); diff --git a/app/server/modules/agents/agents-manager.ts b/app/server/modules/agents/agents-manager.ts index 893ccd1a..ab8048cb 100644 --- a/app/server/modules/agents/agents-manager.ts +++ b/app/server/modules/agents/agents-manager.ts @@ -19,38 +19,6 @@ export type AgentRunBackupRequest = { onProgress: (progress: BackupExecutionProgress) => void; }; -type AgentRuntimeState = { - agentManager: AgentManagerRuntime | null; - localAgent: ChildProcess | null; - isStoppingLocalAgent: boolean; - localAgentRestartTimeout: ReturnType | null; -}; - -type ProcessWithAgentRuntime = NodeJS.Process & { - __zerobyteAgentRuntime?: AgentRuntimeState; -}; - -const getAgentRuntimeState = () => { - const runtimeProcess = process as ProcessWithAgentRuntime; - const existingRuntime = runtimeProcess.__zerobyteAgentRuntime; - - if (existingRuntime) { - return existingRuntime; - } - - const runtime = { - agentManager: null, - localAgent: null, - isStoppingLocalAgent: false, - localAgentRestartTimeout: null, - }; - - runtimeProcess.__zerobyteAgentRuntime = runtime; - return runtime; -}; - -const getAgentManagerRuntime = () => getAgentRuntimeState().agentManager; - type ActiveBackupRun = { scheduleId: number; jobId: string; @@ -60,8 +28,53 @@ type ActiveBackupRun = { cancellationRequested: boolean; }; -const activeBackupsByScheduleId = new Map(); -const activeBackupScheduleIdsByJobId = new Map(); +type AgentRuntimeState = { + agentManager: AgentManagerRuntime | null; + localAgent: ChildProcess | null; + isStoppingLocalAgent: boolean; + localAgentRestartTimeout: ReturnType | null; + activeBackupsByScheduleId: Map; + activeBackupScheduleIdsByJobId: Map; +}; + +type LegacyAgentRuntimeState = Omit & + Partial>; + +export type ProcessWithAgentRuntime = NodeJS.Process & { + __zerobyteAgentRuntime?: LegacyAgentRuntimeState; +}; + +const getAgentRuntimeState = () => { + const runtimeProcess = process as ProcessWithAgentRuntime; + const existingRuntime = runtimeProcess.__zerobyteAgentRuntime; + + if (existingRuntime) { + const runtime: AgentRuntimeState = { + ...existingRuntime, + activeBackupsByScheduleId: existingRuntime.activeBackupsByScheduleId ?? new Map(), + activeBackupScheduleIdsByJobId: existingRuntime.activeBackupScheduleIdsByJobId ?? new Map(), + }; + + runtimeProcess.__zerobyteAgentRuntime = runtime; + return runtime; + } + + const runtime: AgentRuntimeState = { + agentManager: null, + localAgent: null, + isStoppingLocalAgent: false, + localAgentRestartTimeout: null, + activeBackupsByScheduleId: new Map(), + activeBackupScheduleIdsByJobId: new Map(), + }; + + runtimeProcess.__zerobyteAgentRuntime = runtime; + return runtime; +}; + +const getAgentManagerRuntime = () => getAgentRuntimeState().agentManager; +const getActiveBackupsByScheduleId = () => getAgentRuntimeState().activeBackupsByScheduleId; +const getActiveBackupScheduleIdsByJobId = () => getAgentRuntimeState().activeBackupScheduleIdsByJobId; const getUnavailableError = (agentId: string) => { if (agentId === "local") { @@ -72,6 +85,8 @@ const getUnavailableError = (agentId: string) => { }; const clearActiveBackupRun = (scheduleId: number) => { + const activeBackupsByScheduleId = getActiveBackupsByScheduleId(); + const activeBackupScheduleIdsByJobId = getActiveBackupScheduleIdsByJobId(); const activeBackupRun = activeBackupsByScheduleId.get(scheduleId); if (!activeBackupRun) { return null; @@ -93,13 +108,13 @@ const resolveActiveBackupRun = (scheduleId: number, result: BackupExecutionResul }; const getActiveBackupRun = (jobId: string, scheduleId: string, eventName: string, agentId: string) => { - const trackedScheduleId = activeBackupScheduleIdsByJobId.get(jobId); + const trackedScheduleId = getActiveBackupScheduleIdsByJobId().get(jobId); if (trackedScheduleId === undefined) { logger.warn(`Received ${eventName} for unknown job ${jobId} from agent ${agentId}`); return null; } - const activeBackupRun = activeBackupsByScheduleId.get(trackedScheduleId); + const activeBackupRun = getActiveBackupsByScheduleId().get(trackedScheduleId); if (!activeBackupRun) { logger.warn(`Received ${eventName} for inactive job ${jobId} from agent ${agentId}`); return null; @@ -114,7 +129,7 @@ const getActiveBackupRun = (jobId: string, scheduleId: string, eventName: string }; const requestBackupCancellation = async (agentId: string, scheduleId: number) => { - const activeBackupRun = activeBackupsByScheduleId.get(scheduleId); + const activeBackupRun = getActiveBackupsByScheduleId().get(scheduleId); if (!activeBackupRun) { return false; } @@ -223,7 +238,7 @@ export const agentManager = { } const completion = new Promise((resolve) => { - activeBackupsByScheduleId.set(request.scheduleId, { + getActiveBackupsByScheduleId().set(request.scheduleId, { scheduleId: request.scheduleId, jobId: request.payload.jobId, scheduleShortId: request.payload.scheduleId, @@ -231,7 +246,7 @@ export const agentManager = { resolve, cancellationRequested: false, }); - activeBackupScheduleIdsByJobId.set(request.payload.jobId, request.scheduleId); + getActiveBackupScheduleIdsByJobId().set(request.payload.jobId, request.scheduleId); }); try {