@herbertgao/pi-extensions 2026.9.1 → 2026.9.3

package/node_modules/@herbertgao/sol-pi/src/sol-pi/extensions/online-context-compact/extension.ts

@@ -0,0 +1,455 @@
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: MIT
+ */
+import type { AgentMessage, AgentToolResult } from "@earendil-works/pi-agent-core";
+import {
+	buildSessionContext,
+	estimateTokens,
+	findCutPoint,
+	sessionEntryToContextMessages,
+	type ExtensionContext,
+	type ExtensionFactory,
+	type SessionEntry,
+} from "@earendil-works/pi-coding-agent";
+import { formatSavingsCount, showSolPiSavings } from "../../tui.ts";
+import {
+	DEFAULT_COMPACTION_ECONOMICS,
+	decideCompaction,
+	type CompactionDecision,
+} from "./economics.ts";
+import { analyzePlanTransition, formatPlanSnapshot, parsePlanSteps } from "./plan.ts";
+import {
+	appendOnlineState,
+	initialOnlineState,
+	recordBoundary,
+	recordCompaction,
+	recordCorrection,
+	recordProviderRequest,
+	restoreOnlineState,
+	type OnlineState,
+	type ProgressSummary,
+} from "./state.ts";
+import { registerOnlineTools, type PlanUpdateInput } from "./tools.ts";
+
+export const DEFAULT_KEEP_RECENT_TOKENS = 20_000;
+export const DEFAULT_NATIVE_SUMMARY_TOKEN_ESTIMATE = 1_000;
+export const BOUNDARY_COMPACTION_INSTRUCTIONS =
+	"Preserve completed work, verification results, important decisions, and remaining work.";
+export const POST_COMPACTION_PLAN_REMINDER =
+	"Online context compaction finished. The parent task is still active. " +
+	"Before continuing work, call update_plan with a fresh plan for the remaining work.";
+
+export type OnlineContextCompactOptions = {
+	readonly cacheWriteReadRatio?: number | null;
+	readonly keepRecentTokens?: number;
+};
+
+type PendingBoundary = { readonly toolCallId: string };
+type SelectedCompaction = { readonly decision: CompactionDecision };
+type CacheDebt = { readonly debtTokens: number; readonly repaymentTokens: number };
+type PendingContinuation = { readonly promise: Promise<void>; readonly resolve: () => void };
+
+const SUBAGENTS_MANAGER_KEY = Symbol.for("pi-subagents:manager");
+
+export function hasRunningSubagents(): boolean {
+	try {
+		const manager = (globalThis as Record<PropertyKey, unknown>)[SUBAGENTS_MANAGER_KEY];
+		if (!manager || typeof manager !== "object") return false;
+		const hasRunning = (manager as { hasRunning?: unknown }).hasRunning;
+		return typeof hasRunning === "function" && hasRunning.call(manager) === true;
+	} catch {
+		return true;
+	}
+}
+
+export function resolveKeepRecentTokens(value: number | undefined): number {
+	const resolved = value ?? DEFAULT_KEEP_RECENT_TOKENS;
+	if (!Number.isSafeInteger(resolved) || resolved < 1) {
+		throw new Error("Online Context Compact keepRecentTokens must be a positive safe integer");
+	}
+	return resolved;
+}
+
+function resolveCacheWriteReadRatio(value: number | null | undefined): number | null {
+	if (value === undefined || value === null) return null;
+	if (!Number.isFinite(value) || value < 0) {
+		throw new Error("Online Context Compact cacheWriteReadRatio must be finite and non-negative");
+	}
+	return value;
+}
+
+function tokenEstimate(text: string): number {
+	return Math.ceil(Buffer.byteLength(text) / 4);
+}
+
+function result(text: string, details: Readonly<Record<string, unknown>>): AgentToolResult<Readonly<Record<string, unknown>>> {
+	return { content: [{ type: "text", text }], details };
+}
+
+function progressSummary(input: PlanUpdateInput, completedStepId: string): ProgressSummary | undefined {
+	const step = input.steps.find((item) => item.id === completedStepId);
+	if (!step || !input.progress) return;
+	return {
+		stepId: step.id,
+		goal: step.goal,
+		filesChanged: [...input.progress.files_changed],
+		verification: [...input.progress.verification],
+		decisions: [...input.progress.decisions],
+		nextWork: input.steps.filter((item) => item.status !== "completed").map((item) => item.goal),
+	};
+}
+
+function compactionMessageCount(entries: readonly SessionEntry[], startIndex: number, endIndex: number): number {
+	let count = 0;
+	for (let index = startIndex; index < endIndex; index++) {
+		const entry = entries[index];
+		if (entry && entry.type !== "compaction" && sessionEntryToContextMessages(entry).length > 0) count++;
+	}
+	return count;
+}
+
+function branchAfterAbort(entries: readonly SessionEntry[]): SessionEntry[] {
+	const last = entries.at(-1);
+	const markerProvider = ["sol", "pi"].join("-");
+	return [
+		...entries,
+		{
+			type: "message",
+			id: "sol-pi-online-context-compact-abort-marker",
+			parentId: last?.id ?? null,
+			timestamp: new Date(0).toISOString(),
+			message: {
+				role: "assistant",
+				content: [],
+				api: markerProvider,
+				provider: markerProvider,
+				model: "aborted",
+				usage: {
+					input: 0,
+					output: 0,
+					cacheRead: 0,
+					cacheWrite: 0,
+					totalTokens: 0,
+					cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 },
+				},
+				stopReason: "aborted",
+				timestamp: 0,
+			},
+		} as SessionEntry,
+	];
+}
+
+function nativeCompactionFeasible(entries: readonly SessionEntry[], keepRecentTokens: number): boolean {
+	const path = branchAfterAbort(entries);
+	let startIndex = 0;
+	for (let index = path.length - 1; index >= 0; index--) {
+		const entry = path[index];
+		if (entry?.type !== "compaction") continue;
+		const keptIndex = path.findIndex((item) => item.id === entry.firstKeptEntryId);
+		startIndex = keptIndex >= 0 ? keptIndex : index + 1;
+		break;
+	}
+
+	const cut = findCutPoint(path, startIndex, path.length, keepRecentTokens);
+	const historyEnd = cut.isSplitTurn ? cut.turnStartIndex : cut.firstKeptEntryIndex;
+	const historyMessages = historyEnd > startIndex ? compactionMessageCount(path, startIndex, historyEnd) : 0;
+	const prefixMessages =
+		cut.isSplitTurn && cut.turnStartIndex >= 0
+			? compactionMessageCount(path, cut.turnStartIndex, cut.firstKeptEntryIndex)
+			: 0;
+	return historyMessages > 0 || prefixMessages > 0;
+}
+
+function validPositiveInteger(value: unknown): value is number {
+	return typeof value === "number" && Number.isSafeInteger(value) && value > 0;
+}
+
+function releaseParentContinuation(continuation: PendingContinuation | undefined): void {
+	if (continuation) setTimeout(continuation.resolve, 0);
+}
+
+export function createOnlineContextCompactExtension(options: OnlineContextCompactOptions = {}): ExtensionFactory {
+	const keepRecentTokens = resolveKeepRecentTokens(options.keepRecentTokens);
+	const cacheWriteReadRatio = resolveCacheWriteReadRatio(options.cacheWriteReadRatio);
+
+	return (pi) => {
+		let state: OnlineState = initialOnlineState();
+		let restored = false;
+		let observedMessages: readonly AgentMessage[] = [];
+		let pendingBoundary: PendingBoundary | undefined;
+		let selected: SelectedCompaction | undefined;
+		let activeDebt: CacheDebt | undefined;
+		let nextContinuation: PendingContinuation | undefined;
+		let compactionInFlight = false;
+
+		const releaseContinuation = (): void => {
+			const continuation = nextContinuation;
+			nextContinuation = undefined;
+			continuation?.resolve();
+		};
+		const restore = (context: ExtensionContext): void => {
+			releaseContinuation();
+			state = restoreOnlineState(context.sessionManager.getBranch());
+			restored = true;
+			observedMessages = buildSessionContext(
+				context.sessionManager.getEntries(),
+				context.sessionManager.getLeafId(),
+			).messages;
+			pendingBoundary = undefined;
+			selected = undefined;
+			activeDebt = undefined;
+			compactionInFlight = false;
+		};
+		const ensureRestored = (context: ExtensionContext): void => {
+			if (!restored) restore(context);
+		};
+		const save = (): void => appendOnlineState(pi, state);
+		const contextTokens = (context: ExtensionContext): number => {
+			const visible = observedMessages.reduce((total, message) => total + estimateTokens(message), 0);
+			const estimated = visible + tokenEstimate(context.getSystemPrompt());
+			const reported = context.getContextUsage()?.tokens;
+			return validPositiveInteger(reported) ? Math.max(reported, estimated) : estimated;
+		};
+
+		registerOnlineTools(pi, {
+			updatePlan: async (input) => {
+				ensureRestored(input.context);
+				if (input.signal?.aborted) throw new Error("Plan update was aborted");
+				const steps = parsePlanSteps(input.steps);
+				if (!steps || steps.length === 0) throw new Error("Plan must contain at least one valid step");
+
+				const transition = analyzePlanTransition(state.plan, steps);
+				const completedIds = transition.completedSteps.map((step) => step.id);
+				if (completedIds.length > 0) {
+					state = recordBoundary(state, steps, progressSummary(input, completedIds[0] ?? ""));
+					if (!pendingBoundary) pendingBoundary = { toolCallId: input.toolCallId };
+				} else if (JSON.stringify(state.plan) !== JSON.stringify(steps)) {
+					state = { ...state, plan: [...steps] };
+				}
+				save();
+
+				return result(
+					[formatPlanSnapshot(steps), ...transition.advice].join("\n"),
+					{
+						boundary: completedIds.length > 0,
+						completed_step_ids: completedIds,
+						progress_recorded: completedIds.length > 0 && input.progress !== undefined,
+						task_status: "active",
+						plan: steps,
+					},
+				);
+			},
+		});
+
+		pi.on("session_start", (_event, context) => restore(context));
+		pi.on("session_before_tree", () => (compactionInFlight ? { cancel: true } : undefined));
+		pi.on("session_tree", (_event, context) => restore(context));
+
+		pi.on("context", (event, context) => {
+			ensureRestored(context);
+			observedMessages = [...event.messages];
+		});
+
+		pi.on("before_provider_request", (_event, context) => {
+			ensureRestored(context);
+			state = recordProviderRequest(state, contextTokens(context));
+			save();
+		});
+
+		pi.on("input", (event, context) => {
+			if (event.streamingBehavior !== "steer" && !event.text.startsWith("CORRECTION:")) {
+				return { action: "continue" as const };
+			}
+			ensureRestored(context);
+			pendingBoundary = undefined;
+			selected = undefined;
+			activeDebt = undefined;
+			state = recordCorrection(state);
+			save();
+			return { action: "continue" as const };
+		});
+
+		pi.on("turn_end", (event, context) => {
+			const boundary = pendingBoundary;
+			pendingBoundary = undefined;
+			if (!boundary || selected || hasRunningSubagents()) return;
+			const toolResult = event.toolResults.find((item) => item.toolCallId === boundary.toolCallId);
+			if (
+				event.message.role !== "assistant" ||
+				event.message.stopReason === "error" ||
+				event.message.stopReason === "aborted" ||
+				context.signal?.aborted ||
+				!toolResult ||
+				toolResult.isError
+			) {
+				return;
+			}
+
+			const usage = context.getContextUsage();
+			const writeTokens = contextTokens(context);
+			const fixedTokens = tokenEstimate(context.getSystemPrompt());
+			const archiveTokens = Math.max(0, writeTokens - fixedTokens - keepRecentTokens);
+			const contextWindowTokens = validPositiveInteger(usage?.contextWindow)
+				? usage.contextWindow
+				: validPositiveInteger(context.model?.contextWindow)
+					? context.model.contextWindow
+					: null;
+			const averageContextTokenIncrement =
+				state.positiveContextDeltaCount === 0
+					? null
+					: state.positiveContextDeltaTotal / state.positiveContextDeltaCount;
+			const priced = decideCompaction({
+				writeTokens,
+				archiveTokens,
+				memoTokens: DEFAULT_NATIVE_SUMMARY_TOKEN_ESTIMATE,
+				contextTokens: writeTokens,
+				completedBoundaryRequestCounts: state.completedBoundaryRequestCounts,
+				remainingBoundaries: state.plan.filter((step) => step.status !== "completed").length,
+				averageContextTokenIncrement,
+				contextWindowTokens,
+				priorCompactionCount: state.nativeCompactionCount,
+				carriedDebtTokens: state.cacheDebtTokens,
+				cacheDebtRepaymentTokens: state.cacheDebtRepaymentTokens,
+				cacheWriteReadRatio,
+				economics: DEFAULT_COMPACTION_ECONOMICS,
+			});
+			const decision: CompactionDecision =
+				priced.compact && !nativeCompactionFeasible(context.sessionManager.getBranch(), keepRecentTokens)
+					? { ...priced, compact: false, reason: "native_not_compactable" }
+					: priced;
+			if (!decision.compact) return;
+
+			selected = { decision };
+			context.abort();
+		});
+
+		pi.on("agent_settled", async (_event, context) => {
+			// sendMessage() starts a turn without returning its promise. Capture the
+			// child settlement so print/JSON mode cannot dispose while it is running.
+			const parentContinuation = nextContinuation;
+			nextContinuation = undefined;
+			const pending = selected;
+			selected = undefined;
+			if (!context.isIdle()) {
+				selected = pending;
+				nextContinuation = parentContinuation;
+				return;
+			}
+			if (!pending || hasRunningSubagents()) {
+				releaseParentContinuation(parentContinuation);
+				return;
+			}
+
+			activeDebt = {
+				debtTokens: pending.decision.writeTokens * (pending.decision.incrementalCacheCostRatio ?? 0),
+				repaymentTokens: Math.max(0, pending.decision.archiveTokens - pending.decision.memoTokens),
+			};
+			let compacted = false;
+			let compactionError: Error | undefined;
+			try {
+				compactionInFlight = true;
+				await new Promise<void>((resolve) => {
+					let finished = false;
+					const finish = (): void => {
+						if (finished) return;
+						finished = true;
+						resolve();
+					};
+					context.compact({
+						customInstructions: BOUNDARY_COMPACTION_INSTRUCTIONS,
+						onComplete: (compaction) => {
+							try {
+								compacted = true;
+								const removed = Math.max(
+									0,
+									pending.decision.archiveTokens - tokenEstimate(compaction.summary),
+								);
+								if (removed > 0) {
+									showSolPiSavings(
+										context,
+										"Online Context Compact",
+										formatSavingsCount(removed, "context tokens removed"),
+									);
+								}
+							} finally {
+								finish();
+							}
+						},
+						onError: (error) => {
+							compactionError = error;
+							finish();
+						},
+					});
+				});
+				compactionInFlight = false;
+				if (
+					compactionError &&
+					compactionError.name !== "AbortError" &&
+					compactionError.message !== "Compaction cancelled"
+				) {
+					throw compactionError;
+				}
+
+				if (compacted) {
+					let resolveContinuation!: () => void;
+					const continuation: PendingContinuation = {
+						promise: new Promise<void>((resolve) => {
+							resolveContinuation = resolve;
+						}),
+						resolve: () => resolveContinuation(),
+					};
+					nextContinuation = continuation;
+					try {
+						pi.sendMessage(
+							{
+								customType: "sol-pi-online-context-compact",
+								content: POST_COMPACTION_PLAN_REMINDER,
+								display: false,
+							},
+							{ triggerTurn: true },
+						);
+					} catch (error) {
+						if (nextContinuation === continuation) nextContinuation = undefined;
+						continuation.resolve();
+						throw error;
+					}
+					if (context.isIdle() && nextContinuation === continuation) {
+						nextContinuation = undefined;
+						continuation.resolve();
+						throw new Error("Online context compact continuation did not start");
+					}
+					await continuation.promise;
+				}
+			} finally {
+				compactionInFlight = false;
+				activeDebt = undefined;
+				releaseParentContinuation(parentContinuation);
+			}
+		});
+
+		pi.on("session_compact", (event, context) => {
+			ensureRestored(context);
+			state = recordCompaction(
+				state,
+				event.fromExtension || !activeDebt ? { debtTokens: 0, repaymentTokens: 0 } : activeDebt,
+			);
+			save();
+			pendingBoundary = undefined;
+			selected = undefined;
+			activeDebt = undefined;
+			observedMessages = buildSessionContext(
+				context.sessionManager.getEntries(),
+				context.sessionManager.getLeafId(),
+			).messages;
+		});
+
+		pi.on("session_shutdown", () => {
+			releaseContinuation();
+			pendingBoundary = undefined;
+			selected = undefined;
+			activeDebt = undefined;
+			compactionInFlight = false;
+		});
+	};
+}

