@llblab/pi-kit 0.3.2 → 0.5.0

package/node_modules/@llblab/pi-state-flow/lib/extension.ts

@@ -0,0 +1,238 @@
+import type { AgentMessage } from "@earendil-works/pi-agent-core";
+import type { ExtensionAPI, ExtensionContext } from "@earendil-works/pi-coding-agent";
+import { assistantToolCallCount, finalizedAssistantResponse, parseTerminalPatch, stateFlowProtocol, stripStateComments } from "./terminal.ts";
+import { currentRunTrajectory, runtimeContextMessage, VALIDATION_MESSAGE_TYPE, withoutPrivateValidation } from "./context.ts";
+import { SkillReadTracker } from "./skills.ts";
+import { migrationFailure, type Snapshot } from "./snapshot.ts";
+import { emptyState } from "./state.ts";
+import { commitTransition, stageTransition, type StagedTransition } from "./transition.ts";
+import { MAX_VALIDATION_RETRIES, nextValidation } from "./validation.ts";
+import { discoverSnapshotData, hasPriorConversation, SNAPSHOT_ENTRY_TYPE } from "./session.ts";
+import { compactStatus, detailedStatus, STATUS_KEY } from "./status.ts";
+import { abandonValidation, prepareRun, startEpisode, stopEpisode } from "./episode.ts";
+import { recoverSnapshot } from "./recovery.ts";
+
+export default function stateFlowExtension(pi: ExtensionAPI): void {
+	let snapshot: Snapshot = { enabled: false, state: emptyState(), step: 0 };
+	let stagedFinal: StagedTransition | undefined;
+	let retryQueued = false;
+	let runAnchorTimestamp: number | undefined;
+	const skillReads = new SkillReadTracker();
+
+	function persist(): void {
+		pi.appendEntry(SNAPSHOT_ENTRY_TYPE, structuredClone(snapshot));
+	}
+
+	function updateUi(ctx: ExtensionContext): void {
+		ctx.ui.setStatus(STATUS_KEY, compactStatus(snapshot, (color, text) => ctx.ui.theme.fg(color, text)));
+	}
+
+	function clearRunTransient(): void {
+		stagedFinal = undefined;
+		retryQueued = false;
+		runAnchorTimestamp = undefined;
+		skillReads.clear();
+	}
+
+	function restoreActiveBranch(ctx: ExtensionContext): void {
+		clearRunTransient();
+		try {
+			const discovery = discoverSnapshotData(ctx.sessionManager.getBranch());
+			const recovery = recoverSnapshot(discovery.candidates);
+			const skipped = discovery.errors.length + recovery.skipped.length;
+			snapshot = discovery.candidates.length === 0 && discovery.errors.length > 0
+				? migrationFailure({}, `Snapshot restoration failed: ${discovery.errors[0]}`)
+				: recovery.snapshot;
+			if (snapshot.enabled && skipped > 0) {
+				ctx.ui.notify(`State Flow recovered the previous valid snapshot after skipping ${skipped} malformed newer snapshot(s).`, "warning");
+			}
+		} catch (error) {
+			snapshot = migrationFailure(
+				{},
+				`Snapshot restoration failed: ${error instanceof Error ? error.message : String(error)}`,
+			);
+		}
+		if (!snapshot.enabled && snapshot.validation?.attempt === 0) {
+			ctx.ui.notify(`State Flow restored disabled: ${snapshot.validation.error}`, "error");
+		}
+		updateUi(ctx);
+	}
+
+	function abandonRun(ctx: ExtensionContext): void {
+		const changed = abandonValidation(snapshot);
+		clearRunTransient();
+		if (changed) persist();
+		updateUi(ctx);
+	}
+
+	function queueTerminalRegeneration(error: string, ctx: ExtensionContext): void {
+		const decision = nextValidation(snapshot.validation, error);
+		if (decision.kind === "retry") {
+			snapshot.validation = decision.feedback;
+			persist();
+			retryQueued = true;
+			pi.sendMessage({
+				customType: VALIDATION_MESSAGE_TYPE,
+				content: `State Flow rejected the terminal response (attempt ${decision.feedback.attempt}/${MAX_VALIDATION_RETRIES}). ${decision.feedback.instruction}`,
+				display: false,
+			}, { deliverAs: "steer", triggerTurn: true });
+			return;
+		}
+		retryQueued = false;
+		snapshot.validation = undefined;
+		skillReads.clear();
+		persist();
+		ctx.ui.notify(`State Flow remains enabled after ${MAX_VALIDATION_RETRIES} automatic regeneration attempts; the last committed state was preserved: ${decision.error}`, "error");
+	}
+
+	function rejectTerminal(message: { role: "assistant"; content?: unknown }, error: string, ctx: ExtensionContext) {
+		queueTerminalRegeneration(error, ctx);
+		return {
+			message: {
+				...message,
+				role: "assistant" as const,
+				content: [],
+			},
+		};
+	}
+
+	pi.registerCommand("state-flow-start", {
+		description: "Start State Flow mode",
+		handler: async (_args, ctx) => {
+			snapshot = startEpisode(hasPriorConversation(ctx.sessionManager.getBranch()));
+			stagedFinal = undefined;
+			retryQueued = false;
+			runAnchorTimestamp = undefined;
+			skillReads.clear();
+			persist();
+			updateUi(ctx);
+			ctx.ui.notify(
+				snapshot.bootstrap
+					? "State Flow enabled. The next complete agent run will migrate active context into state."
+					: "State Flow enabled. The next prompt starts a fresh stateful agent run.",
+				"info",
+			);
+		},
+	});
+
+	pi.registerCommand("state-flow-status", {
+		description: "Show State Flow runtime status",
+		handler: async (_args, ctx) => {
+			ctx.ui.notify(detailedStatus(snapshot), "info");
+		},
+	});
+
+	pi.registerCommand("state-flow-stop", {
+		description: "Stop State Flow and clear its complete episode",
+		handler: async (_args, ctx) => {
+			snapshot = stopEpisode();
+			stagedFinal = undefined;
+			retryQueued = false;
+			runAnchorTimestamp = undefined;
+			skillReads.clear();
+			persist();
+			updateUi(ctx);
+		},
+	});
+
+	pi.on("before_agent_start", (event) => {
+		if (!snapshot.enabled) return;
+		if (!retryQueued) {
+			skillReads.clear();
+		}
+		const rotatesRun = snapshot.specification !== undefined && !retryQueued;
+		if (prepareRun(snapshot, event.prompt, retryQueued)) {
+			if (rotatesRun) runAnchorTimestamp = undefined;
+			persist();
+		}
+		return {
+			systemPrompt: `${event.systemPrompt}\n\n${stateFlowProtocol(snapshot.bootstrap === true)}`,
+		};
+	});
+
+	pi.on("context", (event) => {
+		if (!snapshot.enabled || snapshot.specification === undefined) return;
+		if (snapshot.bootstrap) {
+			if (retryQueued) return;
+			const messages = withoutPrivateValidation(event.messages as AgentMessage[]);
+			return messages.length === event.messages.length ? undefined : { messages };
+		}
+		const trajectory = currentRunTrajectory(
+			event.messages as AgentMessage[],
+			snapshot.specification,
+			runAnchorTimestamp,
+		);
+		runAnchorTimestamp = trajectory.anchorTimestamp;
+		return { messages: [runtimeContextMessage(snapshot), ...trajectory.messages] };
+	});
+
+	pi.on("tool_execution_start", (event) => {
+		if (!snapshot.enabled) return;
+		// Pi emits this before tool_call. Keep the argument object as a fallback;
+		// tool_call replaces it with the mutable, post-preflight input reference.
+		skillReads.recordStart(event.toolCallId, event.toolName, event.args);
+	});
+
+	pi.on("tool_call", (event) => {
+		if (!snapshot.enabled) return;
+		skillReads.recordCall(event.toolCallId, event.toolName, event.input);
+	});
+
+	pi.on("tool_execution_end", (event) => {
+		if (!snapshot.enabled) return;
+		skillReads.recordEnd(event.toolCallId, event.toolName, event.isError);
+	});
+
+	pi.on("message_end", (event, ctx): any => {
+		if (!snapshot.enabled || event.message.role !== "assistant") return;
+		stagedFinal = undefined;
+		const message = event.message as unknown as { role: "assistant"; stopReason?: string; content?: unknown };
+		if (message.stopReason === "aborted") {
+			abandonRun(ctx);
+			return;
+		}
+		if (message.stopReason === "length" || message.stopReason === "error") {
+			return rejectTerminal(message, `Assistant response ended with ${message.stopReason}`, ctx);
+		}
+		if (assistantToolCallCount(message.content) > 0 || message.stopReason === "toolUse") {
+			retryQueued = snapshot.validation !== undefined;
+			const cleaned = stripStateComments(message.content);
+			return cleaned.changed
+				? { message: { ...message, role: "assistant" as const, content: cleaned.content } }
+				: undefined;
+		}
+		try {
+			const parsed = parseTerminalPatch(message.content);
+			stagedFinal = stageTransition(snapshot.state, parsed.patch, skillReads.successful);
+			retryQueued = false;
+			return { message: { ...message, role: "assistant" as const, content: parsed.responseContent } };
+		} catch (error) {
+			return rejectTerminal(message, error instanceof Error ? error.message : String(error), ctx);
+		}
+	});
+
+	pi.on("turn_end", (event, ctx) => {
+		if (!snapshot.enabled || !stagedFinal) {
+			updateUi(ctx);
+			return;
+		}
+		try {
+			stagedFinal.nextState.response = finalizedAssistantResponse(event.message);
+			if (commitTransition(snapshot, stagedFinal)) {
+				skillReads.clear();
+				persist();
+			}
+		} catch (error) {
+			queueTerminalRegeneration(error instanceof Error ? error.message : String(error), ctx);
+		}
+		stagedFinal = undefined;
+		updateUi(ctx);
+	});
+
+	pi.on("agent_settled", (_event, ctx) => {
+		if (snapshot.enabled && retryQueued) abandonRun(ctx);
+	});
+
+	pi.on("session_start", (_event, ctx) => restoreActiveBranch(ctx));
+	pi.on("session_tree", (_event, ctx) => restoreActiveBranch(ctx));
+}

