@bastani/atomic 0.9.13 → 0.9.14-alpha.1

package/dist/builtin/workflows/src/extension/extension-runtime-state.ts

@@ -1,3 +1,5 @@
+import { getDurableBackend } from "../durable/factory.js";
+import { readWorkflowHeartbeatAnchor, recordWorkflowHeartbeatAnchor } from "../durable/workflow-heartbeat-anchor.js";
 import { cancellationRegistry } from "../runs/background/cancellation-registry.js";
 import type { StageAdapters } from "../runs/foreground/stage-runner.js";
 import type { SessionManager } from "../shared/persistence-restore.js";
@@ -11,6 +13,7 @@ import type {
 	WorkflowPersistencePort,
 	WorkflowRuntimeConfig,
 } from "../shared/types.js";
+import type { WorkflowHeartbeatIdentity } from "../shared/workflow-heartbeat-contract.js";
 import {
 	type ConfigLoadResult,
 	loadWorkflowConfig,
@@ -35,11 +38,57 @@ import {
 import type { ExtensionAPI, PiModelContext } from "./public-types.js";
 import { createExtensionRuntime, type ExtensionRuntime } from "./runtime.js";
 import { createStatusWriter, type StatusWriter } from "./status-writer.js";
+import { registerWorkflowHeartbeatRenderer } from "./workflow-heartbeat-notice.js";
+import {
+	createWorkflowHeartbeatSchedulerState,
+	installWorkflowHeartbeatScheduler,
+	isWorkflowHeartbeatTerminalRun,
+	resetWorkflowHeartbeatSchedulerState,
+	type WorkflowHeartbeatAnchorStore,
+	type WorkflowHeartbeatScheduler,
+	workflowHeartbeatConsumedIdentity,
+	workflowHeartbeatContextInvalidation,
+} from "./workflow-heartbeat-scheduler.js";
 import { workflowModelCatalogFromContext } from "./workflow-model-catalog.js";
 import { makeMcpPort, makePersistencePort } from "./workflow-ports.js";
 import { createWorkflowReloadCoordinator } from "./workflow-reload-coordinator.js";
 import { type WorkflowReloadReport, workflowReloadDiagnostics } from "./workflow-reload-report.js";
 
+/**
+ * Best-effort durable cadence-anchor store for workflow heartbeats.
+ *
+ * `getDurableBackend()` throws `DbosNotReadyError` until a backend exists, and
+ * the scheduler installs at `session_start`, before any workflow has run — so
+ * both sides swallow failure and fall back to the run's own `startedAt`.
+ */
+function durableWorkflowHeartbeatAnchorStore(): WorkflowHeartbeatAnchorStore {
+	return {
+		readAnchorAt(runId) {
+			try {
+				return readWorkflowHeartbeatAnchor(getDurableBackend(), runId);
+			} catch {
+				return undefined;
+			}
+		},
+		async recordAnchorAt(runId, record) {
+			try {
+				return await recordWorkflowHeartbeatAnchor(getDurableBackend(), {
+					runId,
+					anchorAt: record.anchorAt,
+					// The scheduler only records a positive cadence; a run launched
+					// disabled never reaches this seam.
+					intervalMinutes: record.intervalMinutes ?? 0,
+					now: Date.now(),
+				});
+			} catch {
+				// A backend that is not ready must not break a background pass; the
+				// caller retries on the next schedule pass.
+				return false;
+			}
+		},
+	};
+}
+
 export interface WorkflowExtensionRuntimeState {
 	persistenceRef: { current: WorkflowPersistencePort | undefined };
 	mcpPort: WorkflowMcpPort | undefined;
@@ -48,6 +97,7 @@ export interface WorkflowExtensionRuntimeState {
 	discoveryRef: { current: DiscoveryResult | null };
 	lifecycleNotificationState: ReturnType<typeof createWorkflowLifecycleNotificationState>;
 	hilAnswerNotificationState: ReturnType<typeof createWorkflowHilAnswerNotificationState>;
+	workflowHeartbeatSchedulerState: ReturnType<typeof createWorkflowHeartbeatSchedulerState>;
 	/** Seed lifecycle notification state before completed historical snapshots are inserted. */
 	beforeRestoreCompleted(snapshots: readonly RunSnapshot[]): void;
 	runtimeForContext(ctx?: PiModelContext): ExtensionRuntime;
@@ -83,10 +133,12 @@ export function createWorkflowExtensionRuntimeState(
 	let statusWriterRef: StatusWriter = createStatusWriter(store, runtimeConfigRef.current);
 	let lifecycleNotificationsUnsubscribe: (() => void) | null = null;
 	let hilAnswerNotificationsUnsubscribe: (() => void) | null = null;
+	let workflowHeartbeatScheduler: WorkflowHeartbeatScheduler | null = null;
 	let notificationsActive = false;
 	let notificationGeneration = 0;
 	const lifecycleNotificationState = createWorkflowLifecycleNotificationState();
 	const hilAnswerNotificationState = createWorkflowHilAnswerNotificationState();
+	const workflowHeartbeatSchedulerState = createWorkflowHeartbeatSchedulerState();
 	const beforeRestoreCompleted = (snapshots: readonly RunSnapshot[]): void => {
 		seedWorkflowLifecycleNotificationState(lifecycleNotificationState, {
 			...readGraphStoreSnapshot(store),
@@ -102,6 +154,7 @@ export function createWorkflowExtensionRuntimeState(
 			: undefined;
 	registerLifecycleNoticeRenderer({ rendererHost: pi, registerMessageRenderer });
 	registerHilAnswerNoticeRenderer({ rendererHost: pi, registerMessageRenderer });
+	registerWorkflowHeartbeatRenderer({ rendererHost: pi, registerMessageRenderer });
 	const sendWorkflowNotificationMessage: ExtensionAPI["sendMessage"] | undefined =
 		typeof pi.sendMessage === "function" ? (message, options) => pi.sendMessage!(message, options) : undefined;
 	const reinstallLifecycleNotifications = (): void => {
@@ -127,6 +180,86 @@ export function createWorkflowExtensionRuntimeState(
 			sendMessage: sendWorkflowNotificationMessage,
 		});
 	};
+	// `message_end` is the host's injection signal: agent-core emits it at the
+	// moment a message enters the conversation. That is what releases a held
+	// heartbeat slot.
+	//
+	// The turn-settled event is deliberately not used. The host emits it from the
+	// prompt cycle's `finally` whether or not the queued messages were drained, so
+	// pausing the queue mid-turn would release a slot whose card is still parked
+	// and let the next boundary stack a second card behind it. `message_end` never
+	// fires for a parked card. The visible display card is published on the
+	// session-listener channel only, so an extension sees the hidden
+	// reconciliation at consumption and never the card at admission — there is no
+	// ambiguity to disambiguate. Proven end to end against a real `AgentSession`
+	// in test/unit/workflow-heartbeat-parent-pickup.test.ts.
+	//
+	// Registered exactly once, at construction: `pi.on` has no unsubscribe, so a
+	// per-install registration would accumulate a handler on every notification
+	// cycle. The mutable scheduler reference is what routes the signal to the
+	// current installation, and a `null` reference makes it a no-op.
+	//
+	// The same condition decides whether the scheduler holds a slot at all, so the
+	// registration and the option can never disagree: a host that reports no
+	// consumption would leave a held slot with nothing to release it.
+	const parentAvailabilityReported = typeof pi.on === "function";
+	if (parentAvailabilityReported) {
+		// This handler observes, and for one narrow case replaces.
+		//
+		// The last guard on the heartbeat path (issue #1975): the three before it
+		// all sit before `sendMessage`, so a heartbeat the host has already
+		// admitted is beyond them. Consumption is the final moment before its
+		// steer joins the model's context, and `message_end` may return a
+		// replacement of the same role, so a heartbeat whose run has finished,
+		// vanished, or no longer owns that exact scheduled identity is excluded
+		// from context rather than steering the parent about stale work.
+		//
+		// `absent` counts as much as `terminal`. A card recovered at the restart
+		// door is consumed just after `session_start` cleared the store, so its run
+		// is normally absent. A durable resume may instead reuse the run id, but its
+		// future pending boundary has a different `scheduledAt`. Both cases fail
+		// exact pending ownership and preserve the no-backfill rule.
+		pi.on?.("message_end", (event, ctx) => {
+			// The host hands the session manager to the handler on its context; the
+			// top-level `pi.sessionManager` is not populated in every host, so the
+			// per-call one is preferred and the other is the fallback. Verified
+			// against a real `AgentSession` in
+			// test/unit/workflow-heartbeat-parent-pickup.test.ts.
+			const entries = (ctx?.sessionManager ?? pi.sessionManager)?.getEntries?.();
+			const invalidation = workflowHeartbeatContextInvalidation(event, entries, isWorkflowHeartbeatIdentityOwned);
+			if (invalidation !== undefined) return invalidation;
+			const identity = workflowHeartbeatConsumedIdentity(event, entries);
+			if (identity !== undefined) workflowHeartbeatScheduler?.notifyHeartbeatConsumed(identity);
+			return undefined;
+		});
+	}
+
+	/** Whether this exact identity is still pending for a current nonterminal run. */
+	function isWorkflowHeartbeatIdentityOwned(identity: WorkflowHeartbeatIdentity): boolean {
+		const pending = workflowHeartbeatSchedulerState.pending.get(identity.runId);
+		const run = store.runs().find((candidate) => candidate.id === identity.runId);
+		return pending?.scheduledAt === identity.scheduledAt && run !== undefined && !isWorkflowHeartbeatTerminalRun(run);
+	}
+
+	/**
+	 * Workflow heartbeats share the lifecycle-notice lifetime: they are armed
+	 * while notifications are active and disposed with them. The cadence itself
+	 * comes from the live registry, so a workflow reload is picked up without a
+	 * reinstall.
+	 */
+	const reinstallWorkflowHeartbeatScheduler = (): void => {
+		workflowHeartbeatScheduler?.dispose();
+		workflowHeartbeatScheduler = null;
+		if (!notificationsActive) return;
+		workflowHeartbeatScheduler = installWorkflowHeartbeatScheduler({
+			store,
+			state: workflowHeartbeatSchedulerState,
+			sendMessage: sendWorkflowNotificationMessage,
+			resolveIntervalMinutes: (workflowName) => runtimeProxy.registry.get(workflowName)?.heartbeatIntervalMinutes,
+			parentAvailabilityReported,
+			anchorStore: durableWorkflowHeartbeatAnchorStore(),
+		});
+	};
 
 	const hostStageSessionDir: { current: string | undefined } = { current: undefined };
 	const resolveDefaultStageSessionDir = (): string | undefined => hostStageSessionDir.current;
@@ -425,6 +558,7 @@ export function createWorkflowExtensionRuntimeState(
 		discoveryRef,
 		lifecycleNotificationState,
 		hilAnswerNotificationState,
+		workflowHeartbeatSchedulerState,
 		beforeRestoreCompleted,
 		runtimeForContext,
 		resetWorkflowDiscoveryForSession,
@@ -442,6 +576,10 @@ export function createWorkflowExtensionRuntimeState(
 			notificationsActive = active;
 			reinstallLifecycleNotifications();
 			reinstallHilAnswerNotifications();
+			// A fresh host session restarts the cadence from each run's persisted
+			// start time; prior-session schedule and pending state cannot apply.
+			if (!active) resetWorkflowHeartbeatSchedulerState(workflowHeartbeatSchedulerState);
+			reinstallWorkflowHeartbeatScheduler();
 		},
 		updateHostStageSessionDir(sessionManager) {
 			try {