fix(pi): bind writes to acknowledged runtime sessions

docs/AGENT-SETUP.md

@@ -88,6 +88,8 @@ Other write tools still primarily use cwd/repo detection unless their schema say
 
 OpenCode binds `mem_save`, `mem_save_prompt`, `mem_session_summary`, and `mem_capture_passive` to its confirmed top-level runtime session and maps subagents to their authoritative parent.
 
+Pi's native HTTP wrappers use only `ctx.sessionManager.getSessionId()` for session-attributed writes. Model-supplied session IDs are ignored, and a missing or unconfirmed runtime ID fails safely.
+
 To lock write tools to the canonical project for a repo, add `.engram/config.json` at the repo root:
 
 ```json

plugin/pi/index.ts

@@ -178,10 +178,9 @@ function isTimeoutError(error: unknown): boolean {
   return error instanceof Error && (error.name === "TimeoutError" || error.name === "AbortError");
 }
 
-// engramFetch resolves to null on failure and ~20 call sites depend on that fallthrough —
-// ensureSession in particular must not abort a mem_save just because session creation blipped.
-// So the timeout detail travels out-of-band instead of changing what any caller receives,
-// letting executeMemoryTool tell the truth about an ambiguous write without blast radius.
+// engramFetch resolves to null on transport failure. Session-attributed writes
+// treat a null registration response as unacknowledged and stop before writing;
+// other callers retain the existing null fallthrough contract.
 let lastFetchTimeoutMethod: string | undefined;
 
 function takeLastFetchTimeoutMethod(): string | undefined {
@@ -415,14 +414,33 @@ let projectResolutionError: string | undefined;
 let projectDetectionPending = false;
 
 const knownSessions = new Set<string>();
+const sessionRegistrationsInFlight = new Map<string, Promise<void>>();
 const toolCounts = new Map<string, number>();
 
 async function ensureSession(sessionId: string, sessionProject = project): Promise<void> {
   const key = `${sessionProject}:${sessionId}`;
   if (!sessionId || knownSessions.has(key)) return;
-  knownSessions.add(key);
-  const body: SessionBody = { id: sessionId, project: sessionProject, directory };
-  await engramFetch("/sessions", { method: "POST", body });
+
+  const existingRegistration = sessionRegistrationsInFlight.get(key);
+  if (existingRegistration) return existingRegistration;
+
+  const registration = (async () => {
+    const body: SessionBody = { id: sessionId, project: sessionProject, directory };
+    const acknowledgement = await engramFetch("/sessions", { method: "POST", body });
+    if (acknowledgement === null) {
+      throw new Error(`gentle-engram could not confirm session registration for Pi runtime session ${sessionId}`);
+    }
+    knownSessions.add(key);
+  })();
+  sessionRegistrationsInFlight.set(key, registration);
+
+  try {
+    await registration;
+  } finally {
+    if (sessionRegistrationsInFlight.get(key) === registration) {
+      sessionRegistrationsInFlight.delete(key);
+    }
+  }
 }
 
 async function detectServerProject(cwd: string): Promise<CurrentProjectResponse | undefined> {
@@ -507,6 +525,14 @@ function getSessionId(ctx: SessionContext): string | undefined {
   return ctx.sessionManager.getSessionId();
 }
 
+function requireRuntimeSessionID(ctx: SessionContext): string {
+  const sessionId = ctx.sessionManager.getSessionId()?.trim();
+  if (!sessionId) {
+    throw new Error("Pi runtime session ID is unavailable; session-attributed writes require a native SessionContext ID");
+  }
+  return sessionId;
+}
+
 const optionalString = (description: string) => Type.Optional(Type.String({ description }));
 const optionalNumber = (description: string) => Type.Optional(Type.Number({ description }));
 const optionalBoolean = (description: string) => Type.Optional(Type.Boolean({ description }));
@@ -525,7 +551,6 @@ const MEMORY_TOOL_SCHEMAS: Record<string, ReturnType<typeof Type.Object>> = {
     title: Type.String({ description: "Short, searchable title" }),
     content: Type.String({ description: "Structured memory content" }),
     type: optionalString("Observation type/category"),
-    session_id: optionalString("Session ID to associate with"),
     scope: optionalString("Scope: project or personal"),
     topic_key: optionalString("Stable topic key for upserts"),
     project: optionalString("Optional explicit project"),
@@ -550,12 +575,10 @@ const MEMORY_TOOL_SCHEMAS: Record<string, ReturnType<typeof Type.Object>> = {
   }),
   mem_save_prompt: Type.Object({
     content: Type.String({ description: "The user's prompt text" }),
-    session_id: optionalString("Session ID to associate with"),
     project: optionalString("Optional project"),
   }),
   mem_session_summary: Type.Object({
     content: Type.String({ description: "Full session summary" }),
-    session_id: optionalString("Session ID"),
     project: optionalString("Optional project to use when automatic detection is unavailable"),
   }),
   mem_context: Type.Object({
@@ -591,7 +614,6 @@ const MEMORY_TOOL_SCHEMAS: Record<string, ReturnType<typeof Type.Object>> = {
   }),
   mem_capture_passive: Type.Object({
     content: Type.String({ description: "Text output containing a ## Key Learnings section" }),
-    session_id: optionalString("Session ID to associate with"),
     source: optionalString("Source identifier, e.g. subagent-stop or session-end"),
   }),
   mem_review: Type.Object({
@@ -652,7 +674,7 @@ async function callMemoryTool(toolName: string, params: Record<string, unknown>,
   const sessionId = getSessionId(ctx);
   const requestedProject = typeof params.project === "string" && params.project ? params.project : undefined;
   const activeProject = requestedProject || project;
-  const activeSessionId = String(params.session_id || (requestedProject ? `manual-save-${requestedProject}` : sessionId) || `manual-save-${project}`);
+  const runtimeSessionForWrite = () => requireRuntimeSessionID(ctx);
 
   switch (toolName) {
     case "mem_search":
@@ -674,8 +696,9 @@ async function callMemoryTool(toolName: string, params: Record<string, unknown>,
       return engramFetch(`/timeline${queryString({ observation_id: params.observation_id, before: params.before, after: params.after, project: params.project })}`);
     case "mem_get_observation":
       return engramFetch(`/observations/${encodeURIComponent(String(params.id))}`);
-    case "mem_save":
+    case "mem_save": {
       if (!requestedProject) requireResolvedProject();
+      const activeSessionId = runtimeSessionForWrite();
       await ensureSession(activeSessionId, activeProject);
       return engramFetch("/observations", {
         method: "POST",
@@ -689,6 +712,7 @@ async function callMemoryTool(toolName: string, params: Record<string, unknown>,
           topic_key: params.topic_key,
         },
       });
+    }
     case "mem_update":
       return engramFetch(`/observations/${encodeURIComponent(String(params.id))}`, {
         method: "PATCH",
@@ -704,27 +728,31 @@ async function callMemoryTool(toolName: string, params: Record<string, unknown>,
       return engramFetch(`/observations/${encodeURIComponent(String(params.id))}${queryString({ hard: params.hard_delete })}`, { method: "DELETE" });
     case "mem_suggest_topic_key":
       return { topic_key: slugifyTopicKey(params) };
-    case "mem_save_prompt":
+    case "mem_save_prompt": {
       if (!requestedProject) requireResolvedProject();
-      await ensureSession(activeSessionId, activeProject);
+      const promptSessionId = runtimeSessionForWrite();
+      await ensureSession(promptSessionId, activeProject);
       return engramFetch("/prompts", {
         method: "POST",
-        body: { session_id: activeSessionId, content: params.content, project: activeProject },
+        body: { session_id: promptSessionId, content: params.content, project: activeProject },
       });
-    case "mem_session_summary":
+    }
+    case "mem_session_summary": {
       if (!requestedProject) requireResolvedProject();
-      await ensureSession(activeSessionId, activeProject);
+      const summarySessionId = runtimeSessionForWrite();
+      await ensureSession(summarySessionId, activeProject);
       return engramFetch("/observations", {
         method: "POST",
         body: {
-          session_id: activeSessionId,
+          session_id: summarySessionId,
           type: "session_summary",
           title: "Session summary",
           content: params.content,
           project: activeProject,
           scope: "project",
         },
       });
+    }
     case "mem_session_start":
       requireResolvedProject();
       return engramFetch("/sessions", {
@@ -749,18 +777,20 @@ async function callMemoryTool(toolName: string, params: Record<string, unknown>,
     }
     case "mem_doctor":
       return engramFetch(`/doctor${queryString({ project: params.project, check: params.check, cwd: params.project ? undefined : ctx.cwd })}`);
-    case "mem_capture_passive":
+    case "mem_capture_passive": {
       requireResolvedProject();
-      await ensureSession(activeSessionId);
+      const passiveSessionId = runtimeSessionForWrite();
+      await ensureSession(passiveSessionId);
       return engramFetch("/observations/passive", {
         method: "POST",
         body: {
-          session_id: activeSessionId,
+          session_id: passiveSessionId,
           content: params.content,
           project,
           source: params.source || "pi-tool",
         },
       });
+    }
     case "mem_review": {
       const action = String(params.action || "").trim();
       if (action === "list") {
@@ -884,7 +914,7 @@ export default function registerEngram(pi: ExtensionAPI) {
     await refreshProjectDetection(ctx.cwd);
     if (projectDetectionPending || projectResolutionError) return;
     const sessionId = getSessionId(ctx);
-    if (sessionId) await ensureSessionBestEffort(sessionId);
+    if (sessionId) await ensureSession(sessionId);
 
     const summary = extractCompactedSummary(event);
     if (sessionId && summary) {

plugin/pi/test/index-source.test.mjs

@@ -2,7 +2,7 @@ import assert from "node:assert/strict";
 import { readFileSync } from "node:fs";
 import { test } from "node:test";
 
-const source = readFileSync(new URL("../index.ts", import.meta.url), "utf8");
+const source = readFileSync(new URL("../index.ts", import.meta.url), "utf8").replaceAll("\r\n", "\n");
 
 function extractFunctionBody(name, marker) {
   const signatureIndex = source.indexOf(`function ${name}`);
@@ -103,6 +103,22 @@ function buildScheduleEngramSelfHealForTest({ waitUnref, isEngramRunning, maxAtt
   return factory(waitUnref, isEngramRunning, 1, maxAttempts);
 }
 
+function buildEnsureSessionForTest(engramFetch) {
+  const body = extractFunctionBody("ensureSession", "{\n  const key")
+    .replace("const body: SessionBody", "const body");
+  const factory = new Function("knownSessions", "sessionRegistrationsInFlight", "engramFetch", "project", "directory", `
+    return async function ensureSession(sessionId, sessionProject = project) {
+      ${body}
+    };
+  `);
+  const knownSessions = new Set();
+  const sessionRegistrationsInFlight = new Map();
+  return {
+    ensureSession: factory(knownSessions, sessionRegistrationsInFlight, engramFetch, "engram", "/work/engram"),
+    knownSessions,
+  };
+}
+
 function sessionCtx(id, sink) {
   return {
     sessionManager: { getSessionId: () => id },
@@ -112,7 +128,7 @@ function sessionCtx(id, sink) {
 
 test("mem_session_summary accepts explicit project fallback", () => {
   assert.match(source, /mem_session_summary: Type\.Object\(\{[\s\S]*project: optionalString\("Optional project to use when automatic detection is unavailable"\)/);
-  assert.match(source, /case "mem_session_summary":[\s\S]*if \(!requestedProject\) requireResolvedProject\(\);[\s\S]*ensureSession\(activeSessionId, activeProject\)[\s\S]*project: activeProject/);
+  assert.match(source, /case "mem_session_summary":[\s\S]*if \(!requestedProject\) requireResolvedProject\(\);[\s\S]*ensureSession\(summarySessionId, activeProject\)[\s\S]*project: activeProject/);
 });
 
 test("mem_search exposes and forwards match_mode and all_projects", () => {
@@ -286,35 +302,47 @@ test("the tool layer reports unknown write outcome instead of inviting a blind r
   assert.doesNotMatch(unreachable, /timed out/);
 });
 
-test("a session-creation timeout still lets the observation write through", async () => {
-  // Regression: when engramFetch threw on timeout, the unguarded ensureSession call in
-  // mem_save aborted the whole tool call before /observations was ever attempted, silently
-  // dropping the user's memory while telling the agent not to retry.
-  assert.match(source, /await ensureSession\(activeSessionId, activeProject\);/);
-  assert.doesNotMatch(source, /throw new EngramTimeoutError/);
+test("session registration requires acknowledgement and failed acknowledgement remains retryable", async () => {
+  let calls = 0;
+  const { ensureSession, knownSessions } = buildEnsureSessionForTest(async () => {
+    calls += 1;
+    return calls === 1 ? null : { status: "created" };
+  });
 
-  const originalFetch = globalThis.fetch;
-  const paths = [];
-  globalThis.fetch = async (url, init) => {
-    const path = new URL(url).pathname;
-    paths.push(path);
-    if (path === "/sessions") {
-      const timeout = new Error("The operation was aborted due to timeout");
-      timeout.name = "TimeoutError";
-      throw timeout;
-    }
-    return { ok: true, async json() { return { id: 1 }; } };
-  };
-  try {
-    const { engramFetch } = buildEngramFetchForTest();
-    // ensureSession's own call fails soft...
-    assert.equal(await engramFetch("/sessions", { method: "POST", body: { id: "s" } }), null);
-    // ...and the observation write that follows it still lands.
-    assert.deepEqual(await engramFetch("/observations", { method: "POST", body: { title: "t" } }), { id: 1 });
-    assert.deepEqual(paths, ["/sessions", "/observations"]);
-  } finally {
-    globalThis.fetch = originalFetch;
+  await assert.rejects(ensureSession("runtime"), /could not confirm session registration/);
+  assert.equal(knownSessions.has("engram:runtime"), false);
+  await ensureSession("runtime");
+  assert.equal(knownSessions.has("engram:runtime"), true);
+  await ensureSession("runtime");
+  assert.equal(calls, 2);
+});
+
+test("session compaction strictly registers before forwarding its summary", () => {
+  const compactStart = source.indexOf('pi.on("session_compact"');
+  const compactEnd = source.indexOf('\n  pi.on("before_agent_start"', compactStart);
+  assert.notEqual(compactStart, -1, "session_compact handler not found");
+  assert.notEqual(compactEnd, -1, "session_compact handler end not found");
+  const compactHandler = source.slice(compactStart, compactEnd);
+
+  const registration = compactHandler.indexOf("if (sessionId) await ensureSession(sessionId);");
+  const summaryPost = compactHandler.indexOf('bestEffortEngramFetch("/observations"');
+  assert.notEqual(registration, -1, "session_compact must await strict session registration");
+  assert.notEqual(summaryPost, -1, "session_compact summary post not found");
+  assert.ok(registration < summaryPost, "strict registration must precede summary forwarding");
+  assert.doesNotMatch(compactHandler, /ensureSessionBestEffort/, "session_compact must not hide registration failure");
+});
+
+test("four session-attributed writes ignore model session_id and require the Pi runtime ID", () => {
+  for (const tool of ["mem_save", "mem_save_prompt", "mem_session_summary", "mem_capture_passive"]) {
+    const schema = source.match(new RegExp(`${tool}: Type\\.Object\\(\\{([\\s\\S]*?)\\n  \\}\\),`));
+    assert.ok(schema, `${tool} schema not found`);
+    assert.doesNotMatch(schema[1], /session_id:/, `${tool} must not invite model-supplied session identity`);
   }
+  assert.match(source, /function requireRuntimeSessionID/);
+  assert.match(source, /ctx\.sessionManager\.getSessionId\(\)/);
+  assert.match(source, /Pi runtime session ID is unavailable/);
+  assert.doesNotMatch(source, /const activeSessionId = String\(params\.session_id/);
+  assert.doesNotMatch(source, /manual-save-\$\{requestedProject\}/);
 });
 
 test("a timeout on the session leg does not mislabel an unrelated failure on the write leg", async () => {