package/node_modules/@llblab/pi-state-flow/lib/json.ts

@@ -0,0 +1,78 @@
+import { createHash } from "node:crypto";
+
+export type JsonValue = null | boolean | number | string | JsonValue[] | JsonObject;
+export interface JsonObject { [key: string]: JsonValue }
+
+export function applyPatch(state: JsonObject, patch: JsonObject): JsonObject {
+	const next: JsonObject = structuredClone(state);
+	for (const [key, value] of Object.entries(patch)) {
+		if (value === null) {
+			delete next[key];
+			continue;
+		}
+		const current = next[key];
+		const materialized = isObject(current) && isObject(value)
+			? applyPatch(current, value)
+			: structuredClone(value);
+		Object.defineProperty(next, key, {
+			value: materialized,
+			enumerable: true,
+			configurable: true,
+			writable: true,
+		});
+	}
+	return next;
+}
+
+export function isObject(value: JsonValue | unknown): value is JsonObject {
+	return typeof value === "object" && value !== null && !Array.isArray(value);
+}
+
+export function validatePatch(value: unknown): asserts value is JsonObject {
+	if (!isObject(value)) throw new Error("State patch must be a JSON object");
+	if (!isJsonValue(value)) throw new Error("State patch must contain finite, acyclic JSON data");
+}
+
+export function canonicalJson(value: JsonValue | unknown): string {
+	if (!isJsonValue(value)) throw new Error("Value must be finite, acyclic JSON data");
+	return JSON.stringify(orderValue(value));
+}
+
+export function hashJson(value: JsonValue | unknown): string {
+	return createHash("sha256").update(canonicalJson(value)).digest("hex");
+}
+
+function orderValue(value: unknown): unknown {
+	if (Array.isArray(value)) return value.map((item) => orderValue(item));
+	if (!isObject(value)) return value;
+	return Object.fromEntries(
+		Object.keys(value).sort().map((key) => [key, orderValue(value[key])]),
+	);
+}
+
+export function containsNull(value: unknown): boolean {
+	if (value === null) return true;
+	if (Array.isArray(value)) return value.some((item) => containsNull(item));
+	if (!isObject(value)) return false;
+	return Object.values(value).some((item) => containsNull(item));
+}
+
+export function isJsonValue(value: unknown, ancestors = new WeakSet<object>()): value is JsonValue {
+	if (value === null || typeof value === "string" || typeof value === "boolean") return true;
+	if (typeof value === "number") return Number.isFinite(value);
+	if (typeof value !== "object") return false;
+	if (ancestors.has(value)) return false;
+	ancestors.add(value);
+	try {
+		if (Array.isArray(value)) return value.every((item) => isJsonValue(item, ancestors));
+		if (!isObject(value)) return false;
+		const prototype = Object.getPrototypeOf(value);
+		if (prototype !== Object.prototype && prototype !== null) return false;
+		if (Reflect.ownKeys(value).some((key) => typeof key === "symbol")) return false;
+		return Object.values(value).every((item) => isJsonValue(item, ancestors));
+	} catch {
+		return false;
+	} finally {
+		ancestors.delete(value);
+	}
+}

