Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
9 changes: 5 additions & 4 deletions README.md
Original file line number Diff line number Diff line change
Expand Up @@ -130,8 +130,9 @@ To keep the prefix byte-stable, the extension snapshots the memory context at de

- **`session_start`** — fresh snapshot per session
- **`session_before_compact`** — handoff is written then snapshot refreshes (one intentional cache boundary at compaction)
- **`memory_write` with `target: long_term`** — marks the snapshot dirty so the next turn refreshes (long-term writes are rare, intentional, and the user expects them to stick as ambient context)
- **Day rollover** — snapshot's captured date no longer matches today
- **`session_start`** is the only checkpoint in `stable` mode. Long-term writes and day rollovers do **not** re-render the block: a refresh rewrites the tail of the system prompt and voids the prefix cache for the whole conversation, which is the cost the snapshot exists to avoid, paid on the most common in-session event. The written fact is already in tool-call history, and `memory_read` / `memory_search` reach the files directly.
- **Deletions and restores** are delivered as an injected session message, not by re-rendering the block and not by appending to the system prompt. A forgotten memory must stop being authoritative, but the system prompt is a cache prefix: appending one line to it reprocesses the whole conversation on a recurrent/hybrid model (measured 15.5k tokens / 17.8 s). A message lands at the tail of the history, which is free, and pi persists it for later turns.
- Set `PI_MEMORY_SNAPSHOT=refresh` for the old behaviour (refresh on long-term write and day rollover, with a `Snapshot <reason> at <hh:mm:ss>` caveat line).

`memory_write` with `target: daily` and `scratchpad` writes do **not** mark dirty — they're high-frequency and the write content is already echoed via tool-call args. The model can always call `memory_read` / `memory_search` for the authoritative latest state.

Expand Down Expand Up @@ -186,7 +187,7 @@ This ensures in-progress context survives compaction and is visible in the next
| Variable | Values | Default | Description |
|----------|--------|---------|-------------|
| `PI_MEMORY_DIR` | path | `~/.pi/agent/memory` | Override the memory storage directory |
| `PI_MEMORY_SNAPSHOT` | `stable`, `per-turn` | `stable` | `stable` snapshots memory at checkpoints for KV cache stability; `per-turn` rebuilds every turn (legacy behavior) |
| `PI_MEMORY_SNAPSHOT` | `stable`, `refresh`, `per-turn` | `stable` | `stable` snapshots once at session start and never re-renders it (deletions append a correction); `refresh` also re-renders on long-term writes and day rollover; `per-turn` rebuilds every turn (legacy behavior) |
| `PI_MEMORY_QMD_UPDATE` | `background`, `manual`, `off` | `background` | Controls automatic `qmd update` + `qmd embed` after writes |
| `PI_MEMORY_QMD_SEARCH_TIMEOUT_MS` | positive integer (milliseconds) | `60000` | Sets the timeout for explicit `memory_search` qmd queries |
| `PI_MEMORY_NO_SEARCH` | `1` | unset | Disable selective injection in `per-turn` mode (no effect in `stable` mode) |
Expand All @@ -206,7 +207,7 @@ Run the `memory_status` tool first — it reports most of these at a glance.
| “need embeddings” on semantic/deep search | Vectors not built yet | Embedding starts automatically in the background — retry shortly. If `PI_MEMORY_QMD_UPDATE` is `manual`/`off`, run `qmd embed` yourself |
| Collection `pi-memory` missing | Auto-setup didn't run (qmd installed mid-session) | Run any `memory_search` (auto-creates it) or `qmd collection add ~/.pi/agent/memory --name pi-memory` |
| qmd works in the shell but not from pi on Windows | Broken `.cmd`/`.ps1` shims | The extension bypasses them by invoking qmd's JS entry with `node`; make sure the npm global `node_modules` dir is on `PATH` |
| Memory isn't being injected after a write | Cache-stable snapshot only refreshes at checkpoints | Long-term writes refresh next turn; for daily/scratchpad use `memory_read`, or set `PI_MEMORY_SNAPSHOT=per-turn` |
| Memory isn't being injected after a write | The snapshot is taken once per session and deliberately not re-rendered | The write is visible in tool-call history; use `memory_read` / `memory_search` for the current state, or set `PI_MEMORY_SNAPSHOT=refresh` (costs a full prompt reprocess per write) |