package/node_modules/@herbertgao/sol-pi/src/sol-pi/extensions/online-context-compact/index.ts

@@ -0,0 +1,49 @@
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: MIT
+ */
+import type { ExtensionAPI } from "@earendil-works/pi-coding-agent";
+import { DEFAULT_CACHE_WRITE_READ_RATIO } from "../../config.ts";
+import { createOnlineContextCompactExtension } from "./extension.ts";
+
+export {
+	DEFAULT_COMPACTION_ECONOMICS,
+	decideCompaction,
+	estimateRemainingRequests,
+	type CompactionDecision,
+	type CompactionEconomics,
+	type CompactionReason,
+} from "./economics.ts";
+export {
+	BOUNDARY_COMPACTION_INSTRUCTIONS,
+	createOnlineContextCompactExtension,
+	DEFAULT_KEEP_RECENT_TOKENS,
+	DEFAULT_NATIVE_SUMMARY_TOKEN_ESTIMATE,
+	POST_COMPACTION_PLAN_REMINDER,
+	type OnlineContextCompactOptions,
+	resolveKeepRecentTokens,
+} from "./extension.ts";
+export {
+	analyzePlanTransition,
+	formatPlanSnapshot,
+	parsePlanSteps,
+	type PlanStatus,
+	type PlanStep,
+} from "./plan.ts";
+export {
+	initialOnlineState,
+	ONLINE_STATE_ENTRY,
+	restoreOnlineState,
+	type OnlineState,
+	type ProgressSummary,
+} from "./state.ts";
+export type { PlanProgress, PlanUpdateInput } from "./tools.ts";
+
+export function registerOnlineContextCompact(
+	pi: ExtensionAPI,
+	cacheWriteReadRatio = DEFAULT_CACHE_WRITE_READ_RATIO,
+): void {
+	createOnlineContextCompactExtension({ cacheWriteReadRatio })(pi);
+}
+
+export default registerOnlineContextCompact;

