Merge pull request #926 from alvinunreal/fix/issue-908-session-map-cleanup

fix: bound session metadata maps (#908)

src/config/constants.ts

@@ -94,6 +94,13 @@ export const DEFAULT_READ_CONTEXT_MIN_LINES = 10;
 export const DEFAULT_READ_CONTEXT_MAX_FILES = 8;
 export const DEFAULT_MAX_RETAINED_SNAPSHOTS = 20;
 
+/**
+ * Maximum session metadata entries retained per plugin instance.
+ * Prevents unbounded growth when session.deleted events are missed.
+ * Oldest entries are evicted first when this threshold is reached.
+ */
+export const DEFAULT_MAX_SESSION_METADATA_ENTRIES = 1000;
+
 export type ImageRouting = 'auto' | 'direct';
 
 /**

src/index.ts

@@ -19,6 +19,7 @@ import {
   AGENT_ALIASES,
   DEFAULT_MAX_CONTEXT_LINES,
   DEFAULT_MAX_RETAINED_SNAPSHOTS,
+  DEFAULT_MAX_SESSION_METADATA_ENTRIES,
   DEFAULT_MAX_SESSIONS_PER_AGENT,
   DEFAULT_READ_CONTEXT_MAX_FILES,
   DEFAULT_READ_CONTEXT_MIN_LINES,
@@ -75,6 +76,7 @@ import {
 } from './utils';
 import { isPluginDisabledByEnv } from './utils/env';
 import { initLogger, log } from './utils/logger';
+import { SessionMetadataStore } from './utils/session-metadata';
 import { collapseSystemInPlace } from './utils/system-collapse';
 
 /**
@@ -147,10 +149,15 @@ const OhMyOpenCodeLite: Plugin = async (ctx) => {
   let multiplexerEnabled: boolean;
   let multiplexerSessionManager: MultiplexerSessionManager;
   let autoUpdateChecker: ReturnType<typeof createAutoUpdateCheckerHook>;
-  let sessionAgentMap: Map<string, string>;
-  // ponytail: cache sessionID -> project directory so TUI model writes
-  // land in the right per-project file after a project switch (ctx.directory is stale)
-  const sessionDirectories = new Map<string, string>();
+  const sessionMetadata = new SessionMetadataStore({
+    maxEntries: DEFAULT_MAX_SESSION_METADATA_ENTRIES,
+    onEvict: (sessionID) => {
+      log('[session] evicted oldest session metadata', {
+        threshold: DEFAULT_MAX_SESSION_METADATA_ENTRIES,
+        droppedSessionId: sessionID,
+      });
+    },
+  });
   let sessionLifecycle: SessionLifecycle;
 
   let chatHeadersHook: ReturnType<typeof createChatHeadersHook>;
@@ -296,9 +303,6 @@ const OhMyOpenCodeLite: Plugin = async (ctx) => {
       companion: config.companion,
     });
 
-    // Track session → agent mapping for serve-mode system prompt injection
-    sessionAgentMap = new Map<string, string>();
-
     chatHeadersHook = createChatHeadersHook(ctx);
 
     // Initialize foreground fallback manager for runtime model switching.
@@ -332,9 +336,9 @@ const OhMyOpenCodeLite: Plugin = async (ctx) => {
       continueOnIdle: config.backgroundJobs?.continueOnIdle === true,
       backgroundJobBoard: backgroundJobCoordinator,
       shouldManageSession: (sessionID) =>
-        sessionAgentMap.get(sessionID) === 'orchestrator',
+        sessionMetadata.getAgent(sessionID) === 'orchestrator',
       registerSessionAsOrchestrator: (sessionID) => {
-        sessionAgentMap.set(sessionID, 'orchestrator');
+        sessionMetadata.setAgent(sessionID, 'orchestrator');
       },
       isFallbackInProgress: (sessionID) =>
         foregroundFallback.isFallbackInProgress(sessionID),
@@ -373,7 +377,7 @@ const OhMyOpenCodeLite: Plugin = async (ctx) => {
     // Both message transforms share this gate so a rejected nudge cannot be
     // followed by a phase reminder in the same outgoing turn.
     const shouldInjectOrchestratorReminder = (sessionID: string) =>
-      sessionAgentMap.get(sessionID) === 'orchestrator';
+      sessionMetadata.getAgent(sessionID) === 'orchestrator';
 
     phaseReminder = createPhaseReminderHook({
       shouldInject: shouldInjectOrchestratorReminder,
@@ -415,14 +419,14 @@ const OhMyOpenCodeLite: Plugin = async (ctx) => {
       client: ctx.client,
       backgroundJobBoard: backgroundJobCoordinator,
       shouldManageSession: (sessionID) =>
-        sessionAgentMap.get(sessionID) === 'orchestrator',
+        sessionMetadata.getAgent(sessionID) === 'orchestrator',
     });
     waitForUserTools = createWaitForUserTool({
       shouldManageSession: (sessionID) =>
-        sessionAgentMap.get(sessionID) === 'orchestrator',
+        sessionMetadata.getAgent(sessionID) === 'orchestrator',
       resolveAgentName: (agent) => resolveRuntimeAgentName(config, agent),
       registerSessionAsOrchestrator: (sessionID) => {
-        sessionAgentMap.set(sessionID, 'orchestrator');
+        sessionMetadata.setAgent(sessionID, 'orchestrator');
       },
       beginUserWait: (sessionID) =>
         taskSessionManagerHook.beginUserWait(sessionID),
@@ -918,6 +922,24 @@ const OhMyOpenCodeLite: Plugin = async (ctx) => {
         };
       };
 
+      const eventSessionID =
+        event.properties?.info?.id ?? event.properties?.sessionID;
+      const statusType = event.properties?.status?.type;
+      if (eventSessionID) {
+        if (
+          event.type === 'session.status' &&
+          (statusType === 'busy' || statusType === 'retry')
+        ) {
+          sessionMetadata.markOrchestratorActive(eventSessionID);
+        } else if (
+          event.type === 'session.idle' ||
+          (event.type === 'session.status' && statusType === 'idle') ||
+          event.type === 'session.deleted'
+        ) {
+          sessionMetadata.markOrchestratorIdle(eventSessionID);
+        }
+      }
+
       if (event.type === 'message.updated') {
         const info = event.properties?.info;
         const providerID =
@@ -942,7 +964,7 @@ const OhMyOpenCodeLite: Plugin = async (ctx) => {
               model,
               variant: variant ?? null,
             },
-            (info?.sessionID && sessionDirectories.get(info.sessionID)) ??
+            (info?.sessionID && sessionMetadata.getDirectory(info.sessionID)) ??
               ctx.directory,
           );
         }
@@ -952,7 +974,7 @@ const OhMyOpenCodeLite: Plugin = async (ctx) => {
         const createdSessionId = event.properties?.info?.id;
         const createdSessionDir = event.properties?.info?.directory;
         if (createdSessionId && createdSessionDir) {
-          sessionDirectories.set(createdSessionId, createdSessionDir);
+          sessionMetadata.setDirectory(createdSessionId, createdSessionDir);
         }
       }
 
@@ -1014,7 +1036,7 @@ const OhMyOpenCodeLite: Plugin = async (ctx) => {
         const sessionID = props?.sessionID;
         companionManager.onSessionStatus({
           sessionId: sessionID,
-          agent: sessionID ? sessionAgentMap.get(sessionID) : undefined,
+          agent: sessionID ? sessionMetadata.getAgent(sessionID) : undefined,
           status: props?.status?.type,
         });
       }
@@ -1030,8 +1052,7 @@ const OhMyOpenCodeLite: Plugin = async (ctx) => {
         }
         companionManager.onSessionDeleted(sessionID);
         if (sessionID) {
-          sessionAgentMap.delete(sessionID);
-          sessionDirectories.delete(sessionID);
+          sessionMetadata.delete(sessionID);
         }
       }
     },
@@ -1119,7 +1140,7 @@ const OhMyOpenCodeLite: Plugin = async (ctx) => {
 
       if (agent) {
         foregroundFallback.registerSessionAgent(input.sessionID, agent);
-        sessionAgentMap.set(input.sessionID, agent);
+        sessionMetadata.setAgent(input.sessionID, agent);
         // A chat message means this session is actively working. This also
         // covers the race where session.status busy fires before the
         // session's agent is known.
@@ -1143,7 +1164,7 @@ const OhMyOpenCodeLite: Plugin = async (ctx) => {
       output: { system: string[] },
     ): Promise<void> => {
       const agentName = input.sessionID
-        ? sessionAgentMap.get(input.sessionID)
+        ? sessionMetadata.getAgent(input.sessionID)
         : undefined;
       if (agentName === 'orchestrator') {
         const alreadyInjected = output.system.some(

src/utils/session-metadata.test.ts

@@ -0,0 +1,63 @@
+import { describe, expect, test } from 'bun:test';
+import { SessionMetadataStore } from './session-metadata';
+
+describe('SessionMetadataStore', () => {
+  test('keeps two active orchestrators through metadata overflow', () => {
+    const store = new SessionMetadataStore({ maxEntries: 3 });
+
+    store.setAgent('orchestrator-a', 'orchestrator');
+    store.setAgent('orchestrator-b', 'orchestrator');
+    store.setAgent('old-specialist', 'explore');
+    store.setDirectory('new-session', '/tmp/project');
+
+    expect(store.size).toBe(3);
+    expect(store.getAgent('orchestrator-a')).toBe('orchestrator');
+    expect(store.getAgent('orchestrator-b')).toBe('orchestrator');
+    expect(store.hasAgent('old-specialist')).toBe(false);
+  });
+
+  test('makes an idle orchestrator evictable without dropping another active one', () => {
+    const store = new SessionMetadataStore({ maxEntries: 3 });
+
+    store.setAgent('orchestrator-a', 'orchestrator');
+    store.setAgent('orchestrator-b', 'orchestrator');
+    store.setAgent('old-specialist', 'explore');
+    store.markOrchestratorIdle('orchestrator-a');
+    store.setDirectory('new-session', '/tmp/project');
+
+    expect(store.size).toBe(3);
+    expect(store.hasAgent('orchestrator-a')).toBe(false);
+    expect(store.getAgent('orchestrator-b')).toBe('orchestrator');
+    expect(store.hasAgent('old-specialist')).toBe(true);
+  });
+
+  test('bounds agent-only metadata', () => {
+    const store = new SessionMetadataStore({ maxEntries: 2 });
+
+    store.setAgent('agent-a', 'explore');
+    store.setAgent('agent-b', 'oracle');
+    store.setAgent('agent-c', 'fixer');
+
+    expect(store.size).toBe(2);
+    expect(store.hasAgent('agent-a')).toBe(false);
+    expect(store.hasAgent('agent-b')).toBe(true);
+    expect(store.hasAgent('agent-c')).toBe(true);
+  });
+
+  test('eviction removes directory and agent metadata for one session', () => {
+    const evicted: string[] = [];
+    const store = new SessionMetadataStore({
+      maxEntries: 1,
+      onEvict: (sessionID) => evicted.push(sessionID),
+    });
+
+    store.setDirectory('old-session', '/tmp/project');
+    store.setAgent('old-session', 'explore');
+    store.setDirectory('new-session', '/tmp/project');
+
+    expect(store.size).toBe(1);
+    expect(store.hasDirectory('old-session')).toBe(false);
+    expect(store.hasAgent('old-session')).toBe(false);
+    expect(evicted).toEqual(['old-session']);
+  });
+});

src/utils/session-metadata.ts

@@ -0,0 +1,90 @@
+type SessionMetadataEviction = (sessionID: string) => void;
+
+export class SessionMetadataStore {
+  readonly #agents = new Map<string, string>();
+  readonly #directories = new Map<string, string>();
+  readonly #insertionOrder = new Map<string, undefined>();
+  readonly #activeOrchestratorSessionIDs = new Set<string>();
+  readonly #maxEntries: number;
+  readonly #onEvict?: SessionMetadataEviction;
+
+  constructor(options: {
+    maxEntries: number;
+    onEvict?: SessionMetadataEviction;
+  }) {
+    this.#maxEntries = options.maxEntries;
+    this.#onEvict = options.onEvict;
+  }
+
+  getAgent(sessionID: string): string | undefined {
+    return this.#agents.get(sessionID);
+  }
+
+  getDirectory(sessionID: string): string | undefined {
+    return this.#directories.get(sessionID);
+  }
+
+  setAgent(sessionID: string, agent: string): void {
+    this.#agents.set(sessionID, agent);
+
+    if (agent === 'orchestrator') {
+      this.#activeOrchestratorSessionIDs.add(sessionID);
+    } else {
+      this.#activeOrchestratorSessionIDs.delete(sessionID);
+    }
+
+    this.#track(sessionID);
+  }
+
+  setDirectory(sessionID: string, directory: string): void {
+    this.#directories.set(sessionID, directory);
+    this.#track(sessionID);
+  }
+
+  markOrchestratorActive(sessionID: string): void {
+    if (this.#agents.get(sessionID) === 'orchestrator') {
+      this.#activeOrchestratorSessionIDs.add(sessionID);
+    }
+  }
+
+  markOrchestratorIdle(sessionID: string): void {
+    this.#activeOrchestratorSessionIDs.delete(sessionID);
+  }
+
+  delete(sessionID: string): void {
+    this.#agents.delete(sessionID);
+    this.#directories.delete(sessionID);
+    this.#insertionOrder.delete(sessionID);
+    this.#activeOrchestratorSessionIDs.delete(sessionID);
+  }
+
+  get size(): number {
+    return this.#insertionOrder.size;
+  }
+
+  hasAgent(sessionID: string): boolean {
+    return this.#agents.has(sessionID);
+  }
+
+  hasDirectory(sessionID: string): boolean {
+    return this.#directories.has(sessionID);
+  }
+
+  #track(sessionID: string): void {
+    if (!this.#insertionOrder.has(sessionID)) {
+      this.#insertionOrder.set(sessionID, undefined);
+    }
+
+    while (this.#insertionOrder.size > this.#maxEntries) {
+      const evictableSessionID = [...this.#insertionOrder.keys()].find(
+        (candidate) => !this.#activeOrchestratorSessionIDs.has(candidate),
+      );
+      if (evictableSessionID === undefined) return;
+
+      this.#insertionOrder.delete(evictableSessionID);
+      this.#agents.delete(evictableSessionID);
+      this.#directories.delete(evictableSessionID);
+      this.#onEvict?.(evictableSessionID);
+    }
+  }
+}