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

refactor: improve error handling

Browse source
This commit is contained in:
Nicolas Meienberger 2026-05-05 19:07:22 +02:00
parent 33c7942fb2
commit 7623a1710b
No known key found for this signature in database
2 changed files with 52 additions and 7 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

26
app/server/modules/agents/__tests__/controller-runtime.test.ts
View file

@ -24,12 +24,12 @@ vi.mock("../helpers/tokens", () => ({
validateAgentToken: tokenMocks.validateAgentToken,
}));
const createSocket = (id: string) => ({
const createSocket = (id: string, agentId = LOCAL_AGENT_ID) => ({
data: {
id,
agentId: LOCAL_AGENT_ID,
agentId,
organizationId: null,
agentName: LOCAL_AGENT_NAME,
agentName: agentId === LOCAL_AGENT_ID ? LOCAL_AGENT_NAME : `${LOCAL_AGENT_NAME} ${agentId}`,
agentKind: LOCAL_AGENT_KIND,
},
send: vi.fn(() => 1),
@ -162,6 +162,26 @@ test("websocket lifecycle updates agent connection status", async () => {
expect(stop).toHaveBeenCalledWith(true);
});
test("shutdown closes all sessions and stops the server when marking one agent offline fails", async () => {
agentsServiceMocks.markAgentOffline.mockRejectedValueOnce(new Error("db unavailable"));
const stop = vi.fn(() => Promise.resolve());
const serve = vi.spyOn(Bun, "serve").mockReturnValue(fromPartial({ port: 3001, stop }));
const { runtime, onEvent } = await startRuntime(vi.fn());
const websocket = serve.mock.calls[0]?.[0].websocket;
const firstSocket = createSocket("connection-1", "agent-1");
const secondSocket = createSocket("connection-2", "agent-2");
await websocket?.open?.(fromPartial(firstSocket));
await websocket?.open?.(fromPartial(secondSocket));
await Effect.runPromise(runtime.stop);
expect(agentsServiceMocks.markAgentOffline).toHaveBeenCalledWith("agent-1");
expect(agentsServiceMocks.markAgentOffline).toHaveBeenCalledWith("agent-2");
expect(onEvent).toHaveBeenCalledWith(expect.objectContaining({ type: "agent.disconnected", agentId: "agent-1" }));
expect(onEvent).toHaveBeenCalledWith(expect.objectContaining({ type: "agent.disconnected", agentId: "agent-2" }));
expect(stop).toHaveBeenCalledWith(true);
});
test("closing a replaced connection does not report the active agent as disconnected", async () => {
const serve = vi
.spyOn(Bun, "serve")

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

@ -51,10 +51,20 @@ export function createAgentManagerRuntime(onEvent: (event: AgentManagerEvent) =>
});
});
const markAgentOfflineForShutdown = (agentId: string) =>
Effect.tryPromise({
try: () => agentsService.markAgentOffline(agentId),
catch: (error) => new StopAgentManagerServerError({ cause: error }),
}).pipe(
Effect.catchAll((error) =>
logger.effect.error(`Failed to mark agent ${agentId} offline during shutdown: ${toMessage(error)}`),
),
);
const closeAllSessions = Effect.gen(function* () {
const currentSessions = [...sessions.entries()];
for (const [agentId, sessionHandle] of currentSessions) {
yield* Effect.promise(() => agentsService.markAgentOffline(agentId));
yield* markAgentOfflineForShutdown(agentId);
yield* closeSession(sessionHandle);
}
sessions = new Map();
@ -174,6 +184,21 @@ export function createAgentManagerRuntime(onEvent: (event: AgentManagerEvent) =>
yield* logger.effect.info(`Agent "${ws.data.agentName}" (${ws.data.agentId}) disconnected`);
});
const runWebSocketHandler = (
ws: Bun.ServerWebSocket<AgentConnectionData>,
event: string,
effect: Effect.Effect<void>,
) =>
Effect.runPromise(
effect.pipe(
Effect.catchAllCause((cause) =>
logger.effect.error(
`Agent websocket ${event} failed for ${ws.data.agentId} on ${ws.data.id}: ${toMessage(cause)}`,
),
),
),
);
const acquireServer = Effect.acquireRelease(
Effect.sync(() =>
Bun.serve<AgentConnectionData>({
@ -205,13 +230,13 @@ export function createAgentManagerRuntime(onEvent: (event: AgentManagerEvent) =>
},
websocket: {
open: async (ws) => {
await Effect.runPromise(handleOpen(ws));
await runWebSocketHandler(ws, "open", handleOpen(ws));
},
message: async (ws, data) => {
await Effect.runPromise(handleMessage(ws, data));
await runWebSocketHandler(ws, "message", handleMessage(ws, data));
},
close: async (ws) => {
await Effect.runPromise(handleClose(ws));
await runWebSocketHandler(ws, "close", handleClose(ws));
},
},
}),