fix(agents): await controller shutdown and guard backup dispatch

This commit is contained in:
Nicolas Meienberger 2026-04-09 19:06:24 +02:00
parent d1f872573e
commit d769e08c7b
No known key found for this signature in database
4 changed files with 46 additions and 18 deletions

View file

@ -43,6 +43,6 @@ export const stopLocalAgent = async () => {
export const agentManager = getAgentManagerRuntime();
export const stopAgentRuntime = async () => {
getAgentManagerRuntime().stop();
await getAgentManagerRuntime().stop();
await stopLocalAgent();
};

View file

@ -1,5 +1,6 @@
import { Effect, Exit, Fiber, Scope } from "effect";
import { Data, Effect, Exit, Fiber, Scope } from "effect";
import { logger } from "@zerobyte/core/node";
import { toMessage } from "@zerobyte/core/utils";
import type {
BackupCancelPayload,
BackupCancelledPayload,
@ -37,6 +38,10 @@ type ControllerAgentSessionHandle = {
scope: Scope.CloseableScope;
};
class StopAgentManagerServerError extends Data.TaggedError("StopAgentManagerServerError")<{
cause: unknown;
}> {}
export function createAgentManagerRuntime() {
let sessions = new Map<string, ControllerAgentSessionHandle>();
let backupHandlers: AgentBackupEventHandlers = {};
@ -171,13 +176,22 @@ export function createAgentManagerRuntime() {
}),
),
(server) =>
Effect.sync(() => {
closeAllSessions();
void server.stop(true);
}),
Effect.sync(closeAllSessions).pipe(
Effect.andThen(
Effect.tryPromise({
try: () => server.stop(true),
catch: (error) => new StopAgentManagerServerError({ cause: error }),
}),
),
Effect.catchAll((error) =>
Effect.sync(() => {
logger.error(`Failed to stop Agent Manager server: ${toMessage(error.cause)}`);
}),
),
),
);
const stop = () => {
const stop = async () => {
if (!runtimeScope) {
return;
}
@ -185,12 +199,12 @@ export function createAgentManagerRuntime() {
logger.info("Stopping Agent Manager...");
const scope = runtimeScope;
runtimeScope = null;
Effect.runSync(Scope.close(scope, Exit.succeed(undefined)));
await Effect.runPromise(Scope.close(scope, Exit.succeed(undefined)));
};
const start = () => {
const start = async () => {
if (runtimeScope) {
stop();
await stop();
}
logger.info("Starting Agent Manager...");
@ -201,7 +215,7 @@ export function createAgentManagerRuntime() {
runtimeScope = scope;
logger.info(`Agent Manager listening on port ${server.port}`);
} catch (error) {
Effect.runSync(Scope.close(scope, Exit.fail(error)));
await Effect.runPromise(Scope.close(scope, Exit.fail(error)));
throw error;
}
};
@ -221,7 +235,11 @@ export function createAgentManagerRuntime() {
return false;
}
Effect.runSync(session.sendBackup(payload));
if (!Effect.runSync(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;
},
@ -233,7 +251,11 @@ export function createAgentManagerRuntime() {
return false;
}
Effect.runSync(session.sendBackupCancel(payload));
if (!Effect.runSync(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;
},

View file

@ -46,8 +46,8 @@ type ControllerAgentSessionHandlers = {
export type ControllerAgentSession = {
readonly connectionId: string;
handleMessage: (data: string) => Effect.Effect<void>;
sendBackup: (payload: BackupRunPayload) => Effect.Effect<void>;
sendBackupCancel: (payload: BackupCancelPayload) => Effect.Effect<void>;
sendBackup: (payload: BackupRunPayload) => Effect.Effect<boolean>;
sendBackupCancel: (payload: BackupCancelPayload) => Effect.Effect<boolean>;
isReady: () => Effect.Effect<boolean>;
run: Effect.Effect<void, never, Scope.Scope>;
};
@ -70,6 +70,7 @@ export const createControllerAgentSession = (
Effect.catchAllCause((cause) =>
Effect.sync(() => {
logger.error(`Failed to queue outbound message for agent ${socket.data.agentId}: ${toMessage(cause)}`);
return false;
}),
),
);
@ -232,8 +233,13 @@ export const createControllerAgentSession = (
},
sendBackup: (payload) => {
return Effect.gen(function* () {
yield* setTrackedBackupJob(payload.jobId, { scheduleId: payload.scheduleId, state: "pending" });
yield* offerOutbound(createControllerMessage("backup.run", payload));
const queued = yield* offerOutbound(createControllerMessage("backup.run", payload));
if (queued) {
yield* setTrackedBackupJob(payload.jobId, { scheduleId: payload.scheduleId, state: "pending" });
}
return queued;
});
},
sendBackupCancel: (payload) => offerOutbound(createControllerMessage("backup.cancel", payload)),

View file

@ -24,7 +24,7 @@ export async function spawnLocalAgentProcess(runtime: LocalAgentState) {
const args = config.__prod__ ? ["run", agentEntryPoint] : ["run", "--watch", agentEntryPoint];
const agentProcess = spawn("bun", args, {
env: {
PATH: process.env.PATH,
...process.env,
ZEROBYTE_CONTROLLER_URL: "ws://localhost:3001",
ZEROBYTE_AGENT_TOKEN: agentToken,
},