danvics/zerobyte
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions 1

refactor: simplify session management

Browse source
This commit is contained in:
Nicolas Meienberger 2026-05-05 17:42:00 +02:00
parent 91286eb84f
commit 88fe3abc11
No known key found for this signature in database
3 changed files with 92 additions and 90 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

16
app/server/modules/agents/__tests__/session.test.ts
View file

@ -1,4 +1,4 @@
import { Effect, Exit, Scope } from "effect"; import { Effect, Exit, Fiber, Scope } from "effect";
import { expect, test, vi } from "vitest"; import { expect, test, vi } from "vitest";
import waitForExpect from "wait-for-expect"; import waitForExpect from "wait-for-expect";
import { fromPartial } from "@total-typescript/shoehorn"; import { fromPartial } from "@total-typescript/shoehorn";
@ -45,7 +45,9 @@ const createSession = (
return { return {
session, session,
run: () => { run: () => {
Effect.runFork(Scope.extend(session.run, scope)); const fiber = Effect.runFork(Scope.extend(session.run, scope));
Effect.runSync(Scope.addFinalizer(scope, Fiber.interrupt(fiber)));
return fiber;
}, },
socket, socket,
close: () => { close: () => {
@ -61,6 +63,16 @@ const createSession = (
} }
}; };
test("closing the session scope interrupts the session runner", async () => {
const { run, closeAsync } = createSession();
const fiber = run();
await closeAsync();
const exit = await Effect.runPromise(Fiber.await(fiber).pipe(Effect.timeout("100 millis")));
expect(Exit.isInterrupted(exit)).toBe(true);
});
test("close reports a transport disconnect", () => { test("close reports a transport disconnect", () => {
const onDisconnect = vi.fn(() => Effect.void); const onDisconnect = vi.fn(() => Effect.void);
const { close } = createSession({ onDisconnect }); const { close } = createSession({ onDisconnect });

4
app/server/modules/agents/agents-manager.ts
View file

@ -172,9 +172,7 @@ export const startAgentController = async () => {
runtime.agentManager = null; runtime.agentManager = null;
} }
const nextAgentManager = createAgentManagerRuntime(); const nextAgentManager = createAgentManagerRuntime(backupEventHandlers);
nextAgentManager.setBackupEventHandlers(backupEventHandlers);
await nextAgentManager.start(); await nextAgentManager.start();
runtime.agentManager = nextAgentManager; runtime.agentManager = nextAgentManager;
}; };

162
app/server/modules/agents/controller/server.ts
View file

@ -35,8 +35,8 @@ export type AgentBackupEventHandlers = {
}; };
type ControllerAgentSessionHandle = { type ControllerAgentSessionHandle = {
agentId: string;
session: ControllerAgentSession; session: ControllerAgentSession;
runFiber: Fiber.RuntimeFiber<void, never>;
scope: Scope.CloseableScope; scope: Scope.CloseableScope;
}; };
@ -44,15 +44,14 @@ class StopAgentManagerServerError extends Data.TaggedError("StopAgentManagerServ
cause: unknown; cause: unknown;
}> {} }> {}
export function createAgentManagerRuntime() { export function createAgentManagerRuntime(handlers: AgentBackupEventHandlers) {
let sessions = new Map<string, ControllerAgentSessionHandle>(); let sessions = new Map<string, ControllerAgentSessionHandle>();
let backupHandlers: AgentBackupEventHandlers = {};
let runtimeScope: Scope.CloseableScope | null = null; let runtimeScope: Scope.CloseableScope | null = null;
const closeSession = (sessionHandle: ControllerAgentSessionHandle) => const closeSession = (sessionHandle: ControllerAgentSessionHandle) =>
Effect.gen(function* () { Effect.gen(function* () {
yield* Fiber.interrupt(sessionHandle.runFiber);
yield* Scope.close(sessionHandle.scope, Exit.succeed(undefined)); yield* Scope.close(sessionHandle.scope, Exit.succeed(undefined));
yield* Effect.sync(() => sessions.delete(sessionHandle.agentId));
}); });
const closeAllSessions = Effect.gen(function* () { const closeAllSessions = Effect.gen(function* () {
@ -65,17 +64,14 @@ export function createAgentManagerRuntime() {
}); });
const getSessionHandle = (agentId: string) => sessions.get(agentId); const getSessionHandle = (agentId: string) => sessions.get(agentId);
const getSession = (agentId: string) => getSessionHandle(agentId)?.session; const getSession = (agentId: string) => getSessionHandle(agentId)?.session;
const createSessionHandlers = (ws: Bun.ServerWebSocket<AgentConnectionData>, markConnecting: Promise<unknown>) => { const createSessionHandlers = (params: { agentId: string; agentName: string; sessionId: string }) => {
const agentId = ws.data.agentId; const { agentId, agentName, sessionId } = params;
const agentName = ws.data.agentName;
return { return {
onReady: ({ at }: { at: number }) => { onReady: ({ at }: { at: number }) => {
return Effect.promise(async () => { return Effect.promise(async () => {
await markConnecting;
await agentsService.markAgentOnline(agentId, at); await agentsService.markAgentOnline(agentId, at);
}); });
}, },
@ -83,70 +79,89 @@ export function createAgentManagerRuntime() {
return Effect.promise(() => agentsService.markAgentSeen(agentId, at)); return Effect.promise(() => agentsService.markAgentSeen(agentId, at));
}, },
onDisconnect: () => { onDisconnect: () => {
if (getSession(ws.data.agentId)?.connectionId !== ws.data.id) { if (getSession(agentId)?.connectionId !== sessionId) {
return Effect.void; return Effect.void;
} }
return Effect.sync(() => backupHandlers.onAgentDisconnected?.({ agentId, agentName })); return Effect.sync(() => handlers.onAgentDisconnected?.({ agentId, agentName }));
}, },
onBackupStarted: (payload: BackupStartedPayload) => { onBackupStarted: (payload: BackupStartedPayload) => {
return Effect.sync(() => backupHandlers.onBackupStarted?.({ agentId, agentName, payload })); return Effect.sync(() => handlers.onBackupStarted?.({ agentId, agentName, payload }));
}, },
onBackupProgress: (payload: BackupProgressPayload) => { onBackupProgress: (payload: BackupProgressPayload) => {
return Effect.sync(() => backupHandlers.onBackupProgress?.({ agentId, agentName, payload })); return Effect.sync(() => handlers.onBackupProgress?.({ agentId, agentName, payload }));
}, },
onBackupCompleted: (payload: BackupCompletedPayload) => { onBackupCompleted: (payload: BackupCompletedPayload) => {
return Effect.sync(() => backupHandlers.onBackupCompleted?.({ agentId, agentName, payload })); return Effect.sync(() => handlers.onBackupCompleted?.({ agentId, agentName, payload }));
}, },
onBackupFailed: (payload: BackupFailedPayload) => { onBackupFailed: (payload: BackupFailedPayload) => {
return Effect.sync(() => backupHandlers.onBackupFailed?.({ agentId, agentName, payload })); return Effect.sync(() => handlers.onBackupFailed?.({ agentId, agentName, payload }));
}, },
onBackupCancelled: (payload: BackupCancelledPayload) => { onBackupCancelled: (payload: BackupCancelledPayload) => {
return Effect.sync(() => backupHandlers.onBackupCancelled?.({ agentId, agentName, payload })); return Effect.sync(() => handlers.onBackupCancelled?.({ agentId, agentName, payload }));
}, },
}; };
}; };
const createSession = (ws: Bun.ServerWebSocket<AgentConnectionData>, markConnecting: Promise<unknown>) => { const createSession = (ws: Bun.ServerWebSocket<AgentConnectionData>) =>
// Manual scope management because we are out of Effect Effect.gen(function* () {
const scope = Effect.runSync(Scope.make()); const scope = yield* Scope.make();
try { const session = yield* Scope.extend(
const session = Effect.runSync( createControllerAgentSession(
Scope.extend(createControllerAgentSession(ws, createSessionHandlers(ws, markConnecting)), scope), ws,
createSessionHandlers({
agentId: ws.data.agentId,
agentName: ws.data.agentName,
sessionId: ws.data.id,
}),
),
scope,
); );
const runFiber = Effect.runFork(Scope.extend(session.run, scope)); const runFiber = yield* Effect.fork(Scope.extend(session.run, scope));
yield* Scope.addFinalizer(scope, Fiber.interrupt(runFiber));
return { session, runFiber, scope }; return { agentId: ws.data.agentId, session, scope };
} catch (error) { });
Effect.runSync(Scope.close(scope, Exit.fail(error)));
throw error; const setSession = (sessionHandle: ControllerAgentSessionHandle) =>
} Effect.gen(function* () {
}; const existingSession = sessions.get(sessionHandle.agentId);
sessions.set(sessionHandle.agentId, sessionHandle);
const setSession = (agentId: string, sessionHandle: ControllerAgentSessionHandle) => {
const existingSession = getSessionHandle(agentId); if (existingSession) {
sessions.set(agentId, sessionHandle); yield* closeSession(existingSession);
}
if (existingSession) { });
void Effect.runPromise(closeSession(existingSession)).catch((error) => {
logger.error(`Failed to close existing agent session for ${agentId}: ${toMessage(error)}`); const removeSession = (agentId: string, connectionId: string) =>
}); Effect.gen(function* () {
} const handle = sessions.get(agentId);
}; if (!handle || handle.session.connectionId !== connectionId) {
return false;
const removeSession = (agentId: string, connectionId: string) => { }
const sessionHandle = getSessionHandle(agentId);
if (!sessionHandle || sessionHandle.session.connectionId !== connectionId) { yield* closeSession(handle);
return false;
} yield* Effect.promise(() => agentsService.markAgentOffline(agentId));
return true;
sessions.delete(agentId); });
void Effect.runPromise(closeSession(sessionHandle)).catch((error) => {
logger.error(`Failed to close agent session for ${agentId}: ${toMessage(error)}`); const handleMessage = (ws: Bun.ServerWebSocket<AgentConnectionData>, data: unknown) =>
Effect.gen(function* () {
if (typeof data !== "string") {
yield* logger.effect.warn(`Ignoring non-text message from agent ${ws.data.agentId}`);
return;
}
const session = getSession(ws.data.agentId);
if (!session || session.connectionId !== ws.data.id) {
yield* logger.effect.warn(`No active session for agent ${ws.data.agentId} on ${ws.data.id}`);
return;
}
yield* session.handleMessage(data);
}); });
return true;
};
const acquireServer = Effect.acquireRelease( const acquireServer = Effect.acquireRelease(
Effect.sync(() => Effect.sync(() =>
@ -178,43 +193,24 @@ export function createAgentManagerRuntime() {
return new Response("WebSocket upgrade failed", { status: 400 }); return new Response("WebSocket upgrade failed", { status: 400 });
}, },
websocket: { websocket: {
open: (ws) => { open: async (ws) => {
const markConnecting = agentsService.markAgentConnecting({ await agentsService.markAgentConnecting({
agentId: ws.data.agentId, agentId: ws.data.agentId,
organizationId: ws.data.organizationId, organizationId: ws.data.organizationId,
agentName: ws.data.agentName, agentName: ws.data.agentName,
agentKind: ws.data.agentKind, agentKind: ws.data.agentKind,
}); });
void markConnecting.catch((error) => {
logger.error(`Failed to mark agent ${ws.data.agentId} as connecting: ${toMessage(error)}`); const sessionHandle = await Effect.runPromise(createSession(ws));
}); await Effect.runPromise(setSession(sessionHandle));
setSession(ws.data.agentId, createSession(ws, markConnecting));
logger.info(`Agent "${ws.data.agentName}" (${ws.data.agentId}) connected on ${ws.data.id}`); logger.info(`Agent "${ws.data.agentName}" (${ws.data.agentId}) connected on ${ws.data.id}`);
}, },
message: (ws, data) => { message: async (ws, data) => {
if (typeof data !== "string") { await Effect.runPromise(handleMessage(ws, data));
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) => { close: async (ws) => {
if (removeSession(ws.data.agentId, ws.data.id)) { await Effect.runPromise(removeSession(ws.data.agentId, ws.data.id));
void agentsService.markAgentOffline(ws.data.agentId).catch((error) => {
logger.error(`Failed to mark agent ${ws.data.agentId} as offline: ${toMessage(error)}`);
});
}
logger.info(`Agent "${ws.data.agentName}" (${ws.data.agentId}) disconnected`); logger.info(`Agent "${ws.data.agentName}" (${ws.data.agentId}) disconnected`);
}, },
}, },
@ -303,10 +299,6 @@ export function createAgentManagerRuntime() {
logger.info(`Sent backup cancel for command ${payload.jobId} to agent ${agentId}`); logger.info(`Sent backup cancel for command ${payload.jobId} to agent ${agentId}`);
return true; return true;
}, },
setBackupEventHandlers: (handlers: AgentBackupEventHandlers) => {
backupHandlers = handlers;
},
getBackupEventHandlers: () => backupHandlers,
stop, stop,
}; };
} }