fix: mark agent offline only if the session was removed properly

app/server/modules/agents/agents-manager.ts

@@ -1,7 +1,7 @@
 import { logger } from "@zerobyte/core/node";
 import type { BackupRunPayload } from "@zerobyte/contracts/agent-protocol";
 import { config } from "../../core/config";
-import type { AgentBackupEventHandlers } from "./controller/server";
+import { createAgentManagerRuntime, type AgentBackupEventHandlers } from "./controller/server";
 import { spawnLocalAgentProcess, stopLocalAgentProcess } from "./local/process";
 import type { BackupExecutionProgress, BackupExecutionResult } from "./helpers/runtime-state";
 import { createAgentRuntimeState } from "./helpers/runtime-state";
@@ -155,7 +155,6 @@ export const startAgentController = async () => {
 		runtime.agentManager = null;
 	}
 
-	const { createAgentManagerRuntime } = await import("./controller/server");
 	const nextAgentManager = createAgentManagerRuntime();
 	nextAgentManager.setBackupEventHandlers(backupEventHandlers);
 

app/server/modules/agents/controller/server.ts

@@ -197,10 +197,11 @@ export function createAgentManagerRuntime() {
 						});
 					},
 					close: (ws) => {
-						removeSession(ws.data.agentId, ws.data.id);
+						if (removeSession(ws.data.agentId, ws.data.id)) {
-						void agentsService.markAgentOffline(ws.data.agentId).catch((error) => {
+							void agentsService.markAgentOffline(ws.data.agentId).catch((error) => {
-							logger.error(`Failed to mark agent ${ws.data.agentId} as offline: ${toMessage(error)}`);
+								logger.error(`Failed to mark agent ${ws.data.agentId} as offline: ${toMessage(error)}`);
-						});
+							});
+						}
 						logger.info(`Agent "${ws.data.agentName}" (${ws.data.agentId}) disconnected`);
 					},
 				},
@@ -214,11 +215,9 @@ export function createAgentManagerRuntime() {
 						catch: (error) => new StopAgentManagerServerError({ cause: error }),
 					}),
 				),
-				Effect.catchAll((error) =>
+				Effect.catchAll((error) => {
-					Effect.sync(() => {
+					return logger.effect.error(`Failed to stop Agent Manager server: ${toMessage(error.cause)}`);
-						logger.error(`Failed to stop Agent Manager server: ${toMessage(error.cause)}`);
+				}),
-					}),
-				),
 			),
 	);
 
@@ -233,6 +232,7 @@ export function createAgentManagerRuntime() {
 		await Effect.runPromise(Scope.close(scope, Exit.succeed(undefined)));
 	};
 
