@davideasden/pi-undo 0.2.24 → 0.2.25

package/src/pi-runtime.ts

@@ -1,4 +1,5 @@
 import { join, resolve } from "node:path";
+import { writeJsonAtomic } from "./atomic-fs.ts";
 
 import type {
 	ExtensionAPI,
@@ -28,11 +29,17 @@ import { DurableCursorWriter, SessionState, type SessionEntrySource } from "./se
 import { SnapshotStore } from "./snapshot-store.ts";
 import { StatusReporter } from "./status-reporter.ts";
 import { WorkspaceLock } from "./workspace-lock.ts";
+import { allCompleted, checkOperation, configuredTimeout, currentOperationContext, isUnconfirmedExit, operationHasUnconfirmedExit, runWithOperationContext, withRecoveryBudget, type ProcessDiagnostic } from "./operation-context.ts";
 
 type ReadonlySessionManager = ExtensionContext["sessionManager"];
 
-export async function createPiUndoRuntime(context: ExtensionContext, pi: ExtensionAPI) {
+export function createPiUndoRuntime(context: ExtensionContext, pi: ExtensionAPI) {
+	return withRecoveryBudget(() => initializePiUndoRuntime(context, pi));
+}
+
+async function initializePiUndoRuntime(context: ExtensionContext, pi: ExtensionAPI) {
 	const manager = context.sessionManager;
+	const reporter = new StatusReporter(context);
 	const sessionState = sessionStateFor(manager);
 	const sessionIdentity = await sessionState.getSessionIdentity() ?? volatileSessionIdentity(manager, context.cwd);
 	const privateRoot = join(manager.getSessionDir(), ".pi-undo");
@@ -66,11 +73,15 @@ export async function createPiUndoRuntime(context: ExtensionContext, pi: Extensi
 		loadPending: () => journal.loadPending(),
 		assessForeignTransaction: (pending) => journal.isInertForeignPrepared(pending),
 		assessCompensatedTransaction: (pending) => journal.isFullyCompensated(pending),
-		inspectCursor: (pending) => inspectCursorMarkers(pending.descriptor.sessionIdentity.path, pending.descriptor),
+		inspectCursor: (pending) => inspectCursorMarkers(pending.descriptor.sessionIdentity.path, pending.descriptor,
+			pending.descriptor.action === "tree" && pending.state.observedLogicalLeaf !== undefined
+				? pending.state.observedLogicalLeaf : pending.descriptor.toLogicalLeaf),
 		finalizeCursor: (pending, inspection) => finalizeCursorMarker(
 			pending.descriptor.sessionIdentity.path,
 			pending.descriptor,
 			inspection,
+			pending.descriptor.action === "tree" && pending.state.observedLogicalLeaf !== undefined
+				? pending.state.observedLogicalLeaf : pending.descriptor.toLogicalLeaf,
 		),
 		recoverMutations: async (pending, decision) => {
 			const mutationJournal = journal.mutationJournal(pending.descriptor.opId);
@@ -141,7 +152,7 @@ export async function createPiUndoRuntime(context: ExtensionContext, pi: Extensi
 		if (finalizationFailure !== undefined) throw finalizationFailure;
 	};
 	const scheduleFinalization = (opId: string): void => {
-		finalizationQueue = finalizationQueue.then(async () => {
+		finalizationQueue = runWithOperationContext(undefined, () => finalizationQueue.then(() => withRecoveryBudget(async () => {
 			const lease = await workspaceLock.acquire(initialTopology.workspaceIdentity);
 			try {
 				const mutationJournal = journal.mutationJournal(opId);
@@ -161,9 +172,9 @@ export async function createPiUndoRuntime(context: ExtensionContext, pi: Extensi
 				});
 				await journal.markCommitted(opId);
 			} finally {
-				await lease.release();
+				if (!operationHasUnconfirmedExit()) await lease.release();
 			}
-		}).catch((error: unknown) => {
+		}))).catch((error: unknown) => {
 			finalizationFailure = error;
 		});
 	};
@@ -183,6 +194,9 @@ export async function createPiUndoRuntime(context: ExtensionContext, pi: Extensi
 		loadPending: () => journal.loadPending(),
 	};
 
+	const diagnostics: ProcessDiagnostic[] = [];
+	const phases: Array<{ phase: string; elapsedMs: number }> = [];
+	let operationStarted = 0;
 	const dependencies: ControllerDependencies = {
 		workspaceIdentity: initialTopology.workspaceIdentity,
 		sessionIdentity,
@@ -210,10 +224,12 @@ export async function createPiUndoRuntime(context: ExtensionContext, pi: Extensi
 			}
 		},
 		restoreSessionLeaf: async (logicalLeafId) => {
-			if (commandContext === undefined || logicalLeafId === null) return false;
+			if (commandContext === undefined) return false;
+			const targetId = sessionNavigationEntry(manager, logicalLeafId);
+			if (targetId === undefined) return false;
 			internalNavigation = true;
 			try {
-				const result = await commandContext.navigateTree(logicalLeafId, { summarize: false });
+				const result = await commandContext.navigateTree(targetId, { summarize: false });
 				return !result.cancelled && sessionStateFor(manager).getLogicalLeafId() === logicalLeafId;
 			} finally {
 				internalNavigation = false;
@@ -225,7 +241,7 @@ export async function createPiUndoRuntime(context: ExtensionContext, pi: Extensi
 		capture,
 		captureBaseline,
 		captureSafety: async (referenceManifestId, targetManifestId, scopePaths) => {
-			const [reference, target] = await Promise.all([
+			const [reference, target] = await allCompleted([
 				store.loadManifest(referenceManifestId),
 				store.loadManifest(targetManifestId),
 			]);
@@ -253,6 +269,30 @@ export async function createPiUndoRuntime(context: ExtensionContext, pi: Extensi
 		},
 		journal: transactionJournal,
 		clock: Date.now,
+		operationTimeoutMs: configuredTimeout("PI_UNDO_OPERATION_TIMEOUT_MS", 300_000),
+		onOperationStart: (opId) => {
+			operationStarted = performance.now();
+			diagnostics.length = 0;
+			phases.length = 0;
+			reporter.startOperation(opId);
+		},
+		onProgress: (phase) => {
+			reporter.setOperationPhase(phase);
+			if (phases.at(-1)?.phase === phase) return;
+			if (phases.length === 128) phases.shift();
+			phases.push({ phase, elapsedMs: Math.round(performance.now() - operationStarted) });
+		},
+		onProcess: (diagnostic) => {
+			if (diagnostics.length === 128) diagnostics.shift();
+			diagnostics.push(diagnostic);
+		},
+		onOperationEnd: async (opId, result) => {
+			const totalMs = Math.round(performance.now() - operationStarted);
+			if (totalMs < 1_000 && (result.code === "ok" || result.code === "noop")) return;
+			await writeJsonAtomic(join(privateRoot, "diagnostics", `${manager.getSessionId()}-latest.json`), {
+				schemaVersion: 1, opId, code: result.code, totalMs, phases, processes: diagnostics,
+			});
+		},
 	};
 	const startupRecovery = await workspaceLock.withLock(
 		initialTopology.workspaceIdentity,
@@ -266,17 +306,30 @@ export async function createPiUndoRuntime(context: ExtensionContext, pi: Extensi
 	});
 	return {
 		controller,
-		reporter: new StatusReporter(context),
+		reporter,
 		diffSource: store,
 		recovery: startupRecovery.kind === "locked"
 			? { reason: startupRecovery.reason, files: startupRecovery.files, opId: startupRecovery.opId }
 			: undefined,
+		async dispose(): Promise<void> {
+			await controller.dispose();
+			await finalizationQueue;
+			if (isUnconfirmedExit(finalizationFailure)) throw finalizationFailure;
+			reporter.endOperation();
+		},
 		setCommandContext(next: ExtensionCommandContext | undefined): void {
 			commandContext = next;
 		},
 		isInternalNavigation(): boolean {
 			return internalNavigation;
 		},
+		normalizeTreeEvent(event: { newLeafId: string | null; summaryEntry?: { parentId: string | null } }) {
+			return {
+				newLeafId: logicalLeafAt(manager, event.newLeafId),
+				navigationTargetLeafId: logicalLeafAt(manager,
+					event.summaryEntry === undefined ? event.newLeafId : event.summaryEntry.parentId),
+			};
+		},
 	};
 }
 
@@ -401,6 +454,21 @@ function logicalLeafAt(manager: ReadonlySessionManager, leafId: string | null):
 	return sessionStateFor(manager, leafId).getLogicalLeafId();
 }
 
+function sessionNavigationEntry(manager: ReadonlySessionManager, logicalLeafId: string | null): string | undefined {
+	if (logicalLeafId !== null) {
+		const entry = manager.getEntry(logicalLeafId) as unknown;
+		if (isRecord(entry) && entry.type !== "custom_message" &&
+			!(entry.type === "message" && isRecord(entry.message) && entry.message.role === "user")) return logicalLeafId;
+	}
+	// Pi 的公开导航 API 不接受 null；选择 parent 对应逻辑位置的用户消息。
+	for (const entry of manager.getEntries() as unknown[]) {
+		if (!isRecord(entry) || typeof entry.id !== "string" || entry.type !== "message" ||
+			!isRecord(entry.message) || entry.message.role !== "user") continue;
+		if (logicalLeafAt(manager, entryParent(manager, entry.id)) === logicalLeafId) return entry.id;
+	}
+	return undefined;
+}
+
 async function resolveTreeTarget(
 	manager: ReadonlySessionManager,
 	identity: SessionFileIdentity,
@@ -421,12 +489,42 @@ async function resolveTreeTarget(
 			undoStack: sessionStateFor(manager, physicalLeaf).getCheckpoints(identity),
 		};
 	}
-	let checkpoints = [...sessionStateFor(manager, physicalLeaf).getCheckpoints(identity)];
 	const exact = findCheckpointByEndLeaf(manager, identity, logicalLeafId);
-	if (exact !== undefined) checkpoints = checkpointFrontierById(manager, identity, exact.checkpointId);
-	const checkpoint = exact ?? checkpoints.at(-1);
-	if (checkpoint === undefined) throw new Error("tree target 缺少完整 checkpoint");
-	return { logicalLeafId, targetManifestId: checkpoint.afterManifestId, undoStack: checkpoints };
+	if (exact !== undefined) {
+		return {
+			logicalLeafId,
+			targetManifestId: exact.afterManifestId,
+			undoStack: checkpointFrontierById(manager, identity, exact.checkpointId),
+		};
+	}
+	// cursor 可以证明 summary/撤回后逻辑边界对应的文件状态；不能退回任意旧栈顶。
+	const cursor = sessionStateFor(manager, physicalLeaf).getCursor(identity);
+	if (cursor !== null && cursor.toLogicalLeaf === logicalLeafId) {
+		return {
+			logicalLeafId,
+			targetManifestId: cursor.targetManifestId,
+			undoStack: cursor.undoHead === null ? [] : checkpointFrontierById(manager, identity, cursor.undoHead),
+		};
+	}
+	// 用户输入前的控制条目/根叶可由可信 before checkpoint 证明。
+	const boundaries = new Map<string, { checkpoint: CheckpointRecord; physicalLeaf: string | null }>();
+	for (const value of manager.getEntries() as unknown[]) {
+		if (!isRecord(value) || value.type !== "custom" || value.customType !== "pi-undo:checkpoint" || typeof value.id !== "string") continue;
+		for (const checkpoint of sessionStateFor(manager, value.id).getCheckpoints(identity)) {
+			const beforeLeaf = entryParent(manager, checkpoint.userEntryId);
+			if (logicalLeafAt(manager, beforeLeaf) === logicalLeafId) boundaries.set(checkpoint.checkpointId, { checkpoint, physicalLeaf: beforeLeaf });
+		}
+	}
+	const beforeStates = [...boundaries.values()];
+	if (beforeStates.length > 0 && new Set(beforeStates.map(({ checkpoint }) => checkpoint.beforeManifestId)).size === 1) {
+		const boundary = beforeStates[0]!;
+		return {
+			logicalLeafId,
+			targetManifestId: boundary.checkpoint.beforeManifestId,
+			undoStack: sessionStateFor(manager, boundary.physicalLeaf).getCheckpoints(identity),
+		};
+	}
+	throw new Error("tree target 缺少精确的 checkpoint 边界");
 }
 
 function findCheckpointByEndLeaf(
@@ -538,17 +636,23 @@ function volatileSessionIdentity(manager: ReadonlySessionManager, cwd: string):
 
 async function waitForIdle(context: ExtensionCommandContext | undefined, deadlineMs: number): Promise<boolean> {
 	if (context === undefined) return false;
-	const remaining = Math.max(0, deadlineMs - Date.now());
+	checkOperation();
+	const operation = currentOperationContext();
+	const remaining = Math.max(0, Math.min(deadlineMs, operation?.deadline ?? Infinity) - Date.now());
 	let timeout: ReturnType<typeof setTimeout> | undefined;
+	let onAbort: (() => void) | undefined;
 	try {
 		return await Promise.race([
 			context.waitForIdle().then(() => true, () => false),
 			new Promise<boolean>((resolveTimeout) => {
 				timeout = setTimeout(() => resolveTimeout(false), remaining);
+				onAbort = () => resolveTimeout(false);
+				operation?.signal.addEventListener("abort", onAbort, { once: true });
 			}),
 		]);
 	} finally {
 		if (timeout !== undefined) clearTimeout(timeout);
+		if (onAbort !== undefined) operation?.signal.removeEventListener("abort", onAbort);
 	}
 }
 

package/src/quarantine.ts

@@ -7,6 +7,7 @@ import { fsyncDirectory, writeBytesExclusive } from "./atomic-fs.ts";
 import { canonicalJson, checksum } from "./encoding.ts";
 import type { MutationJournal } from "./mutation-journal.ts";
 import type { MutationRecord } from "./model.ts";
+import { checkOperation } from "./operation-context.ts";
 import { assertNoSymlinkEscape, relativeSafePath } from "./path-safety.ts";
 
 const BATCH_FILE_IO_CONCURRENCY = 32;
@@ -828,6 +829,7 @@ export class QuarantineManager {
 	}
 
 	private async assertWorkspaceIdentity(): Promise<void> {
+		checkOperation();
 		const identity = await realpath(this.requestedWorkspaceRoot);
 		if (identity !== this.workspaceRoot) {
 			throw new QuarantineError("unsafe_artifact", "workspace root identity 已变化");
@@ -925,6 +927,7 @@ async function mapConcurrentFailClosed<T>(
 			const index = nextIndex;
 			nextIndex += 1;
 			try {
+				checkOperation();
 				await operation(values[index]!);
 			} catch (error) {
 				if (!failed) failure = error;

package/src/recovery.ts

@@ -94,9 +94,9 @@ export class JournalRecovery {
 			if (inspection.kind === "conflict") {
 				return { kind: "locked", reason: "cursor_conflict", operations: recovered };
 			}
-			const expectedLeaf = inspection.kind === "match"
-				? journal.descriptor.toLogicalLeaf
-				: journal.descriptor.fromLogicalLeaf;
+			const committedLeaf = journal.descriptor.action === "tree" && journal.state.observedLogicalLeaf !== undefined
+				? journal.state.observedLogicalLeaf : journal.descriptor.toLogicalLeaf;
+			const expectedLeaf = inspection.kind === "match" ? committedLeaf : journal.descriptor.fromLogicalLeaf;
 			if (this.dependencies.getLogicalLeafId() !== expectedLeaf) {
 				return { kind: "locked", reason: "session_leaf_mismatch", operations: recovered };
 			}