package/node_modules/@llblab/pi-state-flow/lib/recovery.ts

@@ -0,0 +1,38 @@
+import { emptyState } from "./state.ts";
+import { migrateSnapshot, migrationFailure, type Snapshot } from "./snapshot.ts";
+
+export interface SnapshotRecovery {
+	snapshot: Snapshot;
+	skipped: string[];
+}
+
+function failureMessage(snapshot: Snapshot): string | undefined {
+	return !snapshot.enabled && snapshot.validation?.attempt === 0
+		? snapshot.validation.error
+		: undefined;
+}
+
+/** Recover the newest valid snapshot, falling back through the active branch. */
+export function recoverSnapshot(candidates: readonly unknown[]): SnapshotRecovery {
+	const skipped: string[] = [];
+	let newestFailure: Snapshot | undefined;
+	for (const candidate of candidates) {
+		let migrated: Snapshot;
+		try {
+			migrated = migrateSnapshot(candidate);
+		} catch (error) {
+			migrated = migrationFailure(
+				{},
+				`Snapshot restoration failed: ${error instanceof Error ? error.message : String(error)}`,
+			);
+		}
+		const failure = failureMessage(migrated);
+		if (failure === undefined) return { snapshot: migrated, skipped };
+		newestFailure ??= migrated;
+		skipped.push(failure);
+	}
+	return {
+		snapshot: newestFailure ?? { enabled: false, state: emptyState(), step: 0 },
+		skipped,
+	};
+}

