Browse Source

Merge pull request #615 from alvinunreal/fix/companion-exit-listener-leak

fix(companion): dedup process exit listener across re-inits
Alvin 1 month ago
parent
commit
1d84a8c3f0
2 changed files with 112 additions and 3 deletions
  1. 66 1
      src/companion/manager.test.ts
  2. 46 2
      src/companion/manager.ts

+ 66 - 1
src/companion/manager.test.ts

@@ -11,6 +11,8 @@ import {
 // Point writes at a temp dir so tests don't touch the real state file.
 const TEST_DIR = path.join(os.tmpdir(), `companion-test-${process.pid}`);
 const XDG_DIR = path.join(TEST_DIR, 'xdg');
+const managers: CompanionManager[] = [];
+
 function readState() {
   return JSON.parse(readFileSync(stateFilePath(), 'utf8'));
 }
@@ -21,6 +23,9 @@ beforeEach(() => {
 });
 
 afterEach(() => {
+  for (const manager of managers.splice(0)) {
+    manager.onExit();
+  }
   rmSync(TEST_DIR, { recursive: true, force: true });
   delete process.env.XDG_DATA_HOME;
 });
@@ -30,7 +35,23 @@ function make(
   cwd = '/home/user/myproject',
   config: any = { enabled: true, position: 'bottom-right', size: 'medium' },
 ) {
-  return new CompanionManager(id, cwd, config);
+  const manager = new CompanionManager(id, cwd, config);
+  managers.push(manager);
+  return manager;
+}
+
+function attachFakeChild(manager: CompanionManager): { killed: () => boolean } {
+  let killed = false;
+  (
+    manager as unknown as {
+      companionProcess: { kill: () => void } | null;
+    }
+  ).companionProcess = {
+    kill: () => {
+      killed = true;
+    },
+  };
+  return { killed: () => killed };
 }
 
 describe('CompanionManager', () => {
@@ -210,6 +231,50 @@ describe('CompanionManager', () => {
     expect(readState().sessions[0].active_agents).toEqual(['fixer', 'fixer']);
   });
 
+  it('keeps at most one process exit listener across reloads', () => {
+    const baseline = process.listenerCount('exit');
+    for (let i = 0; i < 5; i++) {
+      const m = make('reload-session');
+      m.onLoad();
+    }
+    // Re-inits must dedup the exit listener rather than stacking one each time.
+    expect(process.listenerCount('exit')).toBeLessThanOrEqual(baseline + 1);
+    // onExit releases the live listener again.
+    managers.at(-1)?.onExit();
+    expect(process.listenerCount('exit')).toBeLessThanOrEqual(baseline);
+  });
+
+  it('cleans up a superseded manager for the same session on reload', () => {
+    const first = make('reload-session');
+    first.onLoad();
+    const firstChild = attachFakeChild(first);
+
+    const second = make('reload-session');
+    second.onLoad();
+
+    expect(firstChild.killed()).toBe(true);
+    expect(readState().sessions).toHaveLength(1);
+    expect(readState().sessions[0].session_id).toBe('reload-session');
+
+    second.onExit();
+  });
+
+  it('cleans up active managers when companion is disabled on reload', () => {
+    const enabled = make('disable-session');
+    enabled.onLoad();
+    const child = attachFakeChild(enabled);
+
+    const disabled = new CompanionManager('disable-session', '/path', {
+      enabled: false,
+      position: 'bottom-right',
+      size: 'medium',
+    });
+    disabled.onLoad();
+
+    expect(child.killed()).toBe(true);
+    expect(readState().sessions).toEqual([]);
+  });
+
   it('removes its entry on exit', () => {
     const m = make('sess-a', '/a');
     const m2 = make('sess-b', '/b');

+ 46 - 2
src/companion/manager.ts

@@ -12,6 +12,16 @@ import * as path from 'node:path';
 import type { CompanionConfig } from '../config/schema';
 import { log } from '../utils/logger';
 
+// Only one companion `process.on('exit')` listener should be live per process.
+// The plugin function can re-run (config.update() → Instance.dispose()),
+// constructing fresh CompanionManager instances; without deduping, every
+// re-init would leak another exit listener. Track live managers separately so
+// replacing the listener never drops cleanup for detached companion children.
+// Module-level state survives re-inits because the module itself is not
+// re-evaluated.
+let activeExitListener: (() => void) | null = null;
+const activeManagers = new Set<CompanionManager>();
+
 interface CompanionSession {
   session_id: string;
   cwd: string;
@@ -153,6 +163,7 @@ export class CompanionManager {
 
   onLoad(): void {
     if (this.config?.enabled !== true) {
+      CompanionManager.disposeActiveManagers(this.id);
       try {
         if (!existsSync(stateFilePath())) return;
         writeState((state) => {
@@ -163,11 +174,37 @@ export class CompanionManager {
       } catch {}
       return;
     }
-    process.on('exit', () => this.onExit());
+    this.registerActiveManager();
     this.flush();
     this.spawnIfAvailable();
   }
 
+  /**
+   * Register this manager behind a single process `exit` listener. Re-inits for
+   * the same OpenCode session dispose the superseded manager immediately so its
+   * detached child does not survive until process exit.
+   */
+  private registerActiveManager(): void {
+    for (const manager of [...activeManagers]) {
+      if (manager !== this && manager.id === this.id) {
+        manager.onExit();
+      }
+    }
+
+    activeManagers.add(this);
+    if (!activeExitListener) {
+      activeExitListener = () => CompanionManager.disposeActiveManagers();
+      process.on('exit', activeExitListener);
+    }
+  }
+
+  private static disposeActiveManagers(sessionId?: string): void {
+    for (const manager of [...activeManagers]) {
+      if (sessionId && manager.id !== sessionId) continue;
+      manager.onExit();
+    }
+  }
+
   /**
    * Feed every session.status event here, with the agent name resolved
    * from sessionAgentMap. Orchestrator sessions drive overall status;
@@ -223,13 +260,20 @@ export class CompanionManager {
   }
 
   onExit(): void {
-    if (this.config?.enabled !== true) return;
+    activeManagers.delete(this);
     if (this.companionProcess) {
       try {
         this.companionProcess.kill();
       } catch {}
       this.companionProcess = null;
     }
+    if (activeManagers.size === 0 && activeExitListener) {
+      try {
+        process.removeListener('exit', activeExitListener);
+      } catch {}
+      activeExitListener = null;
+    }
+    if (this.config?.enabled !== true) return;
     writeState((state) => {
       state.sessions = state.sessions.filter((s) => s.session_id !== this.id);
     });