Browse Source

fix tmux pane lifecycle across plugin instances

Alvin Unreal 2 months ago
parent
commit
1b844f0798
3 changed files with 224 additions and 17 deletions
  1. 10 9
      src/index.ts
  2. 115 1
      src/multiplexer/session-manager.test.ts
  3. 99 7
      src/multiplexer/session-manager.ts

+ 10 - 9
src/index.ts

@@ -784,6 +784,16 @@ const OhMyOpenCodeLite: Plugin = async (ctx) => {
         }
       }
 
+      // Handle multiplexer pane spawning for OpenCode's Task tool sessions
+      await multiplexerSessionManager.onSessionCreated(event);
+
+      // Handle session status/idle events for pane cleanup early so child panes
+      // close promptly even if later hooks do additional work on idle.
+      await multiplexerSessionManager.onSessionStatus(event);
+
+      // Handle session.deleted events for pane cleanup
+      await multiplexerSessionManager.onSessionDeleted(event);
+
       // Runtime model fallback for foreground agents (rate-limit detection)
       await foregroundFallback.handleEvent(input.event);
 
@@ -799,15 +809,6 @@ const OhMyOpenCodeLite: Plugin = async (ctx) => {
       // Handle auto-update checking
       await autoUpdateChecker.event(input);
 
-      // Handle multiplexer pane spawning for OpenCode's Task tool sessions
-      await multiplexerSessionManager.onSessionCreated(event);
-
-      // Handle session.status events for pane cleanup
-      await multiplexerSessionManager.onSessionStatus(event);
-
-      // Handle session.deleted events for pane cleanup
-      await multiplexerSessionManager.onSessionDeleted(event);
-
       await interviewManager.handleEvent(
         input as {
           event: { type: string; properties?: Record<string, unknown> };

+ 115 - 1
src/multiplexer/session-manager.test.ts

@@ -1,5 +1,8 @@
 import { afterEach, beforeEach, describe, expect, mock, test } from 'bun:test';
-import { MultiplexerSessionManager } from './session-manager';
+import {
+  MultiplexerSessionManager,
+  resetMultiplexerSessionManagerState,
+} from './session-manager';
 
 const originalFetch = globalThis.fetch;
 let mockSessionStatuses: Record<string, { type: string }> = {};
@@ -75,6 +78,7 @@ function createDeferred<T>() {
 
 describe('MultiplexerSessionManager', () => {
   beforeEach(() => {
+    resetMultiplexerSessionManagerState();
     mockSessionStatuses = {};
     mockFetch.mockClear();
     globalThis.fetch = mockFetch as typeof fetch;
@@ -258,6 +262,116 @@ describe('MultiplexerSessionManager', () => {
       expect(mockMultiplexer.closePane).toHaveBeenCalledWith('p-1');
     });
 
+    test('closes pane immediately on session.idle event', async () => {
+      const ctx = createMockContext();
+      mockMultiplexer.spawnPane.mockResolvedValue({
+        success: true,
+        paneId: 'p-idle-event',
+      });
+
+      const manager = new MultiplexerSessionManager(
+        ctx,
+        defaultMultiplexerConfig,
+      );
+
+      await manager.onSessionCreated({
+        type: 'session.created',
+        properties: { info: { id: 'idle-event-child', parentID: 'parent' } },
+      });
+
+      await manager.onSessionStatus({
+        type: 'session.idle',
+        properties: { sessionID: 'idle-event-child' },
+      });
+
+      expect(mockMultiplexer.closePane).toHaveBeenCalledWith('p-idle-event');
+      expect(mockFetch).not.toHaveBeenCalled();
+    });
+
+    test('closes pane when idle event reaches a different manager instance', async () => {
+      const ctx = createMockContext();
+      mockMultiplexer.spawnPane.mockResolvedValue({
+        success: true,
+        paneId: 'p-shared-idle',
+      });
+
+      const spawningManager = new MultiplexerSessionManager(
+        ctx,
+        defaultMultiplexerConfig,
+      );
+      const idleManager = new MultiplexerSessionManager(
+        ctx,
+        defaultMultiplexerConfig,
+      );
+
+      await spawningManager.onSessionCreated({
+        type: 'session.created',
+        properties: { info: { id: 'shared-child', parentID: 'parent' } },
+      });
+
+      await idleManager.onSessionStatus({
+        type: 'session.idle',
+        properties: { sessionID: 'shared-child' },
+      });
+
+      expect(mockMultiplexer.closePane).toHaveBeenCalledWith('p-shared-idle');
+    });
+
+    test('respawns resumed known session from a different manager instance', async () => {
+      const ctx = createMockContext();
+      mockMultiplexer.spawnPane
+        .mockResolvedValueOnce({
+          success: true,
+          paneId: 'p-first',
+        })
+        .mockResolvedValueOnce({
+          success: true,
+          paneId: 'p-resumed',
+        });
+
+      const firstManager = new MultiplexerSessionManager(
+        ctx,
+        defaultMultiplexerConfig,
+      );
+      const secondManager = new MultiplexerSessionManager(
+        ctx,
+        defaultMultiplexerConfig,
+      );
+
+      await firstManager.onSessionCreated({
+        type: 'session.created',
+        properties: {
+          info: {
+            id: 'resumed-child',
+            parentID: 'parent',
+            title: 'Resumed Worker',
+            directory: '/resumed/dir',
+          },
+        },
+      });
+
+      await secondManager.onSessionStatus({
+        type: 'session.idle',
+        properties: { sessionID: 'resumed-child' },
+      });
+
+      await secondManager.onSessionStatus({
+        type: 'session.status',
+        properties: {
+          sessionID: 'resumed-child',
+          status: { type: 'busy' },
+        },
+      });
+
+      expect(mockMultiplexer.spawnPane).toHaveBeenCalledTimes(2);
+      expect(mockMultiplexer.spawnPane).toHaveBeenLastCalledWith(
+        'resumed-child',
+        'Resumed Worker',
+        `http://localhost:${process.env.OPENCODE_PORT ?? '4096'}/`,
+        '/resumed/dir',
+      );
+    });
+
     test('does not close on transient status absence', async () => {
       const ctx = createMockContext();
       const manager = new MultiplexerSessionManager(

+ 99 - 7
src/multiplexer/session-manager.ts

@@ -25,6 +25,13 @@ interface KnownSession {
   directory: string;
 }
 
+interface SharedSessionState {
+  sessions: Map<string, TrackedSession>;
+  knownSessions: Map<string, KnownSession>;
+  spawningSessions: Set<string>;
+  closingSessions: Map<string, Promise<void>>;
+}
+
 interface SessionEvent {
   type: string;
   properties?: {
@@ -43,6 +50,32 @@ type CloseReason = 'idle' | 'deleted' | 'missing' | 'timeout';
 
 const SESSION_TIMEOUT_MS = 10 * 60 * 1000;
 const SESSION_MISSING_GRACE_MS = POLL_INTERVAL_BACKGROUND_MS * 3;
+const SHARED_STATE_KEY = Symbol.for(
+  'oh-my-opencode-slim.multiplexer-session-manager.state',
+);
+
+function getSharedState(): SharedSessionState {
+  const globalWithState = globalThis as typeof globalThis & {
+    [SHARED_STATE_KEY]?: SharedSessionState;
+  };
+
+  globalWithState[SHARED_STATE_KEY] ??= {
+    sessions: new Map(),
+    knownSessions: new Map(),
+    spawningSessions: new Set(),
+    closingSessions: new Map(),
+  };
+
+  return globalWithState[SHARED_STATE_KEY];
+}
+
+export function resetMultiplexerSessionManagerState(): void {
+  const state = getSharedState();
+  state.sessions.clear();
+  state.knownSessions.clear();
+  state.spawningSessions.clear();
+  state.closingSessions.clear();
+}
 
 /**
  * Tracks child sessions and spawns/closes multiplexer panes for them.
@@ -51,17 +84,24 @@ const SESSION_MISSING_GRACE_MS = POLL_INTERVAL_BACKGROUND_MS * 3;
  * with polling kept as a fallback for reliability.
  */
 export class MultiplexerSessionManager {
+  private instanceId = Math.random().toString(36).slice(2, 8);
   private serverUrl: string;
   private directory: string;
   private multiplexer: Multiplexer | null = null;
-  private sessions = new Map<string, TrackedSession>();
-  private knownSessions = new Map<string, KnownSession>();
-  private spawningSessions = new Set<string>();
-  private closingSessions = new Map<string, Promise<void>>();
+  private sessions: SharedSessionState['sessions'];
+  private knownSessions: SharedSessionState['knownSessions'];
+  private spawningSessions: SharedSessionState['spawningSessions'];
+  private closingSessions: SharedSessionState['closingSessions'];
   private pollInterval?: ReturnType<typeof setInterval>;
   private enabled = false;
 
   constructor(ctx: PluginInput, config: MultiplexerConfig) {
+    const sharedState = getSharedState();
+    this.sessions = sharedState.sessions;
+    this.knownSessions = sharedState.knownSessions;
+    this.spawningSessions = sharedState.spawningSessions;
+    this.closingSessions = sharedState.closingSessions;
+
     this.directory = ctx.directory;
     const defaultPort = process.env.OPENCODE_PORT ?? '4096';
     this.serverUrl =
@@ -74,9 +114,12 @@ export class MultiplexerSessionManager {
       this.multiplexer.isInsideSession();
 
     log('[multiplexer-session-manager] initialized', {
+      instanceId: this.instanceId,
       enabled: this.enabled,
       type: config.type,
       serverUrl: this.serverUrl,
+      trackedSessions: this.sessions.size,
+      knownSessions: this.knownSessions.size,
     });
   }
 
@@ -96,6 +139,7 @@ export class MultiplexerSessionManager {
 
     if (this.isTrackedOrSpawning(sessionId)) {
       log('[multiplexer-session-manager] session already tracked or spawning', {
+        instanceId: this.instanceId,
         sessionId,
       });
       return;
@@ -118,6 +162,7 @@ export class MultiplexerSessionManager {
       const serverRunning = await isServerRunning(this.serverUrl);
       if (!serverRunning) {
         log('[multiplexer-session-manager] server not running, skipping', {
+          instanceId: this.instanceId,
           serverUrl: this.serverUrl,
         });
         return;
@@ -133,6 +178,7 @@ export class MultiplexerSessionManager {
           sessionId,
           parentId,
           title,
+          instanceId: this.instanceId,
         },
       );
 
@@ -140,6 +186,7 @@ export class MultiplexerSessionManager {
         .spawnPane(sessionId, title, this.serverUrl, directory)
         .catch((err) => {
           log('[multiplexer-session-manager] failed to spawn pane', {
+            instanceId: this.instanceId,
             error: String(err),
           });
           return { success: false, paneId: undefined };
@@ -157,6 +204,7 @@ export class MultiplexerSessionManager {
             {
               sessionId,
               paneId: paneResult.paneId,
+              instanceId: this.instanceId,
               error: String(err),
             },
           ),
@@ -176,6 +224,7 @@ export class MultiplexerSessionManager {
       });
 
       log('[multiplexer-session-manager] pane spawned', {
+        instanceId: this.instanceId,
         sessionId,
         paneId: paneResult.paneId,
       });
@@ -188,6 +237,22 @@ export class MultiplexerSessionManager {
 
   async onSessionStatus(event: SessionEvent): Promise<void> {
     if (!this.enabled) return;
+
+    if (event.type === 'session.idle') {
+      const sessionId = event.properties?.sessionID;
+      if (!sessionId) return;
+
+      log('[multiplexer-session-manager] session idle event received', {
+        instanceId: this.instanceId,
+        sessionId,
+        tracked: this.sessions.has(sessionId),
+        known: this.knownSessions.has(sessionId),
+      });
+
+      await this.closeSession(sessionId, 'idle');
+      return;
+    }
+
     if (event.type !== 'session.status') return;
 
     const sessionId = event.properties?.sessionID;
@@ -199,6 +264,12 @@ export class MultiplexerSessionManager {
     }
 
     if (event.properties?.status?.type === 'busy') {
+      log('[multiplexer-session-manager] session busy event received', {
+        instanceId: this.instanceId,
+        sessionId,
+        tracked: this.sessions.has(sessionId),
+        known: this.knownSessions.has(sessionId),
+      });
       await this.respawnIfKnown(sessionId);
     }
   }
@@ -211,6 +282,7 @@ export class MultiplexerSessionManager {
     if (!sessionId) return;
 
     log('[multiplexer-session-manager] session deleted, closing pane', {
+      instanceId: this.instanceId,
       sessionId,
     });
 
@@ -224,14 +296,18 @@ export class MultiplexerSessionManager {
       () => this.pollSessions(),
       POLL_INTERVAL_BACKGROUND_MS,
     );
-    log('[multiplexer-session-manager] polling started');
+    log('[multiplexer-session-manager] polling started', {
+      instanceId: this.instanceId,
+    });
   }
 
   private stopPolling(): void {
     if (this.pollInterval) {
       clearInterval(this.pollInterval);
       this.pollInterval = undefined;
-      log('[multiplexer-session-manager] polling stopped');
+      log('[multiplexer-session-manager] polling stopped', {
+        instanceId: this.instanceId,
+      });
     }
   }
 
@@ -307,11 +383,21 @@ export class MultiplexerSessionManager {
     if (existingClose) return existingClose;
 
     const tracked = this.sessions.get(sessionId);
-    if (!tracked || !this.multiplexer) return;
+    if (!tracked || !this.multiplexer) {
+      log('[multiplexer-session-manager] close skipped; session not tracked', {
+        instanceId: this.instanceId,
+        sessionId,
+        reason,
+        tracked: !!tracked,
+        hasMultiplexer: !!this.multiplexer,
+      });
+      return;
+    }
 
     this.sessions.delete(sessionId);
 
     log('[multiplexer-session-manager] closing session pane', {
+      instanceId: this.instanceId,
       sessionId,
       paneId: tracked.paneId,
       reason,
@@ -322,6 +408,7 @@ export class MultiplexerSessionManager {
       .then(() => undefined)
       .catch((err) =>
         log('[multiplexer-session-manager] failed to close session pane', {
+          instanceId: this.instanceId,
           sessionId,
           paneId: tracked.paneId,
           reason,
@@ -357,6 +444,7 @@ export class MultiplexerSessionManager {
         log(
           '[multiplexer-session-manager] server not running, skipping busy respawn',
           {
+            instanceId: this.instanceId,
             serverUrl: this.serverUrl,
             sessionId,
           },
@@ -371,6 +459,7 @@ export class MultiplexerSessionManager {
       log(
         '[multiplexer-session-manager] child session busy again, respawning pane',
         {
+          instanceId: this.instanceId,
           sessionId,
           parentId: known.parentId,
           title: known.title,
@@ -381,6 +470,7 @@ export class MultiplexerSessionManager {
         .spawnPane(sessionId, known.title, this.serverUrl, known.directory)
         .catch((err) => {
           log('[multiplexer-session-manager] failed to respawn pane', {
+            instanceId: this.instanceId,
             error: String(err),
           });
           return { success: false, paneId: undefined };
@@ -396,6 +486,7 @@ export class MultiplexerSessionManager {
           log(
             '[multiplexer-session-manager] closing stale respawned pane failed',
             {
+              instanceId: this.instanceId,
               sessionId,
               paneId: paneResult.paneId,
               error: String(err),
@@ -417,6 +508,7 @@ export class MultiplexerSessionManager {
       });
 
       log('[multiplexer-session-manager] pane respawned on busy', {
+        instanceId: this.instanceId,
         sessionId,
         paneId: paneResult.paneId,
       });