From 4c79695dd5b885265dd37207285b12f99f8195af Mon Sep 17 00:00:00 2001 From: gnezim Date: Wed, 29 Apr 2026 21:35:35 +0300 Subject: [PATCH] chore: automate agent memory capture --- .pi/extensions/agent-memory.ts | 319 +++++++++++++++++++++++++ docs/agent-memory/README.md | 39 ++- docs/agent-memory/log.md | 5 + docs/agent-memory/prompt-change-log.md | 7 + 4 files changed, 369 insertions(+), 1 deletion(-) create mode 100644 .pi/extensions/agent-memory.ts diff --git a/.pi/extensions/agent-memory.ts b/.pi/extensions/agent-memory.ts new file mode 100644 index 00000000..71038d60 --- /dev/null +++ b/.pi/extensions/agent-memory.ts @@ -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; +}; + +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): 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).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}`); + }, + }); +} + diff --git a/docs/agent-memory/README.md b/docs/agent-memory/README.md index a2fc7a2c..9fbe1221 100644 --- a/docs/agent-memory/README.md +++ b/docs/agent-memory/README.md @@ -53,6 +53,44 @@ Daily entries live in `docs/agent-memory/daily/YYYY-MM-DD.md` when they are safe - `prompt-evolution/` stores proposed prompt and workflow changes before they are applied. - `prompt-change-log.md` records accepted prompt, agent, and workflow changes. +## Automation + +The project-local Pi extension lives at `.pi/extensions/agent-memory.ts`. + +It does four things automatically: + +- injects the reviewed memory index and prompt change log into each agent turn through `before_agent_start` +- records private raw metrics and session snapshots under `.agent-memory/raw/` +- records provider header latency, agent duration, turn duration, tool counts, prompt submission gaps, and compaction/shutdown events +- exposes `/memory-status`, `/memory-capture`, and `/memory-compile` + +Raw metrics are intentionally gitignored. They are evidence for later compilation, not reviewed memory. + +The extension does not infer true active typing time. It records: + +- `pauseInclusiveGapMs`: time between previous agent completion and next prompt submission +- `idleExcludedMs`: the portion above the 5-minute idle cap +- `agentDurationMs`: prompt submission to `agent_end` +- `turnDurationMs`: one LLM/tool turn duration +- `headerLatencyMs`: provider request to HTTP response headers + +Use `npx @ccusage/pi@latest session` and LiteLLM metrics for token/cost and provider-side inference timing. + +## Commands + +Inside Pi: + +```text +/memory-status +/memory-capture [note] +/memory-compile [goal] +/pi-remember +/pi-memory +/pi-evolve +``` + +`/memory-compile` captures a private snapshot and then asks Pi to run `/pi-evolve` against the latest evidence. + ## Guardrails - Memory can suggest prompt changes; it must not silently rewrite prompts. @@ -61,4 +99,3 @@ Daily entries live in `docs/agent-memory/daily/YYYY-MM-DD.md` when they are safe - Prefer small, testable prompt changes over broad rewrites. - If a lesson is only true for one feature, store it with that scope. - If evidence is weak, classify it as `hypothesis`, not `rule`. - diff --git a/docs/agent-memory/log.md b/docs/agent-memory/log.md index 1565200b..bedb4fce 100644 --- a/docs/agent-memory/log.md +++ b/docs/agent-memory/log.md @@ -5,3 +5,8 @@ - Created shared memory schema and guarded prompt-evolution process. - Added commands for manual capture, memory query, and prompt evolution. +## [2026-04-29] automation | Pi extension + +- Added `.pi/extensions/agent-memory.ts`. +- Verified extension command `/memory-status` in Pi print mode. +- Metrics and raw snapshots are private runtime artifacts under `.agent-memory/raw/`. diff --git a/docs/agent-memory/prompt-change-log.md b/docs/agent-memory/prompt-change-log.md index 4e28a5a6..2f477298 100644 --- a/docs/agent-memory/prompt-change-log.md +++ b/docs/agent-memory/prompt-change-log.md @@ -6,3 +6,10 @@ - Added memory and prompt-evolution scaffold. - Prompt changes must include evidence, validation commands, reviewer notes, and rollback guidance. +## [2026-04-29] automation | Agent memory extension + +- Added project-local Pi extension `.pi/extensions/agent-memory.ts`. +- Automatically injects reviewed memory index and prompt change log into future turns. +- Captures private raw metrics and session snapshots under `.agent-memory/raw/`. +- Added `/memory-status`, `/memory-capture`, and `/memory-compile` extension commands. +- Raw telemetry remains gitignored; reviewed memory remains under `docs/agent-memory/`.