@lll9p/pi-better-compaction 0.2.1 → 0.5.0

package/src/extension-runtime.ts

@@ -6,15 +6,18 @@ import type {
 	SessionBeforeCompactEvent,
 } from "@earendil-works/pi-coding-agent";
 import { executeNativeCompaction } from "./compact-client";
+import { executeV2Compaction } from "./compact-client-v2";
 import { loadExtensionConfig } from "./config";
 import { writeDebugArtifact } from "./debug";
 import { resolveLatestNativeCompactionEntry } from "./details-store";
+import { registerMidRunGuard } from "./midrun";
 import { runNativeFallbackCompaction } from "./native-fallback";
 import {
 	rewriteResponsesPayloadWithNativeReplay,
 	serializeLiveTailToResponsesInput,
 } from "./payload-rewrite";
 import { getCompactionRequestExtras, rememberRequestContext } from "./request-context-cache";
+import { buildRetainedMessages } from "./retained-messages";
 import {
 	isResponsesCompatiblePayload,
 	resolveNativeCompactionEnvironment,
@@ -26,6 +29,8 @@ import {
 	createNativeCompactionResult,
 	EXTENSION_ID,
 	isNativeCompactionDetails,
+	NATIVE_COMPACTION_STRATEGY,
+	NATIVE_COMPACTION_STRATEGY_V2,
 	type ExtensionConfig,
 	type NativeCompactionDetails,
 	type NativeCompactionRequestMeta,
@@ -36,6 +41,20 @@ type ResponsesCompactOutcome =
 	| { outcome: "aborted" }
 	| { outcome: "failed" };
 
+export type ExtensionRuntimeDependencies = {
+	loadExtensionConfig: typeof loadExtensionConfig;
+	executeNativeCompaction: typeof executeNativeCompaction;
+	executeV2Compaction: typeof executeV2Compaction;
+	runNativeFallbackCompaction: typeof runNativeFallbackCompaction;
+};
+
+const DEFAULT_DEPENDENCIES: ExtensionRuntimeDependencies = {
+	loadExtensionConfig,
+	executeNativeCompaction,
+	executeV2Compaction,
+	runNativeFallbackCompaction,
+};
+
 function buildCompactionRequestMeta(event: SessionBeforeCompactEvent): NativeCompactionRequestMeta {
 	return {
 		tokensBefore: event.preparation.tokensBefore,
@@ -87,14 +106,15 @@ function buildCompactionInstructions(systemPrompt: string, customInstructions?:
 		return systemPrompt;
 	}
 
-	return `${systemPrompt}\n\nAdditional user guidance for this manual /compact request:\n${guidance}`;
+	return `${systemPrompt}\n\nAdditional compaction guidance:\n${guidance}`;
 }
 
-async function runResponsesNativeCompact(
+async function runResponsesV1Compact(
 	event: SessionBeforeCompactEvent,
 	ctx: ExtensionContext,
 	config: ExtensionConfig,
 	runtime: NativeCompactionRuntime,
+	dependencies: ExtensionRuntimeDependencies,
 ): Promise<ResponsesCompactOutcome> {
 	const instructions = buildCompactionInstructions(ctx.getSystemPrompt(), event.customInstructions);
 	const branchEntries = ctx.sessionManager.getBranch();
@@ -135,7 +155,7 @@ async function runResponsesNativeCompact(
 		writeDebugArtifact(
 			"compaction-event",
 			{
-				event: "session_before_compact.responses-compact-skip",
+				event: "session_before_compact.v1-compact-skip",
 				reason: latestNativeCompaction.reason,
 				provider: runtime.provider,
 				api: runtime.api,
@@ -157,7 +177,7 @@ async function runResponsesNativeCompact(
 		request = { ...request, ...extras };
 	}
 
-	const compactResult = await executeNativeCompaction({
+	const compactResult = await dependencies.executeNativeCompaction({
 		runtime,
 		request,
 		signal: event.signal,
@@ -169,7 +189,7 @@ async function runResponsesNativeCompact(
 		writeDebugArtifact(
 			"compaction-event",
 			{
-				event: "session_before_compact.responses-compact-failure",
+				event: "session_before_compact.v1-compact-failure",
 				reason: compactResult.reason,
 				status: compactResult.status,
 				errorMessage: compactResult.errorMessage,
@@ -196,7 +216,7 @@ async function runResponsesNativeCompact(
 		writeDebugArtifact(
 			"compaction-event",
 			{
-				event: "session_before_compact.invalid-native-details",
+				event: "session_before_compact.v1-invalid-native-details",
 				reason: error instanceof Error ? error.message : String(error),
 				provider: runtime.provider,
 				api: runtime.api,
@@ -219,7 +239,7 @@ async function runResponsesNativeCompact(
 	writeDebugArtifact(
 		"compaction-event",
 		{
-			event: "session_before_compact.responses-compact-success",
+			event: "session_before_compact.v1-compact-success",
 			provider: runtime.provider,
 			api: runtime.api,
 			model: runtime.model,
@@ -238,8 +258,171 @@ async function runResponsesNativeCompact(
 	return { outcome: "success", compaction };
 }
 
-async function handleSessionBeforeCompact(event: SessionBeforeCompactEvent, ctx: ExtensionContext) {
-	const { config } = loadExtensionConfig();
+/**
+ * V2 compaction: stream a Responses request with compaction_trigger appended.
+ * On success, returns retained messages + encrypted compaction blob.
+ */
+async function runResponsesV2Compact(
+	event: SessionBeforeCompactEvent,
+	ctx: ExtensionContext,
+	config: ExtensionConfig,
+	runtime: NativeCompactionRuntime,
+	dependencies: ExtensionRuntimeDependencies,
+): Promise<ResponsesCompactOutcome> {
+	const instructions = buildCompactionInstructions(ctx.getSystemPrompt(), event.customInstructions);
+	const branchEntries = ctx.sessionManager.getBranch();
+	const latestNativeCompaction = resolveLatestNativeCompactionEntry(branchEntries, {
+		provider: runtime.provider,
+		api: runtime.api,
+		model: runtime.model,
+		baseUrl: runtime.baseUrl,
+	});
+
+	let requestSource: "session-context" | "non-native-session-context" | "latest-native-replay";
+	let request: NativeCompactionRequestBody;
+	if (latestNativeCompaction.ok) {
+		const liveTailEntries = branchEntries.slice(latestNativeCompaction.index + 1);
+		requestSource = "latest-native-replay";
+		const input: ResponsesInputItem[] = [
+			...(cloneOpaqueWindow(latestNativeCompaction.entry.details.compactedWindow) as ResponsesInputItem[]),
+			...serializeLiveTailToResponsesInput({ model: runtime.currentModel, entries: liveTailEntries }),
+		];
+		request = {
+			model: runtime.currentModel.id,
+			input,
+			instructions,
+		};
+	} else if (
+		latestNativeCompaction.reason === "no-compaction" ||
+		(latestNativeCompaction.reason === "latest-compaction-not-native" &&
+			config.allowCompactionContinuityBreak)
+	) {
+		requestSource =
+			latestNativeCompaction.reason === "no-compaction" ? "session-context" : "non-native-session-context";
+		request = serializeMessagesToCompactRequest({
+			model: runtime.currentModel,
+			messages: ctx.sessionManager.buildSessionContext().messages,
+			instructions,
+		});
+	} else {
+		writeDebugArtifact(
+			"compaction-event",
+			{
+				event: "session_before_compact.v2-compact-skip",
+				reason: latestNativeCompaction.reason,
+				provider: runtime.provider,
+				api: runtime.api,
+				model: runtime.model,
+				baseUrl: runtime.baseUrl,
+				latestCompactionIndex: latestNativeCompaction.latestCompactionIndex,
+				latestCompactionIdentity: getCompactionIdentityDebugInfo(latestNativeCompaction.latestCompaction),
+			},
+			config,
+			ctx,
+		);
+		return { outcome: "failed" };
+	}
+
+	const extras = getCompactionRequestExtras(runtime.model, getSessionId(ctx));
+	if (extras) {
+		request = { ...request, ...extras };
+	}
+
+	const v2Result = await dependencies.executeV2Compaction({
+		runtime,
+		request,
+		signal: event.signal,
+		settings: config,
+		context: ctx,
+	});
+
+	if (!v2Result.ok) {
+		writeDebugArtifact(
+			"compaction-event",
+			{
+				event: "session_before_compact.v2-compact-failure",
+				reason: v2Result.reason,
+				status: v2Result.status,
+				errorMessage: v2Result.errorMessage,
+			},
+			config,
+			ctx,
+		);
+		return v2Result.reason === "aborted" ? { outcome: "aborted" } : { outcome: "failed" };
+	}
+
+	// Build compacted window: retained messages + compaction blob.
+	const retainedMessages = buildRetainedMessages(request.input);
+	const compactedWindow = [...retainedMessages, v2Result.compactionItem];
+
+	let details: NativeCompactionDetails;
+	try {
+		details = createNativeCompactionDetails(
+			{
+				provider: runtime.provider,
+				api: runtime.api,
+				model: runtime.model,
+				baseUrl: runtime.baseUrl,
+				compactedWindow,
+				compactResponseId: v2Result.responseId,
+				createdAt: v2Result.createdAt,
+				requestMeta: buildCompactionRequestMeta(event),
+			},
+			NATIVE_COMPACTION_STRATEGY_V2,
+		);
+	} catch (error) {
+		writeDebugArtifact(
+			"compaction-event",
+			{
+				event: "session_before_compact.v2-invalid-native-details",
+				reason: error instanceof Error ? error.message : String(error),
+				provider: runtime.provider,
+				api: runtime.api,
+				model: runtime.model,
+				baseUrl: runtime.baseUrl,
+			},
+			config,
+			ctx,
+		);
+		return { outcome: "failed" };
+	}
+
+	// V2 blob is encrypted; no summary text can be extracted.
+	const compaction = createNativeCompactionResult({
+		firstKeptEntryId: event.preparation.firstKeptEntryId,
+		tokensBefore: event.preparation.tokensBefore,
+		details,
+	});
+
+	writeDebugArtifact(
+		"compaction-event",
+		{
+			event: "session_before_compact.v2-compact-success",
+			provider: runtime.provider,
+			api: runtime.api,
+			model: runtime.model,
+			requestSource,
+			requestInputItems: request.input.length,
+			requestExtras: extras ? Object.keys(extras) : [],
+			compactResponseId: v2Result.responseId,
+			retainedMessageCount: retainedMessages.length,
+			compactedItems: compactedWindow.length,
+			usage: v2Result.usage,
+			firstKeptEntryId: event.preparation.firstKeptEntryId,
+		},
+		config,
+		ctx,
+	);
+
+	return { outcome: "success", compaction };
+}
+
+async function handleSessionBeforeCompact(
+	event: SessionBeforeCompactEvent,
+	ctx: ExtensionContext,
+	dependencies: ExtensionRuntimeDependencies,
+) {
+	const { config } = dependencies.loadExtensionConfig();
 	if (!config.enabled) {
 		return undefined;
 	}
@@ -271,7 +454,14 @@ async function handleSessionBeforeCompact(event: SessionBeforeCompactEvent, ctx:
 		responsesCompactApis: config.responsesCompactApis,
 	});
 	if (resolution.ok) {
-		const responsesOutcome = await runResponsesNativeCompact(event, ctx, config, resolution.runtime);
+		let responsesOutcome: ResponsesCompactOutcome;
+
+		if (config.compactionVersion === "v2") {
+			responsesOutcome = await runResponsesV2Compact(event, ctx, config, resolution.runtime, dependencies);
+		} else {
+			responsesOutcome = await runResponsesV1Compact(event, ctx, config, resolution.runtime, dependencies);
+		}
+
 		if (responsesOutcome.outcome === "success") {
 			return { compaction: responsesOutcome.compaction };
 		}
@@ -296,7 +486,12 @@ async function handleSessionBeforeCompact(event: SessionBeforeCompactEvent, ctx:
 	}
 
 	// Branch 2: run pi's native compaction method with the configured model.
-	const fallback = await runNativeFallbackCompaction({ ctx, event, config });
+	const fallback = await dependencies.runNativeFallbackCompaction({
+		ctx,
+		event,
+		config,
+		sessionId: getSessionId(ctx),
+	});
 	if (fallback.ok) {
 		if (ctx.hasUI) {
 			ctx.ui.notify(
@@ -309,6 +504,7 @@ async function handleSessionBeforeCompact(event: SessionBeforeCompactEvent, ctx:
 			{
 				event: "session_before_compact.fallback-success",
 				model: fallback.model,
+				usage: fallback.usage,
 			},
 			config,
 			ctx,
@@ -344,8 +540,12 @@ async function handleSessionBeforeCompact(event: SessionBeforeCompactEvent, ctx:
 	return undefined;
 }
 
-async function handleBeforeProviderRequest(event: BeforeProviderRequestEvent, ctx: ExtensionContext) {
-	const { config } = loadExtensionConfig();
+async function handleBeforeProviderRequest(
+	event: BeforeProviderRequestEvent,
+	ctx: ExtensionContext,
+	dependencies: ExtensionRuntimeDependencies,
+) {
+	const { config } = dependencies.loadExtensionConfig();
 	if (!config.enabled) {
 		return undefined;
 	}
@@ -468,9 +668,14 @@ async function handleBeforeProviderRequest(event: BeforeProviderRequestEvent, ct
 	return rewrite.rewrittenPayload;
 }
 
-export default function (pi: ExtensionAPI) {
+export function registerExtensionRuntime(
+	pi: ExtensionAPI,
+	dependencies: ExtensionRuntimeDependencies = DEFAULT_DEPENDENCIES,
+): void {
+	registerMidRunGuard(pi, dependencies.loadExtensionConfig);
+
 	pi.on("session_start", (_event, ctx) => {
-		const { config, source, warnings } = loadExtensionConfig();
+		const { config, source, warnings } = dependencies.loadExtensionConfig();
 		if (!config.enabled) return;
 
 		if (warnings.length > 0 && ctx.hasUI && config.debug) {
@@ -499,6 +704,31 @@ export default function (pi: ExtensionAPI) {
 		}
 	});
 
-	pi.on("session_before_compact", handleSessionBeforeCompact);
-	pi.on("before_provider_request", handleBeforeProviderRequest);
+	pi.on("session_before_compact", (event, ctx) =>
+		handleSessionBeforeCompact(event, ctx, dependencies),
+	);
+	pi.on("before_provider_request", (event, ctx) =>
+		handleBeforeProviderRequest(event, ctx, dependencies),
+	);
+
+	pi.on("session_compact_failed", (event, ctx) => {
+		const { config } = dependencies.loadExtensionConfig();
+		if (!config.enabled) return;
+
+		writeDebugArtifact(
+			"compaction-event",
+			{
+				event: "session_compact_failed",
+				reason: event.reason,
+				errorMessage: event.errorMessage,
+				aborted: event.aborted,
+				willRetry: event.willRetry,
+				fromExtension: event.fromExtension,
+			},
+			config,
+			ctx,
+		);
+	});
 }
+
+export default registerExtensionRuntime;

package/src/midrun.ts

@@ -0,0 +1,229 @@
+import type { ExtensionAPI, ExtensionContext } from "@earendil-works/pi-coding-agent";
+import { loadExtensionConfig } from "./config";
+import { writeDebugArtifact } from "./debug";
+import { EXTENSION_ID, type LoadedExtensionConfig } from "./types";
+
+type MidRunPhase = "idle" | "abort-pending" | "compacting" | "resume-pending" | "failed";
+
+type MidRunState = {
+	phase: MidRunPhase;
+	generation: number;
+	sessionId?: string;
+	baselineCompactionId?: string;
+	triggerTokens?: number;
+	triggerPercent?: number;
+	triggerContextWindow?: number;
+};
+
+type ConfigLoader = () => LoadedExtensionConfig;
+
+const RESUME_CUSTOM_TYPE = "pi-better-compaction-midrun-resume";
+const RESUME_PROMPT = `[pi-better-compaction/midrun]
+Context compaction completed. Continue the interrupted task from the exact point where execution stopped.`;
+const MIDRUN_COMPACTION_INSTRUCTIONS =
+	"Preserve the active task, completed work, decisions, changed files, failures, current tool-loop state, and exact next steps so execution can resume immediately after compaction.";
+
+function getSessionId(ctx: ExtensionContext): string | undefined {
+	try {
+		return ctx.sessionManager.getSessionId();
+	} catch {
+		return undefined;
+	}
+}
+
+function getLatestCompactionId(ctx: ExtensionContext): string | undefined {
+	const branch = ctx.sessionManager.getBranch();
+	for (let index = branch.length - 1; index >= 0; index--) {
+		const entry = branch[index];
+		if (entry?.type === "compaction") return entry.id;
+	}
+	return undefined;
+}
+
+function resetState(state: MidRunState): void {
+	state.phase = "idle";
+	state.generation += 1;
+	state.sessionId = undefined;
+	state.baselineCompactionId = undefined;
+	state.triggerTokens = undefined;
+	state.triggerPercent = undefined;
+	state.triggerContextWindow = undefined;
+}
+
+function sameSession(state: MidRunState, ctx: ExtensionContext): boolean {
+	return state.sessionId !== undefined && state.sessionId === getSessionId(ctx);
+}
+
+function notifyFailure(ctx: ExtensionContext, message: string): void {
+	if (ctx.hasUI) {
+		ctx.ui.notify(`${EXTENSION_ID}: ${message}`, "error");
+	}
+}
+
+function scheduleResume(
+	pi: ExtensionAPI,
+	ctx: ExtensionContext,
+	state: MidRunState,
+	generation: number,
+): void {
+	state.phase = "resume-pending";
+
+	setImmediate(() => {
+		if (
+			state.generation !== generation ||
+			state.phase !== "resume-pending" ||
+			!sameSession(state, ctx)
+		) {
+			return;
+		}
+
+		if (!ctx.isIdle() || ctx.hasPendingMessages()) {
+			resetState(state);
+			return;
+		}
+
+		resetState(state);
+		pi.sendMessage(
+			{
+				customType: RESUME_CUSTOM_TYPE,
+				content: RESUME_PROMPT,
+				display: false,
+				details: { source: "midrun-compaction" },
+			},
+			{ triggerTurn: true },
+		);
+	});
+}
+
+export function registerMidRunGuard(
+	pi: ExtensionAPI,
+	loadConfig: ConfigLoader = loadExtensionConfig,
+): void {
+	const state: MidRunState = { phase: "idle", generation: 0 };
+
+	pi.on("turn_end", (event, ctx) => {
+		const { config } = loadConfig();
+		if (!config.enabled || !config.midRun.enabled || state.phase !== "idle") return;
+		if (event.toolResults.length === 0 || ctx.hasPendingMessages()) return;
+
+		const usage = ctx.getContextUsage();
+		if (!usage || usage.tokens == null || usage.percent == null) return;
+		if (usage.percent < config.midRun.thresholdPercent) return;
+
+		const sessionId = getSessionId(ctx);
+		if (!sessionId) return;
+
+		state.phase = "abort-pending";
+		state.generation += 1;
+		state.sessionId = sessionId;
+		state.baselineCompactionId = getLatestCompactionId(ctx);
+		state.triggerTokens = usage.tokens;
+		state.triggerPercent = usage.percent;
+		state.triggerContextWindow = usage.contextWindow;
+
+		writeDebugArtifact(
+			"lifecycle",
+			{
+				event: "midrun.threshold",
+				turnIndex: event.turnIndex,
+				tokens: usage.tokens,
+				contextWindow: usage.contextWindow,
+				percent: usage.percent,
+				thresholdPercent: config.midRun.thresholdPercent,
+				baselineCompactionId: state.baselineCompactionId,
+			},
+			config,
+			ctx,
+		);
+
+		// Never compact while the agent run is active; let Pi settle first.
+		ctx.abort();
+	});
+
+	pi.on("agent_settled", (_event, ctx) => {
+		if (state.phase !== "abort-pending") return;
+		if (!sameSession(state, ctx)) {
+			resetState(state);
+			return;
+		}
+
+		// An earlier agent_settled handler may already have started or queued another run.
+		if (!ctx.isIdle() || ctx.hasPendingMessages()) {
+			resetState(state);
+			return;
+		}
+
+		const generation = state.generation;
+		const latestCompactionId = getLatestCompactionId(ctx);
+		if (latestCompactionId !== state.baselineCompactionId) {
+			const { config } = loadConfig();
+			writeDebugArtifact(
+				"lifecycle",
+				{
+					event: "midrun.coalesced",
+					reason: "compaction-already-occurred-during-abort",
+					baselineCompactionId: state.baselineCompactionId,
+					latestCompactionId,
+				},
+				config,
+				ctx,
+			);
+			scheduleResume(pi, ctx, state, generation);
+			return;
+		}
+
+		const { config } = loadConfig();
+		if (!config.enabled || !config.midRun.enabled) {
+			scheduleResume(pi, ctx, state, generation);
+			return;
+		}
+
+		const usage = ctx.getContextUsage();
+		if (usage?.percent != null && usage.percent < config.midRun.thresholdPercent) {
+			scheduleResume(pi, ctx, state, generation);
+			return;
+		}
+
+		state.phase = "compacting";
+		ctx.compact({
+			customInstructions: MIDRUN_COMPACTION_INSTRUCTIONS,
+			onComplete: () => {
+				if (
+					state.generation !== generation ||
+					state.phase !== "compacting" ||
+					!sameSession(state, ctx)
+				) {
+					return;
+				}
+				scheduleResume(pi, ctx, state, generation);
+			},
+			onError: (error) => {
+				if (state.generation !== generation || !sameSession(state, ctx)) return;
+
+				const latest = getLatestCompactionId(ctx);
+				if (/Already compacted/i.test(error.message) && latest !== state.baselineCompactionId) {
+					scheduleResume(pi, ctx, state, generation);
+					return;
+				}
+
+				state.phase = "failed";
+				writeDebugArtifact(
+					"lifecycle",
+					{
+						event: "midrun.failed",
+						errorMessage: error.message,
+						triggerTokens: state.triggerTokens,
+						triggerPercent: state.triggerPercent,
+						triggerContextWindow: state.triggerContextWindow,
+					},
+					config,
+					ctx,
+				);
+				notifyFailure(ctx, `mid-run compaction failed: ${error.message}`);
+			},
+		});
+	});
+
+	pi.on("session_start", () => resetState(state));
+	pi.on("session_shutdown", () => resetState(state));
+}

package/src/native-fallback.ts

@@ -26,6 +26,7 @@ export type NativeFallbackResult =
 			ok: true;
 			result: CompactionResult;
 			model: { provider: string; id: string };
+			usage?: CompactionResult["usage"];
 	  }
 	| {
 			ok: false;
@@ -38,9 +39,21 @@ export type NativeFallbackResult =
 export type NativeCompactFn = typeof compact;
 
 type ResolvedAuth =
-	| { ok: true; apiKey?: string; headers?: Record<string, string>; env?: Record<string, string> }
+	| { ok: true; apiKey?: string; headers?: Record<string, string | null>; env?: Record<string, string> }
 	| { ok: false; error: string };
 
+/** Strip null-valued entries so downstream consumers receive a clean Record<string, string>. */
+function filterNullHeaders(headers: Record<string, string | null> | undefined): Record<string, string> | undefined {
+	if (!headers) return undefined;
+	const filtered: Record<string, string> = {};
+	for (const [key, value] of Object.entries(headers)) {
+		if (value !== null) {
+			filtered[key] = value;
+		}
+	}
+	return Object.keys(filtered).length > 0 ? filtered : undefined;
+}
+
 /** Parse "provider/model-id" (model ids may themselves contain slashes). */
 export function parseModelSpec(spec: string): ParsedModelSpec | undefined {
 	const trimmed = spec.trim();
@@ -82,6 +95,7 @@ export async function runNativeFallbackCompaction(args: {
 	event: SessionBeforeCompactEvent;
 	config: ExtensionConfig;
 	compactFn?: NativeCompactFn;
+	sessionId?: string;
 }): Promise<NativeFallbackResult> {
 	const { ctx, event, config } = args;
 	const compactFn = args.compactFn ?? compact;
@@ -120,12 +134,15 @@ export async function runNativeFallbackCompaction(args: {
 			event.preparation,
 			model,
 			auth.apiKey,
-			auth.headers,
+			filterNullHeaders(auth.headers),
 			event.customInstructions,
 			event.signal,
 			config.compactionThinkingLevel,
-			undefined,
+			undefined, // streamFn
 			auth.env,
+			undefined, // retry (use pi defaults)
+			undefined, // callbacks
+			args.sessionId,
 		);
 
 		if (event.signal.aborted) {
@@ -139,6 +156,7 @@ export async function runNativeFallbackCompaction(args: {
 			ok: true,
 			result,
 			model: { provider: model.provider, id: model.id },
+			usage: result.usage,
 		};
 	} catch (error) {
 		if (event.signal.aborted || isAbortError(error)) {