@zhushanwen/pi-context-engineering 0.1.0

package/src/config.ts

@@ -0,0 +1,172 @@
+import { readFileSync } from "node:fs";
+import { homedir } from "node:os";
+import { join } from "node:path";
+
+// ── 配置嵌套类型 ──
+
+export interface L0Config {
+  enabled: boolean;
+  expireMinutes: number;
+  bashTruncateChars: number;
+  thinkingExpireMinutes: number;
+  protectRecentTurns: number;
+  keepRecent: number;
+}
+
+export interface L1Config {
+  enabled: boolean;
+  summaryThresholdChars: number;
+  keepHeadLines: number;
+  keepTailLines: number;
+  protectRecentTurns: number;
+}
+
+export interface L2Config {
+  enabled: boolean;
+  emergencyThreshold: number;
+  protectRecentTurns: number;
+}
+
+export interface McConfig {
+  enabled: boolean;
+  gapThresholdMinutes: number;
+  keepRecent: number;
+}
+
+export interface BudgetConfig {
+  enabled: boolean;
+  maxToolResultCharsPerMessage: number;
+  previewSize: number;
+}
+
+export interface ContextEngineeringConfig {
+  enabled: boolean;
+  l0: L0Config;
+  l1: L1Config;
+  l2: L2Config;
+  mc: McConfig;
+  budget: BudgetConfig;
+}
+
+// ── 默认配置 ──
+
+export const DEFAULT_CONFIG: ContextEngineeringConfig = {
+  enabled: true,
+  l0: {
+    enabled: true,
+    expireMinutes: 30,
+    bashTruncateChars: 4000,
+    thinkingExpireMinutes: 5,
+    protectRecentTurns: 2,
+    keepRecent: 5,
+  },
+  l1: {
+    enabled: true,
+    summaryThresholdChars: 8000,
+    keepHeadLines: 10,
+    keepTailLines: 5,
+    protectRecentTurns: 2,
+  },
+  l2: {
+    enabled: true,
+    emergencyThreshold: 0.9,
+    protectRecentTurns: 3,
+  },
+  mc: {
+    enabled: true,
+    gapThresholdMinutes: 60,
+    keepRecent: 5,
+  },
+  budget: {
+    enabled: true,
+    maxToolResultCharsPerMessage: 200_000,
+    previewSize: 2000,
+  },
+};
+
+// ── 深合并工具 ──
+
+function deepMerge<T>(
+  base: T,
+  override: Record<string, unknown>,
+): T {
+  const result = { ...base } as Record<string, unknown>;
+  for (const key of Object.keys(override)) {
+    const baseVal = result[key];
+    const overVal = override[key];
+    if (
+      baseVal != null &&
+      overVal != null &&
+      typeof baseVal === "object" &&
+      typeof overVal === "object" &&
+      !Array.isArray(baseVal) &&
+      !Array.isArray(overVal)
+    ) {
+      result[key] = deepMerge(
+        baseVal as Record<string, unknown>,
+        overVal as Record<string, unknown>,
+      );
+    } else {
+      result[key] = overVal;
+    }
+  }
+  return result as T;
+}
+
+// ── 读取配置 ──
+
+export function loadConfig(
+  settingsPath?: string,
+): ContextEngineeringConfig {
+  const filePath =
+    settingsPath ?? join(homedir(), ".pi", "agent", "settings.json");
+
+  let raw: string;
+  try {
+    raw = readFileSync(filePath, "utf-8");
+  } catch {
+    return { ...DEFAULT_CONFIG };
+  }
+
+  let parsed: Record<string, unknown>;
+  try {
+    parsed = JSON.parse(raw) as Record<string, unknown>;
+  } catch {
+    return { ...DEFAULT_CONFIG };
+  }
+
+  const override = parsed["context-engineering"];
+  if (override == null || typeof override !== "object") {
+    return { ...DEFAULT_CONFIG };
+  }
+
+  return deepMerge<ContextEngineeringConfig>(
+    DEFAULT_CONFIG,
+    override as Record<string, unknown>,
+  );
+}
+
+// ── 命令参数解析 ──
+
+export function parseLevelArgs(
+  args: string,
+): { target: "global" | "l0" | "l1" | "l2" | "mc" | "budget"; action: "on" | "off" } | null {
+  const tokens = args.trim().split(/\s+/);
+  if (tokens.length < 2) {
+    return null;
+  }
+
+  const [rawTarget, rawAction] = tokens;
+
+  const validTargets = new Set(["global", "l0", "l1", "l2", "mc", "budget"]);
+  const validActions = new Set(["on", "off"]);
+
+  if (!validTargets.has(rawTarget) || !validActions.has(rawAction)) {
+    return null;
+  }
+
+  return {
+    target: rawTarget as "global" | "l0" | "l1" | "l2" | "mc" | "budget",
+    action: rawAction as "on" | "off",
+  };
+}