+	// TODO: Move the effect boundary up
 	const start = async () => {
 		if (runtimeScope) {
 			await stop();

app/server/modules/agents/controller/session.ts

@@ -136,16 +136,13 @@ export const createControllerAgentSession = (
 		yield* Effect.addFinalizer(() => closeSession());
 
 		const handleSendFailure = (reason: string) => {
-			logger.error(
+			return Effect.gen(function* () {
-				`Closing session for agent ${socket.data.agentId} on ${socket.data.id} after an outbound websocket send failed: ${reason}`,
-			);
-
-			socket.close();
-
-			void Effect.runPromise(closeSession()).catch((error) => {
 				logger.error(
-					`Failed to close session for agent ${socket.data.agentId} on ${socket.data.id}: ${toMessage(error)}`,
+					`Closing session for agent ${socket.data.agentId} on ${socket.data.id} after an outbound websocket send failed: ${reason}`,
 				);
+
+				yield* Effect.sync(() => socket.close());
+				yield* closeSession();
 			});
 		};
 
@@ -154,17 +151,16 @@ export const createControllerAgentSession = (
 				Effect.forever(
 					Effect.gen(function* () {
 						const message = yield* Queue.take(outboundQueue);
-						yield* Effect.sync(() => {
+
-							try {
+						const sendResult = yield* Effect.try({
-								const sendResult = socket.send(message);
+							try: () => socket.send(message),
-								if (sendResult === 0) {
+							catch: (error) => toMessage(error),
-									handleSendFailure("connection issue");
-								}
-							} catch (error) {
-								handleSendFailure(toMessage(error));
-							}
 						});
-					}),
+
+						if (sendResult === 0) {
+							yield* handleSendFailure("connection issue");
+						}
+					}).pipe(Effect.catchAll((reason) => handleSendFailure(reason))),
 				),
 			);
 
@@ -200,9 +196,7 @@ export const createControllerAgentSession = (
 							at: readyAt,
 						});
 
-						yield* Effect.sync(() => {
+						yield* logger.effect.info(`Agent "${socket.data.agentName}" (${socket.data.agentId}) is ready`);
-							logger.info(`Agent "${socket.data.agentName}" (${socket.data.agentId}) is ready`);
-						});
 						break;
 					}
 					case "backup.started": {
@@ -210,7 +204,7 @@ export const createControllerAgentSession = (
 							scheduleId: message.payload.scheduleId,
 							state: "active",
 						});
-						logger.info(
+						yield* logger.effect.info(
 							`Backup ${message.payload.jobId} started on agent ${socket.data.agentId} for schedule ${message.payload.scheduleId}`,
 						);
 						yield* handlers.onBackupStarted(message.payload);
@@ -258,16 +252,12 @@ export const createControllerAgentSession = (
 					const parsed = parseAgentMessage(data);
 
 					if (parsed === null) {
-						yield* Effect.sync(() => {
+						yield* logger.effect.warn(`Invalid JSON from agent ${socket.data.agentId}`);
-							logger.warn(`Invalid JSON from agent ${socket.data.agentId}`);
-						});
 						return;
 					}
 
 					if (!parsed.success) {
-						yield* Effect.sync(() => {
+						yield* logger.effect.warn(`Invalid agent message from ${socket.data.agentId}: ${parsed.error.message}`);
-							logger.warn(`Invalid agent message from ${socket.data.agentId}: ${parsed.error.message}`);
-						});
 						return;
 					}
 

apps/agent/src/commands/backup-run.ts

@@ -21,7 +21,7 @@ export const handleBackupRunCommand = (context: ControllerCommandContext, payloa
 			return;
 		}
 
-		logger.info(`Starting backup ${payload.jobId} for schedule ${payload.scheduleId}`);
+		yield* logger.effect.info(`Starting backup ${payload.jobId} for schedule ${payload.scheduleId}`);
 		const abortController = new AbortController();
 		yield* context.setRunningJob(payload.jobId, { scheduleId: payload.scheduleId, abortController });
 

packages/core/src/node/logger.ts

@@ -2,6 +2,7 @@ import { format } from "date-fns";
 import { createConsola, type ConsolaReporter } from "consola";
 import { formatWithOptions } from "node:util";
 import { sanitizeSensitiveData } from "../utils/sanitize";
+import { Effect } from "effect";
 
 type LogLevel = "debug" | "info" | "warn" | "error";
 
@@ -102,4 +103,10 @@ export const logger = {
 	info: (...messages: unknown[]) => consola.info(formatMessages(messages).join(" ")),
 	warn: (...messages: unknown[]) => consola.warn(formatMessages(messages).join(" ")),
 	error: (...messages: unknown[]) => consola.error(formatMessages(messages).join(" ")),
+	effect: {
+		debug: (...messages: unknown[]) => Effect.sync(() => consola.debug(formatMessages(messages).join(" "))),
+		info: (...messages: unknown[]) => Effect.sync(() => consola.info(formatMessages(messages).join(" "))),
+		warn: (...messages: unknown[]) => Effect.sync(() => consola.warn(formatMessages(messages).join(" "))),
+		error: (...messages: unknown[]) => Effect.sync(() => consola.error(formatMessages(messages).join(" "))),
+	},
 };