@agwab/pi-workflow 0.3.0 → 0.4.0

package/src/roles.ts

@@ -1,4 +1,5 @@
-import { AgentDefinition, CompiledRole, RoleSpec } from "./types.js";
+import { compactStrings } from "./strings.js";
+import type { AgentDefinition, CompiledRole, RoleSpec } from "./types.js";
 
 export const DEFAULT_SAFE_SECTIONS = [
   "Core Principles",
@@ -24,14 +25,10 @@ export function compileRole(name: string, spec: RoleSpec, sourceAgent?: AgentDef
   const maxChars = spec.maxChars ?? DEFAULT_MAX_ROLE_CHARS;
   const includeSections = spec.includeSections ?? [...DEFAULT_SAFE_SECTIONS];
   const excludedSections = [...ALWAYS_EXCLUDED_SECTIONS, ...(spec.excludeSections ?? [])];
-  const parts: string[] = [];
-
-  if (sourceAgent) {
-    const extracted = extractMarkdownSections(sourceAgent.body, includeSections, excludedSections);
-    if (extracted.trim() !== "") parts.push(extracted.trim());
-  }
-
-  if (spec.prompt?.trim()) parts.push(spec.prompt.trim());
+  const parts = compactStrings([
+    sourceAgent ? extractMarkdownSections(sourceAgent.body, includeSections, excludedSections) : undefined,
+    spec.prompt,
+  ], { unique: false });
 
   const fullContent = parts.join("\n\n");
   const truncated = fullContent.length > maxChars;

package/src/store.ts

@@ -55,7 +55,25 @@ const runLeaseContext = new AsyncLocalStorage<{
 	cwd: string;
 	runId: string;
 	ownerId: string;
+	abortSignal: AbortSignal;
 }>();
+type RunLeaseTestHooks = {
+	heartbeatIntervalMs?: number;
+	onAfterReclaimRename?: (context: {
+		lockFile: string;
+		reclaimFile: string;
+	}) => void | Promise<void>;
+	onBeforeRestoreReclaimFile?: (context: {
+		lockFile: string;
+		reclaimFile: string;
+	}) => void | Promise<void>;
+	onBeforeReleaseLockRename?: (context: {
+		lockFile: string;
+		releaseFile: string;
+		ownerId: string;
+	}) => void | Promise<void>;
+};
+let runLeaseTestHooks: RunLeaseTestHooks = {};
 const TASK_STATUSES: Array<keyof Omit<TaskSummary, "total">> = [
 	"pending",
 	"running",
@@ -148,10 +166,14 @@ export async function writeJsonAtomic(
 	await rename(temp, file);
 }
 
+export function setRunLeaseTestHooksForTests(hooks?: RunLeaseTestHooks): void {
+	runLeaseTestHooks = hooks ?? {};
+}
+
 export async function withRunLease<T>(
 	cwd: string,
 	runId: string,
-	action: () => Promise<T>,
+	action: (abortSignal: AbortSignal) => Promise<T>,
 ): Promise<T | undefined> {
 	const dir = workflowRunDir(cwd, runId);
 	await ensureDir(dir);
@@ -160,8 +182,14 @@ export async function withRunLease<T>(
 	const lock = await acquireLock(lockFile, ownerId);
 	if (!lock) return undefined;
 
+	const abortController = new AbortController();
+	const abortLease = (error: unknown): void => {
+		if (abortController.signal.aborted) return;
+		abortController.abort(asLeaseError(error));
+	};
 	const supervisorFile = join(dir, "supervisor.json");
 	const heartbeat = async (): Promise<void> => {
+		assertLeaseNotAborted(abortController.signal);
 		await assertLockOwner(lockFile, ownerId);
 		const timestamp = nowIso();
 		const now = new Date();
@@ -176,22 +204,51 @@ export async function withRunLease<T>(
 	};
 
 	await heartbeat();
-	const heartbeatTimer = setInterval(
-		() => {
-			void heartbeat().catch(() => undefined);
-		},
-		Math.max(1000, Math.floor(LEASE_STALE_MS / 3)),
-	);
+	const heartbeatTimer = setInterval(() => {
+		void heartbeat().catch(abortLease);
+	}, runLeaseHeartbeatIntervalMs());
 	heartbeatTimer.unref?.();
 
 	try {
-		return await runLeaseContext.run({ cwd, runId, ownerId }, action);
+		const result = await runLeaseContext.run(
+			{ cwd, runId, ownerId, abortSignal: abortController.signal },
+			() => action(abortController.signal),
+		);
+		assertLeaseNotAborted(abortController.signal);
+		return result;
 	} finally {
 		clearInterval(heartbeatTimer);
 		await releaseLock(lockFile, ownerId);
 	}
 }
 
+function runLeaseHeartbeatIntervalMs(): number {
+	return Math.max(
+		1,
+		Math.floor(
+			runLeaseTestHooks.heartbeatIntervalMs ??
+				Math.max(1000, Math.floor(LEASE_STALE_MS / 3)),
+		),
+	);
+}
+
+function assertLeaseNotAborted(signal: AbortSignal): void {
+	if (signal.aborted) throw abortSignalError(signal);
+}
+
+function abortSignalError(signal: AbortSignal): Error {
+	return asLeaseError((signal as AbortSignal & { reason?: unknown }).reason);
+}
+
+function asLeaseError(error: unknown): Error {
+	if (error instanceof Error) return error;
+	return new Error(
+		error === undefined
+			? "Lost supervisor lease"
+			: `Lost supervisor lease: ${String(error)}`,
+	);
+}
+
 async function acquireLock(
 	lockFile: string,
 	ownerId: string,
@@ -231,6 +288,7 @@ async function reclaimStaleLock(lockFile: string): Promise<boolean> {
 		if ((error as NodeJS.ErrnoException).code === "ENOENT") return true;
 		return false;
 	}
+	await runLeaseTestHooks.onAfterReclaimRename?.({ lockFile, reclaimFile });
 
 	const claimed = await readLockSnapshot(reclaimFile);
 	if (!claimed) return true;
@@ -251,13 +309,22 @@ async function restoreReclaimFile(
 	reclaimFile: string,
 	lockFile: string,
 ): Promise<void> {
+	await runLeaseTestHooks.onBeforeRestoreReclaimFile?.({
+		lockFile,
+		reclaimFile,
+	});
 	try {
 		await link(reclaimFile, lockFile);
 	} catch (error) {
-		if ((error as NodeJS.ErrnoException).code !== "EEXIST") throw error;
-	} finally {
-		await unlink(reclaimFile).catch(() => undefined);
+		if ((error as NodeJS.ErrnoException).code === "EEXIST") {
+			throw new Error(
+				`Could not restore reclaimed lock because another owner acquired ${lockFile}`,
+				{ cause: error },
+			);
+		}
+		throw error;
 	}
+	await unlink(reclaimFile).catch(() => undefined);
 }
 
 function isReclaimableLockSnapshot(snapshot: LockSnapshot): boolean {
@@ -339,8 +406,27 @@ async function acquireLockWithWait(
 }
 
 async function releaseLock(lockFile: string, ownerId: string): Promise<void> {
-	if (await ownsLock(lockFile, ownerId))
-		await unlink(lockFile).catch(() => undefined);
+	const snapshot = await readLockSnapshot(lockFile);
+	if (!snapshot || snapshot.ownerId !== ownerId) return;
+	const releaseFile = `${lockFile}.release-${process.pid}-${randomBytes(3).toString("hex")}`;
+	await runLeaseTestHooks.onBeforeReleaseLockRename?.({
+		lockFile,
+		releaseFile,
+		ownerId,
+	});
+	try {
+		await rename(lockFile, releaseFile);
+	} catch (error) {
+		if ((error as NodeJS.ErrnoException).code === "ENOENT") return;
+		throw error;
+	}
+	const claimed = await readLockSnapshot(releaseFile);
+	if (!claimed) return;
+	if (sameLockOwnerSnapshot(snapshot, claimed)) {
+		await unlink(releaseFile).catch(() => undefined);
+		return;
+	}
+	await restoreReclaimFile(releaseFile, lockFile);
 }
 
 async function assertLockOwner(
@@ -1095,6 +1181,7 @@ async function assertActiveRunLease(cwd: string, runId: string): Promise<void> {
 	const context = runLeaseContext.getStore();
 	if (!context) return;
 	if (context.cwd !== cwd || context.runId !== runId) return;
+	assertLeaseNotAborted(context.abortSignal);
 	await assertLockOwner(
 		join(workflowRunDir(cwd, runId), "supervisor.lock"),
 		context.ownerId,
@@ -1243,6 +1330,7 @@ async function updateIndexIncremental(
 	const changedEntry = buildIndexEntry(cwd, changedRun);
 	const entries = existing.runs
 		.filter((entry) => entry.runId !== changedRun.runId)
+		.map(stripIndexTaskRows)
 		.concat(changedEntry);
 	return {
 		schemaVersion: 1,
@@ -1290,6 +1378,11 @@ function selectIndexEntries(
 	);
 }
 
+function stripIndexTaskRows(entry: WorkflowIndexRunEntry): WorkflowIndexRunEntry {
+	const { tasks: _tasks, ...slim } = entry;
+	return slim;
+}
+
 function buildIndexEntry(
 	cwd: string,
 	run: WorkflowRunRecord,
@@ -1308,17 +1401,6 @@ function buildIndexEntry(
 		round: run.round,
 		fanout: run.fanout,
 		runJson: toProjectPath(cwd, workflowRunPath(cwd, run.runId)),
-		tasks: run.tasks.map((task) => ({
-			taskId: task.taskId,
-			displayName: task.displayName,
-			agent: task.agent,
-			kind: task.kind,
-			stageId: task.stageId,
-			backendHandle: task.backendHandle,
-			status: task.status,
-			statusDetail: task.statusDetail,
-			lastMessage: task.lastMessage,
-		})),
 	};
 }
 
@@ -1328,15 +1410,16 @@ function isIndexRecordLike(
 	return (
 		value?.schemaVersion === 1 &&
 		Array.isArray(value.runs) &&
-		value.runs.every(
-			(entry) =>
-				entry &&
-				typeof entry === "object" &&
+		value.runs.every((entry) => {
+			if (!entry || typeof entry !== "object") return false;
+			const tasks = (entry as { tasks?: unknown }).tasks;
+			return (
 				typeof entry.runId === "string" &&
 				typeof entry.updatedAt === "string" &&
 				typeof entry.status === "string" &&
-				Array.isArray(entry.tasks),
-		)
+				(tasks === undefined || Array.isArray(tasks))
+			);
+		})
 	);
 }
 
@@ -1411,13 +1494,19 @@ const RESUMABLE_BLOCKED_STATUS_DETAILS = new Set([
 	"dynamic_approval_timeout",
 ]);
 
+export function isBlockedTaskResumableForResume(
+	task: Pick<WorkflowTaskRunRecord, "status" | "statusDetail">,
+): boolean {
+	return (
+		task.status === "blocked" &&
+		RESUMABLE_BLOCKED_STATUS_DETAILS.has(task.statusDetail)
+	);
+}
+
 export function resetTaskForResume(task: WorkflowTaskRunRecord): boolean {
 	if (
 		!RESUMABLE_TASK_STATUSES.has(task.status) &&
-		!(
-			task.status === "blocked" &&
-			RESUMABLE_BLOCKED_STATUS_DETAILS.has(task.statusDetail)
-		)
+		!isBlockedTaskResumableForResume(task)
 	) {
 		return false;
 	}

package/src/strings.ts

@@ -0,0 +1,38 @@
+export interface CompactStringsOptions {
+	/** Trim returned strings before filtering. Defaults to true. */
+	trim?: boolean;
+	/** Drop duplicate strings after optional trimming. Defaults to true. */
+	unique?: boolean;
+	/** Drop strings whose raw/trimmed form is empty. Defaults to true. */
+	dropEmpty?: boolean;
+	/** Drop strings whose trimmed form is empty even when trim=false. */
+	dropWhitespaceOnly?: boolean;
+}
+
+export function compactStrings(
+	values: readonly unknown[],
+	options: CompactStringsOptions = {},
+): string[] {
+	const trim = options.trim ?? true;
+	const unique = options.unique ?? true;
+	const dropEmpty = options.dropEmpty ?? true;
+	const dropWhitespaceOnly = options.dropWhitespaceOnly ?? trim;
+	const seen = new Set<string>();
+	const result: string[] = [];
+	for (const value of values) {
+		if (typeof value !== "string") continue;
+		const compacted = trim ? value.trim() : value;
+		if (
+			dropEmpty &&
+			(dropWhitespaceOnly ? value.trim().length === 0 : compacted.length === 0)
+		) {
+			continue;
+		}
+		if (unique) {
+			if (seen.has(compacted)) continue;
+			seen.add(compacted);
+		}
+		result.push(compacted);
+	}
+	return result;
+}