Extract three module globals into RuntimeSingletons container for atomic test swap

Author: ofriw

notebook/store.ts

@@ -10,39 +10,37 @@ import {
 	DEFAULT_MAX_LINES,
 	truncateHead,
 } from "@earendil-works/pi-coding-agent";
-import { AsyncLocalStorage } from "node:async_hooks";
 import type { AgenticodingState } from "../state.js";
-
-/**
- * Module-level write lock state.
- *
- * Concurrent callers serialize by chaining on the prior promise. Reentrancy is
- * tracked per async call chain so a nested saveNotebookPage fails explicitly
- * without rejecting unrelated concurrent writers that happen to overlap.
- */
-let writeLock: Promise<void> = Promise.resolve();
-const writeContext = new AsyncLocalStorage<true>();
+import { createWriteLock, __setSingletons, getSingletons } from "../runtime-singletons.js";
 
 /** Reset write lock state. Only for test cleanup after concurrent runs. */
 export function resetNotebookWriteLock(): void {
-	writeLock = Promise.resolve();
+	__setSingletons(
+		{ ...getSingletons(), writeLock: createWriteLock() },
+		{ forceWriteLock: true },
+	);
 }
 
 async function withWriteLock<T>(fn: () => Promise<T>): Promise<T> {
-	if (writeContext.getStore()) {
+	const s = getSingletons();
+	const lock = s.writeLock;
+	if (s.writeContext.getStore()) {
 		throw new Error(
 			"Notebook write lock is not reentrant — saveNotebookPage called from within its own critical section.",
 		);
 	}
 	let release: () => void;
-	const prev = writeLock;
-	writeLock = new Promise<void>((resolve) => {
+	const prev = lock.tail;
+	const next = new Promise<void>((resolve) => {
 		release = resolve;
 	});
+	lock.pending += 1;
+	lock.tail = next;
 	await prev;
 	try {
-		return await writeContext.run(true, fn);
+		return await s.writeContext.run(true, fn);
 	} finally {
+		lock.pending -= 1;
 		release!();
 	}
 }

runtime-singletons.ts

@@ -0,0 +1,79 @@
+/**
+ * Shared singleton container for the agenticoding extension.
+ *
+ * Allows tests to replace all module-level singletons (write lock, frame
+ * scheduler, etc.) with one atomic swap via __setSingletons(), instead of
+ * patching each singleton individually per test.
+ *
+ * In production the frame scheduler is registered by spawn/renderer.ts at
+ * module import time.  In tests, createTestHarness() provides a fresh
+ * container that tests own and dispose.
+ */
+
+import { AsyncLocalStorage } from "node:async_hooks";
+
+// ── Types ─────────────────────────────────────────────────────────────
+
+/** Minimal frame scheduler interface that the container understands. */
+export interface RuntimeFrameScheduler {
+	markDirty(component: unknown): void;
+	cancelDirty(component: unknown): void;
+	flushNow(): void;
+	clear(): void;
+}
+
+export interface RuntimeWriteLock {
+	pending: number;
+	tail: Promise<void>;
+}
+
+export interface RuntimeSingletons {
+	writeLock: RuntimeWriteLock;
+	writeContext: AsyncLocalStorage<true>;
+	frameScheduler: RuntimeFrameScheduler;
+}
+
+export function createWriteLock(): RuntimeWriteLock {
+	return {
+		pending: 0,
+		tail: Promise.resolve(),
+	};
+}
+
+// ── Pre‑init defaults (overwritten by spawn/renderer.ts at import time) ──
+
+let current: RuntimeSingletons = {
+	writeLock: createWriteLock(),
+	writeContext: new AsyncLocalStorage<true>(),
+	frameScheduler: {
+		markDirty: () => {},
+		cancelDirty: () => {},
+		flushNow: () => {},
+		clear: () => {},
+	},
+};
+
+// ── Public API ────────────────────────────────────────────────────────
+
+/** Atomically replace all singletons.  Test‑only — use __ naming convention. */
+export function __setSingletons(
+	s: RuntimeSingletons,
+	options?: { forceWriteLock?: boolean },
+): void {
+	if (!options?.forceWriteLock && current.writeLock.pending > 0) {
+		console.warn(
+			"[runtime-singletons] writeLock has %d pending operation(s) — " +
+				"preserving existing lock chain to avoid breaking in-flight writes. " +
+				"Use { forceWriteLock: true } to override.",
+			current.writeLock.pending,
+		);
+		current = { ...s, writeLock: current.writeLock };
+		return;
+	}
+	current = s;
+}
+
+/** Read the current singleton container. */
+export function getSingletons(): RuntimeSingletons {
+	return current;
+}

spawn/renderer.ts

@@ -32,6 +32,10 @@ import type { Theme, ThemeColor } from "@earendil-works/pi-coding-agent";
 import { Container, Spacer, Text, truncateToWidth, visibleWidth } from "@earendil-works/pi-tui";
 import type { TUI } from "@earendil-works/pi-tui";
 import type { AgenticodingState } from "../state.js";
+import {
+	__setSingletons,
+	getSingletons,
+} from "../runtime-singletons.js";
 import {
 	getLastAssistantText,
 	type SpawnOutcome,
@@ -249,7 +253,7 @@ interface SpawnFrameTarget {
  * streaming events (50-100+/sec) do not trigger an equal number of heavy
  * component mutations.
  */
-class SpawnFrameScheduler {
+export class SpawnFrameScheduler {
 	private readonly frameMs: number;
 	private dirtyComponents = new Set<SpawnFrameTarget>();
 	private frameTimer: ReturnType<typeof setTimeout> | null = null;
@@ -316,8 +320,20 @@ class SpawnFrameScheduler {
 	}
 }
 
-/** Module-level singleton shared by all NestedAgentSessionComponent instances. */
+/**
+ * Module-level singleton shared by all NestedAgentSessionComponent instances.
+ *
+ * Registered into the RuntimeSingletons container at module evaluation time.
+ * Test harnesses overwrite this with a fresh SpawnFrameScheduler via
+ * createTestHarness().  ESM guarantees all static imports resolve before any
+ * module body runs, so the harness always wins.
+ *
+ * IMPORTANT: never use dynamic import() to load this module *after* a
+ * createTestHarness() call, or the production scheduler will overwrite the
+ * test one.
+ */
 const spawnFrameScheduler = new SpawnFrameScheduler();
+__setSingletons({ ...getSingletons(), frameScheduler: spawnFrameScheduler });
 
 // ── NestedAgentSessionComponent ───────────────────────────────────────
 
@@ -396,7 +412,7 @@ class NestedAgentSessionComponent extends Container implements SpawnFrameTarget
 		this.renderQueued = false;
 		this.queuedRenderToken = undefined;
 		this.renderScheduleToken++;
-		spawnFrameScheduler.cancelDirty(this);
+		getSingletons().frameScheduler.cancelDirty(this);
 	}
 
 	/**
@@ -409,7 +425,7 @@ class NestedAgentSessionComponent extends Container implements SpawnFrameTarget
 		if (this.renderQueued) return;
 		this.renderQueued = true;
 		this.queuedRenderToken = ++this.renderScheduleToken;
-		spawnFrameScheduler.markDirty(this);
+		getSingletons().frameScheduler.markDirty(this);
 	}
 
 	/**
@@ -555,7 +571,7 @@ class NestedAgentSessionComponent extends Container implements SpawnFrameTarget
 	dispose(): void {
 		this.unsubscribe?.();
 		this.unsubscribe = undefined;
-		spawnFrameScheduler.cancelDirty(this);
+		getSingletons().frameScheduler.cancelDirty(this);
 		this.clearPendingState();
 		// Snapshot fields before clearing: if session.abort() triggers re-entrant
 		// dispose, the nulled-out fields prevent double-abort.
@@ -731,7 +747,7 @@ class NestedAgentSessionComponent extends Container implements SpawnFrameTarget
 		if (!this.session) return;
 
 		// Flush any pending state first so accumulated updates don't double-apply
-		spawnFrameScheduler.cancelDirty(this);
+		getSingletons().frameScheduler.cancelDirty(this);
 		this.clearPendingState();
 
 		this.clear();
@@ -1209,16 +1225,20 @@ export { NestedAgentSessionComponent, renderSpawnCall, renderSpawnResult };
  * Synchronously flush all pending spawn frame work.
  * Exported for tests.  Not needed in production — the frame timer handles
  * everything automatically.
+ *
+ * Delegate through getSingletons() so that test harness swaps are respected.
  */
 export function flushSpawnFrameScheduler(): void {
-	spawnFrameScheduler.flushNow();
+	getSingletons().frameScheduler.flushNow();
 }
 
 /**
  * Reset the frame scheduler, discarding any pending dirty markers.
  * Exported for tests.  In production the scheduler lifecycle is tied to
  * component dispose(), so this is never needed.
+ *
+ * Delegate through getSingletons() so that test harness swaps are respected.
  */
 export function resetSpawnFrameScheduler(): void {
-	spawnFrameScheduler.clear();
+	getSingletons().frameScheduler.clear();
 }