Browse Source

fix(companion): clean up superseded managers

Alvin Unreal 1 month ago
parent
commit
1e90e02d5b
2 changed files with 91 additions and 30 deletions
  1. 54 4
      src/companion/manager.test.ts
  2. 37 26
      src/companion/manager.ts

+ 54 - 4
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', () => {
@@ -212,11 +233,9 @@ describe('CompanionManager', () => {
 
   it('keeps at most one process exit listener across reloads', () => {
     const baseline = process.listenerCount('exit');
-    const managers: CompanionManager[] = [];
     for (let i = 0; i < 5; i++) {
-      const m = make(`reload-${i}`);
+      const m = make('reload-session');
       m.onLoad();
-      managers.push(m);
     }
     // Re-inits must dedup the exit listener rather than stacking one each time.
     expect(process.listenerCount('exit')).toBeLessThanOrEqual(baseline + 1);
@@ -225,6 +244,37 @@ describe('CompanionManager', () => {
     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');

+ 37 - 26
src/companion/manager.ts

@@ -14,11 +14,13 @@ 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 a fresh CompanionManager each time; without deduping, every
-// re-init would leak another exit listener and retain the previous manager
-// (and its detached child reference). Module-level state survives re-inits
-// because the module itself is not re-evaluated.
+// 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;
@@ -152,7 +154,6 @@ export class CompanionManager {
   private readonly busyAgentSessions = new Map<string, string>();
   private readonly config?: CompanionConfig;
   private companionProcess: ChildProcess | null = null;
-  private exitListener: (() => void) | null = null;
 
   constructor(sessionId: string, cwd: string, config?: CompanionConfig) {
     this.id = sessionId;
@@ -162,6 +163,7 @@ export class CompanionManager {
 
   onLoad(): void {
     if (this.config?.enabled !== true) {
+      CompanionManager.disposeActiveManagers(this.id);
       try {
         if (!existsSync(stateFilePath())) return;
         writeState((state) => {
@@ -172,26 +174,35 @@ export class CompanionManager {
       } catch {}
       return;
     }
-    this.registerExitListener();
+    this.registerActiveManager();
     this.flush();
     this.spawnIfAvailable();
   }
 
   /**
-   * Register a single process `exit` listener, replacing any previously
-   * registered companion listener so re-inits don't stack listeners (and so
-   * the prior manager + its child reference become collectable).
+   * 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 registerExitListener(): void {
-    if (activeExitListener) {
-      try {
-        process.removeListener('exit', activeExitListener);
-      } catch {}
+  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();
     }
-    const listener = () => this.onExit();
-    process.on('exit', listener);
-    activeExitListener = listener;
-    this.exitListener = listener;
   }
 
   /**
@@ -249,20 +260,20 @@ export class CompanionManager {
   }
 
   onExit(): void {
-    if (this.config?.enabled !== true) return;
-    if (this.exitListener) {
-      try {
-        process.removeListener('exit', this.exitListener);
-      } catch {}
-      if (activeExitListener === this.exitListener) activeExitListener = null;
-      this.exitListener = null;
-    }
+    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);
     });