refactor(pi): share session manager runtime registry

src/agents/pi-extensions/compaction-safeguard-runtime.ts

@@ -1,35 +1,12 @@
+import { createSessionManagerRuntimeRegistry } from "./session-manager-runtime-registry.js";
+
 export type CompactionSafeguardRuntimeValue = {
   maxHistoryShare?: number;
   contextWindowTokens?: number;
 };
 
-// Session-scoped runtime registry keyed by object identity.
+const registry = createSessionManagerRuntimeRegistry<CompactionSafeguardRuntimeValue>();
-// Follows the same WeakMap pattern as context-pruning/runtime.ts.
-const REGISTRY = new WeakMap<object, CompactionSafeguardRuntimeValue>();
 
-export function setCompactionSafeguardRuntime(
+export const setCompactionSafeguardRuntime = registry.set;
-  sessionManager: unknown,
-  value: CompactionSafeguardRuntimeValue | null,
-): void {
-  if (!sessionManager || typeof sessionManager !== "object") {
-    return;
-  }
 
-  const key = sessionManager;
+export const getCompactionSafeguardRuntime = registry.get;
-  if (value === null) {
-    REGISTRY.delete(key);
-    return;
-  }
-
-  REGISTRY.set(key, value);
-}
-
-export function getCompactionSafeguardRuntime(
-  sessionManager: unknown,
-): CompactionSafeguardRuntimeValue | null {
-  if (!sessionManager || typeof sessionManager !== "object") {
-    return null;
-  }
-
-  return REGISTRY.get(sessionManager) ?? null;
-}

src/agents/pi-extensions/context-pruning/runtime.ts

@@ -1,4 +1,5 @@
 import type { EffectiveContextPruningSettings } from "./settings.js";
+import { createSessionManagerRuntimeRegistry } from "../session-manager-runtime-registry.js";
 
 export type ContextPruningRuntimeValue = {
   settings: EffectiveContextPruningSettings;
@@ -7,34 +8,10 @@ export type ContextPruningRuntimeValue = {
   lastCacheTouchAt?: number | null;
 };
 
-// Session-scoped runtime registry keyed by object identity.
 // Important: this relies on Pi passing the same SessionManager object instance into
 // ExtensionContext (ctx.sessionManager) that we used when calling setContextPruningRuntime.
-const REGISTRY = new WeakMap<object, ContextPruningRuntimeValue>();
+const registry = createSessionManagerRuntimeRegistry<ContextPruningRuntimeValue>();
 
-export function setContextPruningRuntime(
+export const setContextPruningRuntime = registry.set;
-  sessionManager: unknown,
-  value: ContextPruningRuntimeValue | null,
-): void {
-  if (!sessionManager || typeof sessionManager !== "object") {
-    return;
-  }
 
-  const key = sessionManager;
+export const getContextPruningRuntime = registry.get;
-  if (value === null) {
-    REGISTRY.delete(key);
-    return;
-  }
-
-  REGISTRY.set(key, value);
-}
-
-export function getContextPruningRuntime(
-  sessionManager: unknown,
-): ContextPruningRuntimeValue | null {
-  if (!sessionManager || typeof sessionManager !== "object") {
-    return null;
-  }
-
-  return REGISTRY.get(sessionManager) ?? null;
-}

src/agents/pi-extensions/session-manager-runtime-registry.test.ts

@@ -0,0 +1,22 @@
+import { describe, expect, it } from "vitest";
+import { createSessionManagerRuntimeRegistry } from "./session-manager-runtime-registry.js";
+
+describe("createSessionManagerRuntimeRegistry", () => {
+  it("stores, reads, and clears values by object identity", () => {
+    const registry = createSessionManagerRuntimeRegistry<{ value: number }>();
+    const key = {};
+    expect(registry.get(key)).toBeNull();
+    registry.set(key, { value: 1 });
+    expect(registry.get(key)).toEqual({ value: 1 });
+    registry.set(key, null);
+    expect(registry.get(key)).toBeNull();
+  });
+
+  it("ignores non-object keys", () => {
+    const registry = createSessionManagerRuntimeRegistry<{ value: number }>();
+    registry.set(null, { value: 1 });
+    registry.set(123, { value: 1 });
+    expect(registry.get(null)).toBeNull();
+    expect(registry.get(123)).toBeNull();
+  });
+});

src/agents/pi-extensions/session-manager-runtime-registry.ts

@@ -0,0 +1,29 @@
+export function createSessionManagerRuntimeRegistry<TValue>() {
+  // Session-scoped runtime registry keyed by object identity.
+  // The SessionManager instance must stay stable across set/get calls.
+  const registry = new WeakMap<object, TValue>();
+
+  const set = (sessionManager: unknown, value: TValue | null): void => {
+    if (!sessionManager || typeof sessionManager !== "object") {
+      return;
+    }
+
+    const key = sessionManager;
+    if (value === null) {
+      registry.delete(key);
+      return;
+    }
+
+    registry.set(key, value);
+  };
+
+  const get = (sessionManager: unknown): TValue | null => {
+    if (!sessionManager || typeof sessionManager !== "object") {
+      return null;
+    }
+
+    return registry.get(sessionManager) ?? null;
+  };
+
+  return { set, get };
+}