package/node_modules/@herbertgao/sol-pi/src/sol-pi/extensions/online-context-compact/plan.ts

@@ -0,0 +1,79 @@
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: MIT
+ */
+
+export const PLAN_STATUSES = ["pending", "in_progress", "completed"] as const;
+
+export type PlanStatus = (typeof PLAN_STATUSES)[number];
+
+export type PlanStep = {
+	readonly id: string;
+	readonly goal: string;
+	readonly status: PlanStatus;
+};
+
+export type PlanTransition = {
+	readonly completedSteps: readonly PlanStep[];
+	readonly advice: readonly string[];
+};
+
+const MAX_PLAN_STEPS = 128;
+const MAX_PLAN_STRING_BYTES = 16_384;
+
+function isRecord(value: unknown): value is Readonly<Record<string, unknown>> {
+	return typeof value === "object" && value !== null && !Array.isArray(value);
+}
+
+function isBoundedString(value: unknown): value is string {
+	return typeof value === "string" && value.length > 0 && Buffer.byteLength(value) <= MAX_PLAN_STRING_BYTES;
+}
+
+function isPlanStatus(value: unknown): value is PlanStatus {
+	return PLAN_STATUSES.some((status) => status === value);
+}
+
+export function parsePlanSteps(value: unknown): readonly PlanStep[] | undefined {
+	if (!Array.isArray(value) || value.length > MAX_PLAN_STEPS) return;
+	const steps: PlanStep[] = [];
+	for (const item of value) {
+		if (
+			!isRecord(item) ||
+			Object.keys(item).length !== 3 ||
+			!isBoundedString(item.id) ||
+			!isBoundedString(item.goal) ||
+			!isPlanStatus(item.status)
+		) {
+			return;
+		}
+		steps.push({ id: item.id, goal: item.goal, status: item.status });
+	}
+	if (new Set(steps.map((step) => step.id)).size !== steps.length) return;
+	return steps;
+}
+
+export function analyzePlanTransition(previous: readonly PlanStep[], next: readonly PlanStep[]): PlanTransition {
+	const previousById = new Map(previous.map((step) => [step.id, step]));
+	const completedSteps: PlanStep[] = [];
+	const advice: string[] = [];
+
+	for (const step of next) {
+		const prior = previousById.get(step.id);
+		if ((!prior || prior.status !== "completed") && step.status === "completed") completedSteps.push(step);
+		if (prior && prior.goal !== step.goal) {
+			advice.push(`Plan step ${JSON.stringify(step.id)} changed goal; reuse an id only for the same goal.`);
+		}
+	}
+
+	const inProgress = next.filter((step) => step.status === "in_progress").length;
+	if (inProgress > 1) advice.push("Keep at most one plan step in_progress.");
+	if (inProgress === 0 && next.some((step) => step.status === "pending")) {
+		advice.push("Mark one pending plan step in_progress before starting it.");
+	}
+
+	return { completedSteps, advice };
+}
+
+export function formatPlanSnapshot(steps: readonly PlanStep[]): string {
+	return `<sol-pi-plan task_status="active">${JSON.stringify({ steps })}</sol-pi-plan>`;
+}