feat(controller): add agent manager and session handling (#763)

app/server/modules/agents/__tests__/session.test.ts

@@ -0,0 +1,146 @@
+import { Effect, Exit, Scope } from "effect";
+import { expect, test, vi } from "vitest";
+import { fromPartial } from "@total-typescript/shoehorn";
+import { createAgentMessage } from "@zerobyte/contracts/agent-protocol";
+import { createControllerAgentSession } from "../controller/session";
+
+const createSocket = () => {
+	return fromPartial<Parameters<typeof createControllerAgentSession>[0]>({
+		data: { id: "connection-1", agentId: "local", organizationId: null, agentName: "Local Agent" },
+		send: vi.fn(),
+	});
+};
+
+const createSession = (handlers: Parameters<typeof createControllerAgentSession>[1] = {}) => {
+	const scope = Effect.runSync(Scope.make());
+
+	try {
+		const session = Effect.runSync(Scope.extend(createControllerAgentSession(createSocket(), handlers), scope));
+
+		return {
+			session,
+			close: () => {
+				Effect.runSync(Scope.close(scope, Exit.succeed(undefined)));
+			},
+		};
+	} catch (error) {
+		Effect.runSync(Scope.close(scope, Exit.fail(error)));
+		throw error;
+	}
+};
+
+test("close emits a synthetic backup.cancelled for a started backup", () => {
+	const onBackupCancelled = vi.fn();
+	const { session, close } = createSession({
+		onBackupCancelled,
+	});
+
+	Effect.runSync(
+		session.handleMessage(
+			createAgentMessage("backup.started", {
+				jobId: "job-1",
+				scheduleId: "schedule-1",
+			}),
+		),
+	);
+
+	close();
+
+	expect(onBackupCancelled).toHaveBeenCalledTimes(1);
+	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.",
+	});
+});
+
+test.each([
+	{
+		name: "backup.completed",
+		jobId: "job-1",
+		scheduleId: "schedule-1",
+		terminalMessage: createAgentMessage("backup.completed", {
+			jobId: "job-1",
+			scheduleId: "schedule-1",
+			exitCode: 0,
+			result: null,
+		}),
+		expectedCancelledCalls: 0,
+	},
+	{
+		name: "backup.failed",
+		jobId: "job-2",
+		scheduleId: "schedule-2",
+		terminalMessage: createAgentMessage("backup.failed", {
+			jobId: "job-2",
+			scheduleId: "schedule-2",
+			error: "backup failed",
+		}),
+		expectedCancelledCalls: 0,
+	},
+	{
+		name: "backup.cancelled",
+		jobId: "job-3",
+		scheduleId: "schedule-3",
+		terminalMessage: createAgentMessage("backup.cancelled", {
+			jobId: "job-3",
+			scheduleId: "schedule-3",
+			message: "Backup was cancelled",
+		}),
+		expectedCancelledCalls: 1,
+	},
+])("close does not emit an extra synthetic backup.cancelled after $name", (testCase) => {
+	const onBackupCancelled = vi.fn();
+	const { session, close } = createSession({
+		onBackupCancelled,
+	});
+
+	Effect.runSync(
+		session.handleMessage(
+			createAgentMessage("backup.started", {
+				jobId: testCase.jobId,
+				scheduleId: testCase.scheduleId,
+			}),
+		),
+	);
+	Effect.runSync(session.handleMessage(testCase.terminalMessage));
+	close();
+
+	expect(onBackupCancelled).toHaveBeenCalledTimes(testCase.expectedCancelledCalls);
+});
+
+test("close emits a synthetic backup.cancelled for a queued backup", () => {
+	const onBackupCancelled = vi.fn();
+	const { session, close } = createSession({
+		onBackupCancelled,
+	});
+
+	Effect.runSync(
+		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: [],
+			},
+		}),
+	);
+
+	close();
+
+	expect(onBackupCancelled).toHaveBeenCalledTimes(1);
+	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.",
+	});
+});

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

