pi-jev-compact 0.1.0

package/src/index.ts

@@ -0,0 +1,162 @@
+import { describeStats, serializePruned, toLibraryMessages } from "./adapter.ts";
+import { envConfig, type FastJevConfig, fileConfig, resolveFastJevConfig } from "./config.ts";
+import { JevClient } from "./jev/client.ts";
+import { compact, reductionRatio } from "./jev/compact.ts";
+import type { CompactResult } from "./jev/types.ts";
+
+export { describeStats, serializePruned, toLibraryMessages } from "./adapter.ts";
+export { envConfig, fileConfig, resolveFastJevConfig } from "./config.ts";
+export * from "./jev/index.ts";
+
+/** The config plus key the compaction handler runs with. */
+export type { FastJevConfig };
+
+export default function (pi: import("@earendil-works/pi-coding-agent").ExtensionAPI) {
+	const WIDGET = "fast-jev";
+
+	pi.on("session_shutdown", async (_event, ctx) => {
+		ctx.ui.setWidget(WIDGET, []);
+	});
+
+	pi.on("session_before_compact", async (event, ctx) => {
+		let config: FastJevConfig;
+		try {
+			// Project file only for trusted projects; project-local extensions are
+			// loaded after trust anyway, but global installs need the explicit check.
+			const trusted = typeof ctx.isProjectTrusted === "function" ? ctx.isProjectTrusted() : false;
+			config = resolveFastJevConfig(envConfig(), trusted ? fileConfig(ctx.cwd) : {});
+		} catch {
+			return undefined;
+		}
+		if (!config.apiKey) return undefined;
+
+		const { preparation } = event;
+		const summarized = [...preparation.messagesToSummarize];
+		if (preparation.isSplitTurn && preparation.turnPrefixMessages.length > 0) {
+			summarized.push(...preparation.turnPrefixMessages);
+		}
+		const transcript = toLibraryMessages(summarized);
+
+		try {
+			const result = await compact(
+				transcript,
+				new JevClient({
+					apiKey: config.apiKey,
+					...(config.model !== undefined && { model: config.model }),
+				}),
+				{
+					...(config.goal !== undefined && { goal: config.goal }),
+					...(config.keepThreshold !== undefined && { keepThreshold: config.keepThreshold }),
+					preserveRecentMessages: config.preserveRecentMessages,
+					...(config.maxStateTokens !== undefined && { maxStateTokens: config.maxStateTokens }),
+					...(config.maxRequestTokens !== undefined && {
+						maxRequestTokens: config.maxRequestTokens,
+					}),
+					...(config.truncateHeadChars !== undefined && {
+						truncateHeadChars: config.truncateHeadChars,
+					}),
+				},
+				event.signal,
+			);
+			const ratio = reductionRatio(result);
+			if (ratio < config.minReductionRatio) {
+				if (ctx.hasUI) {
+					ctx.ui.notify(
+						`fast-jev: fallback to built-in summary (below ${Math.round(
+							config.minReductionRatio * 100,
+						)}% minimum: ${describeStats(result.stats, ratio)})`,
+						"info",
+					);
+				}
+				return undefined;
+			}
+
+			const summary = serializePruned(preparation.previousSummary, result.messages);
+			if (ctx.hasUI) {
+				ctx.ui.notify(
+					`fast-jev: no summary, pruned history kept verbatim (${describeStats(result.stats, ratio)})`,
+					"info",
+				);
+				ctx.ui.setWidget(WIDGET, [
+					describeStats(result.stats, ratio),
+					...result.decisions
+						.filter((decision) => decision.reason !== "pinned" && decision.reason !== "kept")
+						.slice(0, 20)
+						.map(
+							(decision) =>
+								`${decision.id} ${decision.tool}: ${decision.action} (call=${decision.keepCall.toFixed(2)}, result=${decision.keepResult.toFixed(2)})`,
+						),
+				]);
+			}
+			return {
+				compaction: {
+					summary,
+					firstKeptEntryId: preparation.firstKeptEntryId,
+					tokensBefore: preparation.tokensBefore,
+					details: { engine: "fast-jev", decisions: result.decisions, stats: result.stats },
+				},
+			};
+		} catch (error) {
+			if (ctx.hasUI) {
+				const message = error instanceof Error ? error.message : String(error);
+				ctx.ui.notify(`fast-jev: fallback to built-in summary (${message})`, "error");
+			}
+			return undefined;
+		}
+	});
+
+	// On-demand trigger: `/jev-compact [instructions]` runs pi's internal
+	// compaction flow, so the session_before_compact handler above decides
+	// between the fast-jev path and the built-in summary.
+	pi.registerCommand("jev-compact", {
+		description: "Trigger compaction through the fast-jev path (falls back to built-in summary)",
+		handler: async (args, ctx) => {
+			ctx.compact({
+				...(args && args.trim().length > 0 && { customInstructions: args.trim() }),
+				onError: (error) => {
+					ctx.ui.notify(`fast-jev: compaction failed (${error.message})`, "error");
+				},
+			});
+		},
+	});
+
+	// Outcome reporting: fires for every compaction (fast-jev or built-in),
+	// with the saved entry's details carrying our stats when we ran.
+	pi.on("session_compact", async (event, ctx) => {
+		const { compactionEntry, fromExtension, reason } = event;
+		const details =
+			compactionEntry.details && typeof compactionEntry.details === "object"
+				? (compactionEntry.details as Record<string, unknown>)
+				: undefined;
+		const engine = details?.["engine"] === "fast-jev" ? "fast-jev" : "built-in";
+		const stats =
+			details?.["stats"] && typeof details["stats"] === "object"
+				? (details["stats"] as CompactResult["stats"])
+				: undefined;
+		const lines = [
+			`fast-jev: compaction done (${engine}${fromExtension ? ", extension-provided" : ""}, ${reason})`,
+		];
+		if (engine === "fast-jev" && stats) {
+			lines.push(
+				`messages ${stats.messagesAfter}/${stats.messagesBefore}, calls ${stats.kept} kept / ${stats.resultsDropped} truncated / ${stats.callsDropped} dropped / ${stats.pinned} pinned`,
+				`state ~${stats.stateTokens} tokens (${stats.stateStage}), ${stats.requests} request(s), ${stats.ms}ms`,
+			);
+		}
+		if (ctx.hasUI) {
+			ctx.ui.notify(lines.join(" — "), "info");
+			if (engine === "fast-jev") ctx.ui.setWidget(WIDGET, lines);
+		}
+	});
+
+	pi.on("session_compact_failed", async (event, ctx) => {
+		if (!ctx.hasUI) return;
+		if (event.aborted) {
+			ctx.ui.notify("fast-jev: compaction cancelled", "info");
+		} else {
+			ctx.ui.notify(
+				`fast-jev: compaction failed (${event.errorMessage ?? "unknown error"})`,
+				"error",
+			);
+		}
+	});
+}