## Running tests

Expand Down
93 changes: 85 additions & 8 deletions index.ts
Original file line number Diff line number Diff line change
Expand Up @@ -1396,6 +1396,35 @@ let snapshotTakenAt: string | null = null;
let snapshotTakenOnDate: string | null = null;
let snapshotReason: string | null = null;
let snapshotDirty = false;
// Append-only list of things that stopped being true AFTER the snapshot was
// taken (deletions, restores). Re-rendering the whole block to reflect them
// would move every byte of it and cost a full prompt reprocess, so these are
// drained into an injected message instead (see before_agent_start) and still
// stop the model from acting on a memory that was explicitly forgotten. Writes
// are NOT listed here: the written fact is already in the tool-call history.
let snapshotCorrections: string[] = [];

const SNAPSHOT_CORRECTIONS_MAX = 20;

function noteSnapshotCorrection(line: string) {
if (snapshotCorrections.length >= SNAPSHOT_CORRECTIONS_MAX) return;
if (!snapshotCorrections.includes(line)) snapshotCorrections.push(line);
}

// One short, stable line per changed entry: strip the HTML id/timestamp
// comments, take the first line, cap the length.
function correctionPreview(entries: string[]): string[] {
return entries
.map((e) =>
e
.replace(/<!--[\s\S]*?-->/g, "")
.trim()
.split("\n")[0]
.trim()
.slice(0, 160),
)
.filter(Boolean);
}

