From 9084c4e3452d8612004c4f58c22999f01922f7ec Mon Sep 17 00:00:00 2001 From: Peter Steinberger Date: Sun, 15 Feb 2026 17:39:10 +0000 Subject: [PATCH] refactor(pi): share session manager runtime registry --- .../compaction-safeguard-runtime.ts | 33 +++---------------- .../pi-extensions/context-pruning/runtime.ts | 31 +++-------------- .../session-manager-runtime-registry.test.ts | 22 +++++++++++++ .../session-manager-runtime-registry.ts | 29 ++++++++++++++++ 4 files changed, 60 insertions(+), 55 deletions(-) create mode 100644 src/agents/pi-extensions/session-manager-runtime-registry.test.ts create mode 100644 src/agents/pi-extensions/session-manager-runtime-registry.ts diff --git a/src/agents/pi-extensions/compaction-safeguard-runtime.ts b/src/agents/pi-extensions/compaction-safeguard-runtime.ts index bda1b1de638..df3919cf815 100644 --- a/src/agents/pi-extensions/compaction-safeguard-runtime.ts +++ b/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. -// Follows the same WeakMap pattern as context-pruning/runtime.ts. -const REGISTRY = new WeakMap(); +const registry = createSessionManagerRuntimeRegistry(); -export function setCompactionSafeguardRuntime( - sessionManager: unknown, - value: CompactionSafeguardRuntimeValue | null, -): void { - if (!sessionManager || typeof sessionManager !== "object") { - return; - } +export const setCompactionSafeguardRuntime = registry.set; - const key = sessionManager; - 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; -} +export const getCompactionSafeguardRuntime = registry.get; diff --git a/src/agents/pi-extensions/context-pruning/runtime.ts b/src/agents/pi-extensions/context-pruning/runtime.ts index 7780464d1da..6d4fd07a2fb 100644 --- a/src/agents/pi-extensions/context-pruning/runtime.ts +++ b/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(); +const registry = createSessionManagerRuntimeRegistry(); -export function setContextPruningRuntime( - sessionManager: unknown, - value: ContextPruningRuntimeValue | null, -): void { - if (!sessionManager || typeof sessionManager !== "object") { - return; - } +export const setContextPruningRuntime = registry.set; - const key = sessionManager; - 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; -} +export const getContextPruningRuntime = registry.get; diff --git a/src/agents/pi-extensions/session-manager-runtime-registry.test.ts b/src/agents/pi-extensions/session-manager-runtime-registry.test.ts new file mode 100644 index 00000000000..59e004fab7a --- /dev/null +++ b/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(); + }); +}); diff --git a/src/agents/pi-extensions/session-manager-runtime-registry.ts b/src/agents/pi-extensions/session-manager-runtime-registry.ts new file mode 100644 index 00000000000..a23a7385d6a --- /dev/null +++ b/src/agents/pi-extensions/session-manager-runtime-registry.ts @@ -0,0 +1,29 @@ +export function createSessionManagerRuntimeRegistry() { + // Session-scoped runtime registry keyed by object identity. + // The SessionManager instance must stay stable across set/get calls. + const registry = new WeakMap(); + + 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 }; +}