package/src/jev/client.ts

@@ -0,0 +1,50 @@
+import { buildJevRequest, parseJevResponse } from "./request.ts";
+import type { JevAsker, JevQuestions, JevResponse, JevState } from "./types.ts";
+
+const ENV_KEY = "TYPESAFE_API_KEY";
+
+export interface JevClientOptions {
+	/** Defaults to `process.env.TYPESAFE_API_KEY`. */
+	apiKey?: string;
+	/** Defaults to `jev-latest`. */
+	model?: string;
+	/** Defaults to the System One endpoint. */
+	baseUrl?: string;
+	/** Defaults to the global `fetch`. */
+	fetch?: typeof fetch;
+}
+
+/** Asks Jev over HTTP with the global `fetch` (or an injected one). */
+export class JevClient implements JevAsker {
+	private readonly apiKey: string;
+	private readonly model: string | undefined;
+	private readonly baseUrl: string | undefined;
+	private readonly fetcher: typeof fetch;
+
+	constructor(options: JevClientOptions = {}) {
+		this.apiKey = options.apiKey ?? process.env[ENV_KEY] ?? "";
+		this.model = options.model;
+		this.baseUrl = options.baseUrl;
+		this.fetcher = options.fetch ?? fetch;
+	}
+
+	async ask(state: JevState, questions: JevQuestions, signal?: AbortSignal): Promise<JevResponse> {
+		if (!this.apiKey) throw new Error("TYPESAFE_API_KEY is not configured");
+		const request = buildJevRequest(
+			{
+				apiKey: this.apiKey,
+				...(this.model !== undefined && { model: this.model }),
+				...(this.baseUrl !== undefined && { baseUrl: this.baseUrl }),
+			},
+			state,
+			questions,
+		);
+		const response = await this.fetcher(request.url, {
+			method: request.method,
+			headers: request.headers,
+			body: request.body,
+			...(signal !== undefined && { signal }),
+		});
+		return parseJevResponse(response.status, response.ok, await response.text());
+	}
+}