function refreshMemorySnapshot(reason: string) {
memorySnapshot = buildMemoryContext("");
Expand All @@ -1405,9 +1434,11 @@ function refreshMemorySnapshot(reason: string) {
snapshotDirty = false;
}

function getSnapshotMode(): "stable" | "per-turn" {
function getSnapshotMode(): "stable" | "refresh" | "per-turn" {
const mode = (process.env.PI_MEMORY_SNAPSHOT ?? "stable").toLowerCase();
return mode === "per-turn" ? "per-turn" : "stable";
if (mode === "per-turn") return "per-turn";
if (mode === "refresh") return "refresh";
return "stable";
}

/** Reset snapshot state (for testing). */
Expand All @@ -1417,6 +1448,7 @@ export function _resetMemorySnapshot() {
snapshotTakenOnDate = null;
snapshotReason = null;
snapshotDirty = false;
snapshotCorrections = [];
}

// ---------------------------------------------------------------------------
Expand Down Expand Up @@ -1548,18 +1580,33 @@ export default function (pi: ExtensionAPI) {
const searchResults = skipSearch ? "" : await searchRelevantMemories(event.prompt ?? "");
memoryContext = buildMemoryContext(searchResults);
} else {
// "stable" means stable: once taken, the block is emitted byte-for-byte
// for the rest of the session. Refreshing on a long-term write or a
// midnight rollover rewrites the tail of the system prompt and voids the
// whole conversation's prefix cache — the exact cost the snapshot exists
// to avoid, paid on the single most common in-session event. The fresh
// state is not lost: the write is in tool-call history a few messages
// back, deletions are sent as a correction message below, and
// memory_read / memory_search reach the files directly. "refresh" restores the old
// checkpoint behaviour.
const today = todayStr();
const needsRefresh = memorySnapshot === null || snapshotDirty || snapshotTakenOnDate !== today;
if (needsRefresh) {
const stale = mode === "refresh" && (snapshotDirty || snapshotTakenOnDate !== today);
if (memorySnapshot === null || stale) {
const reason =
memorySnapshot === null ? "before_agent_start" : snapshotDirty ? "long_term_write" : "day_rollover";
refreshMemorySnapshot(reason);
}
memoryContext = memorySnapshot ?? "";
// Deliberately carries no timestamp and no reason word: both change
// between turns without the memory itself changing, which is enough on
// its own to invalidate the cache this branch is trying to preserve.
snapshotCaveat =
`Snapshot ${snapshotReason} at ${snapshotTakenAt}. ` +
"Use memory_read / memory_search for the authoritative latest state; " +
"recent writes may also be visible in tool-call history.";
mode === "refresh"
? `Snapshot ${snapshotReason} at ${snapshotTakenAt}. ` +
"Use memory_read / memory_search for the authoritative latest state; " +
"recent writes may also be visible in tool-call history."
: "Loaded once at session start and not re-read since. Use memory_read / memory_search " +
"for the authoritative latest state; anything written this session is in tool-call history.";
}

if (!memoryContext) return;
Expand All @@ -1578,8 +1625,31 @@ export default function (pi: ExtensionAPI) {
memoryContext,
);

// Corrections are delivered as an injected message, NOT appended to the
// system prompt. Appending would still move the prefix boundary, and a
// recurrent/hybrid model (GatedDeltaNet, Mamba) cannot rewind its state to
// a partial match — it reuses zero tokens and reprocesses the entire
// conversation, measured at 15.5k tokens / 17.8 s for one appended line.
// A message lands at the tail of the history, which costs nothing, and pi
// persists it in the session so it stays visible on later turns; that is
// also why the queue is drained after emitting rather than re-sent.
let correctionMessage: { customType: string; content: string; display: boolean } | undefined;
if (mode !== "refresh" && snapshotCorrections.length > 0) {
correctionMessage = {
customType: "pi-memory-correction",
content: [
"Memory corrections - these override the ## Memory block in the system prompt,",
"which was loaded at session start and is not re-read:",
...snapshotCorrections,
].join("\n"),
display: true,
};
snapshotCorrections = [];
}

return {
systemPrompt: event.systemPrompt + headerLines.join("\n"),
...(correctionMessage ? { message: correctionMessage } : {}),
};
});

Expand Down Expand Up @@ -2116,8 +2186,12 @@ export default function (pi: ExtensionAPI) {
fs.writeFileSync(filePath, result.content, "utf-8");
// Deleted facts must leave the injected snapshot too, whichever file
// they lived in — a forgotten-but-still-injected memory defeats the
// point of forgetting.
// point of forgetting. In stable mode that is done by appending a
// correction rather than re-rendering the block.
snapshotDirty = true;
for (const line of correctionPreview(result.removed)) {
noteSnapshotCorrection(`- FORGOTTEN, no longer true: ${line}${target === "daily" ? " (daily log)" : ""}`);
}
await ensureQmdAvailableForUpdate();
scheduleQmdUpdate();

Expand Down Expand Up @@ -2186,6 +2260,9 @@ export default function (pi: ExtensionAPI) {
const separator = existing.trim() ? "\n\n" : "";
fs.writeFileSync(targetPath, `${existing}${separator}${missingEntries.join("\n\n")}\n`, "utf-8");
snapshotDirty = true;
for (const line of correctionPreview(missingEntries)) {
noteSnapshotCorrection(`- RESTORED, true again: ${line}`);
}
await ensureQmdAvailableForUpdate();
scheduleQmdUpdate();
}
Expand Down
58 changes: 53 additions & 5 deletions test/unit.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -1912,7 +1912,29 @@ describe("KV cache stability: memory snapshot", () => {
expect(result2.systemPrompt).not.toBe(result1.systemPrompt);
});

test("memory_write target=long_term marks snapshot dirty so next turn refreshes", async () => {
test("memory_write target=long_term does NOT refresh the snapshot (cache stays warm)", async () => {
fs.writeFileSync(path.join(tmpDir, "MEMORY.md"), "OLD_FACT line", "utf-8");

const result1 = await hooks.before_agent_start({ systemPrompt: "base" }, {});
expect(result1.systemPrompt).toContain("OLD_FACT");

await tools.memory_write.execute(
"tc1",
{ target: "long_term", content: "NEW_FACT_ABOUT_X", mode: "append" },
null,
null,
createMockCtx(),
);

// The write is already in tool-call history; re-rendering the block would
// rewrite the prompt tail and void the whole conversation's prefix cache.
const result2 = await hooks.before_agent_start({ systemPrompt: "base" }, {});
expect(result2.systemPrompt).toBe(result1.systemPrompt);
expect(result2.systemPrompt).not.toContain("NEW_FACT_ABOUT_X");
});

test("PI_MEMORY_SNAPSHOT=refresh restores checkpoint refresh on long_term writes", async () => {
process.env.PI_MEMORY_SNAPSHOT = "refresh";
fs.writeFileSync(path.join(tmpDir, "MEMORY.md"), "OLD_FACT line", "utf-8");

const result1 = await hooks.before_agent_start({ systemPrompt: "base" }, {});
Expand All @@ -1928,10 +1950,32 @@ describe("KV cache stability: memory snapshot", () => {

const result2 = await hooks.before_agent_start({ systemPrompt: "base" }, {});
expect(result2.systemPrompt).toContain("NEW_FACT_ABOUT_X");
// Snapshot did refresh, so previous bytes are no longer identical.
expect(result2.systemPrompt).not.toBe(result1.systemPrompt);
});

test("memory_forget sends a correction message and leaves the systemPrompt untouched", async () => {
fs.writeFileSync(path.join(tmpDir, "MEMORY.md"), "WRONG_FACT_ABOUT_Z\n\nkeep me\n", "utf-8");

const result1 = await hooks.before_agent_start({ systemPrompt: "base" }, {});
expect(result1.systemPrompt).toContain("WRONG_FACT_ABOUT_Z");
expect(result1.message).toBeUndefined();

await tools.memory_forget.execute("tc1", { match: "WRONG_FACT_ABOUT_Z" }, null, null, {});

// A forgotten memory must stop being authoritative — but via a message at
// the tail of the history, not by moving the cached prompt prefix.
const result2 = await hooks.before_agent_start({ systemPrompt: "base" }, {});
expect(result2.systemPrompt).toBe(result1.systemPrompt);
expect(result2.message).toBeDefined();
expect(result2.message.customType).toBe("pi-memory-correction");
expect(result2.message.content).toContain("WRONG_FACT_ABOUT_Z");

// pi persists the injected message, so it is sent once and then drained.
const result3 = await hooks.before_agent_start({ systemPrompt: "base" }, {});
expect(result3.systemPrompt).toBe(result1.systemPrompt);
expect(result3.message).toBeUndefined();
});

test("memory_write target=daily does NOT mark snapshot dirty (cache stays warm)", async () => {
fs.writeFileSync(path.join(tmpDir, "MEMORY.md"), "Stable long-term content", "utf-8");

Expand Down Expand Up @@ -1992,11 +2036,15 @@ describe("KV cache stability: memory snapshot", () => {
}
});

test("snapshot caveat is included in stable mode header", async () => {
test("stable mode header carries a caveat with no volatile timestamp", async () => {
fs.writeFileSync(path.join(tmpDir, "MEMORY.md"), "anything", "utf-8");
const result = await hooks.before_agent_start({ systemPrompt: "base" }, {});
// Reader-facing hint that ambient context may lag behind disk.
expect(result.systemPrompt.toLowerCase()).toContain("snapshot");
// Reader-facing hint that ambient context may lag behind disk...
expect(result.systemPrompt).toContain("not re-read since");
expect(result.systemPrompt).toContain("memory_search");
// ...but no clock and no reason word: either would change the bytes
// between turns without the memory itself changing.
expect(result.systemPrompt).not.toMatch(/\d{2}:\d{2}:\d{2}/);
});
});

Expand Down
Loading