diff --git a/app/server/modules/agents/__tests__/agents-manager.test.ts b/app/server/modules/agents/__tests__/agents-manager.test.ts new file mode 100644 index 00000000..dda50938 --- /dev/null +++ b/app/server/modules/agents/__tests__/agents-manager.test.ts @@ -0,0 +1,87 @@ +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, + }; +}); + +const { spawnLocalAgent, stopLocalAgent } = await import("../agents-manager"); + +const flushMicrotasks = async () => { + await Promise.resolve(); +}; + +type FakeChildProcess = EventEmitter & { + stdout: PassThrough; + stderr: PassThrough; + exitCode: number | null; + signalCode: NodeJS.Signals | null; + kill: ReturnType; +}; + +const createFakeChild = () => { + const child = new EventEmitter() as FakeChildProcess; + + child.stdout = new PassThrough(); + child.stderr = new PassThrough(); + child.exitCode = null; + child.signalCode = null; + child.kill = vi.fn(() => { + child.exitCode = 0; + child.emit("exit", 0, null); + return true; + }); + + return child; +}; + +afterEach(async () => { + await stopLocalAgent(); + spawnMock.mockReset(); + vi.restoreAllMocks(); + vi.useRealTimers(); +}); + +test("respawns the local agent after an unexpected exit", async () => { + vi.useFakeTimers(); + + const firstChild = createFakeChild(); + const secondChild = createFakeChild(); + spawnMock + .mockReturnValueOnce(firstChild as unknown as ReturnType) + .mockReturnValueOnce(secondChild as unknown as ReturnType); + + await spawnLocalAgent(); + + firstChild.exitCode = 1; + firstChild.emit("exit", 1, null); + + vi.advanceTimersByTime(1_000); + await flushMicrotasks(); + + expect(spawnMock).toHaveBeenCalledTimes(2); +}); + +test("does not respawn the local agent after an intentional stop", async () => { + vi.useFakeTimers(); + + const child = createFakeChild(); + spawnMock.mockReturnValue(child as unknown as ReturnType); + + await spawnLocalAgent(); + await stopLocalAgent(); + + vi.advanceTimersByTime(1_000); + await flushMicrotasks(); + + expect(spawnMock).toHaveBeenCalledTimes(1); + expect(child.kill).toHaveBeenCalledTimes(1); +}); diff --git a/app/server/modules/agents/__tests__/controller-agent-session.test.ts b/app/server/modules/agents/__tests__/controller-agent-session.test.ts deleted file mode 100644 index c2a35dd5..00000000 --- a/app/server/modules/agents/__tests__/controller-agent-session.test.ts +++ /dev/null @@ -1,121 +0,0 @@ -import { expect, mock, test } from "bun:test"; -import { createAgentMessage } from "@zerobyte/contracts/agent-protocol"; -import { createControllerAgentSession } from "../controller-agent-session"; - -const createSocket = () => { - return { - data: { id: "connection-1", agentId: "local", organizationId: null, agentName: "Local Agent" }, - send: mock(() => undefined), - } as unknown as Parameters[0]; -}; - -test("close emits a synthetic backup.cancelled for a started backup", () => { - const onBackupCancelled = mock(() => undefined); - const session = createControllerAgentSession(createSocket(), { - onBackupCancelled, - }); - - session.handleMessage( - createAgentMessage("backup.started", { - jobId: "job-1", - scheduleId: "schedule-1", - }), - ); - - session.close(); - - expect(onBackupCancelled).toHaveBeenCalledTimes(1); - expect(onBackupCancelled).toHaveBeenCalledWith({ - jobId: "job-1", - scheduleId: "schedule-1", - message: - "The connection to the backup agent was lost while this backup was running. Restart the backup to ensure it completes.", - }); -}); - -test("close does not emit a synthetic backup.cancelled after a terminal event", () => { - for (const testCase of [ - { - jobId: "job-1", - scheduleId: "schedule-1", - terminalMessage: createAgentMessage("backup.completed", { - jobId: "job-1", - scheduleId: "schedule-1", - exitCode: 0, - result: null, - }), - expectedCancelledCalls: 0, - }, - { - jobId: "job-2", - scheduleId: "schedule-2", - terminalMessage: createAgentMessage("backup.failed", { - jobId: "job-2", - scheduleId: "schedule-2", - error: "backup failed", - }), - expectedCancelledCalls: 0, - }, - { - jobId: "job-3", - scheduleId: "schedule-3", - terminalMessage: createAgentMessage("backup.cancelled", { - jobId: "job-3", - scheduleId: "schedule-3", - message: "Backup was cancelled", - }), - expectedCancelledCalls: 1, - }, - ]) { - const onBackupCancelled = mock(() => undefined); - const session = createControllerAgentSession(createSocket(), { - onBackupCancelled, - }); - - session.handleMessage( - createAgentMessage("backup.started", { - jobId: testCase.jobId, - scheduleId: testCase.scheduleId, - }), - ); - session.handleMessage(testCase.terminalMessage); - session.close(); - - expect(onBackupCancelled).toHaveBeenCalledTimes(testCase.expectedCancelledCalls); - } -}); - -test("close emits a synthetic backup.cancelled for a queued backup", () => { - const onBackupCancelled = mock(() => undefined); - const session = createControllerAgentSession(createSocket(), { - onBackupCancelled, - }); - - session.sendBackup({ - jobId: "job-queued", - scheduleId: "schedule-queued", - organizationId: "org-1", - sourcePath: "/tmp/source", - repositoryConfig: { - backend: "local", - path: "/tmp/repository", - }, - options: {}, - runtime: { - password: "password", - cacheDir: "/tmp/cache", - passFile: "/tmp/pass", - defaultExcludes: [], - }, - }); - - session.close(); - - expect(onBackupCancelled).toHaveBeenCalledTimes(1); - expect(onBackupCancelled).toHaveBeenCalledWith({ - jobId: "job-queued", - scheduleId: "schedule-queued", - message: - "The connection to the backup agent was lost before this backup started. Restart the backup to ensure it completes.", - }); -}); diff --git a/app/server/modules/agents/__tests__/session.test.ts b/app/server/modules/agents/__tests__/session.test.ts index b787e1ea..2dbb33d8 100644 --- a/app/server/modules/agents/__tests__/session.test.ts +++ b/app/server/modules/agents/__tests__/session.test.ts @@ -1,24 +1,36 @@ 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 { createControllerAgentSession } from "../controller/session"; -const createSocket = () => { +const createSocket = ( + overrides: Partial[0]> = {}, +) => { return fromPartial[0]>({ data: { id: "connection-1", agentId: "local", organizationId: null, agentName: "Local Agent" }, - send: vi.fn(), + send: vi.fn(() => 1), + close: vi.fn(), + ...overrides, }); }; -const createSession = (handlers: Parameters[1] = {}) => { +const createSession = ( + handlers: Parameters[1] = {}, + socket = createSocket(), +) => { const scope = Effect.runSync(Scope.make()); try { - const session = Effect.runSync(Scope.extend(createControllerAgentSession(createSocket(), handlers), scope)); + const session = Effect.runSync(Scope.extend(createControllerAgentSession(socket, handlers), scope)); return { session, + run: () => { + Effect.runSync(Scope.extend(Effect.forkScoped(session.run), scope)); + }, + socket, close: () => { Effect.runSync(Scope.close(scope, Exit.succeed(undefined))); }, @@ -50,7 +62,7 @@ test("close emits a synthetic backup.cancelled for a started backup", () => { expect(onBackupCancelled).toHaveBeenCalledWith({ jobId: "job-1", scheduleId: "schedule-1", - message: "The connection to the backup agent was lost. Restart the backup to ensure it completes.", + message: "The connection to the backup agent was lost while this backup was running. Restart the backup to ensure it completes.", }); }); @@ -142,6 +154,45 @@ test("close emits a synthetic backup.cancelled for a queued backup", () => { expect(onBackupCancelled).toHaveBeenCalledWith({ jobId: "job-queued", scheduleId: "schedule-queued", - message: "The connection to the backup agent was lost. Restart the backup to ensure it completes.", + message: "The connection to the backup agent was lost before this backup started. Restart the backup to ensure it completes.", }); }); + +test("a dropped backup.cancel closes the session and emits a synthetic backup.cancelled", async () => { + const send = vi.fn(() => 0); + const socket = createSocket({ send, close: vi.fn() }); + const onBackupCancelled = vi.fn(); + const { session, run, close: closeSession } = createSession({ onBackupCancelled }, socket); + + try { + run(); + Effect.runSync( + session.handleMessage( + createAgentMessage("backup.started", { + jobId: "job-1", + scheduleId: "schedule-1", + }), + ), + ); + Effect.runSync( + session.sendBackupCancel({ + jobId: "job-1", + scheduleId: "schedule-1", + }), + ); + + await waitForExpect(() => { + 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.", + }); + }); + } finally { + closeSession(); + } +}); diff --git a/app/server/modules/agents/agent-tokens.ts b/app/server/modules/agents/agent-tokens.ts deleted file mode 100644 index febb9d94..00000000 --- a/app/server/modules/agents/agent-tokens.ts +++ /dev/null @@ -1,12 +0,0 @@ -import { cryptoUtils } from "~/server/utils/crypto"; - -export const deriveLocalAgentToken = async () => { - return cryptoUtils.deriveSecret("zerobyte:local-agent-token"); -}; - -export const validateAgentToken = async (token: string) => { - const localToken = await deriveLocalAgentToken(); - if (token === localToken) { - return { agentId: "local", organizationId: null, agentName: "local" }; - } -}; diff --git a/app/server/modules/agents/agents-manager.ts b/app/server/modules/agents/agents-manager.ts index 1d912878..b97147b5 100644 --- a/app/server/modules/agents/agents-manager.ts +++ b/app/server/modules/agents/agents-manager.ts @@ -8,6 +8,8 @@ export type { AgentBackupEventHandlers } from "./controller/server"; type AgentRuntimeState = { agentManager: AgentManagerRuntime | null; localAgent: ChildProcess | null; + isStoppingLocalAgent: boolean; + localAgentRestartTimeout: ReturnType | null; }; type ProcessWithAgentRuntime = NodeJS.Process & { @@ -25,6 +27,8 @@ const getAgentRuntimeState = () => { const runtime = { agentManager: null, localAgent: null, + isStoppingLocalAgent: false, + localAgentRestartTimeout: null, }; runtimeProcess.__zerobyteAgentRuntime = runtime; @@ -43,11 +47,11 @@ export const startAgentRuntime = async () => { } const { createAgentManagerRuntime } = await import("./controller/server"); - const agentManager = createAgentManagerRuntime(); - agentManager.setBackupEventHandlers(backupEventHandlers); + const nextAgentManager = createAgentManagerRuntime(); + nextAgentManager.setBackupEventHandlers(backupEventHandlers); - await agentManager.start(); - runtime.agentManager = agentManager; + await nextAgentManager.start(); + runtime.agentManager = nextAgentManager; }; export const agentManager = { diff --git a/app/server/modules/agents/controller-agent-session.ts b/app/server/modules/agents/controller-agent-session.ts deleted file mode 100644 index 900aee72..00000000 --- a/app/server/modules/agents/controller-agent-session.ts +++ /dev/null @@ -1,247 +0,0 @@ -import { Effect, Fiber, Queue, Ref } from "effect"; -import { - createControllerMessage, - parseAgentMessage, - type AgentMessage, - type BackupCancelledPayload, - type BackupCompletedPayload, - type BackupFailedPayload, - type BackupProgressPayload, - type BackupRunPayload, - type BackupCancelPayload, - type BackupStartedPayload, - type ControllerWireMessage, -} from "@zerobyte/contracts/agent-protocol"; -import { logger } from "@zerobyte/core/node"; -import { toMessage } from "@zerobyte/core/utils"; - -export type AgentConnectionData = { - id: string; - agentId: string; - organizationId: string | null; - agentName: string; -}; - -type AgentSocket = Bun.ServerWebSocket; - -type SessionState = { - isReady: boolean; - lastSeenAt: number | null; - lastPongAt: number | null; -}; - -type ControllerAgentSessionHandlers = { - onBackupStarted?: (payload: BackupStartedPayload) => void; - onBackupProgress?: (payload: BackupProgressPayload) => void; - onBackupCompleted?: (payload: BackupCompletedPayload) => void; - onBackupFailed?: (payload: BackupFailedPayload) => void; - onBackupCancelled?: (payload: BackupCancelledPayload) => void; -}; - -export type ControllerAgentSession = { - readonly connectionId: string; - handleMessage: (data: string) => void; - sendBackup: (payload: BackupRunPayload) => string; - sendBackupCancel: (payload: BackupCancelPayload) => void; - isReady: () => boolean; - close: () => void; -}; - -export const createControllerAgentSession = ( - socket: AgentSocket, - handlers: ControllerAgentSessionHandlers = {}, -): ControllerAgentSession => { - const outboundQueue = Effect.runSync(Queue.bounded(64)); - const activeBackupJobs = Effect.runSync(Ref.make>(new Map())); - const pendingBackupJobs = Effect.runSync(Ref.make>(new Map())); - const state = Effect.runSync( - Ref.make({ - isReady: false, - lastSeenAt: null, - lastPongAt: null, - }), - ); - - const offerOutbound = (message: ControllerWireMessage) => { - void Effect.runPromise(Queue.offer(outboundQueue, message)).catch((error) => { - logger.error(`Failed to queue outbound message for agent ${socket.data.agentId}: ${toMessage(error)}`); - }); - }; - - const updateState = (update: (current: SessionState) => SessionState) => { - Effect.runSync(Ref.update(state, update)); - }; - - const setActiveBackupJob = (jobId: string, scheduleId: string) => { - Effect.runSync( - Ref.update(activeBackupJobs, (current) => { - const next = new Map(current); - next.set(jobId, scheduleId); - return next; - }), - ); - }; - - const setPendingBackupJob = (jobId: string, scheduleId: string) => { - Effect.runSync( - Ref.update(pendingBackupJobs, (current) => { - const next = new Map(current); - next.set(jobId, scheduleId); - return next; - }), - ); - }; - - const deleteActiveBackupJob = (jobId: string) => { - Effect.runSync( - Ref.update(activeBackupJobs, (current) => { - const next = new Map(current); - next.delete(jobId); - return next; - }), - ); - }; - - const deletePendingBackupJob = (jobId: string) => { - Effect.runSync( - Ref.update(pendingBackupJobs, (current) => { - const next = new Map(current); - next.delete(jobId); - return next; - }), - ); - }; - - const writerFiber = Effect.runFork( - Effect.forever( - Effect.gen(function* () { - const message = yield* Queue.take(outboundQueue); - yield* Effect.sync(() => { - try { - socket.send(message); - } catch (error) { - logger.error( - `Failed to send message to agent ${socket.data.agentId} on ${socket.data.id}: ${toMessage(error)}`, - ); - } - }); - }), - ), - ); - - const heartbeatFiber = Effect.runFork( - Effect.forever( - Effect.gen(function* () { - yield* Effect.sleep("15 seconds"); - yield* Queue.offer( - outboundQueue, - createControllerMessage("heartbeat.ping", { - sentAt: Date.now(), - }), - ); - }), - ), - ); - - const handleAgentMessage = (message: AgentMessage) => { - updateState((current) => ({ ...current, lastSeenAt: Date.now() })); - - switch (message.type) { - case "agent.ready": { - updateState((current) => ({ ...current, isReady: true })); - logger.info(`Agent "${socket.data.agentName}" (${socket.data.agentId}) is ready`); - break; - } - case "backup.started": { - deletePendingBackupJob(message.payload.jobId); - setActiveBackupJob(message.payload.jobId, message.payload.scheduleId); - logger.info( - `Backup ${message.payload.jobId} started on agent ${socket.data.agentId} for schedule ${message.payload.scheduleId}`, - ); - handlers.onBackupStarted?.(message.payload); - break; - } - case "backup.progress": { - handlers.onBackupProgress?.(message.payload); - break; - } - case "backup.completed": { - deletePendingBackupJob(message.payload.jobId); - deleteActiveBackupJob(message.payload.jobId); - handlers.onBackupCompleted?.(message.payload); - break; - } - case "backup.failed": { - deletePendingBackupJob(message.payload.jobId); - deleteActiveBackupJob(message.payload.jobId); - handlers.onBackupFailed?.(message.payload); - break; - } - case "backup.cancelled": { - deletePendingBackupJob(message.payload.jobId); - deleteActiveBackupJob(message.payload.jobId); - handlers.onBackupCancelled?.(message.payload); - break; - } - case "heartbeat.pong": { - updateState((current) => ({ ...current, lastPongAt: message.payload.sentAt })); - break; - } - } - }; - - return { - connectionId: socket.data.id, - handleMessage: (data: string) => { - const parsed = parseAgentMessage(data); - - if (parsed === null) { - logger.warn(`Invalid JSON from agent ${socket.data.agentId}`); - return; - } - - if (!parsed.success) { - logger.warn(`Invalid agent message from ${socket.data.agentId}: ${parsed.error.message}`); - return; - } - - handleAgentMessage(parsed.data); - }, - sendBackup: (payload) => { - setPendingBackupJob(payload.jobId, payload.scheduleId); - offerOutbound(createControllerMessage("backup.run", payload)); - return payload.jobId; - }, - sendBackupCancel: (payload) => { - offerOutbound(createControllerMessage("backup.cancel", payload)); - }, - isReady: () => Effect.runSync(Ref.get(state)).isReady, - close: () => { - updateState((current) => ({ ...current, isReady: false })); - const pendingJobs = Effect.runSync(Ref.get(pendingBackupJobs)); - Effect.runSync(Ref.set(pendingBackupJobs, new Map())); - for (const [jobId, scheduleId] of pendingJobs) { - handlers.onBackupCancelled?.({ - jobId, - scheduleId, - message: - "The connection to the backup agent was lost before this backup started. Restart the backup to ensure it completes.", - }); - } - - const activeJobs = Effect.runSync(Ref.get(activeBackupJobs)); - Effect.runSync(Ref.set(activeBackupJobs, new Map())); - for (const [jobId, scheduleId] of activeJobs) { - handlers.onBackupCancelled?.({ - jobId, - scheduleId, - message: - "The connection to the backup agent was lost while this backup was running. Restart the backup to ensure it completes.", - }); - } - void Effect.runPromise(Fiber.interrupt(writerFiber)).catch(() => {}); - void Effect.runPromise(Fiber.interrupt(heartbeatFiber)).catch(() => {}); - void Effect.runPromise(Queue.shutdown(outboundQueue)).catch(() => {}); - }, - }; -}; diff --git a/app/server/modules/agents/controller/session.ts b/app/server/modules/agents/controller/session.ts index cae39436..e208e1a5 100644 --- a/app/server/modules/agents/controller/session.ts +++ b/app/server/modules/agents/controller/session.ts @@ -102,7 +102,10 @@ export const createControllerAgentSession = ( yield* updateState((current) => ({ ...current, isReady: false })); const trackedJobs = yield* takeTrackedBackupJobs; for (const [jobId, trackedJob] of trackedJobs) { - let message = "The connection to the backup agent was lost. Restart the backup to ensure it completes."; + const message = + trackedJob.state === "pending" + ? "The connection to the backup agent was lost before this backup started. Restart the backup to ensure it completes." + : "The connection to the backup agent was lost while this backup was running. Restart the backup to ensure it completes."; yield* Effect.sync(() => { handlers.onBackupCancelled?.({ jobId, scheduleId: trackedJob.scheduleId, message }); @@ -114,21 +117,34 @@ export const createControllerAgentSession = ( yield* Effect.addFinalizer(() => releaseSession); + const handleSendFailure = (reason: string) => { + logger.error( + `Closing session for agent ${socket.data.agentId} on ${socket.data.id} after an outbound websocket send failed: ${reason}`, + ); + + try { + socket.close(); + } catch (error) { + logger.error(`Failed to close socket for agent ${socket.data.agentId} on ${socket.data.id}: ${toMessage(error)}`); + } + }; + const run = Effect.gen(function* () { yield* Effect.forkScoped( Effect.forever( - Effect.gen(function* () { - const message = yield* Queue.take(outboundQueue); - yield* Effect.sync(() => { - try { - socket.send(message); - } catch (error) { - logger.error( - `Failed to send message to agent ${socket.data.agentId} on ${socket.data.id}: ${toMessage(error)}`, - ); + Effect.gen(function* () { + const message = yield* Queue.take(outboundQueue); + yield* Effect.sync(() => { + try { + const sendResult = socket.send(message); + if (sendResult <= 0) { + handleSendFailure(sendResult === 0 ? "connection issue" : "backpressure"); } - }); - }), + } catch (error) { + handleSendFailure(toMessage(error)); + } + }); + }), ), ); diff --git a/app/server/modules/agents/local/process.ts b/app/server/modules/agents/local/process.ts index 6c353782..3852b36a 100644 --- a/app/server/modules/agents/local/process.ts +++ b/app/server/modules/agents/local/process.ts @@ -7,6 +7,8 @@ import { deriveLocalAgentToken } from "../helpers/tokens"; type LocalAgentState = { localAgent: ChildProcess | null; + isStoppingLocalAgent: boolean; + localAgentRestartTimeout: ReturnType | null; }; export async function spawnLocalAgentProcess(runtime: LocalAgentState) { @@ -48,27 +50,47 @@ export async function spawnLocalAgentProcess(runtime: LocalAgentState) { }); agentProcess.on("exit", (code, signal) => { + const shouldRestart = runtime.localAgent === agentProcess && !runtime.isStoppingLocalAgent; if (runtime.localAgent === agentProcess) { runtime.localAgent = null; } logger.info(`Agent process exited with code ${code} and signal ${signal}`); + + if (!shouldRestart) { + return; + } + + runtime.localAgentRestartTimeout = setTimeout(() => { + runtime.localAgentRestartTimeout = null; + void spawnLocalAgentProcess(runtime).catch((error) => { + logger.error(`Failed to restart local agent: ${error instanceof Error ? error.message : String(error)}`); + }); + }, 1_000); }); } export async function stopLocalAgentProcess(runtime: LocalAgentState) { + if (runtime.localAgentRestartTimeout) { + clearTimeout(runtime.localAgentRestartTimeout); + runtime.localAgentRestartTimeout = null; + } + if (!runtime.localAgent) { return; } const agentProcess = runtime.localAgent; runtime.localAgent = null; + runtime.isStoppingLocalAgent = true; if (agentProcess.exitCode !== null || agentProcess.signalCode !== null) { + runtime.isStoppingLocalAgent = false; return; } const exited = new Promise((resolve) => { agentProcess.once("exit", () => { + runtime.isStoppingLocalAgent = false; resolve(); }); }); diff --git a/apps/agent/src/__tests__/controller-session.test.ts b/apps/agent/src/__tests__/controller-session.test.ts index e9ff4e83..4f238b8f 100644 --- a/apps/agent/src/__tests__/controller-session.test.ts +++ b/apps/agent/src/__tests__/controller-session.test.ts @@ -70,3 +70,25 @@ test("emits backup.failed when a backup command hits a restic error", async () = session.close(); } }); + +test("closes the websocket when an outbound send throws", async () => { + const close = vi.fn(() => undefined); + const session = createControllerSession( + fromPartial({ + send: () => { + throw new Error("socket write failed"); + }, + close, + }), + ); + + try { + session.onOpen(); + + await waitForExpect(() => { + expect(close).toHaveBeenCalledTimes(1); + }); + } finally { + session.close(); + } +}); diff --git a/apps/agent/src/controller-session.ts b/apps/agent/src/controller-session.ts index 9948d00e..d8f82e24 100644 --- a/apps/agent/src/controller-session.ts +++ b/apps/agent/src/controller-session.ts @@ -17,6 +17,7 @@ export type ControllerSession = { }; export const createControllerSession = (ws: WebSocket): ControllerSession => { + let isClosed = false; const outboundQueue = Effect.runSync(Queue.bounded(64)); const inboundQueue = Effect.runSync(Queue.bounded(64)); const runningJobsRef = Effect.runSync(Ref.make>(new Map())); @@ -63,6 +64,31 @@ export const createControllerSession = (ws: WebSocket): ControllerSession => { offerOutbound, }; + const closeSession = () => { + if (isClosed) { + return; + } + + isClosed = true; + void Effect.runPromise(abortRunningJobs).catch(() => {}); + void Effect.runPromise(Fiber.interrupt(writerFiber)).catch(() => {}); + void Effect.runPromise(Fiber.interrupt(processorFiber)).catch(() => {}); + void Effect.runPromise(Queue.shutdown(outboundQueue)).catch(() => {}); + void Effect.runPromise(Queue.shutdown(inboundQueue)).catch(() => {}); + }; + + const handleSendFailure = (reason: string) => { + logger.error(`Closing agent session after an outbound websocket send failed: ${reason}`); + + try { + ws.close(); + } catch (error) { + logger.error(`Failed to close controller websocket after send failure: ${toMessage(error)}`); + } + + closeSession(); + }; + const writerFiber = Effect.runFork( Effect.forever( Effect.gen(function* () { @@ -71,7 +97,7 @@ export const createControllerSession = (ws: WebSocket): ControllerSession => { try { ws.send(message); } catch (error) { - logger.error(`Failed to send controller message: ${toMessage(error)}`); + handleSendFailure(toMessage(error)); } }); }), @@ -106,16 +132,15 @@ export const createControllerSession = (ws: WebSocket): ControllerSession => { }); }, onMessage: (data) => { + if (typeof data !== "string") { + logger.warn("Agent received a non-text message"); + return; + } + void Effect.runPromise(offerInbound(data as ControllerWireMessage)).catch((error) => { logger.error(`Failed to queue inbound message: ${toMessage(error)}`); }); }, - close: () => { - void Effect.runPromise(abortRunningJobs).catch(() => {}); - void Effect.runPromise(Fiber.interrupt(writerFiber)).catch(() => {}); - void Effect.runPromise(Fiber.interrupt(processorFiber)).catch(() => {}); - void Effect.runPromise(Queue.shutdown(outboundQueue)).catch(() => {}); - void Effect.runPromise(Queue.shutdown(inboundQueue)).catch(() => {}); - }, + close: closeSession, }; }; diff --git a/apps/agent/src/index.ts b/apps/agent/src/index.ts index f9e08b6d..2593792c 100644 --- a/apps/agent/src/index.ts +++ b/apps/agent/src/index.ts @@ -5,7 +5,7 @@ const controllerUrl = process.env.ZEROBYTE_CONTROLLER_URL; const agentToken = process.env.ZEROBYTE_AGENT_TOKEN; const RECONNECT_DELAY_MS = 1000; -class Agent { +export class Agent { private ws: WebSocket | null = null; private controllerSession: ControllerSession | null = null; private reconnectTimeout: ReturnType | null = null; @@ -22,6 +22,10 @@ class Agent { } connect() { + if (this.reconnectTimeout) { + clearTimeout(this.reconnectTimeout); + this.reconnectTimeout = null; + } if (this.ws) { return; }