fix: bound session metadata coherently

src/config/constants.ts

@@ -94,7 +94,7 @@ export const DEFAULT_READ_CONTEXT_MAX_FILES = 8;
 export const DEFAULT_MAX_RETAINED_SNAPSHOTS = 20;
 
 /**
- * Maximum session-directory mappings retained per plugin instance.
+ * 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.
  */

src/index.ts

@@ -74,6 +74,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';
 
 /**
@@ -173,10 +174,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_DIRECTORIES,
+    onEvict: (sessionID) => {
+      log('[session] evicted oldest session metadata', {
+        threshold: DEFAULT_MAX_SESSION_DIRECTORIES,
+        droppedSessionId: sessionID,
+      });
+    },
+  });
   let sessionLifecycle: SessionLifecycle;
 
   let chatHeadersHook: ReturnType<typeof createChatHeadersHook>;
@@ -320,9 +326,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.
@@ -356,9 +359,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),
@@ -397,7 +400,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,
@@ -439,14 +442,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),
@@ -963,7 +966,7 @@ const OhMyOpenCodeLite: Plugin = async (ctx) => {
               model,
               variant: variant ?? null,
             },
-            (info?.sessionID && sessionDirectories.get(info.sessionID)) ??
+            (info?.sessionID && sessionMetadata.getDirectory(info.sessionID)) ??
               ctx.directory,
           );
         }
@@ -973,21 +976,7 @@ const OhMyOpenCodeLite: Plugin = async (ctx) => {
         const createdSessionId = event.properties?.info?.id;
         const createdSessionDir = event.properties?.info?.directory;
         if (createdSessionId && createdSessionDir) {
-          // Guard against unbounded growth when session.deleted events
-          // are missed: evict the oldest entry when the threshold is
-          // reached, keeping both maps aligned.
-          if (sessionDirectories.size >= DEFAULT_MAX_SESSION_DIRECTORIES) {
-            const oldestId = sessionDirectories.keys().next().value;
-            if (oldestId !== undefined) {
-              sessionDirectories.delete(oldestId);
-              sessionAgentMap.delete(oldestId);
-              log('[session] evicted oldest session-directory mapping', {
-                threshold: DEFAULT_MAX_SESSION_DIRECTORIES,
-                droppedSessionId: oldestId,
-              });
-            }
-          }
-          sessionDirectories.set(createdSessionId, createdSessionDir);
+          sessionMetadata.setDirectory(createdSessionId, createdSessionDir);
         }
       }
 
@@ -1047,9 +1036,12 @@ const OhMyOpenCodeLite: Plugin = async (ctx) => {
           | { sessionID?: string; status?: { type?: string } }
           | undefined;
         const sessionID = props?.sessionID;
+        if (sessionID && props?.status?.type === 'busy') {
+          sessionMetadata.markOrchestratorBusy(sessionID);
+        }
         companionManager.onSessionStatus({
           sessionId: sessionID,
-          agent: sessionID ? sessionAgentMap.get(sessionID) : undefined,
+          agent: sessionID ? sessionMetadata.getAgent(sessionID) : undefined,
           status: props?.status?.type,
         });
       }
@@ -1065,8 +1057,7 @@ const OhMyOpenCodeLite: Plugin = async (ctx) => {
         }
         companionManager.onSessionDeleted(sessionID);
         if (sessionID) {
-          sessionAgentMap.delete(sessionID);
-          sessionDirectories.delete(sessionID);
+          sessionMetadata.delete(sessionID);
         }
       }
     },
@@ -1154,7 +1145,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.
@@ -1178,7 +1169,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,65 @@
+import { describe, expect, test } from 'bun:test';
+import { SessionMetadataStore } from './session-metadata';
+
+describe('SessionMetadataStore', () => {
+  test('bounds the union of directory and agent metadata', () => {
+    const evicted: string[] = [];
+    const store = new SessionMetadataStore({
+      maxEntries: 2,
+      onEvict: (sessionID) => evicted.push(sessionID),
+    });
+
+    store.setDirectory('directory-only', '/tmp/project');
+    store.setAgent('agent-only', 'explore');
+    store.setDirectory('newest', '/tmp/other-project');
+
+    expect(store.size).toBe(2);
+    expect(store.hasDirectory('directory-only')).toBe(false);
+    expect(store.hasAgent('directory-only')).toBe(false);
+    expect(store.hasAgent('agent-only')).toBe(true);
+    expect(store.hasDirectory('newest')).toBe(true);
+    expect(evicted).toEqual(['directory-only']);
+  });
+
+  test('retains the active orchestrator while evicting older metadata', () => {
+    const store = new SessionMetadataStore({ maxEntries: 2 });
+
+    store.setAgent('orchestrator-session', 'orchestrator');
+    store.setDirectory('orchestrator-session', '/tmp/project');
+    store.setAgent('older-specialist', 'explore');
+    store.setDirectory('newer-specialist', '/tmp/project');
+
+    expect(store.size).toBe(2);
+    expect(store.getAgent('orchestrator-session')).toBe('orchestrator');
+    expect(store.getDirectory('orchestrator-session')).toBe('/tmp/project');
+    expect(store.hasAgent('older-specialist')).toBe(false);
+    expect(store.hasDirectory('older-specialist')).toBe(false);
+  });
+
+  test('allows a deleted orchestrator to be evicted after cleanup', () => {
+    const store = new SessionMetadataStore({ maxEntries: 2 });
+
+    store.setAgent('orchestrator-session', 'orchestrator');
+    store.setAgent('specialist-session', 'explore');
+    store.delete('orchestrator-session');
+    store.setDirectory('new-session', '/tmp/project');
+
+    expect(store.size).toBe(2);
+    expect(store.hasAgent('orchestrator-session')).toBe(false);
+    expect(store.hasAgent('specialist-session')).toBe(true);
+    expect(store.hasDirectory('new-session')).toBe(true);
+  });
+
+  test('protects the orchestrator reported busy by the session event', () => {
+    const store = new SessionMetadataStore({ maxEntries: 2 });
+
+    store.setAgent('first-orchestrator', 'orchestrator');
+    store.setAgent('second-orchestrator', 'orchestrator');
+    store.markOrchestratorBusy('first-orchestrator');
+    store.setDirectory('new-session', '/tmp/project');
+
+    expect(store.getAgent('first-orchestrator')).toBe('orchestrator');
+    expect(store.hasAgent('second-orchestrator')).toBe(false);
+    expect(store.hasDirectory('new-session')).toBe(true);
+  });
+});

src/utils/session-metadata.ts

@@ -0,0 +1,88 @@
+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 #maxEntries: number;
+  readonly #onEvict?: SessionMetadataEviction;
+  #activeOrchestratorSessionID: string | undefined;
+
+  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.#activeOrchestratorSessionID = sessionID;
+    } else if (this.#activeOrchestratorSessionID === sessionID) {
+      this.#activeOrchestratorSessionID = undefined;
+    }
+
+    this.#track(sessionID);
+  }
+
+  setDirectory(sessionID: string, directory: string): void {
+    this.#directories.set(sessionID, directory);
+    this.#track(sessionID);
+  }
+
+  markOrchestratorBusy(sessionID: string): void {
+    if (this.#agents.get(sessionID) === 'orchestrator') {
+      this.#activeOrchestratorSessionID = sessionID;
+    }
+  }
+
+  delete(sessionID: string): void {
+    this.#agents.delete(sessionID);
+    this.#directories.delete(sessionID);
+    this.#insertionOrder.delete(sessionID);
+    if (this.#activeOrchestratorSessionID === sessionID) {
+      this.#activeOrchestratorSessionID = undefined;
+    }
+  }
+
+  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) => candidate !== this.#activeOrchestratorSessionID,
+      );
+      if (evictableSessionID === undefined) return;
+
+      this.#insertionOrder.delete(evictableSessionID);
+      this.#agents.delete(evictableSessionID);
+      this.#directories.delete(evictableSessionID);
+      this.#onEvict?.(evictableSessionID);
+    }
+  }
+}