@@ -0,0 +1,48 @@
+import type { ChildProcess } from "node:child_process";
+import { createAgentManagerRuntime, type AgentManagerRuntime } from "./controller/server";
+import { spawnLocalAgentProcess, stopLocalAgentProcess } from "./local/process";
+
+export type { AgentBackupEventHandlers } from "./controller/server";
+
+type AgentRuntimeState = {
+	agentManager: AgentManagerRuntime;
+	localAgent: ChildProcess | null;
+};
+
+type ProcessWithAgentRuntime = NodeJS.Process & {
+	__zerobyteAgentRuntime?: AgentRuntimeState;
+};
+
+const getAgentRuntimeState = () => {
+	const runtimeProcess = process as ProcessWithAgentRuntime;
+	const existingRuntime = runtimeProcess.__zerobyteAgentRuntime;
+
+	if (existingRuntime) {
+		return existingRuntime;
+	}
+
+	const runtime = {
+		agentManager: createAgentManagerRuntime(),
+		localAgent: null,
+	};
+
+	runtimeProcess.__zerobyteAgentRuntime = runtime;
+	return runtime;
+};
+
+const getAgentManagerRuntime = () => getAgentRuntimeState().agentManager;
+
+export const spawnLocalAgent = async () => {
+	await spawnLocalAgentProcess(getAgentRuntimeState());
+};
+
+export const stopLocalAgent = async () => {
+	await stopLocalAgentProcess(getAgentRuntimeState());
+};
+
+export const agentManager = getAgentManagerRuntime();
+
+export const stopAgentRuntime = async () => {
+	await getAgentManagerRuntime().stop();
+	await stopLocalAgent();
+};

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