package/node_modules/@llblab/pi-state-flow/lib/session.ts

@@ -0,0 +1,49 @@
+export const SNAPSHOT_ENTRY_TYPE = "state-flow-snapshot";
+
+interface BranchEntry {
+	type?: unknown;
+	customType?: unknown;
+	data?: unknown;
+	message?: { role?: unknown };
+}
+
+export interface SnapshotDiscovery {
+	candidates: unknown[];
+	errors: string[];
+}
+
+/** Enumerate active-branch snapshots newest-first while containing hostile entries. */
+export function discoverSnapshotData(branch: readonly BranchEntry[]): SnapshotDiscovery {
+	const candidates: unknown[] = [];
+	const errors: string[] = [];
+	for (let index = branch.length - 1; index >= 0; index--) {
+		try {
+			const entry = branch[index];
+			if (entry?.type === "custom" && entry.customType === SNAPSHOT_ENTRY_TYPE) candidates.push(entry.data);
+		} catch (error) {
+			errors.push(error instanceof Error ? error.message : String(error));
+		}
+	}
+	return { candidates, errors };
+}
+
+export function snapshotDataNewestFirst(branch: readonly BranchEntry[]): unknown[] {
+	return discoverSnapshotData(branch).candidates;
+}
+
+export function latestSnapshotData(branch: readonly BranchEntry[]): unknown {
+	return snapshotDataNewestFirst(branch)[0];
+}
+
+export function hasPriorConversation(branch: readonly BranchEntry[]): boolean {
+	for (const entry of branch) {
+		try {
+			if (entry.type !== "message") continue;
+			const role = entry.message?.role;
+			if (role === "user" || role === "assistant" || role === "toolResult") return true;
+		} catch {
+			// A hostile unrelated entry must not prevent explicit episode startup.
+		}
+	}
+	return false;
+}

