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

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

@ -99,17 +99,25 @@ const clearActiveExecution = (jobId: string) => {
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);
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;
}
if (activeExecution.scheduleShortId !== scheduleId) {
logger.warn(
`Ignoring ${eventName} for job ${jobId} due to schedule mismatch ${scheduleId} from executor ${executorId}`,
);
logger.warn(`Ignoring ${eventName} for job ${jobId} due to schedule mismatch ${scheduleId} from agent ${agentId}`);
return null;
}
@ -134,9 +142,7 @@ agentManager.setBackupEventHandlers({
return;
}
requestedCancellationsByScheduleId.delete(activeExecution.scheduleId);
clearActiveExecution(payload.jobId);
activeExecution.resolve({
resolveExecution(payload.jobId, activeExecution, {
status: "completed",
exitCode: payload.exitCode,
result: payload.result,
@ -149,9 +155,7 @@ agentManager.setBackupEventHandlers({
return;
}
requestedCancellationsByScheduleId.delete(activeExecution.scheduleId);
clearActiveExecution(payload.jobId);
activeExecution.resolve({
resolveExecution(payload.jobId, activeExecution, {
status: "failed",
error: payload.errorDetails ?? payload.error,
});
@ -163,9 +167,7 @@ agentManager.setBackupEventHandlers({
}
const wasRequested = requestedCancellationsByScheduleId.has(activeExecution.scheduleId);
requestedCancellationsByScheduleId.delete(activeExecution.scheduleId);
clearActiveExecution(payload.jobId);
activeExecution.resolve({
resolveExecution(payload.jobId, activeExecution, {
status: "cancelled",
message: wasRequested ? undefined : payload.message,
});
@ -207,8 +209,7 @@ export const backupExecutor = {
}
if (!agentManager.sendBackup(LOCAL_AGENT_ID, payload)) {
requestedCancellationsByScheduleId.delete(request.scheduleId);
clearActiveExecution(jobId);
clearExecutionState(jobId, request.scheduleId);
return {
status: "unavailable",
error: new Error("Local backup agent is not connected"),
@ -217,8 +218,7 @@ export const backupExecutor = {
return completion;
} catch (error) {
requestedCancellationsByScheduleId.delete(request.scheduleId);
clearActiveExecution(jobId);
clearExecutionState(jobId, request.scheduleId);
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";
export const handleBackupCancelCommand = (context: ControllerCommandContext, payload: BackupCancelPayload) =>
Effect.sync(() => {
const running = context.getRunningJob(payload.jobId);
Effect.gen(function* () {
const running = yield* context.getRunningJob(payload.jobId);
if (!running) {
logger.warn(`Backup ${payload.jobId} is not running`);
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) =>
Effect.fork(
Effect.gen(function* () {
const existing = context.getRunningJob(payload.jobId);
const existing = yield* context.getRunningJob(payload.jobId);
if (existing) {
context.offerOutbound(
yield* context.offerOutbound(
createAgentMessage("backup.failed", {
jobId: payload.jobId,
scheduleId: payload.scheduleId,
@ -23,9 +23,19 @@ export const handleBackupRunCommand = (context: ControllerCommandContext, payloa
logger.info(`Starting backup ${payload.jobId} for schedule ${payload.scheduleId}`);
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", {
jobId: payload.jobId,
scheduleId: payload.scheduleId,
@ -49,70 +59,52 @@ export const handleBackupRunCommand = (context: ControllerCommandContext, payloa
...payload.options,
signal: abortController.signal,
onProgress: (progress) => {
context.offerOutbound(
createAgentMessage("backup.progress", {
jobId: payload.jobId,
scheduleId: payload.scheduleId,
progress,
}),
);
void Effect.runPromise(
context.offerOutbound(
createAgentMessage("backup.progress", {
jobId: payload.jobId,
scheduleId: payload.scheduleId,
progress,
}),
),
).catch((error) => {
logger.error(`Failed to send backup progress update: ${toMessage(error)}`);
});
},
})
.pipe(
Effect.matchEffect({
onSuccess: (result) => {
return Effect.sync(() => {
if (abortController.signal.aborted) {
context.offerOutbound(
createAgentMessage("backup.cancelled", {
jobId: payload.jobId,
scheduleId: payload.scheduleId,
message: "Backup was cancelled",
}),
);
return;
}
if (abortController.signal.aborted) {
return sendCancelled();
}
context.offerOutbound(
createAgentMessage("backup.completed", {
jobId: payload.jobId,
scheduleId: payload.scheduleId,
exitCode: result.exitCode,
result: result.result,
warningDetails: result.warningDetails ?? undefined,
}),
);
});
return context.offerOutbound(
createAgentMessage("backup.completed", {
jobId: payload.jobId,
scheduleId: payload.scheduleId,
exitCode: result.exitCode,
result: result.result,
warningDetails: result.warningDetails ?? undefined,
}),
);
},
onFailure: (error) => {
return Effect.sync(() => {
if (abortController.signal.aborted) {
context.offerOutbound(
createAgentMessage("backup.cancelled", {
jobId: payload.jobId,
scheduleId: payload.scheduleId,
message: "Backup was cancelled",
}),
);
return;
}
if (abortController.signal.aborted) {
return sendCancelled();
}
context.offerOutbound(
createAgentMessage("backup.failed", {
jobId: payload.jobId,
scheduleId: payload.scheduleId,
error: toMessage(error),
errorDetails: toErrorDetails(error),
}),
);
});
return context.offerOutbound(
createAgentMessage("backup.failed", {
jobId: payload.jobId,
scheduleId: payload.scheduleId,
error: toMessage(error),
errorDetails: toErrorDetails(error),
}),
);
},
}),
Effect.ensuring(
Effect.sync(() => {
context.deleteRunningJob(payload.jobId);
}),
),
Effect.ensuring(context.deleteRunningJob(payload.jobId)),
);
}),
).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"];
export const handleHeartbeatPingCommand = (context: ControllerCommandContext, payload: HeartbeatPingPayload) =>
Effect.sync(() => {
context.offerOutbound(
Effect.gen(function* () {
yield* context.offerOutbound(
createAgentMessage("heartbeat.pong", {
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 { Effect } from "effect";
export type RunningJob = {
scheduleId: string;
@ -6,8 +7,8 @@ export type RunningJob = {
};
export type ControllerCommandContext = {
getRunningJob: (jobId: string) => RunningJob | undefined;
setRunningJob: (jobId: string, job: RunningJob) => void;
deleteRunningJob: (jobId: string) => void;
offerOutbound: (message: AgentWireMessage) => void;
getRunningJob: (jobId: string) => Effect.Effect<RunningJob | undefined, never, never>;
setRunningJob: (jobId: string, job: RunningJob) => Effect.Effect<void, never, never>;
deleteRunningJob: (jobId: string) => Effect.Effect<void, never, never>;
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 { toMessage } from "@zerobyte/core/utils";
import { handleControllerCommand } from "./commands";
import type { ControllerCommandContext, RunningJob } from "./context";
export type ControllerSession = {
onOpen: () => void;
@ -18,45 +19,44 @@ export type ControllerSession = {
export const createControllerSession = (ws: WebSocket): ControllerSession => {
const outboundQueue = Effect.runSync(Queue.bounded<AgentWireMessage>(64));
const inboundQueue = Effect.runSync(Queue.bounded<ControllerWireMessage>(64));
const runningJobsRef = Effect.runSync(
Ref.make<Map<string, { scheduleId: string; abortController: AbortController }>>(new Map()),
);
const runningJobsRef = Effect.runSync(Ref.make<Map<string, RunningJob>>(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 }) => {
Effect.runSync(
Ref.update(runningJobsRef, (current) => {
const next = new Map(current);
next.set(jobId, job);
return next;
}),
);
const setRunningJob = (jobId: string, job: RunningJob) => {
return Ref.update(runningJobsRef, (current) => {
const next = new Map(current);
next.set(jobId, job);
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) => {
Effect.runSync(
Ref.update(runningJobsRef, (current) => {
const next = new Map(current);
next.delete(jobId);
return next;
}),
);
return Ref.update(runningJobsRef, (current) => {
const next = new Map(current);
next.delete(jobId);
return next;
});
};
const offerOutbound = (message: AgentWireMessage) => {
void Effect.runPromise(Queue.offer(outboundQueue, message)).catch((error) => {
logger.error(`Failed to queue outbound controller message: ${toMessage(error)}`);
});
return Queue.offer(outboundQueue, message);
};
const offerInbound = (message: ControllerWireMessage) => {
void Effect.runPromise(Queue.offer(inboundQueue, message)).catch((error) => {
logger.error(`Failed to queue inbound controller message: ${toMessage(error)}`);
});
return Queue.offer(inboundQueue, message);
};
const commandContext = {
const commandContext: ControllerCommandContext = {
getRunningJob,
setRunningJob,
deleteRunningJob,
@ -101,7 +101,9 @@ export const createControllerSession = (ws: WebSocket): ControllerSession => {
return {
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) => {
if (typeof data !== "string") {
@ -109,14 +111,12 @@ export const createControllerSession = (ws: WebSocket): ControllerSession => {
return;
}
offerInbound(data as ControllerWireMessage);
void Effect.runPromise(offerInbound(data as ControllerWireMessage)).catch((error) => {
logger.error(`Failed to queue inbound message: ${toMessage(error)}`);
});
},
close: () => {
const runningJobs = Effect.runSync(Ref.get(runningJobsRef));
for (const running of runningJobs.values()) {
running.abortController.abort();
}
Effect.runSync(Ref.set(runningJobsRef, new Map()));
void Effect.runPromise(abortRunningJobs).catch(() => {});
void Effect.runPromise(Fiber.interrupt(writerFiber)).catch(() => {});
void Effect.runPromise(Fiber.interrupt(processorFiber)).catch(() => {});
void Effect.runPromise(Queue.shutdown(outboundQueue)).catch(() => {});