fix(agent): harden disconnect and send-failure handling
This commit is contained in:
Nicolas Meienberger
parent
889865a5e5
commit
6336444e88
11 changed files with 262 additions and 411 deletions
87
app/server/modules/agents/__tests__/agents-manager.test.ts
Normal file
87
app/server/modules/agents/__tests__/agents-manager.test.ts
Normal file
|
|
@ -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<typeof import("node:child_process")>("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<typeof vi.fn>;
|
||||
};
|
||||
|
||||
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<typeof childProcess.spawn>)
|
||||
.mockReturnValueOnce(secondChild as unknown as ReturnType<typeof childProcess.spawn>);
|
||||
|
||||
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<typeof childProcess.spawn>);
|
||||
|
||||
await spawnLocalAgent();
|
||||
await stopLocalAgent();
|
||||
|
||||
vi.advanceTimersByTime(1_000);
|
||||
await flushMicrotasks();
|
||||
|
||||
expect(spawnMock).toHaveBeenCalledTimes(1);
|
||||
expect(child.kill).toHaveBeenCalledTimes(1);
|
||||
});
|
||||
|
|
@ -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<typeof createControllerAgentSession>[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.",
|
||||
});
|
||||
});
|
||||
|
|
@ -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<Parameters<typeof createControllerAgentSession>[0]> = {},
|
||||
) => {
|
||||
return fromPartial<Parameters<typeof createControllerAgentSession>[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<typeof createControllerAgentSession>[1] = {}) => {
|
||||
const createSession = (
|
||||
handlers: Parameters<typeof createControllerAgentSession>[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();
|
||||
}
|
||||
});
|
||||
|
|
|
|||
|
|
@ -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" };
|
||||
}
|
||||
};
|
||||
|
|
@ -8,6 +8,8 @@ export type { AgentBackupEventHandlers } from "./controller/server";
|
|||
type AgentRuntimeState = {
|
||||
agentManager: AgentManagerRuntime | null;
|
||||
localAgent: ChildProcess | null;
|
||||
isStoppingLocalAgent: boolean;
|
||||
localAgentRestartTimeout: ReturnType<typeof setTimeout> | 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 = {
|
||||
|
|
|
|||
|
|
@ -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<AgentConnectionData>;
|
||||
|
||||
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<ControllerWireMessage>(64));
|
||||
const activeBackupJobs = Effect.runSync(Ref.make<Map<string, string>>(new Map()));
|
||||
const pendingBackupJobs = Effect.runSync(Ref.make<Map<string, string>>(new Map()));
|
||||
const state = Effect.runSync(
|
||||
Ref.make<SessionState>({
|
||||
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(() => {});
|
||||
},
|
||||
};
|
||||
};
|
||||
|
|
@ -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));
|
||||
}
|
||||
});
|
||||
}),
|
||||
),
|
||||
);
|
||||
|
||||
|
|
|
|||
|
|
@ -7,6 +7,8 @@ import { deriveLocalAgentToken } from "../helpers/tokens";
|
|||
|
||||
type LocalAgentState = {
|
||||
localAgent: ChildProcess | null;
|
||||
isStoppingLocalAgent: boolean;
|
||||
localAgentRestartTimeout: ReturnType<typeof setTimeout> | 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<void>((resolve) => {
|
||||
agentProcess.once("exit", () => {
|
||||
runtime.isStoppingLocalAgent = false;
|
||||
resolve();
|
||||
});
|
||||
});
|
||||
|
|
|
|||
|
|
@ -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();
|
||||
}
|
||||
});
|
||||
|
|
|
|||
|
|
@ -17,6 +17,7 @@ export type ControllerSession = {
|
|||
};
|
||||
|
||||
export const createControllerSession = (ws: WebSocket): ControllerSession => {
|
||||
let isClosed = false;
|
||||
const outboundQueue = Effect.runSync(Queue.bounded<AgentWireMessage>(64));
|
||||
const inboundQueue = Effect.runSync(Queue.bounded<ControllerWireMessage>(64));
|
||||
const runningJobsRef = Effect.runSync(Ref.make<Map<string, RunningJob>>(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,
|
||||
};
|
||||
};
|
||||
|
|
|
|||
|
|
@ -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<typeof setTimeout> | null = null;
|
||||
|
|
@ -22,6 +22,10 @@ class Agent {
|
|||
}
|
||||
|
||||
connect() {
|
||||
if (this.reconnectTimeout) {
|
||||
clearTimeout(this.reconnectTimeout);
|
||||
this.reconnectTimeout = null;
|
||||
}
|
||||
if (this.ws) {
|
||||
return;
|
||||
}
|
||||
|
|
|
|||
Loading…
Reference in a new issue