chore: automate agent memory capture

2026-04-29 21:35:35 +03:00
parent 54f1ccc80d
commit 4c79695dd5
4 changed files with 369 additions and 1 deletions
@@ -0,0 +1,319 @@
+import { existsSync, mkdirSync, readFileSync, statSync, writeFileSync, appendFileSync } from "node:fs";
+import { dirname, join, relative } from "node:path";
+
+import type { ExtensionAPI } from "@mariozechner/pi-coding-agent";
+
+type SessionMetric = {
+  type: string;
+  timestamp: string;
+  cwd: string;
+  sessionFile?: string;
+  data: Record<string, unknown>;
+};
+
+type MessageLike = {
+  role?: string;
+  content?: unknown;
+};
+
+type EntryLike = {
+  type?: string;
+  message?: MessageLike;
+};
+
+const MEMORY_ROOT = "docs/agent-memory";
+const RAW_ROOT = ".agent-memory/raw";
+const REPORT_ROOT = ".agent-memory/reports";
+const STATE_ROOT = ".agent-memory/state";
+const MAX_INJECT_CHARS = 9000;
+const MAX_RAW_TEXT_CHARS = 1600;
+const IDLE_TIMEOUT_MS = 5 * 60 * 1000;
+
+const isoDate = (date = new Date()) => date.toISOString().slice(0, 10);
+const isoTime = (date = new Date()) => date.toISOString();
+
+const ensureDir = (path: string) => mkdirSync(path, { recursive: true });
+
+const appendJsonl = (path: string, row: unknown) => {
+  ensureDir(dirname(path));
+  appendFileSync(path, `${JSON.stringify(row)}\n`, "utf8");
+};
+
+const readIfExists = (path: string, maxChars = MAX_INJECT_CHARS): string => {
+  if (!existsSync(path) || !statSync(path).isFile()) return "";
+  const text = readFileSync(path, "utf8").trim();
+  if (text.length <= maxChars) return text;
+  return `${text.slice(0, maxChars)}\n\n[truncated by agent-memory extension]`;
+};
+
+const truncate = (text: string, maxChars = MAX_RAW_TEXT_CHARS): string => {
+  const clean = text.replace(/\s+/g, " ").trim();
+  if (clean.length <= maxChars) return clean;
+  return `${clean.slice(0, maxChars)}... [truncated]`;
+};
+
+const extractTextParts = (content: unknown): string[] => {
+  if (typeof content === "string") return [content];
+  if (!Array.isArray(content)) return [];
+
+  const parts: string[] = [];
+  for (const part of content) {
+    if (!part || typeof part !== "object") continue;
+    const block = part as { type?: string; text?: string };
+    if (block.type === "text" && typeof block.text === "string") parts.push(block.text);
+  }
+  return parts;
+};
+
+const branchConversation = (entries: EntryLike[], maxMessages = 24) => {
+  const messages: Array<{ role: string; text: string }> = [];
+
+  for (const entry of entries) {
+    if (entry.type !== "message" || !entry.message?.role) continue;
+    const role = entry.message.role;
+    if (role !== "user" && role !== "assistant") continue;
+
+    const text = extractTextParts(entry.message.content).join("\n").trim();
+    if (!text) continue;
+    messages.push({ role, text: truncate(text) });
+  }
+
+  return messages.slice(-maxMessages);
+};
+
+const writeMetric = (cwd: string, metric: SessionMetric) => {
+  const day = isoDate();
+  appendJsonl(join(cwd, RAW_ROOT, `${day}.jsonl`), metric);
+};
+
+const memoryInjection = (cwd: string): string => {
+  const indexPath = join(cwd, MEMORY_ROOT, "index.md");
+  const changeLogPath = join(cwd, MEMORY_ROOT, "prompt-change-log.md");
+  const index = readIfExists(indexPath, 6000);
+  const changeLog = readIfExists(changeLogPath, 3000);
+
+  if (!index && !changeLog) return "";
+
+  return [
+    "## Project Agent Memory",
+    "",
+    "Use this as a compact index of reviewed project memory. Do not treat it as exhaustive; read cited files when relevant.",
+    "Never store secrets or raw private transcript content in reviewed memory.",
+    "",
+    index ? `### Memory Index\n\n${index}` : "",
+    changeLog ? `### Prompt Change Log\n\n${changeLog}` : "",
+    "",
+    "When the user gives a durable correction, says a prompt pattern worked, or reports an error/fix, suggest `/pi-remember` or `/pi-evolve` instead of relying on chat history.",
+  ]
+    .filter(Boolean)
+    .join("\n");
+};
+
+export default function agentMemoryExtension(pi: ExtensionAPI) {
+  let sessionStartedAt = Date.now();
+  let currentPrompt = "";
+  let promptStartedAt = 0;
+  let lastAgentEndedAt = 0;
+  let turnStartedAt = 0;
+  let providerRequestStartedAt = 0;
+  let providerRequests = 0;
+  let providerResponses = 0;
+  let toolsStarted = 0;
+  let toolsErrored = 0;
+  const submittedPrompts: Array<{ at: string; text: string; pauseInclusiveGapMs?: number }> = [];
+
+  const baseMetric = (ctx: { cwd: string; sessionManager?: { getSessionFile?: () => string | undefined } }, type: string, data: Record<string, unknown>): SessionMetric => ({
+    type,
+    timestamp: isoTime(),
+    cwd: ctx.cwd,
+    sessionFile: ctx.sessionManager?.getSessionFile?.(),
+    data,
+  });
+
+  const captureSnapshot = (ctx: { cwd: string; sessionManager: { getBranch: () => EntryLike[]; getSessionFile?: () => string | undefined } }, note = "") => {
+    ensureDir(join(ctx.cwd, RAW_ROOT));
+    const row = baseMetric(ctx, "memory_snapshot", {
+      note,
+      sessionStartedAt: new Date(sessionStartedAt).toISOString(),
+      capturedAt: isoTime(),
+      prompts: submittedPrompts.slice(-20),
+      conversation: branchConversation(ctx.sessionManager.getBranch()),
+    });
+    writeMetric(ctx.cwd, row);
+    return row;
+  };
+
+  pi.on("session_start", async (_event, ctx) => {
+    sessionStartedAt = Date.now();
+    ensureDir(join(ctx.cwd, RAW_ROOT));
+    ensureDir(join(ctx.cwd, REPORT_ROOT));
+    ensureDir(join(ctx.cwd, STATE_ROOT));
+    writeMetric(ctx.cwd, baseMetric(ctx, "session_start", { pid: process.pid }));
+  });
+
+  pi.on("input", async (event, ctx) => {
+    if (event.source === "extension") return { action: "continue" };
+
+    const now = Date.now();
+    const pauseInclusiveGapMs = lastAgentEndedAt > 0 ? now - lastAgentEndedAt : undefined;
+    promptStartedAt = now;
+    currentPrompt = event.text;
+    submittedPrompts.push({
+      at: isoTime(new Date(now)),
+      text: truncate(event.text, 500),
+      pauseInclusiveGapMs,
+    });
+
+    writeMetric(ctx.cwd, baseMetric(ctx, "prompt_submitted", {
+      promptChars: event.text.length,
+      pauseInclusiveGapMs,
+      idleExcludedMs: pauseInclusiveGapMs && pauseInclusiveGapMs > IDLE_TIMEOUT_MS ? pauseInclusiveGapMs - IDLE_TIMEOUT_MS : 0,
+    }));
+
+    return { action: "continue" };
+  });
+
+  pi.on("before_agent_start", async (event, ctx) => {
+    const injection = memoryInjection(ctx.cwd);
+    if (!injection) return;
+
+    writeMetric(ctx.cwd, baseMetric(ctx, "memory_injected", {
+      promptChars: event.prompt.length,
+      injectionChars: injection.length,
+      indexPath: relative(ctx.cwd, join(ctx.cwd, MEMORY_ROOT, "index.md")),
+    }));
+
+    return {
+      systemPrompt: `${event.systemPrompt}\n\n${injection}`,
+    };
+  });
+
+  pi.on("agent_start", async (_event, ctx) => {
+    toolsStarted = 0;
+    toolsErrored = 0;
+    providerRequests = 0;
+    providerResponses = 0;
+    writeMetric(ctx.cwd, baseMetric(ctx, "agent_start", {
+      promptChars: currentPrompt.length,
+      promptSubmitToStartMs: promptStartedAt > 0 ? Date.now() - promptStartedAt : undefined,
+    }));
+  });
+
+  pi.on("turn_start", async (_event, ctx) => {
+    turnStartedAt = Date.now();
+    writeMetric(ctx.cwd, baseMetric(ctx, "turn_start", {}));
+  });
+
+  pi.on("before_provider_request", (event, ctx) => {
+    providerRequests += 1;
+    providerRequestStartedAt = Date.now();
+    writeMetric(ctx.cwd, baseMetric(ctx, "provider_request_start", {
+      providerRequestIndex: providerRequests,
+      payloadKeys: event.payload && typeof event.payload === "object" ? Object.keys(event.payload as Record<string, unknown>).sort() : [],
+    }));
+  });
+
+  pi.on("after_provider_response", (event, ctx) => {
+    providerResponses += 1;
+    writeMetric(ctx.cwd, baseMetric(ctx, "provider_response_headers", {
+      providerResponseIndex: providerResponses,
+      status: event.status,
+      headerLatencyMs: providerRequestStartedAt > 0 ? Date.now() - providerRequestStartedAt : undefined,
+    }));
+  });
+
+  pi.on("tool_execution_start", async (event, ctx) => {
+    toolsStarted += 1;
+    writeMetric(ctx.cwd, baseMetric(ctx, "tool_start", {
+      toolName: event.toolName,
+    }));
+  });
+
+  pi.on("tool_execution_end", async (event, ctx) => {
+    if (event.isError) toolsErrored += 1;
+    writeMetric(ctx.cwd, baseMetric(ctx, "tool_end", {
+      toolName: event.toolName,
+      isError: event.isError,
+    }));
+  });
+
+  pi.on("turn_end", async (_event, ctx) => {
+    writeMetric(ctx.cwd, baseMetric(ctx, "turn_end", {
+      turnDurationMs: turnStartedAt > 0 ? Date.now() - turnStartedAt : undefined,
+    }));
+  });
+
+  pi.on("agent_end", async (_event, ctx) => {
+    const now = Date.now();
+    lastAgentEndedAt = now;
+    writeMetric(ctx.cwd, baseMetric(ctx, "agent_end", {
+      agentDurationMs: promptStartedAt > 0 ? now - promptStartedAt : undefined,
+      providerRequests,
+      providerResponses,
+      toolsStarted,
+      toolsErrored,
+    }));
+
+    captureSnapshot(ctx, "automatic agent_end snapshot");
+  });
+
+  pi.on("session_before_compact", async (event, ctx) => {
+    writeMetric(ctx.cwd, baseMetric(ctx, "session_before_compact", {
+      tokensBefore: event.preparation?.tokensBefore,
+      firstKeptEntryId: event.preparation?.firstKeptEntryId,
+    }));
+  });
+
+  pi.on("session_shutdown", async (event, ctx) => {
+    writeMetric(ctx.cwd, baseMetric(ctx, "session_shutdown", {
+      reason: event.reason,
+      sessionDurationMs: Date.now() - sessionStartedAt,
+      promptsSeen: submittedPrompts.length,
+    }));
+  });
+
+  pi.registerCommand("memory-status", {
+    description: "Show agent memory automation status",
+    handler: async (_args, ctx) => {
+      const day = isoDate();
+      const rawPath = join(ctx.cwd, RAW_ROOT, `${day}.jsonl`);
+      const indexPath = join(ctx.cwd, MEMORY_ROOT, "index.md");
+      const message = [
+        `Memory index: ${existsSync(indexPath) ? relative(ctx.cwd, indexPath) : "missing"}`,
+        `Raw metrics today: ${existsSync(rawPath) ? relative(ctx.cwd, rawPath) : "none yet"}`,
+        `Prompts observed this session: ${submittedPrompts.length}`,
+        `Provider requests this turn: ${providerRequests}`,
+      ].join("\n");
+
+      if (ctx.hasUI) ctx.ui.notify(message, "info");
+      else console.log(message);
+    },
+  });
+
+  pi.registerCommand("memory-capture", {
+    description: "Capture a private raw snapshot for later memory compilation",
+    handler: async (args, ctx) => {
+      const row = captureSnapshot(ctx, args.trim() || "manual capture");
+      const rawPath = join(ctx.cwd, RAW_ROOT, `${isoDate()}.jsonl`);
+      pi.appendEntry("agent-memory-capture", row);
+      const message = `Captured private memory snapshot: ${relative(ctx.cwd, rawPath)}`;
+      if (ctx.hasUI) ctx.ui.notify(message, "success");
+      else console.log(message);
+    },
+  });
+
+  pi.registerCommand("memory-compile", {
+    description: "Capture current session and ask Pi to compile durable memory/evolution changes",
+    handler: async (args, ctx) => {
+      if (!ctx.isIdle()) {
+        if (ctx.hasUI) ctx.ui.notify("Agent is busy. Run /memory-compile after the current turn finishes.", "warning");
+        return;
+      }
+
+      captureSnapshot(ctx, args.trim() || "compile request");
+      const goal = args.trim() || "Compile durable lessons from the latest private memory snapshot and propose prompt evolution only if evidence is strong.";
+      pi.sendUserMessage(`/pi-evolve ${goal}`);
+    },
+  });
+}
+