@lll9p/pi-better-compaction 0.2.1 → 0.5.0

package/src/retained-messages.ts

@@ -0,0 +1,98 @@
+/**
+ * V2 compaction retained-message filtering.
+ *
+ * After the API returns an encrypted compaction blob, V2 keeps a window of recent
+ * user/developer/system messages alongside the blob so the model retains explicit
+ * user instructions and context anchors. This module mirrors the filtering and
+ * budget logic from codex-rs `compact_remote_v2.rs`.
+ */
+
+/** Token budget for retained messages (matches codex-rs RETAINED_MESSAGE_TOKEN_BUDGET). */
+export const RETAINED_MESSAGE_TOKEN_BUDGET = 65_536;
+
+/** Messages larger than this are excluded even if they are of a retained role. */
+const MAX_SINGLE_ITEM_TOKENS = 10_000;
+
+/** Rough chars-per-token ratio for budget estimation. */
+const CHARS_PER_TOKEN = 4;
+
+function isRecord(value: unknown): value is Record<string, unknown> {
+	return !!value && typeof value === "object" && !Array.isArray(value);
+}
+
+/**
+ * Estimate the token count of an opaque input item by serializing to JSON and
+ * dividing by the chars-per-token ratio. This is intentionally rough — codex-rs
+ * uses a real tokenizer but we don't have one in the extension runtime.
+ */
+export function estimateItemTokens(item: unknown): number {
+	try {
+		const length = JSON.stringify(item).length;
+		return Math.ceil(length / CHARS_PER_TOKEN);
+	} catch {
+		return 0;
+	}
+}
+
+/**
+ * Whether an input item should be retained alongside the compaction blob.
+ *
+ * Retained roles (matching codex-rs `is_retained_for_remote_compaction_v2`):
+ * - `user` messages
+ * - `developer` messages
+ * - `system` messages
+ *
+ * Everything else (assistant, function_call, function_call_output, reasoning,
+ * compaction, etc.) is excluded.
+ */
+export function isRetainedItem(item: unknown): boolean {
+	if (!isRecord(item)) {
+		return false;
+	}
+
+	const role = item.role;
+	if (typeof role === "string") {
+		return role === "user" || role === "developer" || role === "system";
+	}
+
+	return false;
+}
+
+/**
+ * From a list of Responses API input items, select the most recent retained
+ * messages that fit within the token budget. Items are selected newest-first
+ * (reverse order) and returned in their original chronological order.
+ *
+ * Oversized individual items (> MAX_SINGLE_ITEM_TOKENS) are skipped.
+ */
+export function buildRetainedMessages(
+	input: readonly unknown[],
+	budget: number = RETAINED_MESSAGE_TOKEN_BUDGET,
+): unknown[] {
+	// Filter to retained roles first.
+	const retained: Array<{ index: number; item: unknown; tokens: number }> = [];
+	for (let i = 0; i < input.length; i++) {
+		const item = input[i];
+		if (!isRetainedItem(item)) continue;
+
+		const tokens = estimateItemTokens(item);
+		if (tokens > MAX_SINGLE_ITEM_TOKENS) continue;
+
+		retained.push({ index: i, item, tokens });
+	}
+
+	// Select from newest to oldest within the budget.
+	let remaining = budget;
+	const selected: typeof retained = [];
+
+	for (let i = retained.length - 1; i >= 0; i--) {
+		const entry = retained[i]!;
+		if (entry.tokens > remaining) break;
+		remaining -= entry.tokens;
+		selected.push(entry);
+	}
+
+	// Return in original chronological order.
+	selected.reverse();
+	return selected.map((entry) => structuredClone(entry.item));
+}

package/src/runtime.ts

@@ -4,6 +4,8 @@ import { RESPONSES_COMPACT_CAPABLE_APIS } from "./types";
 
 const OPENAI_COMPACT_PATH = "responses/compact";
 const CODEX_COMPACT_PATH = "codex/responses/compact";
+const OPENAI_RESPONSES_PATH = "responses";
+const CODEX_RESPONSES_PATH = "codex/responses";
 
 type ResponsesCompactApi = (typeof RESPONSES_COMPACT_CAPABLE_APIS)[number];
 