package/src/frozen-fresh.ts

@@ -0,0 +1,36 @@
+// frozen-fresh.ts — Frozen/Fresh 状态跟踪，用于 Budget 工具结果管理
+
+export interface FrozenFreshState {
+  isFrozen(toolUseId: string): boolean;
+  markFrozen(toolUseId: string, replacement: string): void;
+  getReplacement(toolUseId: string): string | undefined;
+  getAllFrozenIds(): Set<string>;
+  reset(): void;
+}
+
+export function createFrozenFreshState(): FrozenFreshState {
+  // toolCallId → replacement text
+  const frozen = new Map<string, string>();
+
+  return {
+    isFrozen(toolUseId: string): boolean {
+      return frozen.has(toolUseId);
+    },
+
+    markFrozen(toolUseId: string, replacement: string): void {
+      frozen.set(toolUseId, replacement);
+    },
+
+    getReplacement(toolUseId: string): string | undefined {
+      return frozen.get(toolUseId);
+    },
+
+    getAllFrozenIds(): Set<string> {
+      return new Set(frozen.keys());
+    },
+
+    reset(): void {
+      frozen.clear();
+    },
+  };
+}

package/src/index.ts

@@ -0,0 +1,105 @@
+import type {
+  ExtensionAPI,
+  ExtensionCommandContext,
+} from "@mariozechner/pi-coding-agent";
+import { Type } from "typebox";
+
+import { loadConfig, type ContextEngineeringConfig } from "./config";
+import { createRecallStore, type RecallStore } from "./recall-store";
+import { createFrozenFreshState, type FrozenFreshState } from "./frozen-fresh";
+import {
+  compressContext,
+  type CompressionStats,
+  type AgentMessage as CompressorMessage,
+} from "./compressor";
+import { handleContextEngineeringCommand, handleContextStatsCommand } from "./commands";
+
+const RecallParams = Type.Object({
+  id: Type.String({ description: "Context ID (ctx-xxxxxxxx) to recall" }),
+});
+
+function zeroStats(): CompressionStats {
+  return { l0Expired: 0, l0Truncated: 0, l0ThinkingCleared: 0, l1Condensed: 0, l2Triggered: false, validationFailed: false, mcTriggered: false, mcCleared: 0, budgetPersisted: 0 };
+}
+
+function addStats(target: CompressionStats, delta: CompressionStats): void {
+  target.l0Expired += delta.l0Expired;
+  target.l0Truncated += delta.l0Truncated;
+  target.l0ThinkingCleared += delta.l0ThinkingCleared;
+  target.l1Condensed += delta.l1Condensed;
+  target.mcCleared += delta.mcCleared;
+  target.budgetPersisted += delta.budgetPersisted;
+  if (delta.l2Triggered) target.l2Triggered = true;
+  if (delta.validationFailed) target.validationFailed = true;
+  if (delta.mcTriggered) target.mcTriggered = true;
+}
+
+function recallResult(id: string, store: RecallStore) {
+  const stored = store.recall(id);
+  if (!stored) return {
+    content: [{ type: "text" as const, text: `[recall_context] ID "${id}" not found. Content may have been lost on session reload.` }],
+    details: { found: false, id },
+  };
+  return {
+    content: [{ type: "text" as const, text: `[Recalled content (${stored.level}, ${new Date(stored.compressedAt).toISOString()})]\n\n${stored.original}` }],
+    details: { found: true, id, level: stored.level },
+  };
+}
+
+// Extension entry — handlers close over mutable `config`/`store`/`cumulativeStats`
+// so session_start reassignment is visible to every registered handler.
+
+export default function contextEngineeringExtension(pi: ExtensionAPI): void {
+  let config: ContextEngineeringConfig = loadConfig();
+  let store: RecallStore = createRecallStore();
+  let cumulativeStats: CompressionStats = zeroStats();
+  let frozenFreshState: FrozenFreshState = createFrozenFreshState();
+
+  pi.on("session_start", () => {
+    config = loadConfig();
+    store = createRecallStore();
+    cumulativeStats = zeroStats();
+    frozenFreshState = createFrozenFreshState();
+  });
+
+  pi.on("context", (event, ctx) => {
+    try {
+      // Pi Extension API types differ from our internal message types.
+      // Both sides define the same shape but TypeScript can't verify across packages.
+      // If Pi's message format changes, compressContext will gracefully fail via the catch below.
+      const msgs = event.messages as unknown as CompressorMessage[];
+      const result = compressContext(msgs, config, store, ctx.getContextUsage() as unknown as Parameters<typeof compressContext>[3], frozenFreshState);
+      addStats(cumulativeStats, result.stats);
+      return { messages: result.messages as unknown as (typeof event.messages)[number][] };
+    } catch (err) {
+      // Silently degrade to original messages, but log for debuggability
+      if (process.env.DEBUG_CONTEXT_ENGINEERING) {
+        console.error("[context-engineering] compressContext failed:", err);
+      }
+      return {};
+    }
+  });
+
+  pi.registerTool({
+    name: "recall_context",
+    label: "Recall Compressed Context",
+    description: "Recall original content compressed by context engineering. Use when you need the full content of an expired, truncated, or condensed tool result.",
+    promptSnippet: "recall_context(id) — retrieve original content compressed by context engineering",
+    parameters: RecallParams,
+    execute: async (_tcId, params, _sig, _upd, _ctx) => recallResult(params.id, store),
+  });
+
+  pi.registerCommand("context-engineering", {
+    description: "View/modify context compression settings",
+    handler: async (_args: string, ctx: ExtensionCommandContext) => {
+      ctx.ui.notify(handleContextEngineeringCommand(_args || undefined, config, cumulativeStats), "info");
+    },
+  });
+
+  pi.registerCommand("context-stats", {
+    description: "View context compression statistics",
+    handler: async (_args: string, ctx: ExtensionCommandContext) => {
+      ctx.ui.notify(handleContextStatsCommand(cumulativeStats), "info");
+    },
+  });
+}