package/src/jev/compact.ts

@@ -0,0 +1,298 @@
+import { noulAnswer } from "./request.ts";
+import { collectToolCalls, estimateTokens, fitState } from "./state.ts";
+import type {
+	CallAnswer,
+	CallDecision,
+	CompactionState,
+	CompactOptions,
+	CompactResult,
+	JevAsker,
+	JevQuestions,
+	Message,
+	ResolvedCompactOptions,
+	ToolCall,
+	ToolResult,
+	ToolUse,
+} from "./types.ts";
+
+export const DEFAULT_OPTIONS: ResolvedCompactOptions = {
+	goal: "",
+	keepThreshold: 0.5,
+	preserveRecentMessages: 6,
+	maxStateTokens: 25_000,
+	maxRequestTokens: 30_000,
+	truncateHeadChars: 300,
+};
+
+/** Tokens the request envelope (`model`, key names) adds around state and questions. */
+const REQUEST_OVERHEAD_TOKENS = 20;
+
+function finite(value: number | undefined, fallback: number): number {
+	return typeof value === "number" && Number.isFinite(value) ? value : fallback;
+}
+
+export function resolveOptions(options: CompactOptions = {}): ResolvedCompactOptions {
+	return {
+		goal: options.goal ?? DEFAULT_OPTIONS.goal,
+		keepThreshold: finite(options.keepThreshold, DEFAULT_OPTIONS.keepThreshold),
+		preserveRecentMessages: Math.max(
+			0,
+			Math.floor(finite(options.preserveRecentMessages, DEFAULT_OPTIONS.preserveRecentMessages)),
+		),
+		maxStateTokens: Math.max(1, finite(options.maxStateTokens, DEFAULT_OPTIONS.maxStateTokens)),
+		maxRequestTokens: Math.max(
+			1,
+			finite(options.maxRequestTokens, DEFAULT_OPTIONS.maxRequestTokens),
+		),
+		truncateHeadChars: Math.max(
+			0,
+			Math.floor(finite(options.truncateHeadChars, DEFAULT_OPTIONS.truncateHeadChars)),
+		),
+	};
+}
+
+/** The two `noul` questions asked about one call: keep the call, keep its result. */
+export function questionsFor(call: ToolCall): JevQuestions {
+	return {
+		[`call_${call.id}`]: {
+			type: "noul",
+			instructions: `Tool call ${call.id} (${call.tool}) should stay in the history: knowing this call was made, with its input, still matters for what the assistant does next`,
+		},
+		[`result_${call.id}`]: {
+			type: "noul",
+			instructions: `The full output of tool call ${call.id} (${call.tool}, ${call.resultChars} chars) should stay in the history verbatim: the assistant still needs its contents and re-running the tool would not do`,
+		},
+	};
+}
+
+/**
+ * Splits the candidate calls into batches whose questions, together with the
+ * (always complete) state, fit one request.
+ */
+export function batchCalls(
+	calls: readonly ToolCall[],
+	stateTokens: number,
+	options: Pick<ResolvedCompactOptions, "maxRequestTokens">,
+): ToolCall[][] {
+	const budget = options.maxRequestTokens - stateTokens - REQUEST_OVERHEAD_TOKENS;
+	const batches: ToolCall[][] = [];
+	let current: ToolCall[] = [];
+	let currentTokens = 0;
+	for (const call of calls) {
+		const tokens = estimateTokens(JSON.stringify(questionsFor(call)));
+		if (current.length > 0 && currentTokens + tokens > budget) {
+			batches.push(current);
+			current = [];
+			currentTokens = 0;
+		}
+		if (current.length === 0 && tokens > budget) {
+			throw new Error(
+				`state leaves no room for questions (~${stateTokens} of ${options.maxRequestTokens} tokens)`,
+			);
+		}
+		current.push(call);
+		currentTokens += tokens;
+	}
+	if (current.length > 0) batches.push(current);
+	return batches;
+}
+
+export function decideCall(
+	call: Pick<ToolCall, "id" | "tool" | "pinned">,
+	answer: CallAnswer,
+	options: Pick<ResolvedCompactOptions, "keepThreshold">,
+): CallDecision {
+	const base = { id: call.id, tool: call.tool, ...answer };
+	if (call.pinned) return { ...base, action: "keep", reason: "pinned" };
+	if (answer.keepResult >= options.keepThreshold) {
+		return { ...base, action: "keep", reason: "kept" };
+	}
+	if (answer.keepCall >= options.keepThreshold) {
+		return { ...base, action: "drop_result", reason: "result_dropped" };
+	}
+	return { ...base, action: "drop_call", reason: "call_dropped" };
+}
+
+async function askBatch(
+	asker: JevAsker,
+	state: CompactionState,
+	batch: readonly ToolCall[],
+	signal?: AbortSignal,
+): Promise<Map<string, CallAnswer>> {
+	const questions: JevQuestions = Object.assign({}, ...batch.map(questionsFor));
+	const { answers } = await asker.ask(state, questions, signal);
+	return new Map(
+		batch.map((call) => [
+			call.id,
+			{
+				keepCall: noulAnswer(answers, `call_${call.id}`),
+				keepResult: noulAnswer(answers, `result_${call.id}`),
+			},
+		]),
+	);
+}
+
+function truncatedResultText(text: string, isError: boolean, headChars: number): string {
+	if (text.length <= headChars + 120) return text;
+	const head = headChars > 0 ? `${text.slice(0, headChars)}\n` : "";
+	return `${head}[fast-jev-compaction truncated ${text.length - headChars} chars of this tool result${
+		isError ? " (error)" : ""
+	}; re-run the tool if needed]`;
+}
+
+/**
+ * Rebuilds the conversation from the decisions. A dropped call disappears
+ * together with its result; a dropped result keeps a bounded head and note.
+ * Messages that lose all their content are removed; untouched messages are
+ * returned as the same objects they came in as.
+ */
+export function applyDecisions(
+	messages: readonly Message[],
+	decisions: readonly CallDecision[],
+	calls: readonly ToolCall[],
+	headChars: number,
+): Message[] {
+	const byId = new Map(calls.map((call) => [call.id, call]));
+	const actions = new Map<string, CallDecision["action"]>();
+	for (const decision of decisions) {
+		const call = byId.get(decision.id);
+		if (call && decision.action !== "keep") actions.set(call.tool_use_id, decision.action);
+	}
+	const kept: Message[] = [];
+	for (const message of messages) {
+		const touched =
+			message.toolUses.some((tool) => actions.has(tool.tool_use_id)) ||
+			(message.toolResults ?? []).some((result) => actions.has(result.tool_use_id));
+		if (!touched) {
+			kept.push(message);
+			continue;
+		}
+		const toolUses = message.toolUses
+			.filter((tool) => actions.get(tool.tool_use_id) !== "drop_call")
+			.map((tool) => {
+				if (actions.get(tool.tool_use_id) !== "drop_result") return tool;
+				const text = truncatedResultText(tool.text ?? "", tool.isError ?? false, headChars);
+				if ((tool.text ?? "") === text) return tool;
+				const copy: ToolUse = {
+					tool_use_id: tool.tool_use_id,
+					tool: tool.tool,
+					input: tool.input,
+					text,
+				};
+				if (tool.isError) copy.isError = true;
+				return copy;
+			});
+		const toolResults = (message.toolResults ?? [])
+			.filter((result) => actions.get(result.tool_use_id) !== "drop_call")
+			.map((result) => {
+				if (actions.get(result.tool_use_id) !== "drop_result") return result;
+				const text = truncatedResultText(result.text, result.isError ?? false, headChars);
+				if (text === result.text) return result;
+				const copy: ToolResult = { tool_use_id: result.tool_use_id, text };
+				if (result.isError) copy.isError = result.isError;
+				return copy;
+			});
+		if (
+			!message.toolUses.some((tool) => actions.get(tool.tool_use_id) === "drop_call") &&
+			!(message.toolResults ?? []).some(
+				(result) => actions.get(result.tool_use_id) === "drop_call",
+			) &&
+			toolUses.every((tool, index) => tool === message.toolUses[index]) &&
+			toolResults.every((result, index) => result === message.toolResults?.[index])
+		) {
+			kept.push(message);
+			continue;
+		}
+		if (message.text.trim().length === 0 && toolUses.length === 0 && toolResults.length === 0) {
+			continue;
+		}
+		const rebuilt: Message = {
+			role: message.role,
+			text: message.text,
+			toolUses,
+		};
+		if (toolResults.length > 0) rebuilt.toolResults = toolResults;
+		kept.push(rebuilt);
+	}
+	return kept;
+}
+
+/** Characters of text, tool input and tool output a message holds. */
+export function messageChars(message: Message): number {
+	let total = message.text.length;
+	for (const tool of message.toolUses) {
+		try {
+			total += JSON.stringify(tool.input).length;
+		} catch {
+			total += 20;
+		}
+	}
+	for (const result of message.toolResults ?? []) total += result.text.length;
+	return total;
+}
+
+export function reductionRatio(result: Pick<CompactResult, "stats">): number {
+	const { charsBefore, charsAfter } = result.stats;
+	return charsBefore === 0 ? 0 : (charsBefore - charsAfter) / charsBefore;
+}
+
+function count(decisions: readonly CallDecision[], reason: CallDecision["reason"]): number {
+	return decisions.filter((decision) => decision.reason === reason).length;
+}
+
+/**
+ * Compacts a transcript by asking Jev, for every tool call outside the pinned
+ * first and newest messages, whether the call and whether its result must
+ * stay. The whole history (results omitted, fitted into `maxStateTokens`) is
+ * sent as state with every batch of questions. Throws when Jev fails or the
+ * history cannot be fitted; the caller decides whether to fall back.
+ */
+export async function compact(
+	messages: readonly Message[],
+	asker: JevAsker,
+	options: CompactOptions = {},
+	signal?: AbortSignal,
+): Promise<CompactResult> {
+	const started = Date.now();
+	const resolved = resolveOptions(options);
+	const calls = collectToolCalls(messages, resolved.preserveRecentMessages);
+	const candidates = calls.filter((call) => !call.pinned);
+	const charsBefore = messages.reduce((sum, message) => sum + messageChars(message), 0);
+
+	let fitted: { tokens: number; stage: string } = { tokens: 0, stage: "" };
+	let batches: ToolCall[][] = [];
+	const answers = new Map<string, CallAnswer>();
+	if (candidates.length > 0) {
+		const state = fitState(messages, calls, resolved);
+		fitted = state;
+		batches = batchCalls(candidates, state.tokens, resolved);
+		const answered = await Promise.all(
+			batches.map((batch) => askBatch(asker, state.state, batch, signal)),
+		);
+		for (const map of answered) for (const [id, answer] of map) answers.set(id, answer);
+	}
+
+	const decisions = calls.map((call) =>
+		decideCall(call, answers.get(call.id) ?? { keepCall: 1, keepResult: 1 }, resolved),
+	);
+	const kept = applyDecisions(messages, decisions, calls, resolved.truncateHeadChars);
+	return {
+		messages: kept,
+		decisions,
+		stats: {
+			messagesBefore: messages.length,
+			messagesAfter: kept.length,
+			charsBefore,
+			charsAfter: kept.reduce((sum, message) => sum + messageChars(message), 0),
+			calls: calls.length,
+			kept: count(decisions, "kept"),
+			resultsDropped: count(decisions, "result_dropped"),
+			callsDropped: count(decisions, "call_dropped"),
+			pinned: count(decisions, "pinned"),
+			stateTokens: fitted.tokens,
+			stateStage: fitted.stage,
+			requests: batches.length,
+			ms: Date.now() - started,
+		},
+	};
+}

