pi-observational-memory-extension 0.1.0 → 0.1.2

package/extensions/index.ts

@@ -378,9 +378,28 @@ export default function (pi: ExtensionAPI) {
   pi.on("session_shutdown", async (_event, ctx) => {
     if (runtime.statusTimer) clearInterval(runtime.statusTimer);
     runtime.currentOperation?.abort();
-    if (runtime.state) await saveState(runtime.state);
-    ctx.ui.setStatus("om", undefined);
-    ctx.ui.setWidget("om", undefined);
+    if (runtime.state) {
+      try {
+        let isStale = false;
+        try {
+          const _ = ctx.cwd;
+        } catch {
+          isStale = true;
+        }
+        if (!isStale && runtime.state.enabled && runtime.state.pendingMessageTokens > 0) {
+          await observeNow(ctx, { force: true, reason: "session_shutdown" });
+        }
+      } catch (error) {
+        // Safe catch on shutdown
+      }
+      await saveState(runtime.state);
+    }
+    try {
+      ctx.ui.setStatus("om", undefined);
+      ctx.ui.setWidget("om", undefined);
+    } catch (e) {
+      // Ignore if UI or context is already stale
+    }
   });
 
   pi.on("model_select", async (event, ctx) => {
@@ -559,6 +578,14 @@ export default function (pi: ExtensionAPI) {
 }
 
 async function safeBoundary(ctx: any, boundary: string): Promise<void> {
+  let isStale = false;
+  try {
+    const _ = ctx.cwd;
+  } catch {
+    isStale = true;
+  }
+  if (isStale) return;
+
   const state = await ensureState(ctx);
   if (!state.enabled || state.status === "failed") return;
   await refreshCounts(ctx);
@@ -577,17 +604,40 @@ async function safeBoundary(ctx: any, boundary: string): Promise<void> {
   }
   if (shouldBuffer(state)) {
     void bufferObservation(ctx, `${boundary}:buffer`).catch(async (error) => {
-      const s = await ensureState(ctx);
-      s.status = "failed";
-      s.lastError = `OM buffer failed: ${errorMessage(error)}`;
-      await saveState(s);
-      updateStatus(ctx);
-      ctx.ui.notify(s.lastError, "error");
+      try {
+        let innerStale = false;
+        try {
+          const _ = ctx.cwd;
+        } catch {
+          innerStale = true;
+        }
+        if (innerStale) return;
+
+        const s = await ensureState(ctx);
+        s.status = "failed";
+        s.lastError = `OM buffer failed: ${errorMessage(error)}`;
+        await saveState(s);
+        updateStatus(ctx);
+        ctx.ui.notify(s.lastError, "error");
+      } catch (innerError) {
+        // Prevent any uncaught exceptions from background handler
+      }
     });
   }
 }
 
 async function ensureState(ctx: any): Promise<PiOMRecord> {
+  let isStale = false;
+  try {
+    const _ = ctx.cwd;
+  } catch {
+    isStale = true;
+  }
+  if (isStale) {
+    if (runtime.state) return runtime.state;
+    throw new Error("Cannot ensure state because the extension context is stale.");
+  }
+
   if (runtime.state && runtime.state.cwd === ctx.cwd && runtime.state.sessionFile === ctx.sessionManager.getSessionFile()) return runtime.state;
   const sessionFile = ctx.sessionManager.getSessionFile?.();
   const sessionId = ctx.sessionManager.getSessionId?.() || (sessionFile ? basename(sessionFile, ".jsonl") : "in-memory");
@@ -642,6 +692,14 @@ async function saveState(state: PiOMRecord): Promise<void> {
 }
 
 async function refreshCounts(ctx: any): Promise<void> {
+  let isStale = false;
+  try {
+    const _ = ctx.cwd;
+  } catch {
+    isStale = true;
+  }
+  if (isStale) return;
+
   const state = await ensureState(ctx);
   const branch = ctx.sessionManager.getBranch() as SessionEntry[];
   const pending = entriesAfter(branch, state.lastObservedEntryId).filter(isMessageLikeEntry);
@@ -653,11 +711,23 @@ async function refreshCounts(ctx: any): Promise<void> {
 function updateStatus(ctx: any): void {
   const state = runtime.state;
   if (!state) return;
-  ctx.ui.setStatus("om", formatShortStatusColored(state));
+  try {
+    ctx.ui.setStatus("om", formatShortStatusColored(state));
+  } catch (e) {
+    // Context might be stale or UI disposed
+  }
   runtime.overlayHandle?.requestRender();
 }
 
 async function observeNow(ctx: any, opts: { force: boolean; reason: string; signal?: AbortSignal; manualText?: string }): Promise<void> {
+  let isStale = false;
+  try {
+    const _ = ctx.cwd;
+  } catch {
+    isStale = true;
+  }
+  if (isStale) return;
+
   const state = await ensureState(ctx);
   assertNoOperation(state);
   const branch = ctx.sessionManager.getBranch() as SessionEntry[];
@@ -749,6 +819,14 @@ async function bufferObservation(ctx: any, reason: string): Promise<void> {
 }
 
 async function activateBuffered(ctx: any, reason: string): Promise<void> {
+  let isStale = false;
+  try {
+    const _ = ctx.cwd;
+  } catch {
+    isStale = true;
+  }
+  if (isStale) return;
+
   const state = await ensureState(ctx);
   if (state.buffered.observations.length === 0) return;
   assertNoOperation(state);
@@ -785,6 +863,14 @@ async function activateBuffered(ctx: any, reason: string): Promise<void> {
 }
 
 async function reflectNow(ctx: any, opts: { reason: string; signal?: AbortSignal; manualPrompt?: string }): Promise<void> {
+  let isStale = false;
+  try {
+    const _ = ctx.cwd;
+  } catch {
+    isStale = true;
+  }
+  if (isStale) return;
+
   const state = await ensureState(ctx);
   assertNoOperation(state);
   if (!state.observations.trim()) throw new Error("No observations to reflect");
@@ -1311,11 +1397,15 @@ function formatTokens(tokens: number): string {
 }
 
 async function writeDebug(ctx: any, name: string, payload: unknown): Promise<void> {
-  const state = await ensureState(ctx);
-  const dir = runtime.debugDir ?? join(ctx.cwd, CONFIG_DIR_NAME, "om", "debug");
-  await mkdir(dir, { recursive: true });
-  const file = join(dir, `${new Date().toISOString().replace(/[:.]/g, "-")}-${sanitizeFileName(name)}.json`);
-  await writeFile(file, JSON.stringify({ extension: EXTENSION_ID, sessionId: state.sessionId, ...payload as any }, null, 2) + "\n", "utf8");
+  try {
+    const state = await ensureState(ctx);
+    const dir = runtime.debugDir || join(ctx.cwd, CONFIG_DIR_NAME, "om", "debug");
+    await mkdir(dir, { recursive: true });
+    const file = join(dir, `${new Date().toISOString().replace(/[:.]/g, "-")}-${sanitizeFileName(name)}.json`);
+    await writeFile(file, JSON.stringify({ extension: EXTENSION_ID, sessionId: state.sessionId, ...payload as any }, null, 2) + "\n", "utf8");
+  } catch (error) {
+    // Ignore debug write errors if context is stale
+  }
 }
 
 function sanitizeFileName(name: string): string {

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "pi-observational-memory-extension",
-  "version": "0.1.0",
+  "version": "0.1.2",
   "description": "Mastra-style Observational Memory extension for Pi compaction and runtime context.",
   "license": "MIT",
   "author": "Nikita Nosov <20nik.nosov21@gmail.com>",