package/node_modules/@llblab/pi-state-flow/lib/skills.ts

@@ -0,0 +1,58 @@
+import { isObject, type JsonObject } from "./json.ts";
+
+function hasContent(value: unknown): boolean {
+	if (typeof value === "string") return value.trim().length > 0;
+	if (Array.isArray(value)) return value.length > 0;
+	if (isObject(value)) return Object.keys(value).length > 0;
+	return value !== undefined && value !== null;
+}
+
+export function hasCompiledSkill(contract: JsonObject, source: string): boolean {
+	if (!isObject(contract.compiled_skills)) return false;
+	return hasContent(contract.compiled_skills[source]);
+}
+
+export function skillPathFromRead(toolName: unknown, args: unknown): string | undefined {
+	if (toolName !== "read" || !isObject(args) || typeof args.path !== "string") return undefined;
+	return /(^|[\\/])SKILL\.md$/.test(args.path) ? args.path : undefined;
+}
+
+interface PendingRead {
+	toolName: string;
+	args: unknown;
+}
+
+/** Correlates Pi's mutable tool lifecycle without exposing runtime event objects. */
+export class SkillReadTracker {
+	readonly successful = new Set<string>();
+	readonly #pending = new Map<string, PendingRead>();
+
+	clear(): void {
+		this.successful.clear();
+		this.#pending.clear();
+	}
+
+	recordStart(toolCallId: string, toolName: string, args: unknown): void {
+		this.#record(toolCallId, toolName, args);
+	}
+
+	recordCall(toolCallId: string, toolName: string, input: unknown): void {
+		this.#record(toolCallId, toolName, input);
+	}
+
+	recordEnd(toolCallId: string, toolName: string, isError: boolean): void {
+		const pending = this.#pending.get(toolCallId);
+		this.#pending.delete(toolCallId);
+		if (isError || !pending || toolName !== pending.toolName) return;
+		const source = skillPathFromRead(pending.toolName, pending.args);
+		if (source) this.successful.add(source);
+	}
+
+	#record(toolCallId: string, toolName: string, args: unknown): void {
+		if (toolName !== "read") {
+			this.#pending.delete(toolCallId);
+			return;
+		}
+		this.#pending.set(toolCallId, { toolName, args });
+	}
+}

package/node_modules/@llblab/pi-state-flow/lib/snapshot.ts

