256 lines
8.5 KiB
TypeScript
256 lines
8.5 KiB
TypeScript
/**
|
|
* Fast Jev Compaction Extension for Pi Coding Agent
|
|
*
|
|
* Uses Jev System One judgments to prune stale tool calls and tool results
|
|
* during session compaction without rewriting user or assistant messages.
|
|
*/
|
|
|
|
import type { AgentMessage } from "@earendil-works/pi-agent-core";
|
|
import type { ExtensionAPI } from "@earendil-works/pi-coding-agent";
|
|
import { readFileSync } from "node:fs";
|
|
import { homedir } from "node:os";
|
|
import { compactMessages } from "./src/messages.js";
|
|
import { reductionRatio } from "./src/compact.js";
|
|
import type { Message, ToolResult, ToolUse } from "./src/types.js";
|
|
|
|
/** Extract text content from a Pi content string or ContentBlock array */
|
|
function extractText(content: unknown): string {
|
|
if (typeof content === "string") return content;
|
|
if (Array.isArray(content)) {
|
|
return content
|
|
.map((block) => {
|
|
if (typeof block === "string") return block;
|
|
if (block && typeof block === "object") {
|
|
if ("text" in block && typeof block.text === "string") return block.text;
|
|
if ("thinking" in block && typeof block.thinking === "string") return block.thinking;
|
|
}
|
|
return "";
|
|
})
|
|
.filter(Boolean)
|
|
.join("\n");
|
|
}
|
|
return "";
|
|
}
|
|
|
|
/** Convert Pi AgentMessages to fast-jev-compaction Message array */
|
|
export function convertPiToJevMessages(messages: AgentMessage[]): Message[] {
|
|
const result: Message[] = [];
|
|
|
|
for (const msg of messages) {
|
|
if (msg.role === "user") {
|
|
result.push({
|
|
role: "user",
|
|
text: extractText((msg as any).content),
|
|
toolUses: [],
|
|
});
|
|
} else if (msg.role === "assistant") {
|
|
const assistantMsg = msg as any;
|
|
const content = Array.isArray(assistantMsg.content) ? assistantMsg.content : [];
|
|
|
|
const textParts: string[] = [];
|
|
const toolUses: ToolUse[] = [];
|
|
|
|
for (const block of content) {
|
|
if (!block || typeof block !== "object") continue;
|
|
if (block.type === "text" && typeof block.text === "string") {
|
|
textParts.push(block.text);
|
|
} else if (block.type === "thinking" && typeof block.thinking === "string") {
|
|
textParts.push(`[thinking: ${block.thinking.slice(0, 200)}...]`);
|
|
} else if (block.type === "toolCall") {
|
|
toolUses.push({
|
|
tool_use_id: block.id || `tool_${Math.random().toString(36).slice(2, 9)}`,
|
|
tool: block.name || "unknown",
|
|
input: block.arguments || {},
|
|
});
|
|
}
|
|
}
|
|
|
|
result.push({
|
|
role: "assistant",
|
|
text: textParts.join("\n"),
|
|
toolUses,
|
|
});
|
|
} else if (msg.role === "toolResult") {
|
|
const toolMsg = msg as any;
|
|
const tool_use_id = toolMsg.toolCallId || "unknown";
|
|
const text = extractText(toolMsg.content);
|
|
const isError = Boolean(toolMsg.isError);
|
|
|
|
result.push({
|
|
role: "user",
|
|
text: "",
|
|
toolUses: [],
|
|
toolResults: [
|
|
{
|
|
tool_use_id,
|
|
text,
|
|
isError,
|
|
},
|
|
],
|
|
});
|
|
} else if (msg.role === "bashExecution") {
|
|
const bashMsg = msg as any;
|
|
const tool_use_id = `bash_${Math.random().toString(36).slice(2, 9)}`;
|
|
const command = bashMsg.command || "";
|
|
const output = bashMsg.output || "";
|
|
const isError = (bashMsg.exitCode ?? 0) !== 0;
|
|
|
|
// Associate as assistant call + user result
|
|
result.push({
|
|
role: "assistant",
|
|
text: "",
|
|
toolUses: [
|
|
{
|
|
tool_use_id,
|
|
tool: "bash",
|
|
input: { command },
|
|
},
|
|
],
|
|
});
|
|
result.push({
|
|
role: "user",
|
|
text: "",
|
|
toolUses: [],
|
|
toolResults: [
|
|
{
|
|
tool_use_id,
|
|
text: output,
|
|
isError,
|
|
},
|
|
],
|
|
});
|
|
} else if (msg.role === "branchSummary" || msg.role === "compactionSummary") {
|
|
const sumMsg = msg as any;
|
|
result.push({
|
|
role: "user",
|
|
text: `[Previous Summary]: ${sumMsg.summary || ""}`,
|
|
toolUses: [],
|
|
});
|
|
}
|
|
}
|
|
|
|
return result;
|
|
}
|
|
|
|
/** Serialize the Jev-pruned message list into an exact, verbatim transcript.
|
|
* Dropped calls/results vanish; kept content stays exact (no LLM summarization). */
|
|
export function serializePrunedTranscript(messages: readonly Message[]): string {
|
|
const lines: string[] = [];
|
|
for (const m of messages) {
|
|
if (m.role === "user") {
|
|
const text = m.text.trim();
|
|
const results = (m.toolResults ?? []).map((r) =>
|
|
` [tool_result${r.isError ? " ERROR" : ""}] ${r.text}`
|
|
);
|
|
if (text) lines.push(`USER: ${text}`);
|
|
if (results.length) lines.push(...results);
|
|
} else if (m.role === "assistant") {
|
|
const text = m.text.trim();
|
|
if (text) lines.push(`ASSISTANT: ${text}`);
|
|
for (const t of m.toolUses) {
|
|
const input = JSON.stringify(t.input ?? {});
|
|
const clipped = input.length > 300 ? `${input.slice(0, 300)}…` : input;
|
|
lines.push(` [tool_call ${t.tool}] ${clipped}`);
|
|
}
|
|
}
|
|
}
|
|
return lines.join("\n");
|
|
}
|
|
|
|
/** Format decisions + stats as a compact header (audit trail only, content is the transcript). */
|
|
function formatJevHeader(result: any): string {
|
|
const { decisions, stats } = result;
|
|
const kept = decisions.filter((d: any) => d.action === "keep").length;
|
|
const truncated = decisions.filter((d: any) => d.action === "drop_result").length;
|
|
const dropped = decisions.filter((d: any) => d.action === "drop_call").length;
|
|
return (
|
|
`## Context pruned by fast-jev-compaction (Jev System One)\n\n` +
|
|
`Dropped ${dropped} stale tool call(s), truncated ${truncated} result(s), kept ${kept} intact. ` +
|
|
`Messages ${stats.messagesBefore}\u2192${stats.messagesAfter}, chars ${stats.charsBefore}\u2192${stats.charsAfter} ` +
|
|
`(${(reductionRatio(result) * 100).toFixed(1)}% reduction). User/assistant text is verbatim; nothing was summarized by an LLM.`
|
|
);
|
|
}
|
|
|
|
export default function fastJevCompactionExtension(pi: ExtensionAPI) {
|
|
/** Resolve the TypeSafe/Jev API key from env, falling back to the secrets file. */
|
|
function resolveJevKey(): string | undefined {
|
|
const fromEnv = process.env.JEV_API_KEY || process.env.TYPESAFE_API_KEY;
|
|
if (fromEnv) return fromEnv;
|
|
// Fallback: read the NixOS secrets file directly (works without a re-login).
|
|
try {
|
|
const path = `${homedir()}/.config/environment.d/10-secrets.conf`;
|
|
const content = readFileSync(path, "utf8");
|
|
const m = content.match(/^\s*(?:export\s+)?(?:JEV_API_KEY|TYPESAFE_API_KEY)\s*=\s*"?([^"\n]+)"?/m);
|
|
if (m?.[1]) return m[1].trim();
|
|
} catch {
|
|
/* ignore */
|
|
}
|
|
return undefined;
|
|
}
|
|
|
|
pi.on("session_before_compact", async (event, ctx) => {
|
|
const apiKey = resolveJevKey();
|
|
|
|
if (!apiKey) {
|
|
ctx.ui.notify("Fast-Jev: JEV_API_KEY not set, using default compaction", "warning");
|
|
return;
|
|
}
|
|
|
|
const { preparation, signal } = event;
|
|
const { messagesToSummarize, turnPrefixMessages, tokensBefore, firstKeptEntryId } = preparation;
|
|
|
|
const allPiMessages = [...messagesToSummarize, ...(turnPrefixMessages || [])];
|
|
|
|
if (allPiMessages.length === 0) {
|
|
return;
|
|
}
|
|
|
|
ctx.ui.notify(
|
|
`Fast-Jev Compaction: evaluating ${allPiMessages.length} messages with Jev System One...`,
|
|
"info"
|
|
);
|
|
|
|
try {
|
|
const jevMessages = convertPiToJevMessages(allPiMessages);
|
|
|
|
const result = await compactMessages(jevMessages, {
|
|
apiKey,
|
|
preserveRecentMessages: 4,
|
|
keepThreshold: 0.5,
|
|
});
|
|
|
|
const savedRatio = reductionRatio(result);
|
|
|
|
if (savedRatio < 0.15) {
|
|
ctx.ui.notify(
|
|
`Fast-Jev: low reduction (${(savedRatio * 100).toFixed(1)}%), using default compaction`,
|
|
"info"
|
|
);
|
|
return;
|
|
}
|
|
|
|
// The summary holds the EXACT pruned transcript (kept content verbatim,
|
|
// dropped tool calls/results removed) — not an LLM summary.
|
|
const transcript = serializePrunedTranscript(result.messages);
|
|
const summary = `${formatJevHeader(result)}\n\n---\n\n${transcript}`;
|
|
|
|
ctx.ui.notify(
|
|
`Fast-Jev Compaction complete: reduced ${(savedRatio * 100).toFixed(1)}% of tool content`,
|
|
"info"
|
|
);
|
|
|
|
return {
|
|
compaction: {
|
|
summary,
|
|
firstKeptEntryId,
|
|
tokensBefore,
|
|
},
|
|
};
|
|
} catch (error: any) {
|
|
const message = error instanceof Error ? error.message : String(error);
|
|
ctx.ui.notify(`Fast-Jev Compaction failed: ${message}, falling back to default`, "warning");
|
|
return;
|
|
}
|
|
});
|
|
}
|