@@ -0,0 +1,274 @@
+import { Data, Effect, Exit, Fiber, Scope } from "effect";
+import { logger } from "@zerobyte/core/node";
+import { toMessage } from "@zerobyte/core/utils";
+import type {
+	BackupCancelPayload,
+	BackupCancelledPayload,
+	BackupCompletedPayload,
+	BackupFailedPayload,
+	BackupProgressPayload,
+	BackupRunPayload,
+	BackupStartedPayload,
+} from "@zerobyte/contracts/agent-protocol";
+import { createControllerAgentSession, type AgentConnectionData, type ControllerAgentSession } from "./session";
+import { validateAgentToken } from "../helpers/tokens";
+
+type AgentBackupEventContext = {
+	agentId: string;
+	agentName: string;
+	payload:
+		| BackupStartedPayload
+		| BackupProgressPayload
+		| BackupCompletedPayload
+		| BackupFailedPayload
+		| BackupCancelledPayload;
+};
+
+export type AgentBackupEventHandlers = {
+	onBackupStarted?: (context: AgentBackupEventContext & { payload: BackupStartedPayload }) => void;
+	onBackupProgress?: (context: AgentBackupEventContext & { payload: BackupProgressPayload }) => void;
+	onBackupCompleted?: (context: AgentBackupEventContext & { payload: BackupCompletedPayload }) => void;
+	onBackupFailed?: (context: AgentBackupEventContext & { payload: BackupFailedPayload }) => void;
+	onBackupCancelled?: (context: AgentBackupEventContext & { payload: BackupCancelledPayload }) => void;
+};
+
+type ControllerAgentSessionHandle = {
+	session: ControllerAgentSession;
+	runFiber: Fiber.RuntimeFiber<void, never>;
+	scope: Scope.CloseableScope;
+};
+
+class StopAgentManagerServerError extends Data.TaggedError("StopAgentManagerServerError")<{
+	cause: unknown;
+}> {}
+
+export function createAgentManagerRuntime() {
+	let sessions = new Map<string, ControllerAgentSessionHandle>();
+	let backupHandlers: AgentBackupEventHandlers = {};
+	let runtimeScope: Scope.CloseableScope | null = null;
+
+	const closeSession = (sessionHandle: ControllerAgentSessionHandle) => {
+		Effect.runSync(Fiber.interrupt(sessionHandle.runFiber));
+		Effect.runSync(Scope.close(sessionHandle.scope, Exit.succeed(undefined)));
+	};
+
+	const closeAllSessions = () => {
+		const currentSessions = sessions;
+		sessions = new Map();
+		for (const sessionHandle of currentSessions.values()) {
+			closeSession(sessionHandle);
+		}
+	};
+
+	const getSessionHandle = (agentId: string) => sessions.get(agentId);
+
+	const getSession = (agentId: string) => getSessionHandle(agentId)?.session;
+
+	const createSessionHandlers = (ws: Bun.ServerWebSocket<AgentConnectionData>) => {
+		const agentId = ws.data.agentId;
+		const agentName = ws.data.agentName;
+
+		return {
+			onBackupStarted: (payload: BackupStartedPayload) => {
+				backupHandlers.onBackupStarted?.({ agentId, agentName, payload });
+			},
+			onBackupProgress: (payload: BackupProgressPayload) => {
+				backupHandlers.onBackupProgress?.({ agentId, agentName, payload });
+			},
+			onBackupCompleted: (payload: BackupCompletedPayload) => {
+				backupHandlers.onBackupCompleted?.({ agentId, agentName, payload });
+			},
+			onBackupFailed: (payload: BackupFailedPayload) => {
+				backupHandlers.onBackupFailed?.({ agentId, agentName, payload });
+			},
+			onBackupCancelled: (payload: BackupCancelledPayload) => {
+				backupHandlers.onBackupCancelled?.({ agentId, agentName, payload });
+			},
+		};
+	};
+
+	const createSession = (ws: Bun.ServerWebSocket<AgentConnectionData>) => {
+		// Manual scope management because we are out of Effect
+		const scope = Effect.runSync(Scope.make());
+
+		try {
+			const session = Effect.runSync(Scope.extend(createControllerAgentSession(ws, createSessionHandlers(ws)), scope));
+			const runFiber = Effect.runFork(Scope.extend(session.run, scope));
+
+			return { session, runFiber, scope };
+		} catch (error) {
+			Effect.runSync(Scope.close(scope, Exit.fail(error)));
+			throw error;
+		}
+	};
+
+	const setSession = (agentId: string, sessionHandle: ControllerAgentSessionHandle) => {
+		const existingSession = getSessionHandle(agentId);
+		if (existingSession) {
+			closeSession(existingSession);
+		}
+
+		sessions.set(agentId, sessionHandle);
+	};
+
+	const removeSession = (agentId: string, connectionId: string) => {
+		const sessionHandle = getSessionHandle(agentId);
+		if (!sessionHandle || sessionHandle.session.connectionId !== connectionId) {
+			return;
+		}
+
+		sessions.delete(agentId);
+		closeSession(sessionHandle);
+	};
+
+	const acquireServer = Effect.acquireRelease(
+		Effect.sync(() =>
+			Bun.serve<AgentConnectionData>({
+				port: 3001,
+				async fetch(req, srv) {
+					const authorizationHeader = req.headers.get("authorization");
+					const token = authorizationHeader?.slice("Bearer ".length);
+
+					if (!token) {
+						return new Response("Missing token", { status: 401 });
+					}
+
+					const result = await validateAgentToken(token);
+					if (!result) {
+						return new Response("Invalid or revoked token", { status: 401 });
+					}
+
+					const upgraded = srv.upgrade(req, {
+						data: {
+							id: Bun.randomUUIDv7(),
+							agentId: result.agentId,
+							organizationId: result.organizationId,
+							agentName: result.agentName,
+						},
+					});
+					if (upgraded) return undefined;
+					return new Response("WebSocket upgrade failed", { status: 400 });
+				},
+				websocket: {
+					open: (ws) => {
+						setSession(ws.data.agentId, createSession(ws));
+						logger.info(`Agent "${ws.data.agentName}" (${ws.data.agentId}) connected on ${ws.data.id}`);
+					},
+					message: (ws, data) => {
+						if (typeof data !== "string") {
+							logger.warn(`Ignoring non-text message from agent ${ws.data.agentId}`);
+							return;
+						}
+
+						const session = getSession(ws.data.agentId);
+						if (!session || session.connectionId !== ws.data.id) {
+							logger.warn(`No active session for agent ${ws.data.agentId} on ${ws.data.id}`);
+							return;
+						}
+
+						void Effect.runPromise(session.handleMessage(data)).catch((error) => {
+							logger.error(
+								`Failed to handle message from agent ${ws.data.agentId} on ${ws.data.id}: ${toMessage(error)}`,
+							);
+						});
+					},
+					close: (ws) => {
+						removeSession(ws.data.agentId, ws.data.id);
+						logger.info(`Agent "${ws.data.agentName}" (${ws.data.agentId}) disconnected`);
+					},
+				},
+			}),
+		),
+		(server) =>
+			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 = async () => {
+		if (!runtimeScope) {
+			return;
+		}
+
+		logger.info("Stopping Agent Manager...");
+		const scope = runtimeScope;
+		runtimeScope = null;
+		await Effect.runPromise(Scope.close(scope, Exit.succeed(undefined)));
+	};
+
+	const start = async () => {
+		if (runtimeScope) {
+			await stop();
+		}
+
+		logger.info("Starting Agent Manager...");
+		const scope = Effect.runSync(Scope.make());
+
+		try {
+			const server = Effect.runSync(Scope.extend(acquireServer, scope));
+			runtimeScope = scope;
+			logger.info(`Agent Manager listening on port ${server.port}`);
+		} catch (error) {
+			await Effect.runPromise(Scope.close(scope, Exit.fail(error)));
+			throw error;
+		}
+	};
+
+	return {
+		start,
+		sendBackup: async (agentId: string, payload: BackupRunPayload) => {
+			const session = getSession(agentId);
+
+			if (!session) {
+				logger.warn(`Cannot send backup command. Agent ${agentId} is not connected.`);
+				return false;
+			}
+
+			if (!Effect.runSync(session.isReady())) {
+				logger.warn(`Cannot send backup command. Agent ${agentId} is not ready.`);
+				return false;
+			}
+
+			if (!(await Effect.runPromise(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;
+		},
+		cancelBackup: async (agentId: string, payload: BackupCancelPayload) => {
+			const session = getSession(agentId);
+
+			if (!session) {
+				logger.warn(`Cannot cancel backup command. Agent ${agentId} is not connected.`);
+				return false;
+			}
+
+			if (!(await Effect.runPromise(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;
+		},
+		setBackupEventHandlers: (handlers: AgentBackupEventHandlers) => {
+			backupHandlers = handlers;
+		},
+		getBackupEventHandlers: () => backupHandlers,
+		stop,
+	};
+}
+
+export type AgentManagerRuntime = ReturnType<typeof createAgentManagerRuntime>;

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

@@ -0,0 +1,249 @@
+import { Effect, Queue, Ref, type Scope } from "effect";
+import {
+	createControllerMessage,
+	parseAgentMessage,
+	type AgentMessage,
+	type BackupCancelPayload,
+	type BackupCancelledPayload,
+	type BackupCompletedPayload,
+	type BackupFailedPayload,
+	type BackupProgressPayload,
+	type BackupRunPayload,
+	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 TrackedBackupJob = {
+	scheduleId: string;
+	state: "pending" | "active";
+};
+
+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) => 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>;
+};
+
+export const createControllerAgentSession = (
+	socket: AgentSocket,
+	handlers: ControllerAgentSessionHandlers = {},
+): Effect.Effect<ControllerAgentSession, never, Scope.Scope> =>
+	Effect.gen(function* () {
+		const outboundQueue = yield* Queue.bounded<ControllerWireMessage>(64);
+		const trackedBackupJobs = yield* Ref.make<Map<string, TrackedBackupJob>>(new Map());
+		const state = yield* Ref.make<SessionState>({
+			isReady: false,
+			lastSeenAt: null,
+			lastPongAt: null,
+		});
+
+		const offerOutbound = (message: ControllerWireMessage) =>
+			Queue.offer(outboundQueue, message).pipe(
+				Effect.catchAllCause((cause) =>
+					Effect.sync(() => {
+						logger.error(`Failed to queue outbound message for agent ${socket.data.agentId}: ${toMessage(cause)}`);
+						return false;
+					}),
+				),
+			);
+
+		const updateState = (update: (current: SessionState) => SessionState) => Ref.update(state, update);
+
+		const setTrackedBackupJob = (jobId: string, trackedBackupJob: TrackedBackupJob) => {
+			return Ref.update(trackedBackupJobs, (current) => {
+				const next = new Map(current);
+				next.set(jobId, trackedBackupJob);
+				return next;
+			});
+		};
+
+		const deleteTrackedBackupJob = (jobId: string) => {
+			return Ref.update(trackedBackupJobs, (current) => {
+				const next = new Map(current);
+				next.delete(jobId);
+				return next;
+			});
+		};
+
+		const takeTrackedBackupJobs = Ref.modify(
+			trackedBackupJobs,
+			(current) => [current, new Map<string, TrackedBackupJob>()] as const,
+		);
+
+		const releaseSession = Effect.gen(function* () {
+			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.";
+
+				yield* Effect.sync(() => {
+					handlers.onBackupCancelled?.({ jobId, scheduleId: trackedJob.scheduleId, message });
+				});
+			}
+
+			yield* Queue.shutdown(outboundQueue);
+		});
+
+		yield* Effect.addFinalizer(() => releaseSession);
+
+		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)}`,
+								);
+							}
+						});
+					}),
+				),
+			);
+
+			yield* Effect.forkScoped(
+				Effect.forever(
+					Effect.gen(function* () {
+						yield* Effect.sleep("15 seconds");
+						yield* Queue.offer(
+							outboundQueue,
+							createControllerMessage("heartbeat.ping", {
+								sentAt: Date.now(),
+							}),
+						);
+					}),
+				),
+			);
+
+			return yield* Effect.never;
+		});
+
+		const handleAgentMessage = (message: AgentMessage) =>
+			Effect.gen(function* () {
+				yield* updateState((current) => ({ ...current, lastSeenAt: Date.now() }));
+
+				switch (message.type) {
+					case "agent.ready": {
+						yield* updateState((current) => ({ ...current, isReady: true }));
+						yield* Effect.sync(() => {
+							logger.info(`Agent "${socket.data.agentName}" (${socket.data.agentId}) is ready`);
+						});
+						break;
+					}
+					case "backup.started": {
+						yield* setTrackedBackupJob(message.payload.jobId, {
+							scheduleId: message.payload.scheduleId,
+							state: "active",
+						});
+						yield* Effect.sync(() => {
+							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": {
+						yield* Effect.sync(() => {
+							handlers.onBackupProgress?.(message.payload);
+						});
+						break;
+					}
+					case "backup.completed": {
+						yield* deleteTrackedBackupJob(message.payload.jobId);
+						yield* Effect.sync(() => {
+							handlers.onBackupCompleted?.(message.payload);
+						});
+						break;
+					}
+					case "backup.failed": {
+						yield* deleteTrackedBackupJob(message.payload.jobId);
+						yield* Effect.sync(() => {
+							handlers.onBackupFailed?.(message.payload);
+						});
+						break;
+					}
+					case "backup.cancelled": {
+						yield* deleteTrackedBackupJob(message.payload.jobId);
+						yield* Effect.sync(() => {
+							handlers.onBackupCancelled?.(message.payload);
+						});
+						break;
+					}
+					case "heartbeat.pong": {
+						yield* updateState((current) => ({ ...current, lastPongAt: message.payload.sentAt }));
+						break;
+					}
+				}
+			});
+
+		return {
+			connectionId: socket.data.id,
+			handleMessage: (data: string) => {
+				return Effect.gen(function* () {
+					const parsed = parseAgentMessage(data);
+
+					if (parsed === null) {
+						yield* Effect.sync(() => {
+							logger.warn(`Invalid JSON from agent ${socket.data.agentId}`);
+						});
+						return;
+					}
+
+					if (!parsed.success) {
+						yield* Effect.sync(() => {
+							logger.warn(`Invalid agent message from ${socket.data.agentId}: ${parsed.error.message}`);
+						});
+						return;
+					}
+
+					yield* handleAgentMessage(parsed.data);
+				});
+			},
+			sendBackup: (payload) => {
+				return Effect.gen(function* () {
+					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)),
+			isReady: () => Ref.get(state).pipe(Effect.map((current) => current.isReady)),
+			run,
+		};
+	});

app/server/modules/agents/helpers/tokens.ts

@@ -0,0 +1,12 @@
+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" };
+	}
+};

app/server/modules/agents/local/process.ts

@@ -0,0 +1,74 @@
+import { type ChildProcess, spawn } from "node:child_process";
+import { existsSync } from "node:fs";
+import path from "node:path";
+import { logger } from "@zerobyte/core/node";
+import { config } from "../../../core/config";
+import { deriveLocalAgentToken } from "../helpers/tokens";
+
+type LocalAgentState = {
+	localAgent: ChildProcess | null;
+};
+
+export async function spawnLocalAgentProcess(runtime: LocalAgentState) {
+	await stopLocalAgentProcess(runtime);
+
+	const sourceEntryPoint = path.join(process.cwd(), "apps", "agent", "src", "index.ts");
+	const productionEntryPoint = path.join(process.cwd(), ".output", "agent", "index.mjs");
+
+	if (config.__prod__ && !existsSync(productionEntryPoint)) {
+		throw new Error(`Local agent entrypoint not found at ${productionEntryPoint}`);
+	}
+
+	const agentEntryPoint = config.__prod__ ? productionEntryPoint : sourceEntryPoint;
+	const agentToken = await deriveLocalAgentToken();
+	const args = config.__prod__ ? ["run", agentEntryPoint] : ["run", "--watch", agentEntryPoint];
+	const agentProcess = spawn("bun", args, {
+		env: {
+			...process.env,
+			ZEROBYTE_CONTROLLER_URL: "ws://localhost:3001",
+			ZEROBYTE_AGENT_TOKEN: agentToken,
+		},
+		stdio: ["ignore", "pipe", "pipe"],
+	});
+
+	runtime.localAgent = agentProcess;
+
+	agentProcess.stdout?.on("data", (data: Buffer) => {
+		const line = data.toString().trim();
+		if (line) logger.info(`[agent] ${line}`);
+	});
+
+	agentProcess.stderr?.on("data", (data: Buffer) => {
+		const line = data.toString().trim();
+		if (line) logger.error(`[agent] ${line}`);
+	});
+
+	agentProcess.on("exit", (code, signal) => {
+		if (runtime.localAgent === agentProcess) {
+			runtime.localAgent = null;
+		}
+		logger.info(`Agent process exited with code ${code} and signal ${signal}`);
+	});
+}
+
+export async function stopLocalAgentProcess(runtime: LocalAgentState) {
+	if (!runtime.localAgent) {
+		return;
+	}
+
+	const agentProcess = runtime.localAgent;
+	runtime.localAgent = null;
+
+	if (agentProcess.exitCode !== null || agentProcess.signalCode !== null) {
+		return;
+	}
+
+	const exited = new Promise<void>((resolve) => {
+		agentProcess.once("exit", () => {
+			resolve();
+		});
+	});
+
+	agentProcess.kill();
+	await exited;
+}

apps/agent/src/index.ts

@@ -35,9 +35,11 @@ class Agent {
 		}
 
 		const url = new URL(controllerUrl);
-		url.searchParams.set("token", agentToken);
-
-		this.ws = new WebSocket(url.toString());
+		this.ws = new WebSocket(url.toString(), {
+			headers: {
+				authorization: `Bearer ${agentToken}`,
+			},
+		});
 		this.controllerSession = createControllerSession(this.ws);
 
 		this.ws.onopen = () => {

apps/agent/tsconfig.json

@@ -2,12 +2,12 @@
 	"compilerOptions": {
 		"plugins": [{ "name": "@effect/language-service" }],
 		// Environment setup & latest features
-		"lib": ["DOM", "ESNext"],
+		"lib": ["ESNext"],
 		"target": "ESNext",
 		"module": "Preserve",
 		"moduleDetection": "force",
 		"jsx": "react-jsx",
-		"types": ["bun", "node"],
+		"types": ["node", "bun"],
 		"allowJs": true,
 
 		// Bundler mode