@@ -0,0 +1,103 @@
+import { applyPatch, containsNull, isJsonValue, isObject, type JsonObject } from "./json.ts";
+import { emptyState, isStateDocument, type StateDocument } from "./state.ts";
+import { MAX_VALIDATION_RETRIES, type ValidationFeedback } from "./validation.ts";
+
+const MAX_RESTORED_STEP = Number.MAX_SAFE_INTEGER - 1;
+
+export interface Snapshot {
+	enabled: boolean;
+	specification?: string;
+	state: StateDocument;
+	step: number;
+	validation?: ValidationFeedback;
+	bootstrap?: boolean;
+}
+
+function isLegacyTwoPartState(value: unknown): value is { contract: JsonObject; working: JsonObject } {
+	return isObject(value)
+		&& isObject(value.contract)
+		&& isObject(value.working)
+		&& Object.keys(value).every((key) => key === "contract" || key === "working");
+}
+
+function restoredStep(value: unknown): number {
+	return typeof value === "number"
+		&& Number.isSafeInteger(value)
+		&& value >= 0
+		&& value <= MAX_RESTORED_STEP
+		? value
+		: 0;
+}
+
+function restoredValidation(value: unknown): ValidationFeedback | undefined {
+	if (!isObject(value)
+		|| !Number.isSafeInteger(value.attempt as number)
+		|| (value.attempt as number) < 0
+		|| (value.attempt as number) > MAX_VALIDATION_RETRIES
+		|| typeof value.error !== "string"
+		|| typeof value.instruction !== "string") return undefined;
+	return {
+		attempt: value.attempt as number,
+		error: value.error,
+		instruction: value.instruction,
+	};
+}
+
+export function migrationFailure(data: JsonObject, error: string): Snapshot {
+	return {
+		enabled: false,
+		specification: typeof data.specification === "string" ? data.specification : undefined,
+		state: emptyState(),
+		step: restoredStep(data.step),
+		validation: {
+			attempt: 0,
+			error,
+			instruction: "Start a fresh State Flow episode; null is reserved for patch deletion.",
+		},
+	};
+}
+
+export function migrateSnapshot(value: unknown): Snapshot {
+	if (!isObject(value)) return { enabled: false, state: emptyState(), step: 0 };
+	const base = {
+		enabled: value.enabled === true,
+		specification: typeof value.specification === "string" ? value.specification : undefined,
+		step: restoredStep(value.step),
+		validation: restoredValidation(value.validation),
+		bootstrap: value.bootstrap === true,
+	};
+	if (isStateDocument(value.state)) {
+		if (!isJsonValue(value.state)) return migrationFailure(value, "Restored state contains non-JSON data");
+		if (containsNull(value.state)) return migrationFailure(value, "Restored state contains null data");
+		return { ...base, state: structuredClone(value.state) };
+	}
+	if (isLegacyTwoPartState(value.state)) {
+		if (!isJsonValue(value.state)) return migrationFailure(value, "Restored state contains non-JSON data");
+		if (containsNull(value.state)) return migrationFailure(value, "Restored state contains null data");
+		return {
+			...base,
+			state: {
+				contract: structuredClone(value.state.contract),
+				working: structuredClone(value.state.working),
+				response: "",
+			},
+		};
+	}
+	if (isObject(value.state)
+		&& (Object.hasOwn(value.state, "response")
+			|| (Object.hasOwn(value.state, "contract") && Object.hasOwn(value.state, "working")))) {
+		return migrationFailure(value, "Restored state has an invalid materialized-state schema");
+	}
+	const legacyBasis = isObject(value.stateBasis)
+		? value.stateBasis
+		: isObject(value.state)
+			? value.state
+			: {};
+	const legacyPatch = isObject(value.previousStatePatch) ? value.previousStatePatch : undefined;
+	if (!isJsonValue(legacyBasis) || (legacyPatch !== undefined && !isJsonValue(legacyPatch))) {
+		return migrationFailure(value, "Legacy state contains non-JSON data");
+	}
+	const legacyState = legacyPatch === undefined ? legacyBasis : applyPatch(legacyBasis, legacyPatch);
+	if (containsNull(legacyState)) return migrationFailure(value, "Legacy state contains null data");
+	return { ...base, state: { contract: {}, working: structuredClone(legacyState), response: "" } };
+}

package/node_modules/@llblab/pi-state-flow/lib/state.ts

@@ -0,0 +1,20 @@
+import { isObject, type JsonObject } from "./json.ts";
+
+/** The only materialized memory shape persisted between State Flow runs. */
+export interface StateDocument extends JsonObject {
+	contract: JsonObject;
+	working: JsonObject;
+	response: string;
+}
+
+export function emptyState(): StateDocument {
+	return { contract: {}, working: {}, response: "" };
+}
+
+export function isStateDocument(value: unknown): value is StateDocument {
+	return isObject(value)
+		&& isObject(value.contract)
+		&& isObject(value.working)
+		&& typeof value.response === "string"
+		&& Object.keys(value).every((key) => key === "contract" || key === "working" || key === "response");
+}

package/node_modules/@llblab/pi-state-flow/lib/status.ts

@@ -0,0 +1,16 @@
+import type { Snapshot } from "./snapshot.ts";
+
+export const STATUS_KEY = "state-flow";
+
+export type Colorize = (color: "accent" | "dim", text: string) => string;
+
+export function compactStatus(snapshot: Snapshot, colorize: Colorize): string | undefined {
+	if (!snapshot.enabled) return undefined;
+	return `${colorize("accent", "state-flow")} ${colorize("dim", `#${snapshot.step}`)}`;
+}
+
+export function detailedStatus(snapshot: Snapshot): string {
+	const stateJson = JSON.stringify(snapshot.state, null, 2);
+	const stateBytes = Buffer.byteLength(stateJson, "utf8");
+	return `State Flow ${snapshot.enabled ? "enabled" : "disabled"}; iteration #${snapshot.step}; state ${stateBytes} bytes; validation attempts ${snapshot.validation?.attempt ?? 0}.\n\n${stateJson}`;
+}