fix(pi-memory): harden dream completion tracking

Author: HenryQW

package-lock.json

@@ -1,4 +1,4 @@
-import { lstat, mkdir, open, readFile, readdir, realpath, writeFile } from "node:fs/promises";
+import { lstat, mkdir, open, readFile, readdir, realpath, rename, unlink } from "node:fs/promises";
 import { join, sep } from "node:path";
 import { StringEnum } from "@earendil-works/pi-ai";
 import { getAgentDir, withFileMutationQueue, type ExtensionAPI } from "@earendil-works/pi-coding-agent";
@@ -78,6 +78,29 @@ async function loadLastDreamAt(): Promise<number | undefined> {
 	}
 }
 
+async function saveLastDreamAt(): Promise<void> {
+	const path = LAST_DREAM_PATH();
+	const tempPath = `${path}.${process.pid}.${Date.now()}.${Math.random().toString(36).slice(2)}`;
+	let created = false;
+	try {
+		const handle = await open(tempPath, "wx", 0o600);
+		created = true;
+		try {
+			await handle.writeFile(`${new Date().toISOString()}\n`);
+		} finally {
+			await handle.close();
+		}
+		// rename replaces a destination symlink rather than following it.
+		await rename(tempPath, path);
+	} finally {
+		if (created) {
+			await unlink(tempPath).catch((error: NodeJS.ErrnoException) => {
+				if (error.code !== "ENOENT") throw error;
+			});
+		}
+	}
+}
+
 function sanitizeEntry(entry: string): string {
 	return entry.split("\n").map((line) => isReservedFrameLine(line) ? FRAME_TOKEN_REPLACEMENT : line).join("\n");
 }
@@ -249,7 +272,7 @@ export default function memoryExtension(pi: ExtensionAPI): void {
 		for (let index = event.messages.length - 1; index >= 0; index--) {
 			const message = event.messages[index];
 			if (message?.role !== "assistant") continue;
-			state.dreamSucceeded = message.stopReason !== "error" && message.stopReason !== "aborted" && message.stopReason !== "length";
+			state.dreamSucceeded = message.stopReason === "stop";
 			break;
 		}
 	});
@@ -264,7 +287,7 @@ export default function memoryExtension(pi: ExtensionAPI): void {
 			return;
 		}
 		try {
-			await writeFile(LAST_DREAM_PATH(), `${new Date().toISOString()}\n`, { mode: 0o600 });
+			await saveLastDreamAt();
 		} catch (error) {
 			ctx.ui.notify(`Dream completed, but its timestamp could not be recorded: ${error instanceof Error ? error.message : String(error)}`, "warning");
 		}