@@ -40,6 +42,7 @@ export type NativeCompactionRuntime = {
 	headers?: Record<string, string>;
 	compactPath: string;
 	compactUrl: string;
+	responsesUrl: string;
 	payload?: ResponsesCompatibleRequestPayload;
 	currentModel: RuntimeModel;
 };
@@ -74,6 +77,29 @@ export function normalizeBaseUrl(baseUrl: string | undefined | null): string | u
 	return normalized ? normalized : undefined;
 }
 
+function buildOpenAIResponsesUrl(baseUrl: string): string {
+	const normalized = normalizeBaseUrl(baseUrl) ?? baseUrl;
+	if (normalized.endsWith("/responses")) {
+		return normalized;
+	}
+	return `${normalized}/${OPENAI_RESPONSES_PATH}`;
+}
+
+function buildCodexResponsesUrl(baseUrl: string): string {
+	const normalized = normalizeBaseUrl(baseUrl) ?? baseUrl;
+	if (normalized.endsWith("/codex/responses")) {
+		return normalized;
+	}
+	if (normalized.endsWith("/codex")) {
+		return `${normalized}/responses`;
+	}
+	return `${normalized}/${CODEX_RESPONSES_PATH}`;
+}
+
+export function buildResponsesUrl(baseUrl: string, api: ResponsesCompactApi): string {
+	return api === "openai-codex-responses" ? buildCodexResponsesUrl(baseUrl) : buildOpenAIResponsesUrl(baseUrl);
+}
+
 function buildOpenAICompactUrl(baseUrl: string): string {
 	const normalized = normalizeBaseUrl(baseUrl) ?? baseUrl;
 	if (normalized.endsWith("/responses")) {
@@ -101,13 +127,25 @@ export function buildCompactPath(api: ResponsesCompactApi): string {
 	return api === "openai-codex-responses" ? CODEX_COMPACT_PATH : OPENAI_COMPACT_PATH;
 }
 
+/** 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;
+}
+
 async function resolveRequestAuth(
 	ctx: ExtensionContext,
 	model: RuntimeModel,
 ): Promise<{ apiKey?: string; headers?: Record<string, string> }> {
 	const modelRegistry = ctx.modelRegistry as {
 		getApiKeyAndHeaders?: (currentModel: RuntimeModel) => Promise<
-			| { ok: true; apiKey?: string; headers?: Record<string, string> }
+			| { ok: true; apiKey?: string; headers?: Record<string, string | null> }
 			| { ok: false; error: string }
 		>;
 	};
@@ -117,7 +155,7 @@ async function resolveRequestAuth(
 	}
 
 	const auth = await modelRegistry.getApiKeyAndHeaders(model);
-	return auth.ok ? { apiKey: auth.apiKey, headers: auth.headers } : {};
+	return auth.ok ? { apiKey: auth.apiKey, headers: filterNullHeaders(auth.headers) } : {};
 }
 
 export function isSupportedApi(api: string): api is ResponsesCompactApi {
@@ -234,6 +272,7 @@ export async function resolveNativeCompactionEnvironment(
 			headers,
 			compactPath: buildCompactPath(descriptor.api),
 			compactUrl: buildCompactUrl(descriptor.baseUrl, descriptor.api),
+			responsesUrl: buildResponsesUrl(descriptor.baseUrl, descriptor.api),
 			payload: requestPayload,
 			currentModel,
 		},

package/src/shared-headers.ts

@@ -0,0 +1,103 @@
+/**
+ * Shared HTTP header construction for V1 and V2 compaction clients.
+ *
+ * Extracts the header building logic that was previously private in
+ * compact-client.ts so both compact-client.ts and compact-client-v2.ts
+ * can share it without duplication.
+ */
+
+import type { NativeCompactionRuntime } from "./runtime";
+
+const JSON_CONTENT_TYPE = "application/json";
+
+function decodeJwtPayload(token: string): Record<string, unknown> | undefined {
+	const parts = token.split(".");
+	if (parts.length !== 3) {
+		return undefined;
+	}
+
+	try {
+		const payloadText = Buffer.from(parts[1]!, "base64url").toString("utf8");
+		const payload = JSON.parse(payloadText);
+		return payload && typeof payload === "object" && !Array.isArray(payload)
+			? (payload as Record<string, unknown>)
+			: undefined;
+	} catch {
+		return undefined;
+	}
+}
+
+function isRecord(value: unknown): value is Record<string, unknown> {
+	return !!value && typeof value === "object" && !Array.isArray(value);
+}
+
+function extractCodexAccountId(token: string): string | undefined {
+	const payload = decodeJwtPayload(token);
+	const authClaims = payload?.["https://api.openai.com/auth"];
+	if (!isRecord(authClaims)) {
+		return undefined;
+	}
+
+	const accountId = authClaims.chatgpt_account_id;
+	return typeof accountId === "string" && accountId.trim().length > 0 ? accountId.trim() : undefined;
+}
+
+function buildCodexUserAgent(): string {
+	const platform = typeof process !== "undefined" ? process.platform : "browser";
+	const arch = typeof process !== "undefined" ? process.arch : "unknown";
+	return `pi (${platform}; ${arch})`;
+}
+
+/**
+ * Build HTTP headers for a compaction request from the resolved runtime.
+ *
+ * Handles model-level headers, extension-resolved headers, authorization,
+ * and Codex-specific headers (account ID, originator, user-agent, beta flag).
+ *
+ * @param accept - The Accept header value. Defaults to `application/json`.
+ */
+export function toHeaders(
+	runtime: NativeCompactionRuntime,
+	accept: string = JSON_CONTENT_TYPE,
+): Record<string, string> {
+	const headers = new Headers();
+	// Model-level headers may contain null values (ProviderHeaders); null means "unset".
+	for (const [key, value] of Object.entries(runtime.currentModel.headers ?? {})) {
+		if (value != null) {
+			headers.set(key, String(value));
+		}
+	}
+	// Extension-resolved headers (already filtered by resolveRequestAuth, but defensive).
+	for (const [key, value] of Object.entries(runtime.headers ?? {})) {
+		if (value == null) {
+			headers.delete(key);
+		} else {
+			headers.set(key, value);
+		}
+	}
+	headers.set("accept", accept);
+	headers.set("content-type", JSON_CONTENT_TYPE);
+	if (!headers.has("authorization")) {
+		headers.set("authorization", `Bearer ${runtime.apiKey}`);
+	}
+
+	if (runtime.api === "openai-codex-responses") {
+		const accountId = extractCodexAccountId(runtime.apiKey);
+		if (accountId) {
+			headers.set("chatgpt-account-id", accountId);
+		}
+		headers.set("originator", "pi");
+		headers.set("user-agent", buildCodexUserAgent());
+		headers.set("openai-beta", "responses=experimental");
+	}
+
+	return Object.fromEntries(headers.entries());
+}
+
+/** Check whether an error represents an intentional abort (AbortController / AbortSignal). */
+export function isAbortError(error: unknown): boolean {
+	return (
+		(error instanceof DOMException && error.name === "AbortError") ||
+		(error instanceof Error && (error.name === "AbortError" || error.name === "ABORT_ERR"))
+	);
+}

package/src/types.ts

@@ -10,6 +10,7 @@ export const REDACTED_VALUE = "[REDACTED]";
  */
 export const RESPONSES_COMPACT_CAPABLE_APIS = ["openai-responses", "openai-codex-responses"] as const;
 export const NATIVE_COMPACTION_STRATEGY = "openai-native-compact-v1";
+export const NATIVE_COMPACTION_STRATEGY_V2 = "openai-native-compact-v2";
 /** Used as CompactionEntry.summary only when no summary text could be extracted from the compact response. */
 export const NATIVE_COMPACTION_FALLBACK_SUMMARY = "[OpenAI native compaction checkpoint]";
 
@@ -23,14 +24,23 @@ export const THINKING_LEVELS: readonly ThinkingLevel[] = [
 	"max",
 ];
 
+export type CompactionVersion = "v1" | "v2";
+export const COMPACTION_VERSIONS: readonly CompactionVersion[] = ["v1", "v2"];
+
 export type DebugArtifactKind =
 	| "provider-request"
 	| "compact-response"
 	| "compaction-event"
 	| "lifecycle";
 
+export type MidRunConfig = {
+	enabled: boolean;
+	thresholdPercent: number;
+};
+
 export type ExtensionConfig = {
 	enabled: boolean;
+	midRun: MidRunConfig;
 	/**
 	 * Allow a Responses session whose latest compaction was not created by this extension
 	 * to restart native compaction from Pi's current serialized session context.
@@ -45,6 +55,13 @@ export type ExtensionConfig = {
 	compactionThinkingLevel: ThinkingLevel;
 	/** Subset of RESPONSES_COMPACT_CAPABLE_APIS that should use the compact endpoint. */
 	responsesCompactApis: string[];
+	/**
+	 * Which compaction protocol to use for Responses-family APIs.
+	 * - "v2" (default): streaming CompactionTrigger via /responses endpoint.
+	 * - "v1": POST /responses/compact endpoint.
+	 * V2 failures automatically fall back to V1.
+	 */
+	compactionVersion: CompactionVersion;
 	notifyOnLoad: boolean;
 	debug: boolean;
 	logProviderPayloads: boolean;
@@ -97,6 +114,7 @@ export type RedactOptions = {
 };
 
 export type NativeCompactionStrategy = typeof NATIVE_COMPACTION_STRATEGY;
+export type NativeCompactionStrategyV2 = typeof NATIVE_COMPACTION_STRATEGY_V2;
 
 export type NativeCompactionRequestMeta = {
 	tokensBefore?: number;
@@ -111,7 +129,7 @@ export type NativeCompactionIdentity = {
 };
 
 export type NativeCompactionDetails = NativeCompactionIdentity & {
-	strategy: NativeCompactionStrategy;
+	strategy: NativeCompactionStrategy | NativeCompactionStrategyV2;
 	compactedWindow: unknown[];
 	compactResponseId?: string;
 	createdAt: string;
@@ -240,7 +258,7 @@ export function isNativeCompactionDetails(value: unknown): value is NativeCompac
 	}
 
 	return (
-		value.strategy === NATIVE_COMPACTION_STRATEGY &&
+		(value.strategy === NATIVE_COMPACTION_STRATEGY || value.strategy === NATIVE_COMPACTION_STRATEGY_V2) &&
 		isNativeCompactionIdentity(value) &&
 		Array.isArray(value.compactedWindow) &&
 		value.compactedWindow.every(isCompactedWindowItem) &&
@@ -254,9 +272,12 @@ export function isNativeCompactionEntry(value: unknown): value is NativeCompacti
 	return isRecord(value) && value.type === "compaction" && isNativeCompactionDetails(value.details);
 }
 
-export function createNativeCompactionDetails(input: CreateNativeCompactionDetailsInput): NativeCompactionDetails {
+export function createNativeCompactionDetails(
+	input: CreateNativeCompactionDetailsInput,
+	strategy: NativeCompactionStrategy | NativeCompactionStrategyV2 = NATIVE_COMPACTION_STRATEGY,
+): NativeCompactionDetails {
 	return {
-		strategy: NATIVE_COMPACTION_STRATEGY,
+		strategy,
 		provider: normalizeString(input.provider),
 		api: normalizeString(input.api),
 		model: normalizeString(input.model),
@@ -289,10 +310,15 @@ export function createNativeCompactionResult(
 
 export const DEFAULT_EXTENSION_CONFIG: ExtensionConfig = {
 	enabled: true,
+	midRun: {
+		enabled: false,
+		thresholdPercent: 80,
+	},
 	allowCompactionContinuityBreak: false,
 	compactionModel: undefined,
 	compactionThinkingLevel: "off",
 	responsesCompactApis: [...RESPONSES_COMPACT_CAPABLE_APIS],
+	compactionVersion: "v2",
 	notifyOnLoad: false,
 	debug: false,
 	logProviderPayloads: false,