From 5bcc3acc56dc965deac83918f8a866b5150bb663 Mon Sep 17 00:00:00 2001 From: Nico <47644445+nicotsx@users.noreply.github.com> Date: Thu, 9 Apr 2026 20:00:54 +0200 Subject: [PATCH] feat(controller): add agent manager and session handling (#763) --- .../modules/agents/__tests__/session.test.ts | 146 ++++++++++ app/server/modules/agents/agents-manager.ts | 48 +++ .../modules/agents/controller/server.ts | 274 ++++++++++++++++++ .../modules/agents/controller/session.ts | 249 ++++++++++++++++ app/server/modules/agents/helpers/tokens.ts | 12 + app/server/modules/agents/local/process.ts | 74 +++++ apps/agent/src/index.ts | 8 +- apps/agent/tsconfig.json | 4 +- 8 files changed, 810 insertions(+), 5 deletions(-) create mode 100644 app/server/modules/agents/__tests__/session.test.ts create mode 100644 app/server/modules/agents/agents-manager.ts create mode 100644 app/server/modules/agents/controller/server.ts create mode 100644 app/server/modules/agents/controller/session.ts create mode 100644 app/server/modules/agents/helpers/tokens.ts create mode 100644 app/server/modules/agents/local/process.ts diff --git a/app/server/modules/agents/__tests__/session.test.ts b/app/server/modules/agents/__tests__/session.test.ts new file mode 100644 index 00000000..87671f4e --- /dev/null +++ b/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[0]>({ + data: { id: "connection-1", agentId: "local", organizationId: null, agentName: "Local Agent" }, + send: vi.fn(), + }); +}; + +const createSession = (handlers: Parameters[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.", + }); +}); diff --git a/app/server/modules/agents/agents-manager.ts b/app/server/modules/agents/agents-manager.ts new file mode 100644 index 00000000..72e900c2 --- /dev/null +++ b/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(); +}; diff --git a/app/server/modules/agents/controller/server.ts b/app/server/modules/agents/controller/server.ts new file mode 100644 index 00000000..3bb025d7 --- /dev/null +++ b/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; + scope: Scope.CloseableScope; +}; + +class StopAgentManagerServerError extends Data.TaggedError("StopAgentManagerServerError")<{ + cause: unknown; +}> {} + +export function createAgentManagerRuntime() { + let sessions = new Map(); + 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) => { + 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) => { + // 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({ + 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; diff --git a/app/server/modules/agents/controller/session.ts b/app/server/modules/agents/controller/session.ts new file mode 100644 index 00000000..cae39436 --- /dev/null +++ b/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; + +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; + sendBackup: (payload: BackupRunPayload) => Effect.Effect; + sendBackupCancel: (payload: BackupCancelPayload) => Effect.Effect; + isReady: () => Effect.Effect; + run: Effect.Effect; +}; + +export const createControllerAgentSession = ( + socket: AgentSocket, + handlers: ControllerAgentSessionHandlers = {}, +): Effect.Effect => + Effect.gen(function* () { + const outboundQueue = yield* Queue.bounded(64); + const trackedBackupJobs = yield* Ref.make>(new Map()); + const state = yield* Ref.make({ + 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()] 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, + }; + }); diff --git a/app/server/modules/agents/helpers/tokens.ts b/app/server/modules/agents/helpers/tokens.ts new file mode 100644 index 00000000..febb9d94 --- /dev/null +++ b/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" }; + } +}; diff --git a/app/server/modules/agents/local/process.ts b/app/server/modules/agents/local/process.ts new file mode 100644 index 00000000..c844a65d --- /dev/null +++ b/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((resolve) => { + agentProcess.once("exit", () => { + resolve(); + }); + }); + + agentProcess.kill(); + await exited; +} diff --git a/apps/agent/src/index.ts b/apps/agent/src/index.ts index 1b545406..f9e08b6d 100644 --- a/apps/agent/src/index.ts +++ b/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 = () => { diff --git a/apps/agent/tsconfig.json b/apps/agent/tsconfig.json index a39e28eb..e30ff81c 100644 --- a/apps/agent/tsconfig.json +++ b/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