diff --git a/app/server/modules/agents/agents-manager.ts b/app/server/modules/agents/agents-manager.ts index 9ca619b2..72e900c2 100644 --- a/app/server/modules/agents/agents-manager.ts +++ b/app/server/modules/agents/agents-manager.ts @@ -43,6 +43,6 @@ export const stopLocalAgent = async () => { export const agentManager = getAgentManagerRuntime(); export const stopAgentRuntime = async () => { - getAgentManagerRuntime().stop(); + await getAgentManagerRuntime().stop(); await stopLocalAgent(); }; diff --git a/app/server/modules/agents/controller/server.ts b/app/server/modules/agents/controller/server.ts index 63b3f1d4..67a2e808 100644 --- a/app/server/modules/agents/controller/server.ts +++ b/app/server/modules/agents/controller/server.ts @@ -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(); 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; }, diff --git a/app/server/modules/agents/controller/session.ts b/app/server/modules/agents/controller/session.ts index 19805aea..cae39436 100644 --- a/app/server/modules/agents/controller/session.ts +++ b/app/server/modules/agents/controller/session.ts @@ -46,8 +46,8 @@ type ControllerAgentSessionHandlers = { export type ControllerAgentSession = { readonly connectionId: string; handleMessage: (data: string) => Effect.Effect; - sendBackup: (payload: BackupRunPayload) => Effect.Effect; - sendBackupCancel: (payload: BackupCancelPayload) => Effect.Effect; + sendBackup: (payload: BackupRunPayload) => Effect.Effect; + sendBackupCancel: (payload: BackupCancelPayload) => Effect.Effect; isReady: () => Effect.Effect; run: Effect.Effect; }; @@ -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)), diff --git a/app/server/modules/agents/local/process.ts b/app/server/modules/agents/local/process.ts index d480cd14..c844a65d 100644 --- a/app/server/modules/agents/local/process.ts +++ b/app/server/modules/agents/local/process.ts @@ -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, },