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

refactor: context as effectful callbacks

Browse source
This commit is contained in:
Nicolas Meienberger 2026-03-31 18:50:05 +02:00
parent 7d61e7d465
commit 7ea7fe783c
7 changed files with 141 additions and 160 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

76
app/server/modules/agents/controller-agent-session.ts
View file

@ -24,12 +24,6 @@ export type AgentConnectionData = {
type AgentSocket = Bun.ServerWebSocket<AgentConnectionData>; type AgentSocket = Bun.ServerWebSocket<AgentConnectionData>;
type SessionState = {
isReady: boolean;
lastSeenAt: number | null;
lastPongAt: number | null;
};
type ControllerAgentSessionHandlers = { type ControllerAgentSessionHandlers = {
onBackupStarted?: (payload: BackupStartedPayload) => void; onBackupStarted?: (payload: BackupStartedPayload) => void;
onBackupProgress?: (payload: BackupProgressPayload) => void; onBackupProgress?: (payload: BackupProgressPayload) => void;
@ -55,13 +49,7 @@ export const createControllerAgentSession = (
const outboundQueue = Effect.runSync(Queue.bounded<ControllerWireMessage>(64)); const outboundQueue = Effect.runSync(Queue.bounded<ControllerWireMessage>(64));
const activeBackupJobs = Effect.runSync(Ref.make<Map<string, string>>(new Map())); const activeBackupJobs = Effect.runSync(Ref.make<Map<string, string>>(new Map()));
const pendingBackupJobs = Effect.runSync(Ref.make<Map<string, string>>(new Map())); const pendingBackupJobs = Effect.runSync(Ref.make<Map<string, string>>(new Map()));
const state = Effect.runSync( const isReadyRef = Effect.runSync(Ref.make(false));
Ref.make<SessionState>({
isReady: false,
lastSeenAt: null,
lastPongAt: null,
}),
);
const offerOutbound = (message: ControllerWireMessage) => { const offerOutbound = (message: ControllerWireMessage) => {
void Effect.runPromise(Queue.offer(outboundQueue, message)).catch((error) => { void Effect.runPromise(Queue.offer(outboundQueue, message)).catch((error) => {
@ -69,8 +57,8 @@ export const createControllerAgentSession = (
}); });
}; };
const updateState = (update: (current: SessionState) => SessionState) => { const setIsReady = (isReady: boolean) => {
Effect.runSync(Ref.update(state, update)); Effect.runSync(Ref.set(isReadyRef, isReady));
}; };
const setActiveBackupJob = (jobId: string, scheduleId: string) => { const setActiveBackupJob = (jobId: string, scheduleId: string) => {
@ -113,35 +101,41 @@ export const createControllerAgentSession = (
); );
}; };
const clearTrackedBackupJob = (jobId: string) => {
deletePendingBackupJob(jobId);
deleteActiveBackupJob(jobId);
};
const cancelTrackedJobs = (jobs: Map<string, string>, message: string) => {
for (const [jobId, scheduleId] of jobs) {
handlers.onBackupCancelled?.({
jobId,
scheduleId,
message,
});
}
};
const closeSession = () => { const closeSession = () => {
if (isClosed) { if (isClosed) {
return; return;
} }
isClosed = true; isClosed = true;
updateState((current) => ({ ...current, isReady: false })); setIsReady(false);
const pendingJobs = Effect.runSync(Ref.get(pendingBackupJobs)); const pendingJobs = Effect.runSync(Ref.get(pendingBackupJobs));
Effect.runSync(Ref.set(pendingBackupJobs, new Map())); Effect.runSync(Ref.set(pendingBackupJobs, new Map()));
cancelTrackedJobs(
for (const [jobId, scheduleId] of pendingJobs) { pendingJobs,
handlers.onBackupCancelled?.({ "The connection to the backup agent was lost before this backup started. Restart the backup to ensure it completes.",
jobId, );
scheduleId,
message:
"The connection to the backup agent was lost before this backup started. Restart the backup to ensure it completes.",
});
}
const activeJobs = Effect.runSync(Ref.get(activeBackupJobs)); const activeJobs = Effect.runSync(Ref.get(activeBackupJobs));
Effect.runSync(Ref.set(activeBackupJobs, new Map())); Effect.runSync(Ref.set(activeBackupJobs, new Map()));
for (const [jobId, scheduleId] of activeJobs) { cancelTrackedJobs(
handlers.onBackupCancelled?.({ activeJobs,
jobId, "The connection to the backup agent was lost while this backup was running. Restart the backup to ensure it completes.",
scheduleId, );
message:
"The connection to the backup agent was lost while this backup was running. Restart the backup to ensure it completes.",
});
}
void Effect.runPromise(Fiber.interrupt(writerFiber)).catch(() => {}); void Effect.runPromise(Fiber.interrupt(writerFiber)).catch(() => {});
void Effect.runPromise(Fiber.interrupt(heartbeatFiber)).catch(() => {}); void Effect.runPromise(Fiber.interrupt(heartbeatFiber)).catch(() => {});
void Effect.runPromise(Queue.shutdown(outboundQueue)).catch(() => {}); void Effect.runPromise(Queue.shutdown(outboundQueue)).catch(() => {});
@ -190,11 +184,9 @@ export const createControllerAgentSession = (
); );
const handleAgentMessage = (message: AgentMessage) => { const handleAgentMessage = (message: AgentMessage) => {
updateState((current) => ({ ...current, lastSeenAt: Date.now() }));
switch (message.type) { switch (message.type) {
case "agent.ready": { case "agent.ready": {
updateState((current) => ({ ...current, isReady: true })); setIsReady(true);
logger.info(`Agent "${socket.data.agentName}" (${socket.data.agentId}) is ready`); logger.info(`Agent "${socket.data.agentName}" (${socket.data.agentId}) is ready`);
break; break;
} }
@ -212,25 +204,21 @@ export const createControllerAgentSession = (
break; break;
} }
case "backup.completed": { case "backup.completed": {
deletePendingBackupJob(message.payload.jobId); clearTrackedBackupJob(message.payload.jobId);
deleteActiveBackupJob(message.payload.jobId);
handlers.onBackupCompleted?.(message.payload); handlers.onBackupCompleted?.(message.payload);
break; break;
} }
case "backup.failed": { case "backup.failed": {
deletePendingBackupJob(message.payload.jobId); clearTrackedBackupJob(message.payload.jobId);
deleteActiveBackupJob(message.payload.jobId);
handlers.onBackupFailed?.(message.payload); handlers.onBackupFailed?.(message.payload);
break; break;
} }
case "backup.cancelled": { case "backup.cancelled": {
deletePendingBackupJob(message.payload.jobId); clearTrackedBackupJob(message.payload.jobId);
deleteActiveBackupJob(message.payload.jobId);
handlers.onBackupCancelled?.(message.payload); handlers.onBackupCancelled?.(message.payload);
break; break;
} }
case "heartbeat.pong": { case "heartbeat.pong": {
updateState((current) => ({ ...current, lastPongAt: message.payload.sentAt }));
break; break;
} }
} }
@ -261,7 +249,7 @@ export const createControllerAgentSession = (
sendBackupCancel: (payload) => { sendBackupCancel: (payload) => {
offerOutbound(createControllerMessage("backup.cancel", payload)); offerOutbound(createControllerMessage("backup.cancel", payload));
}, },
isReady: () => Effect.runSync(Ref.get(state)).isReady, isReady: () => Effect.runSync(Ref.get(isReadyRef)),
close: closeSession, close: closeSession,
}; };
}; };

36
app/server/modules/backups/backup-executor.ts
View file

@ -99,17 +99,25 @@ const clearActiveExecution = (jobId: string) => {
return activeExecution; return activeExecution;
}; };
const getActiveExecution = (jobId: string, scheduleId: string, eventName: string, executorId: string) => { const clearExecutionState = (jobId: string, scheduleId: number) => {
requestedCancellationsByScheduleId.delete(scheduleId);
clearActiveExecution(jobId);
};
const resolveExecution = (jobId: string, activeExecution: ActiveBackupExecution, result: BackupExecutionResult) => {
clearExecutionState(jobId, activeExecution.scheduleId);
activeExecution.resolve(result);
};
const getActiveExecution = (jobId: string, scheduleId: string, eventName: string, agentId: string) => {
const activeExecution = activeExecutionsByJobId.get(jobId); const activeExecution = activeExecutionsByJobId.get(jobId);
if (!activeExecution) { if (!activeExecution) {
logger.warn(`Received ${eventName} for unknown job ${jobId} from executor ${executorId}`); logger.warn(`Received ${eventName} for unknown job ${jobId} from agent ${agentId}`);
return null; return null;
} }
if (activeExecution.scheduleShortId !== scheduleId) { if (activeExecution.scheduleShortId !== scheduleId) {
logger.warn( logger.warn(`Ignoring ${eventName} for job ${jobId} due to schedule mismatch ${scheduleId} from agent ${agentId}`);
`Ignoring ${eventName} for job ${jobId} due to schedule mismatch ${scheduleId} from executor ${executorId}`,
);
return null; return null;
} }
@ -134,9 +142,7 @@ agentManager.setBackupEventHandlers({
return; return;
} }
requestedCancellationsByScheduleId.delete(activeExecution.scheduleId); resolveExecution(payload.jobId, activeExecution, {
clearActiveExecution(payload.jobId);
activeExecution.resolve({
status: "completed", status: "completed",
exitCode: payload.exitCode, exitCode: payload.exitCode,
result: payload.result, result: payload.result,
@ -149,9 +155,7 @@ agentManager.setBackupEventHandlers({
return; return;
} }
requestedCancellationsByScheduleId.delete(activeExecution.scheduleId); resolveExecution(payload.jobId, activeExecution, {
clearActiveExecution(payload.jobId);
activeExecution.resolve({
status: "failed", status: "failed",
error: payload.errorDetails ?? payload.error, error: payload.errorDetails ?? payload.error,
}); });
@ -163,9 +167,7 @@ agentManager.setBackupEventHandlers({
} }
const wasRequested = requestedCancellationsByScheduleId.has(activeExecution.scheduleId); const wasRequested = requestedCancellationsByScheduleId.has(activeExecution.scheduleId);
requestedCancellationsByScheduleId.delete(activeExecution.scheduleId); resolveExecution(payload.jobId, activeExecution, {
clearActiveExecution(payload.jobId);
activeExecution.resolve({
status: "cancelled", status: "cancelled",
message: wasRequested ? undefined : payload.message, message: wasRequested ? undefined : payload.message,
}); });
@ -207,8 +209,7 @@ export const backupExecutor = {
} }
if (!agentManager.sendBackup(LOCAL_AGENT_ID, payload)) { if (!agentManager.sendBackup(LOCAL_AGENT_ID, payload)) {
requestedCancellationsByScheduleId.delete(request.scheduleId); clearExecutionState(jobId, request.scheduleId);
clearActiveExecution(jobId);
return { return {
status: "unavailable", status: "unavailable",
error: new Error("Local backup agent is not connected"), error: new Error("Local backup agent is not connected"),
@ -217,8 +218,7 @@ export const backupExecutor = {
return completion; return completion;
} catch (error) { } catch (error) {
requestedCancellationsByScheduleId.delete(request.scheduleId); clearExecutionState(jobId, request.scheduleId);
clearActiveExecution(jobId);
throw error; throw error;
} }
}, },

4
apps/agent/src/commands/backup-cancel.ts
View file

@ -4,8 +4,8 @@ import { logger } from "@zerobyte/core/node";
import type { ControllerCommandContext } from "../context"; import type { ControllerCommandContext } from "../context";
export const handleBackupCancelCommand = (context: ControllerCommandContext, payload: BackupCancelPayload) => export const handleBackupCancelCommand = (context: ControllerCommandContext, payload: BackupCancelPayload) =>
Effect.sync(() => { Effect.gen(function* () {
const running = context.getRunningJob(payload.jobId); const running = yield* context.getRunningJob(payload.jobId);
if (!running) { if (!running) {
logger.warn(`Backup ${payload.jobId} is not running`); logger.warn(`Backup ${payload.jobId} is not running`);
return; return;

106
apps/agent/src/commands/backup-run.ts
View file

@ -9,9 +9,9 @@ import type { ControllerCommandContext } from "../context";
export const handleBackupRunCommand = (context: ControllerCommandContext, payload: BackupRunPayload) => export const handleBackupRunCommand = (context: ControllerCommandContext, payload: BackupRunPayload) =>
Effect.fork( Effect.fork(
Effect.gen(function* () { Effect.gen(function* () {
const existing = context.getRunningJob(payload.jobId); const existing = yield* context.getRunningJob(payload.jobId);
if (existing) { if (existing) {
context.offerOutbound( yield* context.offerOutbound(
createAgentMessage("backup.failed", { createAgentMessage("backup.failed", {
jobId: payload.jobId, jobId: payload.jobId,
scheduleId: payload.scheduleId, scheduleId: payload.scheduleId,
@ -23,9 +23,19 @@ export const handleBackupRunCommand = (context: ControllerCommandContext, payloa
logger.info(`Starting backup ${payload.jobId} for schedule ${payload.scheduleId}`); logger.info(`Starting backup ${payload.jobId} for schedule ${payload.scheduleId}`);
const abortController = new AbortController(); const abortController = new AbortController();
context.setRunningJob(payload.jobId, { scheduleId: payload.scheduleId, abortController }); yield* context.setRunningJob(payload.jobId, { scheduleId: payload.scheduleId, abortController });
context.offerOutbound( const sendCancelled = () => {
return context.offerOutbound(
createAgentMessage("backup.cancelled", {
jobId: payload.jobId,
scheduleId: payload.scheduleId,
message: "Backup was cancelled",
}),
);
};
yield* context.offerOutbound(
createAgentMessage("backup.started", { createAgentMessage("backup.started", {
jobId: payload.jobId, jobId: payload.jobId,
scheduleId: payload.scheduleId, scheduleId: payload.scheduleId,
@ -49,70 +59,52 @@ export const handleBackupRunCommand = (context: ControllerCommandContext, payloa
...payload.options, ...payload.options,
signal: abortController.signal, signal: abortController.signal,
onProgress: (progress) => { onProgress: (progress) => {
context.offerOutbound( void Effect.runPromise(
createAgentMessage("backup.progress", { context.offerOutbound(
jobId: payload.jobId, createAgentMessage("backup.progress", {
scheduleId: payload.scheduleId, jobId: payload.jobId,
progress, scheduleId: payload.scheduleId,
}), progress,
); }),
),
).catch((error) => {
logger.error(`Failed to send backup progress update: ${toMessage(error)}`);
});
}, },
}) })
.pipe( .pipe(
Effect.matchEffect({ Effect.matchEffect({
onSuccess: (result) => { onSuccess: (result) => {
return Effect.sync(() => { if (abortController.signal.aborted) {
if (abortController.signal.aborted) { return sendCancelled();
context.offerOutbound( }
createAgentMessage("backup.cancelled", {
jobId: payload.jobId,
scheduleId: payload.scheduleId,
message: "Backup was cancelled",
}),
);
return;
}
context.offerOutbound( return context.offerOutbound(
createAgentMessage("backup.completed", { createAgentMessage("backup.completed", {
jobId: payload.jobId, jobId: payload.jobId,
scheduleId: payload.scheduleId, scheduleId: payload.scheduleId,
exitCode: result.exitCode, exitCode: result.exitCode,
result: result.result, result: result.result,
warningDetails: result.warningDetails ?? undefined, warningDetails: result.warningDetails ?? undefined,
}), }),
); );
});
}, },
onFailure: (error) => { onFailure: (error) => {
return Effect.sync(() => { if (abortController.signal.aborted) {
if (abortController.signal.aborted) { return sendCancelled();
context.offerOutbound( }
createAgentMessage("backup.cancelled", {
jobId: payload.jobId,
scheduleId: payload.scheduleId,
message: "Backup was cancelled",
}),
);
return;
}
context.offerOutbound( return context.offerOutbound(
createAgentMessage("backup.failed", { createAgentMessage("backup.failed", {
jobId: payload.jobId, jobId: payload.jobId,
scheduleId: payload.scheduleId, scheduleId: payload.scheduleId,
error: toMessage(error), error: toMessage(error),
errorDetails: toErrorDetails(error), errorDetails: toErrorDetails(error),
}), }),
); );
});
}, },
}), }),
Effect.ensuring( Effect.ensuring(context.deleteRunningJob(payload.jobId)),
Effect.sync(() => {
context.deleteRunningJob(payload.jobId);
}),
),
); );
}), }),
).pipe(Effect.asVoid); ).pipe(Effect.asVoid);

4
apps/agent/src/commands/heartbeat-ping.ts
View file

@ -5,8 +5,8 @@ import type { ControllerCommandContext } from "../context";
type HeartbeatPingPayload = Extract<ControllerMessage, { type: "heartbeat.ping" }>["payload"]; type HeartbeatPingPayload = Extract<ControllerMessage, { type: "heartbeat.ping" }>["payload"];
export const handleHeartbeatPingCommand = (context: ControllerCommandContext, payload: HeartbeatPingPayload) => export const handleHeartbeatPingCommand = (context: ControllerCommandContext, payload: HeartbeatPingPayload) =>
Effect.sync(() => { Effect.gen(function* () {
context.offerOutbound( yield* context.offerOutbound(
createAgentMessage("heartbeat.pong", { createAgentMessage("heartbeat.pong", {
sentAt: payload.sentAt, sentAt: payload.sentAt,
}), }),

9
apps/agent/src/context.ts
View file

@ -1,4 +1,5 @@
import type { AgentWireMessage } from "@zerobyte/contracts/agent-protocol"; import type { AgentWireMessage } from "@zerobyte/contracts/agent-protocol";
import type { Effect } from "effect";
export type RunningJob = { export type RunningJob = {
scheduleId: string; scheduleId: string;
@ -6,8 +7,8 @@ export type RunningJob = {
}; };
export type ControllerCommandContext = { export type ControllerCommandContext = {
getRunningJob: (jobId: string) => RunningJob | undefined; getRunningJob: (jobId: string) => Effect.Effect<RunningJob | undefined, never, never>;
setRunningJob: (jobId: string, job: RunningJob) => void; setRunningJob: (jobId: string, job: RunningJob) => Effect.Effect<void, never, never>;
deleteRunningJob: (jobId: string) => void; deleteRunningJob: (jobId: string) => Effect.Effect<void, never, never>;
offerOutbound: (message: AgentWireMessage) => void; offerOutbound: (message: AgentWireMessage) => Effect.Effect<boolean, never, never>;
}; };

66
apps/agent/src/controller-session.ts
View file

@ -8,6 +8,7 @@ import {
import { logger } from "@zerobyte/core/node"; import { logger } from "@zerobyte/core/node";
import { toMessage } from "@zerobyte/core/utils"; import { toMessage } from "@zerobyte/core/utils";
import { handleControllerCommand } from "./commands"; import { handleControllerCommand } from "./commands";
import type { ControllerCommandContext, RunningJob } from "./context";
export type ControllerSession = { export type ControllerSession = {
onOpen: () => void; onOpen: () => void;
@ -18,45 +19,44 @@ export type ControllerSession = {
export const createControllerSession = (ws: WebSocket): ControllerSession => { export const createControllerSession = (ws: WebSocket): ControllerSession => {
const outboundQueue = Effect.runSync(Queue.bounded<AgentWireMessage>(64)); const outboundQueue = Effect.runSync(Queue.bounded<AgentWireMessage>(64));
const inboundQueue = Effect.runSync(Queue.bounded<ControllerWireMessage>(64)); const inboundQueue = Effect.runSync(Queue.bounded<ControllerWireMessage>(64));
const runningJobsRef = Effect.runSync( const runningJobsRef = Effect.runSync(Ref.make<Map<string, RunningJob>>(new Map()));
Ref.make<Map<string, { scheduleId: string; abortController: AbortController }>>(new Map()),
);
const getRunningJob = (jobId: string) => Effect.runSync(Ref.get(runningJobsRef)).get(jobId); const getRunningJob = (jobId: string) => Ref.get(runningJobsRef).pipe(Effect.map((map) => map.get(jobId)));
const setRunningJob = (jobId: string, job: { scheduleId: string; abortController: AbortController }) => { const setRunningJob = (jobId: string, job: RunningJob) => {
Effect.runSync( return Ref.update(runningJobsRef, (current) => {
Ref.update(runningJobsRef, (current) => { const next = new Map(current);
const next = new Map(current); next.set(jobId, job);
next.set(jobId, job); return next;
return next; });
}),
);
}; };
const abortRunningJobs = Effect.gen(function* () {
const runningJobs = yield* Ref.modify(runningJobsRef, (current) => [current, new Map()]);
yield* Effect.sync(() => {
for (const runningJob of runningJobs.values()) {
runningJob.abortController.abort();
}
});
});
const deleteRunningJob = (jobId: string) => { const deleteRunningJob = (jobId: string) => {
Effect.runSync( return Ref.update(runningJobsRef, (current) => {
Ref.update(runningJobsRef, (current) => { const next = new Map(current);
const next = new Map(current); next.delete(jobId);
next.delete(jobId); return next;
return next; });
}),
);
}; };
const offerOutbound = (message: AgentWireMessage) => { const offerOutbound = (message: AgentWireMessage) => {
void Effect.runPromise(Queue.offer(outboundQueue, message)).catch((error) => { return Queue.offer(outboundQueue, message);
logger.error(`Failed to queue outbound controller message: ${toMessage(error)}`);
});
}; };
const offerInbound = (message: ControllerWireMessage) => { const offerInbound = (message: ControllerWireMessage) => {
void Effect.runPromise(Queue.offer(inboundQueue, message)).catch((error) => { return Queue.offer(inboundQueue, message);
logger.error(`Failed to queue inbound controller message: ${toMessage(error)}`);
});
}; };
const commandContext = { const commandContext: ControllerCommandContext = {
getRunningJob, getRunningJob,
setRunningJob, setRunningJob,
deleteRunningJob, deleteRunningJob,
@ -101,7 +101,9 @@ export const createControllerSession = (ws: WebSocket): ControllerSession => {
return { return {
onOpen: () => { onOpen: () => {
offerOutbound(createAgentMessage("agent.ready", { agentId: "" })); void Effect.runPromise(offerOutbound(createAgentMessage("agent.ready", { agentId: "" }))).catch((error) => {
logger.error(`Failed to queue ready message: ${toMessage(error)}`);
});
}, },
onMessage: (data) => { onMessage: (data) => {
if (typeof data !== "string") { if (typeof data !== "string") {
@ -109,14 +111,12 @@ export const createControllerSession = (ws: WebSocket): ControllerSession => {
return; return;
} }
offerInbound(data as ControllerWireMessage); void Effect.runPromise(offerInbound(data as ControllerWireMessage)).catch((error) => {
logger.error(`Failed to queue inbound message: ${toMessage(error)}`);
});
}, },
close: () => { close: () => {
const runningJobs = Effect.runSync(Ref.get(runningJobsRef)); void Effect.runPromise(abortRunningJobs).catch(() => {});
for (const running of runningJobs.values()) {
running.abortController.abort();
}
Effect.runSync(Ref.set(runningJobsRef, new Map()));
void Effect.runPromise(Fiber.interrupt(writerFiber)).catch(() => {}); void Effect.runPromise(Fiber.interrupt(writerFiber)).catch(() => {});
void Effect.runPromise(Fiber.interrupt(processorFiber)).catch(() => {}); void Effect.runPromise(Fiber.interrupt(processorFiber)).catch(() => {});
void Effect.runPromise(Queue.shutdown(outboundQueue)).catch(() => {}); void Effect.runPromise(Queue.shutdown(outboundQueue)).catch(() => {});