@@ -323,12 +346,14 @@ export default function memoryExtension(pi: ExtensionAPI): void {
 			state.snapshotSanitized = rendered.some(({ sanitized }) => sanitized);
 			state.conflictWarnings = conflictWarnings;
 
-			if (!process.argv.includes(BTW_CHILD_PAYLOAD_ARG) && (memory.entries.length || user.entries.length)) {
+			const memoryChars = memory.entries.join(ENTRY_DELIMITER).length;
+			const userChars = user.entries.join(ENTRY_DELIMITER).length;
+			const validWithinCap = !memory.status && !user.status
+				&& memoryChars <= config.memoryCharLimit && userChars <= config.userCharLimit;
+			if (!process.argv.includes(BTW_CHILD_PAYLOAD_ARG) && validWithinCap && (memory.entries.length || user.entries.length)) {
 				try {
 					const lastDreamAt = await loadLastDreamAt();
 					const age = lastDreamAt === undefined ? undefined : Date.now() - lastDreamAt;
-					const memoryChars = memory.entries.join(ENTRY_DELIMITER).length;
-					const userChars = user.entries.join(ENTRY_DELIMITER).length;
 					const full = memoryChars * 100 >= config.memoryCharLimit * DREAM_USAGE_PERCENT
 						|| userChars * 100 >= config.userCharLimit * DREAM_USAGE_PERCENT;
 					if (age === undefined || age >= DREAM_AFTER_MS || (full && age >= DREAM_FULL_COOLDOWN_MS)) {

packages/pi-memory/extensions/memory.ts

@@ -1,6 +1,6 @@
 {
 	"name": "@henryqw/pi-memory",
-	"version": "1.2.1",
+	"version": "1.3.0",
 	"description": "Auto-managed markdown memory for Pi: capped MEMORY.md/USER.md entry stores with frozen session snapshots.",
 	"keywords": [
 		"pi-package",

packages/pi-memory/package.json

@@ -1,5 +1,5 @@
 import assert from "node:assert/strict";
-import { mkdir, mkdtemp, readdir, readFile, rm, symlink, writeFile } from "node:fs/promises";
+import { lstat, mkdir, mkdtemp, readdir, readFile, rm, symlink, writeFile } from "node:fs/promises";
 import { tmpdir } from "node:os";
 import { join } from "node:path";
 import test from "node:test";
@@ -101,7 +101,7 @@ test("/remember validates input, rejects busy agents, and sends live state", asy
 	}
 });
 
-test("session start recommends /dream when memory is new, stale, or 70% full", async () => {
+test("session start recommends /dream only for valid stores within their caps", async () => {
 	const root = await mkdtemp(join(tmpdir(), "pi-memory-dream-reminder-"));
 	const agentDir = join(root, "agent");
 	const memoryDir = join(root, "memory");
@@ -147,6 +147,28 @@ test("session start recommends /dream when memory is new, stale, or 70% full", a
 		assert.deepEqual(notifications, ["Memory dream recommended; run /dream."]);
 
 		notifications.length = 0;
+		await rm(statePath);
+		await writeFile(join(memoryDir, "MEMORY.md"), "x".repeat(11));
+		await handlers.get("session_start")!({ type: "session_start" }, ctx);
+		assert.deepEqual(notifications, []);
+
+		await writeFile(join(memoryDir, "MEMORY.md"), "123456");
+		await writeFile(join(memoryDir, "USER.md"), "x".repeat(11));
+		await handlers.get("session_start")!({ type: "session_start" }, ctx);
+		assert.deepEqual(notifications, []);
+
+		await writeFile(join(memoryDir, "USER.md"), "");
+		await writeFile(join(memoryDir, "MEMORY.md"), "x".repeat(MAX_FILE_BYTES + 1));
+		await handlers.get("session_start")!({ type: "session_start" }, ctx);
+		assert.deepEqual(notifications, []);
+
+		await rm(join(memoryDir, "MEMORY.md"));
+		await symlink(join(root, "missing-MEMORY.md"), join(memoryDir, "MEMORY.md"));
+		await handlers.get("session_start")!({ type: "session_start" }, ctx);
+		assert.deepEqual(notifications, []);
+
+		await rm(join(memoryDir, "MEMORY.md"));
+		await writeFile(join(memoryDir, "MEMORY.md"), "123456");
 		await writeFile(statePath, "x".repeat(65));
 		await handlers.get("session_start")!({ type: "session_start" }, ctx);
 		assert.match(notifications[0]!, /Timestamp file is too large/);
@@ -248,19 +270,36 @@ test("/dream reuses unchanged memory snapshots and guards the agent-global SYSTE
 		assert.match(messages[0]!, /one memory batch per affected target/);
 		assert.match(messages[0]!, /no memory call if none/);
 
+		await rm(lastDreamPath);
+		await dream.handler("", context(true));
+		await handlers.get("agent_end")!({ type: "agent_end", messages: [{ role: "assistant", stopReason: "toolUse" }] });
+		await handlers.get("agent_settled")!({ type: "agent_settled" }, context(true));
+		await assert.rejects(readFile(lastDreamPath), /ENOENT/);
+		assert.equal(notifications.at(-1), "Dream did not complete; its timestamp was not updated.");
+
+		const dreamTarget = join(root, "dream-target.txt");
+		await writeFile(dreamTarget, "keep this target");
+		await symlink(dreamTarget, lastDreamPath);
+		await dream.handler("", context(true));
+		await handlers.get("agent_end")!({ type: "agent_end", messages: [{ role: "assistant", stopReason: "stop" }] });
+		await handlers.get("agent_settled")!({ type: "agent_settled" }, context(true));
+		assert.equal(await readFile(dreamTarget, "utf8"), "keep this target");
+		assert.equal((await lstat(lastDreamPath)).isSymbolicLink(), false);
+		assert.ok(Number.isFinite(Date.parse((await readFile(lastDreamPath, "utf8")).trim())));
+
 		process.argv.push(CHILD_PAYLOAD_ARG);
 		try {
 			await dream.handler("", context(true));
-			assert.ok(messages[1]!.includes(JSON.stringify({ memory: ["stable fact"], user: ["likes concise replies"] })));
-			assert.doesNotMatch(messages[1]!, /do not reread those files/);
-			assert.ok(messages[1]!.includes(`Read ${JSON.stringify(systemPath)} before semantic deduplication or editing.`));
+			assert.ok(messages.at(-1)!.includes(JSON.stringify({ memory: ["stable fact"], user: ["likes concise replies"] })));
+			assert.doesNotMatch(messages.at(-1)!, /do not reread those files/);
+			assert.ok(messages.at(-1)!.includes(`Read ${JSON.stringify(systemPath)} before semantic deduplication or editing.`));
 		} finally {
 			process.argv.pop();
 		}
 
 		await writeFile(join(memoryDir, "MEMORY.md"), "changed fact");
 		await dream.handler("", context(true));
-		assert.ok(messages[2]!.includes(JSON.stringify({ memory: ["changed fact"], user: ["likes concise replies"] })));
+		assert.ok(messages.at(-1)!.includes(JSON.stringify({ memory: ["changed fact"], user: ["likes concise replies"] })));
 
 		await rm(lastDreamPath);
 		await dream.handler("", context(true));