Initial: pi-skill — 68 skills, 43 extensions, 11 themes for Pi
This commit is contained in:
Kunthawat Greethong
541
extensions/memory-cycle.ts
Normal file
541
extensions/memory-cycle.ts
Normal file
@@ -0,0 +1,541 @@
|
||||
// ABOUTME: Memory-aware compaction extension — hooks into pi's native compaction to save/restore context.
|
||||
// ABOUTME: Writes daily logs, session state, and optionally updates MEMORY.md during every compaction cycle.
|
||||
/**
|
||||
* Memory Cycle — Automatic memory-aware compaction with seamless restore
|
||||
*
|
||||
* Hooks into pi's native compaction system to:
|
||||
* 1. BEFORE compact: Extract session insights (daily log, session state, stable facts)
|
||||
* 2. AFTER compact: Inject restored memory context so agent continues seamlessly
|
||||
*
|
||||
* Also provides:
|
||||
* /cycle [instructions] — Manual command to trigger compact → new session → restore
|
||||
* cycle_memory — LLM-callable tool for the same workflow
|
||||
*
|
||||
* The agent gets a clean context window but retains full awareness of
|
||||
* everything that happened before.
|
||||
*/
|
||||
|
||||
import type { ExtensionAPI } from "@mariozechner/pi-coding-agent";
|
||||
// convertToLlm and serializeConversation available if needed for custom summary generation
|
||||
import { Type } from "@sinclair/typebox";
|
||||
import { Box, Text } from "@mariozechner/pi-tui";
|
||||
import {
|
||||
getProjectName,
|
||||
getTimestamp,
|
||||
extractFileOps,
|
||||
writeDailyLog,
|
||||
writeSessionState,
|
||||
readRecentLogs,
|
||||
readSessionState,
|
||||
extractCompactionContext,
|
||||
buildRestorationContent,
|
||||
buildCycleMemoryInjection,
|
||||
} from "./lib/memory-cycle-helpers.ts";
|
||||
import { getProactiveCompactionPhase } from "./lib/context-gate.ts";
|
||||
|
||||
// ── Tool Parameters ──────────────────────────────────────────────────
|
||||
|
||||
const CycleParams = Type.Object({
|
||||
instructions: Type.Optional(
|
||||
Type.String({ description: "Custom instructions for what to focus on in the summary" }),
|
||||
),
|
||||
});
|
||||
|
||||
// ── Compaction Card Details ──────────────────────────────────────────
|
||||
|
||||
interface CompactionCardDetails {
|
||||
/** "cycle" for cycle_memory, "auto" for footer auto-compact, "manual" for /compact */
|
||||
source: "cycle" | "auto" | "manual";
|
||||
/** Context percentage after compaction */
|
||||
postPercent: number;
|
||||
/** Recent session task, if available */
|
||||
task?: string;
|
||||
/** Recently edited files */
|
||||
recentFiles?: string[];
|
||||
}
|
||||
|
||||
// ── Compaction Card Renderer ─────────────────────────────────────────
|
||||
// Renders a minimal, elegant dark-themed status card when compaction
|
||||
// completes. Appears for cycle_memory, auto-compact, and manual /compact.
|
||||
|
||||
function renderCompactionCard(
|
||||
message: any,
|
||||
_options: any,
|
||||
theme: any,
|
||||
) {
|
||||
const details = message.details;
|
||||
const percent = details?.postPercent ?? 0;
|
||||
const source = details?.source ?? "cycle";
|
||||
|
||||
// ── Title ───────────────────────────────────────────────────
|
||||
const label = source === "auto"
|
||||
? "Context Compacted"
|
||||
: source === "manual"
|
||||
? "Context Compacted"
|
||||
: "Memory Cycle Complete";
|
||||
const title = theme.fg("muted", label);
|
||||
|
||||
// ── Percentage — color-coded by health ──────────────────────
|
||||
const pctColor = percent <= 30 ? "success" : percent <= 60 ? "muted" : "warning";
|
||||
const pctText = theme.fg(pctColor as any, `${percent}%`) +
|
||||
theme.fg("dim", " context used");
|
||||
|
||||
// ── Detail lines (task + files) ─────────────────────────────
|
||||
const detailLines: string[] = [];
|
||||
|
||||
if (details?.task) {
|
||||
const truncated = details.task.length > 72
|
||||
? details.task.slice(0, 69) + "..."
|
||||
: details.task;
|
||||
detailLines.push(
|
||||
theme.fg("dim", "task ") + theme.fg("muted", truncated),
|
||||
);
|
||||
}
|
||||
|
||||
if (details?.recentFiles?.length) {
|
||||
const shown = details.recentFiles.slice(0, 3);
|
||||
const names = shown.map((f: string) => {
|
||||
const parts = f.split("/");
|
||||
return parts.length > 1 ? parts.slice(-2).join("/") : parts[0];
|
||||
});
|
||||
const more = details.recentFiles.length > 3
|
||||
? theme.fg("dim", ` +${details.recentFiles.length - 3}`)
|
||||
: "";
|
||||
detailLines.push(
|
||||
theme.fg("dim", "files ") +
|
||||
theme.fg("muted", names.join(theme.fg("dim", " / "))) + more,
|
||||
);
|
||||
}
|
||||
|
||||
// ── Assemble card body ──────────────────────────────────────
|
||||
const lines: string[] = [
|
||||
title,
|
||||
pctText,
|
||||
];
|
||||
|
||||
if (detailLines.length > 0) {
|
||||
lines.push(""); // blank separator line
|
||||
for (const dl of detailLines) lines.push(dl);
|
||||
}
|
||||
|
||||
const body = lines.join("\n");
|
||||
|
||||
// Custom dark-charcoal background — distinct from the ocean-blue theme
|
||||
// Neutral gray so it reads as a "system" card, not success/error
|
||||
const cardBg = (text: string) => `\x1b[48;2;30;36;42m${text}\x1b[49m`;
|
||||
const box = new Box(
|
||||
3, // generous horizontal padding
|
||||
1, // vertical breathing room
|
||||
cardBg,
|
||||
);
|
||||
box.addChild(new Text(body, 0, 0));
|
||||
return box;
|
||||
}
|
||||
|
||||
// ── Extension ────────────────────────────────────────────────────────
|
||||
|
||||
export default function (pi: ExtensionAPI) {
|
||||
// ── Message Renderers ────────────────────────────────────────
|
||||
// Register custom renderers for compaction status cards.
|
||||
// These render in the chat when display:true is set on sendMessage.
|
||||
|
||||
pi.registerMessageRenderer<CompactionCardDetails>("memory-cycle-resume", renderCompactionCard);
|
||||
pi.registerMessageRenderer<CompactionCardDetails>("auto-compact-resume", renderCompactionCard);
|
||||
pi.registerMessageRenderer<CompactionCardDetails>("memory-restored", renderCompactionCard);
|
||||
|
||||
// ── Proactive compaction state ───────────────────────────────
|
||||
// Two-phase: prep at 70% (wrap up work), hard stop at 80% (call cycle_memory).
|
||||
// Flags prevent repeated injection within the same compaction cycle.
|
||||
let prepInjected = false; // true after 70% prep message sent
|
||||
let compactInjected = false; // true after 80% hard-stop message sent
|
||||
|
||||
// ── Hook: before_agent_start — proactive compaction ──────────
|
||||
// Fires before every agent turn. Checks context usage and injects
|
||||
// messages to guide the LLM toward compaction before overflow.
|
||||
pi.on("before_agent_start", async (_event, ctx) => {
|
||||
const usage = ctx.getContextUsage();
|
||||
const { phase, percent } = getProactiveCompactionPhase(usage?.percent);
|
||||
|
||||
if (phase === "compact" && !compactInjected) {
|
||||
compactInjected = true;
|
||||
ctx.ui.notify(
|
||||
`Context overflow detected, Auto-compacting... (escape to cancel)`,
|
||||
"info",
|
||||
);
|
||||
return {
|
||||
message: {
|
||||
customType: "auto-compact-gate",
|
||||
content: `URGENT: Context window is at ${Math.round(percent)}% capacity. You MUST call cycle_memory immediately to prevent context overflow. Do not perform any other actions first. Call cycle_memory now.`,
|
||||
display: false,
|
||||
},
|
||||
};
|
||||
}
|
||||
|
||||
if (phase === "prep" && !prepInjected) {
|
||||
prepInjected = true;
|
||||
ctx.ui.notify(
|
||||
`Context at ${Math.round(percent)}% -- wrapping up soon`,
|
||||
"info",
|
||||
);
|
||||
return {
|
||||
message: {
|
||||
customType: "auto-compact-gate",
|
||||
content: `Context window is at ${Math.round(percent)}% capacity. Start wrapping up your current work: commit any in-progress changes, save state, and prepare for a memory cycle. When you finish your current step, call cycle_memory. Do not start any new large operations.`,
|
||||
display: false,
|
||||
},
|
||||
};
|
||||
}
|
||||
|
||||
return {};
|
||||
});
|
||||
|
||||
// Track cwd across compact events (before_compact → compact)
|
||||
let preCompactCwd: string = "";
|
||||
|
||||
// When cycle_memory triggers compaction, suppress redundant UI from
|
||||
// session_before_compact and session_compact — the cycle_memory
|
||||
// onComplete handler shows a single clean card instead.
|
||||
let cycleMemoryActive = false;
|
||||
|
||||
// ── Hook: session_before_compact ──────────────────────────────
|
||||
// Runs as part of pi's native compaction (both auto and manual /compact).
|
||||
// We extract session insights and save them to disk BEFORE the context
|
||||
// is compacted. We do NOT cancel or replace compaction — we let pi's
|
||||
// default compaction run normally.
|
||||
pi.on("session_before_compact", async (event, ctx) => {
|
||||
preCompactCwd = ctx.cwd;
|
||||
const { preparation } = event;
|
||||
|
||||
try {
|
||||
const project = getProjectName(ctx.cwd);
|
||||
const { date, time, iso } = getTimestamp();
|
||||
|
||||
// Use pi's already-extracted file operations from preparation
|
||||
const prepFileOps = preparation.fileOps;
|
||||
const readFiles = prepFileOps?.read ? [...prepFileOps.read] : [];
|
||||
const writtenFiles = prepFileOps?.written ? [...prepFileOps.written] : [];
|
||||
const editedFiles = prepFileOps?.edited ? [...prepFileOps.edited] : [];
|
||||
const modifiedFiles = [...new Set([...writtenFiles, ...editedFiles])];
|
||||
|
||||
// Also supplement with branch-level file ops for completeness
|
||||
const branchOps = extractFileOps(ctx.sessionManager.getBranch());
|
||||
for (const f of branchOps.read) { if (!readFiles.includes(f)) readFiles.push(f); }
|
||||
for (const f of branchOps.modified) { if (!modifiedFiles.includes(f)) modifiedFiles.push(f); }
|
||||
|
||||
// Build a compact summary from the messages being compacted
|
||||
const { summaryText, continueText } = extractCompactionContext(
|
||||
preparation.messagesToSummarize,
|
||||
preparation.previousSummary,
|
||||
);
|
||||
|
||||
// Write daily log entry
|
||||
writeDailyLog({
|
||||
project,
|
||||
summary: summaryText,
|
||||
date,
|
||||
time,
|
||||
keyFiles: [...modifiedFiles, ...readFiles].slice(0, 10),
|
||||
continuePrompt: continueText,
|
||||
});
|
||||
|
||||
// Write session state
|
||||
writeSessionState(ctx.cwd, {
|
||||
project,
|
||||
iso,
|
||||
continuePrompt: continueText,
|
||||
currentTask: summaryText,
|
||||
filesEdited: modifiedFiles.slice(0, 10),
|
||||
filesRead: readFiles.slice(0, 10),
|
||||
});
|
||||
|
||||
// Only show notification for manual /compact — cycle_memory shows its own card
|
||||
if (!cycleMemoryActive) {
|
||||
ctx.ui.notify("Memory saved (daily log + session state)", "info");
|
||||
}
|
||||
} catch (err) {
|
||||
const msg = err instanceof Error ? err.message : String(err);
|
||||
console.error(`[memory-cycle] Pre-compact save failed: ${msg}`);
|
||||
// Don't cancel compaction on save failure
|
||||
}
|
||||
|
||||
// Return nothing = let pi's default compaction proceed normally
|
||||
return;
|
||||
});
|
||||
|
||||
// ── Hook: session_compact ─────────────────────────────────────
|
||||
// Fires AFTER compaction completes (both manual /compact and core auto-compaction).
|
||||
// We inject a memory-restore message so the agent knows what happened
|
||||
// and can continue seamlessly.
|
||||
//
|
||||
// For core auto-compaction: the interactive-mode handles UI rebuild via
|
||||
// auto_compaction_start/end events. We just provide the restoration context.
|
||||
// For manual /compact: we send both a display card and restoration context.
|
||||
pi.on("session_compact", async (event, ctx) => {
|
||||
// Reset proactive compaction flags — allows next cycle to trigger
|
||||
prepInjected = false;
|
||||
compactInjected = false;
|
||||
|
||||
const { compactionEntry } = event;
|
||||
|
||||
const recentLogs = readRecentLogs();
|
||||
const sessionState = readSessionState(preCompactCwd || ctx.cwd);
|
||||
|
||||
// Build restoration context
|
||||
const parts = buildRestorationContent(sessionState);
|
||||
if (recentLogs) parts.push("", recentLogs);
|
||||
|
||||
const postUsage = ctx.getContextUsage();
|
||||
const postPercent = postUsage?.percent ? Math.round(postUsage.percent) : 0;
|
||||
|
||||
// When cycle_memory is driving compaction, skip the display card here —
|
||||
// the cycle_memory onComplete handler shows a single clean card instead.
|
||||
// Only show the card for manual /compact or core auto-compaction.
|
||||
if (!cycleMemoryActive) {
|
||||
// Short card visible to user
|
||||
pi.sendMessage(
|
||||
{
|
||||
customType: "memory-restored",
|
||||
content: `Context compacted -- now at ${postPercent}%.`,
|
||||
display: true,
|
||||
details: {
|
||||
source: "manual",
|
||||
postPercent,
|
||||
task: sessionState?.currentTask,
|
||||
recentFiles: sessionState?.filesEdited,
|
||||
} satisfies CompactionCardDetails,
|
||||
},
|
||||
);
|
||||
}
|
||||
|
||||
// Full restoration context for the agent (not displayed)
|
||||
// Always send this — cycle_memory onComplete will add its own,
|
||||
// but for manual /compact this is the only restoration message.
|
||||
if (!cycleMemoryActive) {
|
||||
pi.sendMessage(
|
||||
{
|
||||
customType: "memory-restored",
|
||||
content: parts.join("\n"),
|
||||
display: false,
|
||||
},
|
||||
{ deliverAs: "nextTurn" },
|
||||
);
|
||||
}
|
||||
});
|
||||
|
||||
|
||||
// ── /cycle command ────────────────────────────────────────────
|
||||
// Manual command: compact → new session → restore (full reset)
|
||||
pi.registerCommand("cycle", {
|
||||
description: "Compact → new session → restore: fresh context with full memory",
|
||||
handler: async (args, ctx) => {
|
||||
const customInstructions = args?.trim() || undefined;
|
||||
|
||||
await ctx.waitForIdle();
|
||||
|
||||
const parentSessionFile = ctx.sessionManager.getSessionFile();
|
||||
const entries = ctx.sessionManager.getBranch();
|
||||
|
||||
if (entries.length < 3) {
|
||||
ctx.ui.notify("Session too short to cycle — nothing to compact.", "warning");
|
||||
return;
|
||||
}
|
||||
|
||||
ctx.ui.notify("Memory Cycle: Step 1/3 — Compacting...", "info");
|
||||
|
||||
// Step 1: Compact and capture summary
|
||||
const compactionSummary = await new Promise<string | null>((resolve) => {
|
||||
ctx.compact({
|
||||
customInstructions: customInstructions
|
||||
?? "Create a comprehensive summary preserving all goals, decisions, progress, file changes, and context needed to continue work seamlessly in a fresh session.",
|
||||
onComplete: () => {
|
||||
// The session_before_compact hook already saved memory artifacts.
|
||||
// Extract summary from post-compaction session.
|
||||
const postEntries = ctx.sessionManager.getBranch();
|
||||
for (let i = postEntries.length - 1; i >= 0; i--) {
|
||||
const entry = postEntries[i];
|
||||
if (entry.type === "compaction") {
|
||||
resolve((entry as any).summary ?? null);
|
||||
return;
|
||||
}
|
||||
}
|
||||
resolve(null);
|
||||
},
|
||||
onError: (err) => {
|
||||
ctx.ui.notify(`Compaction failed: ${err.message}`, "error");
|
||||
resolve(null);
|
||||
},
|
||||
});
|
||||
});
|
||||
|
||||
if (!compactionSummary) {
|
||||
ctx.ui.notify("Memory Cycle aborted — compaction produced no summary.", "error");
|
||||
return;
|
||||
}
|
||||
|
||||
ctx.ui.notify("Memory Cycle: Step 2/3 — Creating fresh session...", "info");
|
||||
|
||||
// Gather restoration context
|
||||
const recentLogs = readRecentLogs();
|
||||
const sessionState = readSessionState(ctx.cwd);
|
||||
|
||||
// Step 2: New session with parent link and memory injection
|
||||
const result = await ctx.newSession({
|
||||
parentSession: parentSessionFile,
|
||||
setup: async (sm) => {
|
||||
const memoryText = buildCycleMemoryInjection({
|
||||
compactionSummary,
|
||||
sessionState,
|
||||
recentLogs,
|
||||
});
|
||||
|
||||
sm.appendMessage({
|
||||
role: "user",
|
||||
content: [{ type: "text", text: memoryText }],
|
||||
timestamp: Date.now(),
|
||||
});
|
||||
},
|
||||
});
|
||||
|
||||
if (result.cancelled) {
|
||||
ctx.ui.notify("Memory Cycle cancelled — session switch was blocked.", "warning");
|
||||
return;
|
||||
}
|
||||
|
||||
ctx.ui.notify("Memory Cycle complete — fresh context with full memory.", "success");
|
||||
},
|
||||
});
|
||||
|
||||
// ── Deferred compaction via agent_end hook ────────────────────
|
||||
// The cycle_memory tool CANNOT call ctx.compact() directly because
|
||||
// compact() calls abort() which waits for the agent to be idle,
|
||||
// but the agent is blocked waiting for the tool to return → deadlock.
|
||||
//
|
||||
// Instead: tool sets a flag → returns immediately → agent_end fires
|
||||
// when the agent loop finishes → we compact from there (agent is idle).
|
||||
|
||||
let pendingCycleMemory: { instructions?: string } | null = null;
|
||||
|
||||
pi.on("agent_end", async (_event, ctx) => {
|
||||
if (!pendingCycleMemory) return;
|
||||
|
||||
const request = pendingCycleMemory;
|
||||
pendingCycleMemory = null;
|
||||
|
||||
// Signal to session_before_compact and session_compact hooks
|
||||
// to suppress their redundant UI — we show a single clean card.
|
||||
cycleMemoryActive = true;
|
||||
|
||||
ctx.ui.setStatus("memory-cycle", "Compacting context...");
|
||||
|
||||
ctx.compact({
|
||||
customInstructions: request.instructions
|
||||
?? "Create a comprehensive summary preserving all goals, decisions, progress, file changes, and context needed to continue work seamlessly.",
|
||||
onComplete: () => {
|
||||
cycleMemoryActive = false;
|
||||
|
||||
const postUsage = ctx.getContextUsage();
|
||||
const postPercent = postUsage?.percent ? Math.round(postUsage.percent) : 0;
|
||||
|
||||
// Read restored context for the agent
|
||||
const sessionState = readSessionState(ctx.cwd);
|
||||
const recentLogs = readRecentLogs();
|
||||
const parts = buildRestorationContent(sessionState);
|
||||
if (recentLogs) parts.push("", recentLogs);
|
||||
|
||||
const resumeContent = [
|
||||
"Memory cycle complete — context compacted and restored.",
|
||||
`Context usage now at ${postPercent}%.`,
|
||||
"",
|
||||
...parts,
|
||||
"",
|
||||
"Continue where you left off. Resume the task you were working on before compaction. Do NOT ask the user what to do — just keep working.",
|
||||
].join("\n");
|
||||
|
||||
ctx.ui.setStatus("memory-cycle", undefined);
|
||||
|
||||
// Single clean display card — no separate notify() to avoid
|
||||
// duplicate text noise in the terminal.
|
||||
pi.sendMessage(
|
||||
{
|
||||
customType: "memory-cycle-resume",
|
||||
content: `Memory cycle complete -- context compacted and restored.\nContext usage now at ${postPercent}%.`,
|
||||
display: true,
|
||||
details: {
|
||||
source: "cycle",
|
||||
postPercent,
|
||||
task: sessionState?.currentTask,
|
||||
recentFiles: sessionState?.filesEdited,
|
||||
} satisfies CompactionCardDetails,
|
||||
},
|
||||
);
|
||||
|
||||
// Full restoration context for the agent (not displayed)
|
||||
pi.sendMessage(
|
||||
{
|
||||
customType: "memory-cycle-resume",
|
||||
content: resumeContent,
|
||||
display: false,
|
||||
},
|
||||
{ deliverAs: "followUp", triggerTurn: true },
|
||||
);
|
||||
},
|
||||
onError: (err: Error) => {
|
||||
cycleMemoryActive = false;
|
||||
ctx.ui.setStatus("memory-cycle", undefined);
|
||||
ctx.ui.notify(`Memory Cycle failed: ${err.message}. Try /compact manually.`, "error");
|
||||
},
|
||||
});
|
||||
});
|
||||
|
||||
// ── cycle_memory tool (LLM-callable) ─────────────────────────
|
||||
|
||||
pi.registerTool({
|
||||
name: "cycle_memory",
|
||||
label: "Cycle Memory",
|
||||
description: "Compact current session, start fresh, and restore memory. Use when context is getting large or you want a clean slate while keeping all progress.",
|
||||
promptSnippet: "Compact → clear → restore: fresh context with full memory",
|
||||
promptGuidelines: [
|
||||
"Use cycle_memory when context usage is high (>70%) or the user asks to compact/cycle/refresh memory.",
|
||||
"After cycle_memory completes, you will have a fresh context window with full memory of what happened.",
|
||||
"The tool returns immediately — compaction happens after the current turn ends. You will be resumed automatically with restored context.",
|
||||
],
|
||||
parameters: CycleParams,
|
||||
|
||||
renderCall(args, theme) {
|
||||
const hint = (args as any).instructions as string | undefined;
|
||||
const preview = hint
|
||||
? hint.length > 50 ? hint.slice(0, 47) + "..." : hint
|
||||
: "";
|
||||
const text = theme.fg("dim", "cycle_memory") +
|
||||
(preview ? theme.fg("dim", " ") + theme.fg("muted", preview) : "");
|
||||
return new Text(text, 0, 0);
|
||||
},
|
||||
|
||||
renderResult(result, _options, theme) {
|
||||
const details = result.details as { status?: string } | undefined;
|
||||
const status = details?.status ?? "done";
|
||||
const msg = status === "scheduled"
|
||||
? theme.fg("dim", "Memory cycle scheduled — compacting after this turn")
|
||||
: theme.fg("dim", "Memory cycle complete");
|
||||
return new Text(msg, 0, 0);
|
||||
},
|
||||
|
||||
async execute(_toolCallId, params: { instructions?: string }, _signal, _onUpdate, ctx) {
|
||||
const customInstructions = params.instructions?.trim() || undefined;
|
||||
|
||||
// Schedule compaction for after this agent turn ends (avoids deadlock).
|
||||
// The agent_end hook above picks this up and fires ctx.compact().
|
||||
pendingCycleMemory = { instructions: customInstructions };
|
||||
|
||||
return {
|
||||
content: [
|
||||
{
|
||||
type: "text",
|
||||
text: "Memory cycle scheduled. Compaction will run automatically after this turn completes. You will be resumed with full memory context. Do not call any more tools — just finish this turn.",
|
||||
},
|
||||
],
|
||||
details: { status: "scheduled", instructions: params.instructions },
|
||||
};
|
||||
},
|
||||
});
|
||||
}
|
||||
Reference in New Issue
Block a user