package/src/recall-store.ts

@@ -0,0 +1,63 @@
+import { randomUUID } from "node:crypto";
+
+// ── 类型 ──
+
+export interface StoredContent {
+  id: string;
+  original: string;
+  compressedAt: number;
+  level: "l0-expired" | "l0-truncated" | "l1-condensed" | "l2-emergency" | "mc-cleared" | "budget-persisted";
+}
+
+export interface RecallStore {
+  store: (content: string, level: StoredContent["level"]) => string;
+  recall: (id: string) => StoredContent | undefined;
+  clear: () => void;
+  size: () => number;
+}
+
+// ── 常量 ──
+
+/** UUID 前 12 字符 = 48 bit 熵，碰撞阈值约 16M 条。远超单 session 需求。 */
+const ID_CHARS = 12;
+
+/** 内存保护上限。超过时淘汰最早存入的条目。 */
+const MAX_ENTRIES = 500;
+
+// ── 工厂函数 ──
+
+export function createRecallStore(): RecallStore {
+  const entries = new Map<string, StoredContent>();
+
+  function store(content: string, level: StoredContent["level"]): string {
+    // LRU 淘汰：超过上限时删除最早存入的条目
+    if (entries.size >= MAX_ENTRIES) {
+      const oldest = entries.keys().next().value;
+      if (oldest !== undefined) entries.delete(oldest);
+    }
+
+    const idSuffix = randomUUID().replace(/-/g, '').slice(0, ID_CHARS);
+    const id = `ctx-${idSuffix}`;
+    entries.set(id, {
+      id,
+      original: content,
+      compressedAt: Date.now(),
+      level,
+    });
+    return id;
+  }
+
+  function recall(id: string): StoredContent | undefined {
+    return entries.get(id);
+  }
+
+  function clear(): void {
+    entries.clear();
+  }
+
+  function size(): number {
+    return entries.size;
+  }
+
+  return { store, recall, clear, size };
+}