Browse Source

fix: track all active orchestrator sessions

Alvin Unreal 1 week ago
parent
commit
791895f7f1
4 changed files with 74 additions and 59 deletions
  1. 1 1
      src/config/constants.ts
  2. 21 6
      src/index.ts
  3. 40 42
      src/utils/session-metadata.test.ts
  4. 12 10
      src/utils/session-metadata.ts

+ 1 - 1
src/config/constants.ts

@@ -98,7 +98,7 @@ export const DEFAULT_MAX_RETAINED_SNAPSHOTS = 20;
  * Prevents unbounded growth when session.deleted events are missed.
  * Oldest entries are evicted first when this threshold is reached.
  */
-export const DEFAULT_MAX_SESSION_DIRECTORIES = 1000;
+export const DEFAULT_MAX_SESSION_METADATA_ENTRIES = 1000;
 
 export type ImageRouting = 'auto' | 'direct';
 

+ 21 - 6
src/index.ts

@@ -18,7 +18,7 @@ import { parseList } from './config/agent-mcps';
 import {
   AGENT_ALIASES,
   DEFAULT_MAX_RETAINED_SNAPSHOTS,
-  DEFAULT_MAX_SESSION_DIRECTORIES,
+  DEFAULT_MAX_SESSION_METADATA_ENTRIES,
   DEFAULT_MAX_SESSIONS_PER_AGENT,
   DEFAULT_READ_CONTEXT_MAX_FILES,
   DEFAULT_READ_CONTEXT_MIN_LINES,
@@ -175,10 +175,10 @@ const OhMyOpenCodeLite: Plugin = async (ctx) => {
   let multiplexerSessionManager: MultiplexerSessionManager;
   let autoUpdateChecker: ReturnType<typeof createAutoUpdateCheckerHook>;
   const sessionMetadata = new SessionMetadataStore({
-    maxEntries: DEFAULT_MAX_SESSION_DIRECTORIES,
+    maxEntries: DEFAULT_MAX_SESSION_METADATA_ENTRIES,
     onEvict: (sessionID) => {
       log('[session] evicted oldest session metadata', {
-        threshold: DEFAULT_MAX_SESSION_DIRECTORIES,
+        threshold: DEFAULT_MAX_SESSION_METADATA_ENTRIES,
         droppedSessionId: sessionID,
       });
     },
@@ -942,6 +942,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 =
@@ -1036,9 +1054,6 @@ 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 ? sessionMetadata.getAgent(sessionID) : undefined,

+ 40 - 42
src/utils/session-metadata.test.ts

@@ -2,64 +2,62 @@ 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),
-    });
+  test('keeps two active orchestrators through metadata overflow', () => {
+    const store = new SessionMetadataStore({ maxEntries: 3 });
 
-    store.setDirectory('directory-only', '/tmp/project');
-    store.setAgent('agent-only', 'explore');
-    store.setDirectory('newest', '/tmp/other-project');
+    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(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']);
+    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('retains the active orchestrator while evicting older metadata', () => {
-    const store = new SessionMetadataStore({ maxEntries: 2 });
+  test('makes an idle orchestrator evictable without dropping another active one', () => {
+    const store = new SessionMetadataStore({ maxEntries: 3 });
 
-    store.setAgent('orchestrator-session', 'orchestrator');
-    store.setDirectory('orchestrator-session', '/tmp/project');
-    store.setAgent('older-specialist', 'explore');
-    store.setDirectory('newer-specialist', '/tmp/project');
+    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(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);
+    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('allows a deleted orchestrator to be evicted after cleanup', () => {
+  test('bounds agent-only metadata', () => {
     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');
+    store.setAgent('agent-a', 'explore');
+    store.setAgent('agent-b', 'oracle');
+    store.setAgent('agent-c', 'fixer');
 
     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);
+    expect(store.hasAgent('agent-a')).toBe(false);
+    expect(store.hasAgent('agent-b')).toBe(true);
+    expect(store.hasAgent('agent-c')).toBe(true);
   });
 
-  test('protects the orchestrator reported busy by the session event', () => {
-    const store = new SessionMetadataStore({ maxEntries: 2 });
+  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.setAgent('first-orchestrator', 'orchestrator');
-    store.setAgent('second-orchestrator', 'orchestrator');
-    store.markOrchestratorBusy('first-orchestrator');
+    store.setDirectory('old-session', '/tmp/project');
+    store.setAgent('old-session', 'explore');
     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);
+    expect(store.size).toBe(1);
+    expect(store.hasDirectory('old-session')).toBe(false);
+    expect(store.hasAgent('old-session')).toBe(false);
+    expect(evicted).toEqual(['old-session']);
   });
 });

+ 12 - 10
src/utils/session-metadata.ts

@@ -4,9 +4,9 @@ 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;
-  #activeOrchestratorSessionID: string | undefined;
 
   constructor(options: {
     maxEntries: number;
@@ -28,9 +28,9 @@ export class SessionMetadataStore {
     this.#agents.set(sessionID, agent);
 
     if (agent === 'orchestrator') {
-      this.#activeOrchestratorSessionID = sessionID;
-    } else if (this.#activeOrchestratorSessionID === sessionID) {
-      this.#activeOrchestratorSessionID = undefined;
+      this.#activeOrchestratorSessionIDs.add(sessionID);
+    } else {
+      this.#activeOrchestratorSessionIDs.delete(sessionID);
     }
 
     this.#track(sessionID);
@@ -41,19 +41,21 @@ export class SessionMetadataStore {
     this.#track(sessionID);
   }
 
-  markOrchestratorBusy(sessionID: string): void {
+  markOrchestratorActive(sessionID: string): void {
     if (this.#agents.get(sessionID) === 'orchestrator') {
-      this.#activeOrchestratorSessionID = sessionID;
+      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);
-    if (this.#activeOrchestratorSessionID === sessionID) {
-      this.#activeOrchestratorSessionID = undefined;
-    }
+    this.#activeOrchestratorSessionIDs.delete(sessionID);
   }
 
   get size(): number {
@@ -75,7 +77,7 @@ export class SessionMetadataStore {
 
     while (this.#insertionOrder.size > this.#maxEntries) {
       const evictableSessionID = [...this.#insertionOrder.keys()].find(
-        (candidate) => candidate !== this.#activeOrchestratorSessionID,
+        (candidate) => !this.#activeOrchestratorSessionIDs.has(candidate),
       );
       if (evictableSessionID === undefined) return;