mirror of
https://bitbucket.org/siakitem/my-pi.git
synced 2026-10-26 10:45:22 +00:00
原实现对所有工具结果(含 read)在下一次请求即替换为约 4K 的首尾存根, 不看上下文压力,导致模型看不到大文件中段;后台逐条语义摘要又反复改写 历史、打碎 prompt 缓存。参照 Anthropic 上下文编辑、Manus 与 JetBrains 观察遮蔽研究重构为: - L2:上下文达到阈值时在 turn_end 批量提交 Pi 原生 context_edit,把最近 keepRecentTurns 轮之外的旧 toolResult 换成确定性占位符;trigger/target 回差与 minClearTokens 保证少触发、一次清够;原始条目保留,可经 context_tree_query 取回 - L3:空闲 agent_settled 且超过绝对阈值时调用一次 ctx.compact - 阈值按绝对 token 计并受窗口比例约束,默认针对 1M 窗口(DeepSeek V4.1 Flash):遮蔽 200K→100K、整段摘要 400K - 删除逐条存根、语义 worker、soft/hard/target 比例与自行取消 Pi 压缩的门控; Pi 原生阈值压缩保留为最终安全网 Generated with [Devin](https://devin.ai) Co-Authored-By: Devin <158243242+devin-ai-integration[bot]@users.noreply.github.com>
686 lines
25 KiB
TypeScript
686 lines
25 KiB
TypeScript
import { createHash, randomUUID } from "node:crypto";
|
|
import { Type } from "typebox";
|
|
import {
|
|
buildSessionContext,
|
|
convertToLlm,
|
|
estimateTokens,
|
|
generateSummaryWithUsage,
|
|
sessionEntryToContextMessages,
|
|
truncateHead,
|
|
DEFAULT_MAX_BYTES,
|
|
DEFAULT_MAX_LINES,
|
|
type ExtensionAPI,
|
|
type ExtensionContext,
|
|
type SessionEntry,
|
|
} from "@earendil-works/pi-coding-agent";
|
|
import { DEFAULT_CONFIG, loadConfig } from "./src/config.js";
|
|
import {
|
|
applyContextEdits,
|
|
planMasking,
|
|
rebuildMaskedLedger,
|
|
type MaskedLedgerEntry,
|
|
} from "./src/masking.js";
|
|
import { buildNativeCompactionTransport, isNativeCompactionModel } from "./src/native-provider.js";
|
|
import {
|
|
buildCompactionRequestBody,
|
|
buildReplacementHistory,
|
|
buildToolPayload,
|
|
callRemoteCompaction,
|
|
findNativeCheckpoint,
|
|
isJsonObject,
|
|
mergeFeatureHeader,
|
|
messagesToResponseItems,
|
|
modelKey,
|
|
NATIVE_COMPACTION_KIND,
|
|
NATIVE_COMPACTION_VERSION,
|
|
stripInputFromPayload,
|
|
type JsonObject,
|
|
type NativeCompactionDetails,
|
|
type ResponseItem,
|
|
} from "./src/native-compaction.js";
|
|
import { contentText, SegmentIndex } from "./src/segment-index.js";
|
|
import {
|
|
CompactionServiceRegistry,
|
|
publishCompactionService,
|
|
unpublishCompactionService,
|
|
} from "./src/service.js";
|
|
import { effectiveThreshold } from "./src/token-accounting.js";
|
|
import {
|
|
COMPACTION_CUSTOM_TYPE,
|
|
COMPACTION_VERSION,
|
|
QUERY_TOOL_NAME,
|
|
type MaskedResultRecord,
|
|
type PiCompactionConfig,
|
|
type ProjectionSnapshot,
|
|
} from "./src/types.js";
|
|
|
|
interface SessionRuntime {
|
|
index: SegmentIndex;
|
|
config: PiCompactionConfig;
|
|
codexCache?: { fingerprint: string; modelKey: string; toolsHash: string; items: ResponseItem[] };
|
|
masked: Map<string, MaskedLedgerEntry>;
|
|
lastAutoCompactLeafId: string | null;
|
|
}
|
|
|
|
interface CachedPayload {
|
|
modelKey: string;
|
|
payload: JsonObject;
|
|
}
|
|
|
|
function errorMessage(error: unknown): string {
|
|
return error instanceof Error ? error.message : String(error);
|
|
}
|
|
|
|
interface EnvironmentValue {
|
|
hash: string;
|
|
tokens: number;
|
|
contextWindow: number;
|
|
}
|
|
|
|
interface EnvironmentCacheEntry extends EnvironmentValue {
|
|
prompt: string;
|
|
modelIdentity: string;
|
|
activeNames: string[];
|
|
tools: unknown[];
|
|
}
|
|
|
|
function hash(value: string): string {
|
|
return createHash("sha256").update(value).digest("hex").slice(0, 20);
|
|
}
|
|
|
|
function environment(
|
|
pi: ExtensionAPI,
|
|
ctx: ExtensionContext,
|
|
cache: Map<string, EnvironmentCacheEntry>,
|
|
): EnvironmentValue {
|
|
const sessionId = ctx.sessionManager.getSessionId();
|
|
const prompt = ctx.getSystemPrompt();
|
|
const activeNames = pi.getActiveTools();
|
|
const active = new Set(activeNames);
|
|
const tools = pi.getAllTools().filter((tool) => active.has(tool.name));
|
|
const contextWindow = ctx.model?.contextWindow ?? 0;
|
|
const modelIdentity = ctx.model ? modelKey(ctx.model) : "none";
|
|
const previous = cache.get(sessionId);
|
|
if (
|
|
previous
|
|
&& previous.prompt === prompt
|
|
&& previous.contextWindow === contextWindow
|
|
&& previous.modelIdentity === modelIdentity
|
|
&& previous.activeNames.length === activeNames.length
|
|
&& previous.activeNames.every((name, index) => name === activeNames[index])
|
|
&& previous.tools.length === tools.length
|
|
&& previous.tools.every((tool, index) => tool === tools[index])
|
|
) return previous;
|
|
|
|
const serializedTools = JSON.stringify(tools.map((tool) => ({
|
|
name: tool.name,
|
|
description: tool.description,
|
|
parameters: tool.parameters,
|
|
})));
|
|
const text = `${modelIdentity}\n${prompt}\n${serializedTools}`;
|
|
const value: EnvironmentCacheEntry = {
|
|
prompt,
|
|
modelIdentity,
|
|
activeNames: [...activeNames],
|
|
tools,
|
|
hash: hash(text),
|
|
tokens: Math.ceil(text.length / 4),
|
|
contextWindow,
|
|
};
|
|
cache.set(sessionId, value);
|
|
return value;
|
|
}
|
|
|
|
function portableSummary(messages: readonly any[], maxChars = 24_000): string {
|
|
const sections: string[] = ["Portable context checkpoint generated by pi-compaction."];
|
|
for (const message of messages) {
|
|
if (message?.role === "toolResult") {
|
|
const text = Array.isArray(message.content)
|
|
? message.content.flatMap((part: any) => part?.type === "text" ? [part.text] : []).join("\n")
|
|
: "";
|
|
if (text) sections.push(`[Tool ${message.toolName ?? "unknown"}] ${text}`);
|
|
continue;
|
|
}
|
|
if (message?.role === "user") {
|
|
const text = typeof message.content === "string"
|
|
? message.content
|
|
: Array.isArray(message.content)
|
|
? message.content.flatMap((part: any) => part?.type === "text" ? [part.text] : []).join("\n")
|
|
: "";
|
|
if (text) sections.push(`[User] ${text}`);
|
|
continue;
|
|
}
|
|
if (message?.role === "assistant" && Array.isArray(message.content)) {
|
|
const text = message.content.flatMap((part: any) => part?.type === "text" ? [part.text] : []).join("\n");
|
|
if (text) sections.push(`[Assistant] ${text}`);
|
|
}
|
|
}
|
|
const combined = sections.join("\n\n");
|
|
if (combined.length <= maxChars) return combined;
|
|
const head = Math.floor(maxChars * 0.35);
|
|
const tail = maxChars - head - 32;
|
|
return `${combined.slice(0, head)}\n\n… checkpoint truncated …\n\n${combined.slice(-tail)}`;
|
|
}
|
|
|
|
function asPortableMessage(summary: string): any {
|
|
return {
|
|
role: "user",
|
|
content: [{ type: "text", text: summary }],
|
|
timestamp: Date.now(),
|
|
};
|
|
}
|
|
|
|
function setFeatureHeader(headers: Record<string, string | null>): void {
|
|
const existing = Object.entries(headers).find(([name]) => name.toLowerCase() === "x-codex-beta-features");
|
|
if (existing) headers[existing[0]] = mergeFeatureHeader(existing[1]);
|
|
else headers["x-codex-beta-features"] = mergeFeatureHeader(undefined);
|
|
}
|
|
|
|
function cleanHeaders(headers: Record<string, string | null> | undefined): Record<string, string> | undefined {
|
|
if (!headers) return undefined;
|
|
return Object.fromEntries(Object.entries(headers).filter((entry): entry is [string, string] => typeof entry[1] === "string"));
|
|
}
|
|
|
|
function activeToolsHash(pi: ExtensionAPI): string {
|
|
return hash(pi.getActiveTools().join("\0"));
|
|
}
|
|
|
|
export default function piCompactionExtension(pi: ExtensionAPI): void {
|
|
const runtimes = new Map<string, SessionRuntime>();
|
|
const payloads = new Map<string, CachedPayload>();
|
|
const environments = new Map<string, EnvironmentCacheEntry>();
|
|
const service = new CompactionServiceRegistry();
|
|
publishCompactionService(service);
|
|
|
|
const maskedRecords = (sessionId: string, ids: string[]): MaskedResultRecord[] => {
|
|
const masked = runtimes.get(sessionId)?.masked;
|
|
if (!masked) return [];
|
|
return ids.flatMap((id) => {
|
|
const entry = masked.get(id);
|
|
return entry
|
|
? [{
|
|
toolCallId: id,
|
|
resultTimestamp: entry.resultTimestamp,
|
|
deterministicSummary: entry.placeholder,
|
|
}]
|
|
: [];
|
|
});
|
|
};
|
|
|
|
const runtimeFor = (ctx: ExtensionContext): SessionRuntime => {
|
|
const sessionId = ctx.sessionManager.getSessionId();
|
|
let runtime = runtimes.get(sessionId);
|
|
if (!runtime) {
|
|
runtime = {
|
|
index: new SegmentIndex(sessionId),
|
|
config: structuredClone(DEFAULT_CONFIG),
|
|
masked: rebuildLedger(ctx),
|
|
lastAutoCompactLeafId: null,
|
|
};
|
|
runtimes.set(sessionId, runtime);
|
|
service.setSession(sessionId, (ids) => maskedRecords(sessionId, ids));
|
|
}
|
|
return runtime;
|
|
};
|
|
|
|
const ledgerStats = (runtime: SessionRuntime): { maskedResults: number; reclaimedTokens: number } => {
|
|
let reclaimedTokens = 0;
|
|
for (const entry of runtime.masked.values()) {
|
|
reclaimedTokens += Math.max(0, entry.originalTokens - entry.placeholderTokens);
|
|
}
|
|
return { maskedResults: runtime.masked.size, reclaimedTokens };
|
|
};
|
|
|
|
const projectMessages = (
|
|
ctx: ExtensionContext,
|
|
runtime: SessionRuntime,
|
|
messages: readonly any[],
|
|
): ProjectionSnapshot => {
|
|
const env = environment(pi, ctx, environments);
|
|
const stats = ledgerStats(runtime);
|
|
const snapshot = runtime.index.project(messages, {
|
|
environmentHash: env.hash,
|
|
environmentTokens: env.tokens,
|
|
contextWindow: env.contextWindow,
|
|
maskedResults: stats.maskedResults,
|
|
reclaimedTokens: stats.reclaimedTokens,
|
|
});
|
|
service.setSnapshot(snapshot);
|
|
return snapshot;
|
|
};
|
|
|
|
const rebuildLedger = (ctx: ExtensionContext): Map<string, MaskedLedgerEntry> =>
|
|
rebuildMaskedLedger(
|
|
ctx.sessionManager.getBranch(),
|
|
ctx.sessionManager.getLeafId(),
|
|
(message) => estimateTokens(message as any),
|
|
);
|
|
|
|
const hydrateRuntime = (ctx: ExtensionContext, runtime: SessionRuntime): void => {
|
|
runtime.index.hydrate();
|
|
runtime.masked = rebuildLedger(ctx);
|
|
runtime.lastAutoCompactLeafId = null;
|
|
runtime.codexCache = undefined;
|
|
};
|
|
|
|
const branchMessages = (
|
|
ctx: ExtensionContext,
|
|
runtime: SessionRuntime,
|
|
branch: SessionEntry[],
|
|
excludeLastAssistantError = false,
|
|
): { snapshot: ProjectionSnapshot; checkpoint: ReturnType<typeof findNativeCheckpoint> } => {
|
|
const checkpoint = findNativeCheckpoint(branch);
|
|
let messages: any[];
|
|
if (checkpoint.status === "valid") {
|
|
// sessionEntryToContextMessages sees only raw entries; applyContextEdits
|
|
// additionally replays the latest context_edit per target (omissions and
|
|
// placeholder replacements included).
|
|
messages = applyContextEdits(branch.slice(checkpoint.checkpoint.entryIndex + 1));
|
|
if (checkpoint.checkpoint.details.modelKey !== (ctx.model ? modelKey(ctx.model) : "")) {
|
|
messages = [asPortableMessage(checkpoint.checkpoint.details.portableSummary), ...messages];
|
|
}
|
|
} else {
|
|
messages = buildSessionContext(branch).messages;
|
|
}
|
|
if (excludeLastAssistantError) {
|
|
const index = messages.findLastIndex((message) => message.role === "assistant");
|
|
if (index >= 0 && messages[index]?.stopReason === "error") messages = messages.toSpliced(index, 1);
|
|
}
|
|
return { snapshot: projectMessages(ctx, runtime, messages), checkpoint };
|
|
};
|
|
|
|
const codexInput = (
|
|
piModel: any,
|
|
runtime: SessionRuntime,
|
|
snapshot: ProjectionSnapshot,
|
|
checkpoint?: NativeCompactionDetails,
|
|
): ResponseItem[] => {
|
|
const toolsHash = activeToolsHash(pi);
|
|
const key = modelKey(piModel);
|
|
const cached = runtime.codexCache;
|
|
if (cached && cached.fingerprint === snapshot.fingerprint && cached.modelKey === key && cached.toolsHash === toolsHash) {
|
|
return cached.items;
|
|
}
|
|
const tail = messagesToResponseItems(piModel, convertToLlm(snapshot.messages as any), pi.getAllTools());
|
|
const items = checkpoint && checkpoint.modelKey === key
|
|
? [...structuredClone(checkpoint.replacementHistory), ...tail]
|
|
: tail;
|
|
runtime.codexCache = { fingerprint: snapshot.fingerprint, modelKey: key, toolsHash, items };
|
|
snapshot.codexItems = items;
|
|
return items;
|
|
};
|
|
|
|
pi.on("session_start", async (_event, ctx) => {
|
|
publishCompactionService(service);
|
|
const sessionId = ctx.sessionManager.getSessionId();
|
|
const config = await loadConfig();
|
|
const runtime = runtimes.get(sessionId) ?? {
|
|
index: new SegmentIndex(sessionId),
|
|
config,
|
|
masked: new Map<string, MaskedLedgerEntry>(),
|
|
lastAutoCompactLeafId: null,
|
|
};
|
|
runtime.config = config;
|
|
runtime.index.hydrate();
|
|
runtime.masked = rebuildLedger(ctx);
|
|
runtime.lastAutoCompactLeafId = null;
|
|
runtimes.set(sessionId, runtime);
|
|
service.setSession(sessionId, (ids) => maskedRecords(sessionId, ids));
|
|
});
|
|
|
|
pi.on("session_shutdown", (_event, ctx) => {
|
|
const sessionId = ctx.sessionManager.getSessionId();
|
|
runtimes.delete(sessionId);
|
|
payloads.delete(sessionId);
|
|
environments.delete(sessionId);
|
|
service.deleteSession(sessionId);
|
|
if (runtimes.size === 0) unpublishCompactionService(service);
|
|
});
|
|
|
|
pi.on("session_tree", (_event, ctx) => {
|
|
const runtime = runtimeFor(ctx);
|
|
hydrateRuntime(ctx, runtime);
|
|
service.clearSnapshot(ctx.sessionManager.getSessionId());
|
|
});
|
|
|
|
pi.on("session_compact", (_event, ctx) => {
|
|
const runtime = runtimeFor(ctx);
|
|
hydrateRuntime(ctx, runtime);
|
|
service.clearSnapshot(ctx.sessionManager.getSessionId());
|
|
payloads.delete(ctx.sessionManager.getSessionId());
|
|
});
|
|
|
|
pi.on("model_select", (_event, ctx) => {
|
|
const sessionId = ctx.sessionManager.getSessionId();
|
|
payloads.delete(sessionId);
|
|
environments.delete(sessionId);
|
|
const runtime = runtimes.get(sessionId);
|
|
if (runtime) runtime.codexCache = undefined;
|
|
});
|
|
|
|
pi.on("turn_end", (event, ctx) => {
|
|
const runtime = runtimeFor(ctx);
|
|
if (!runtime.config.enabled) return undefined;
|
|
|
|
const previous = runtime.index.getSnapshot();
|
|
const usage = ctx.getContextUsage();
|
|
if (usage?.tokens && previous) {
|
|
const env = environment(pi, ctx, environments);
|
|
runtime.index.recordUsage(
|
|
usage.tokens,
|
|
previous.messageEstimate + estimateTokens(event.message as any),
|
|
env.hash,
|
|
);
|
|
}
|
|
|
|
const currentTokens = usage?.tokens
|
|
?? (event.context?.contextMessages ?? []).reduce((sum, message) => sum + estimateTokens(message as any), 0)
|
|
+ environment(pi, ctx, environments).tokens;
|
|
const plan = planMasking({
|
|
contextEntries: event.context?.contextEntries ?? [],
|
|
currentTokens,
|
|
contextWindow: ctx.model?.contextWindow ?? 0,
|
|
masking: runtime.config.masking,
|
|
estimate: (message) => estimateTokens(message as any),
|
|
});
|
|
if (!plan.drafts.length) return undefined;
|
|
for (const result of plan.masked) {
|
|
runtime.masked.set(result.toolCallId, {
|
|
placeholder: result.placeholder,
|
|
originalTokens: result.originalTokens,
|
|
placeholderTokens: result.placeholderTokens,
|
|
resultTimestamp: result.resultTimestamp,
|
|
});
|
|
}
|
|
return { entries: [...(event.entries ?? []), ...plan.drafts] };
|
|
});
|
|
|
|
pi.on("agent_settled", (_event, ctx) => {
|
|
const runtime = runtimeFor(ctx);
|
|
if (!runtime.config.enabled || ctx.isIdle?.() === false) return;
|
|
const snapshot = runtime.index.getSnapshot();
|
|
const tokens = ctx.getContextUsage()?.tokens ?? snapshot?.effectiveTokens;
|
|
if (tokens == null) return;
|
|
const contextWindow = ctx.model?.contextWindow ?? snapshot?.contextWindow ?? 0;
|
|
const threshold = effectiveThreshold(
|
|
runtime.config.summary.triggerTokens,
|
|
runtime.config.summary.triggerRatio,
|
|
contextWindow,
|
|
);
|
|
if (tokens < threshold) return;
|
|
const leafId = ctx.sessionManager.getLeafId();
|
|
if (leafId !== null && runtime.lastAutoCompactLeafId === leafId) return;
|
|
runtime.lastAutoCompactLeafId = leafId;
|
|
ctx.compact({
|
|
onError: (error) => {
|
|
if (ctx.hasUI) ctx.ui.notify(`pi-compaction auto summary failed: ${errorMessage(error)}`, "warning");
|
|
},
|
|
});
|
|
});
|
|
|
|
pi.on("context", (event, ctx) => {
|
|
const runtime = runtimeFor(ctx);
|
|
const branch = ctx.sessionManager.getBranch() as SessionEntry[];
|
|
const checkpoint = findNativeCheckpoint(branch);
|
|
if (checkpoint.status === "invalid") {
|
|
ctx.abort();
|
|
if (ctx.hasUI) ctx.ui.notify("pi-compaction checkpoint is malformed; request blocked.", "error");
|
|
return { messages: [] };
|
|
}
|
|
if (checkpoint.status === "valid") {
|
|
const projected = branchMessages(ctx, runtime, branch);
|
|
return { messages: projected.snapshot.messages as any };
|
|
}
|
|
// Non-checkpoint path only refreshes the shared snapshot; the message list
|
|
// is returned to Pi unchanged — masking is append-only via context_edit.
|
|
projectMessages(ctx, runtime, event.messages);
|
|
return undefined;
|
|
});
|
|
|
|
pi.on("before_provider_headers", (event, ctx) => {
|
|
if (isNativeCompactionModel(ctx.model)) setFeatureHeader(event.headers);
|
|
});
|
|
|
|
pi.on("before_provider_request", (event, ctx) => {
|
|
const model = ctx.model;
|
|
if (!model || !isNativeCompactionModel(model) || !isJsonObject(event.payload)) return undefined;
|
|
const runtime = runtimeFor(ctx);
|
|
const sessionId = ctx.sessionManager.getSessionId();
|
|
payloads.set(sessionId, { modelKey: modelKey(model), payload: stripInputFromPayload(event.payload) });
|
|
const checkpoint = findNativeCheckpoint(ctx.sessionManager.getBranch() as SessionEntry[]);
|
|
if (checkpoint.status !== "valid" || checkpoint.checkpoint.details.modelKey !== modelKey(model)) return undefined;
|
|
const snapshot = runtime.index.getSnapshot();
|
|
if (!snapshot) return undefined;
|
|
const payload: JsonObject = {
|
|
...event.payload,
|
|
input: codexInput(model, runtime, snapshot, checkpoint.checkpoint.details),
|
|
};
|
|
delete payload.messages;
|
|
delete payload.previous_response_id;
|
|
return payload;
|
|
});
|
|
|
|
pi.on("session_before_compact", async (event, ctx) => {
|
|
const runtime = runtimeFor(ctx);
|
|
if (!runtime.config.enabled) return undefined;
|
|
const { snapshot, checkpoint } = branchMessages(
|
|
ctx,
|
|
runtime,
|
|
event.branchEntries as SessionEntry[],
|
|
event.reason === "overflow" && event.willRetry,
|
|
);
|
|
|
|
const model = ctx.model;
|
|
if (!model) return undefined;
|
|
if (!runtime.config.native.enabled || !isNativeCompactionModel(model)) {
|
|
try {
|
|
const auth = await ctx.modelRegistry.getApiKeyAndHeaders(model);
|
|
if (!auth.ok) return undefined;
|
|
const result = await generateSummaryWithUsage(
|
|
snapshot.messages as any,
|
|
model,
|
|
16_384,
|
|
auth.apiKey,
|
|
cleanHeaders(auth.headers),
|
|
event.signal,
|
|
event.customInstructions,
|
|
undefined,
|
|
runtime.config.summary.thinking,
|
|
);
|
|
return {
|
|
compaction: {
|
|
summary: result.text,
|
|
firstKeptEntryId: event.preparation.firstKeptEntryId,
|
|
tokensBefore: snapshot.effectiveTokens,
|
|
usage: result.usage,
|
|
details: {
|
|
kind: COMPACTION_CUSTOM_TYPE,
|
|
version: COMPACTION_VERSION,
|
|
checkpoint: "portable",
|
|
sourceFingerprint: snapshot.fingerprint,
|
|
},
|
|
},
|
|
};
|
|
} catch (error) {
|
|
if (!event.signal.aborted) console.error("pi-compaction: portable checkpoint failed:", error);
|
|
return undefined;
|
|
}
|
|
}
|
|
|
|
try {
|
|
const auth = await ctx.modelRegistry.getApiKeyAndHeaders(model);
|
|
if (!auth.ok) throw new Error(auth.error);
|
|
const currentCheckpoint = checkpoint.status === "valid" ? checkpoint.checkpoint.details : undefined;
|
|
const input = codexInput(model, runtime, snapshot, currentCheckpoint);
|
|
const sessionId = ctx.sessionManager.getSessionId();
|
|
const cached = payloads.get(sessionId);
|
|
const body = buildCompactionRequestBody({
|
|
basePayload: cached?.modelKey === modelKey(model) ? cached.payload : undefined,
|
|
model,
|
|
input,
|
|
instructions: ctx.getSystemPrompt(),
|
|
tools: buildToolPayload(pi.getAllTools(), pi.getActiveTools()),
|
|
sessionId,
|
|
});
|
|
if (ctx.hasUI) ctx.ui.notify(`pi-compaction: creating ${model.provider} native checkpoint...`, "info");
|
|
const remote = await callRemoteCompaction({
|
|
...buildNativeCompactionTransport({ model, apiKey: auth.apiKey, headers: cleanHeaders(auth.headers), sessionId }),
|
|
body,
|
|
signal: event.signal,
|
|
});
|
|
const details: NativeCompactionDetails = {
|
|
kind: NATIVE_COMPACTION_KIND,
|
|
version: NATIVE_COMPACTION_VERSION,
|
|
checkpoint: "native",
|
|
modelKey: modelKey(model),
|
|
replacementHistory: buildReplacementHistory(
|
|
input,
|
|
remote.compactionItem,
|
|
Math.min(64_000, Math.max(4_096, Math.floor(snapshot.contextWindow * 0.25))),
|
|
),
|
|
portableSummary: portableSummary(snapshot.messages),
|
|
sourceFingerprint: snapshot.fingerprint,
|
|
};
|
|
return {
|
|
compaction: {
|
|
summary: `Pi native checkpoint ${randomUUID()}`,
|
|
firstKeptEntryId: event.preparation.firstKeptEntryId,
|
|
tokensBefore: snapshot.effectiveTokens,
|
|
usage: remote.usage,
|
|
details,
|
|
},
|
|
};
|
|
} catch (error) {
|
|
if (!event.signal.aborted && ctx.hasUI) {
|
|
ctx.ui.notify(`pi-compaction native checkpoint failed; falling back to Pi: ${errorMessage(error)}`, "warning");
|
|
}
|
|
return undefined;
|
|
}
|
|
});
|
|
|
|
pi.on("session_before_tree", async (event, ctx) => {
|
|
const model = ctx.model;
|
|
if (!event.preparation.userWantsSummary || !model) return undefined;
|
|
const runtime = runtimeFor(ctx);
|
|
if (!runtime.config.enabled) return undefined;
|
|
const messages = event.preparation.entriesToSummarize.flatMap((entry) => sessionEntryToContextMessages(entry));
|
|
const snapshot = projectMessages(ctx, runtime, messages);
|
|
try {
|
|
const auth = await ctx.modelRegistry.getApiKeyAndHeaders(model);
|
|
if (!auth.ok) return undefined;
|
|
const result = await generateSummaryWithUsage(
|
|
snapshot.messages as any,
|
|
model,
|
|
16_384,
|
|
auth.apiKey,
|
|
cleanHeaders(auth.headers),
|
|
event.signal,
|
|
event.preparation.customInstructions,
|
|
undefined,
|
|
runtime.config.summary.thinking,
|
|
undefined,
|
|
undefined,
|
|
undefined,
|
|
undefined,
|
|
);
|
|
return {
|
|
summary: {
|
|
summary: result.text,
|
|
usage: result.usage,
|
|
details: {
|
|
kind: COMPACTION_CUSTOM_TYPE,
|
|
version: COMPACTION_VERSION,
|
|
checkpoint: "portable",
|
|
sourceFingerprint: snapshot.fingerprint,
|
|
},
|
|
},
|
|
};
|
|
} catch (error) {
|
|
if (!event.signal.aborted) console.error("pi-compaction: branch summary failed:", error);
|
|
return undefined;
|
|
}
|
|
});
|
|
|
|
pi.registerTool({
|
|
name: QUERY_TOOL_NAME,
|
|
label: "Query Compacted Tool History",
|
|
description: "Recover the original output of a tool result that pi-compaction cleared from context, by tool call id.",
|
|
promptSnippet: "Recover original tool output cleared from context",
|
|
promptGuidelines: [
|
|
"Use context_tree_query when you see a \"[pi-compaction] Earlier\" placeholder and still need the original tool output.",
|
|
],
|
|
parameters: Type.Object({
|
|
toolCallIds: Type.Array(Type.String(), { minItems: 1, maxItems: 20 }),
|
|
}),
|
|
async execute(_id, params, _signal, _update, ctx) {
|
|
const branch = ctx.sessionManager.getBranch() as SessionEntry[];
|
|
const calls = new Map<string, { toolName: string; args: unknown }>();
|
|
const results = new Map<string, any>();
|
|
for (const entry of branch) {
|
|
const message = (entry as any)?.message;
|
|
if ((entry as any)?.type !== "message" || !message) continue;
|
|
if (message.role === "assistant" && Array.isArray(message.content)) {
|
|
for (const block of message.content) {
|
|
if (block?.type === "toolCall" && typeof block.id === "string") {
|
|
calls.set(block.id, {
|
|
toolName: String(block.name ?? "unknown"),
|
|
args: block.arguments ?? block.input ?? block.args ?? {},
|
|
});
|
|
}
|
|
}
|
|
} else if (message.role === "toolResult" && typeof message.toolCallId === "string") {
|
|
results.set(message.toolCallId, message);
|
|
}
|
|
}
|
|
|
|
const found: string[] = [];
|
|
const blocks: string[] = [];
|
|
for (const id of params.toolCallIds) {
|
|
const result = results.get(id);
|
|
if (!result) continue;
|
|
found.push(id);
|
|
const call = calls.get(id) ?? { toolName: String(result.toolName ?? "unknown"), args: {} };
|
|
const perResultBytes = Math.max(1_024, Math.floor(DEFAULT_MAX_BYTES / Math.max(1, params.toolCallIds.length)));
|
|
const args = truncateHead(JSON.stringify(call.args ?? {}), { maxLines: 20, maxBytes: 2_048 }).content;
|
|
const rawText = contentText(result.content);
|
|
const truncated = truncateHead(rawText, {
|
|
maxLines: Math.max(20, Math.floor(DEFAULT_MAX_LINES / Math.max(1, params.toolCallIds.length))),
|
|
maxBytes: perResultBytes,
|
|
});
|
|
blocks.push([
|
|
`## ${id}`,
|
|
`Tool: ${call.toolName}`,
|
|
`Args: ${args}`,
|
|
`Status: ${result.isError ? "ERROR" : "OK"}`,
|
|
"",
|
|
truncated.content,
|
|
...(truncated.truncated ? [`[truncated: ${truncated.outputLines}/${truncated.totalLines} lines]`] : []),
|
|
].join("\n"));
|
|
}
|
|
const missing = params.toolCallIds.filter((id: string) => !results.has(id));
|
|
for (const id of missing) blocks.push(`## ${id}\n(not found)`);
|
|
return {
|
|
content: [{ type: "text", text: blocks.join("\n\n---\n\n") }],
|
|
details: { toolCallIds: found, missing },
|
|
};
|
|
},
|
|
});
|
|
|
|
pi.registerCommand("compaction-status", {
|
|
description: "Show effective/raw context usage and pi-compaction state",
|
|
handler: async (_args, ctx) => {
|
|
const metrics = service.getMetrics(ctx.sessionManager.getSessionId());
|
|
const text = metrics
|
|
? [
|
|
`Effective: ${metrics.effectiveTokens}/${metrics.contextWindow} (${metrics.percent.toFixed(1)}%, ${metrics.confidence})`,
|
|
`Raw: ${metrics.rawTokens}/${metrics.contextWindow}`,
|
|
`Reclaimed: ${metrics.reclaimedTokens}`,
|
|
`Segments: ${metrics.segmentCount}; masked tool results: ${metrics.maskedResults}; generation: ${metrics.generation}`,
|
|
].join("\n")
|
|
: "pi-compaction has no projection snapshot yet.";
|
|
if (ctx.hasUI) ctx.ui.notify(text, "info");
|
|
else process.stderr.write(`${text}\n`);
|
|
},
|
|
});
|
|
}
|