package/src/jev/index.ts

@@ -0,0 +1,6 @@
+export * from "./client.ts";
+export * from "./compact.ts";
+export * from "./messages.ts";
+export * from "./request.ts";
+export * from "./state.ts";
+export * from "./types.ts";

package/src/jev/messages.ts

@@ -0,0 +1,14 @@
+import { JevClient, type JevClientOptions } from "./client.ts";
+import { compact } from "./compact.ts";
+import type { CompactOptions, CompactResult, Message } from "./types.ts";
+
+export type CompactMessagesOptions = CompactOptions & JevClientOptions;
+
+/** `compact` with a `JevClient` built from the options (key from `TYPESAFE_API_KEY` by default). */
+export function compactMessages(
+	messages: readonly Message[],
+	options: CompactMessagesOptions = {},
+	signal?: AbortSignal,
+): Promise<CompactResult> {
+	return compact(messages, new JevClient(options), options, signal);
+}

package/src/jev/request.ts

@@ -0,0 +1,73 @@
+import type { JevAnswer, JevQuestions, JevResponse, JevState } from "./types.ts";
+
+export const SYSTEM_ONE_URL = "https://api.typesafe.ai/v1/systemone";
+export const DEFAULT_MODEL = "jev-latest";
+
+export interface JevRequest {
+	url: string;
+	method: "POST";
+	headers: Record<string, string>;
+	body: string;
+}
+
+/** The HTTP request for one Jev call, for any fetch-like transport. */
+export function buildJevRequest(
+	params: {
+		apiKey: string;
+		model?: string;
+		baseUrl?: string;
+	},
+	state: JevState,
+	questions: JevQuestions,
+): JevRequest {
+	return {
+		url: params.baseUrl ?? SYSTEM_ONE_URL,
+		method: "POST",
+		headers: {
+			authorization: `Bearer ${params.apiKey}`,
+			"content-type": "application/json",
+		},
+		body: JSON.stringify({
+			model: params.model ?? DEFAULT_MODEL,
+			state,
+			questions,
+		}),
+	};
+}
+
+/** Validates a Jev response body; throws on anything but an `answers` object. */
+export function parseJevResponse(status: number, ok: boolean, text: string): JevResponse {
+	if (!ok) {
+		throw new Error(`Jev request failed (${status}): ${text.slice(0, 200)}`);
+	}
+	let parsed: unknown;
+	try {
+		parsed = JSON.parse(text);
+	} catch {
+		throw new Error("Jev returned malformed JSON");
+	}
+	if (
+		parsed === null ||
+		typeof parsed !== "object" ||
+		!("answers" in parsed) ||
+		parsed.answers === null ||
+		typeof parsed.answers !== "object"
+	) {
+		throw new Error("Jev response is missing answers");
+	}
+	return parsed as JevResponse;
+}
+
+/** The `noul` probability of one answer; throws when it is not there. */
+export function noulAnswer(answers: Record<string, JevAnswer>, name: string): number {
+	const answer = answers[name];
+	if (
+		!answer ||
+		!("noul" in answer) ||
+		typeof answer.noul !== "number" ||
+		!Number.isFinite(answer.noul)
+	) {
+		throw new Error(`Invalid Jev answer for ${name}`);
+	}
+	return answer.noul;
+}