Browse Source

fix(companion): serialize singleton pid handoff

Alvin Unreal 3 weeks ago
parent
commit
2a22ee0ee9
3 changed files with 242 additions and 45 deletions
  1. 26 4
      companion/src/app.rs
  2. 88 4
      src/companion/manager.test.ts
  3. 128 37
      src/companion/manager.ts

+ 26 - 4
companion/src/app.rs

@@ -261,9 +261,12 @@ fn choose_session(sessions: &[SessionInfo]) -> Option<usize> {
 
 fn choose_owned_session(sessions: &[SessionInfo], owner_session_id: Option<&str>) -> Option<usize> {
     if let Some(owner_session_id) = owner_session_id {
-        return sessions
+        if let Some(index) = sessions
             .iter()
-            .position(|session| session.session_id == owner_session_id);
+            .position(|session| session.session_id == owner_session_id)
+        {
+            return Some(index);
+        }
     }
 
     choose_session(sessions)
@@ -856,8 +859,9 @@ fn is_pid_alive(_pid: u32) -> bool {
 #[cfg(test)]
 mod tests {
     use super::{
-        apply_config, choose_session, config_key, grid_dims, place_window, restore_window_position,
-        size_from_config, window_size, ConfigKey, SessionInfo, WindowGeometryKey, GAP,
+        apply_config, choose_owned_session, choose_session, config_key, grid_dims, place_window,
+        restore_window_position, size_from_config, window_size, ConfigKey, SessionInfo,
+        WindowGeometryKey, GAP,
     };
     use crate::state::CompanionConfigState;
 
@@ -909,6 +913,24 @@ mod tests {
         assert_eq!(choose_session(&sessions), Some(1));
     }
 
+    #[test]
+    fn owned_session_wins_when_present() {
+        let sessions = vec![
+            session("first", "waiting-input", &["input"]),
+            session("owner", "idle", &["intro"]),
+        ];
+        assert_eq!(choose_owned_session(&sessions, Some("owner")), Some(1));
+    }
+
+    #[test]
+    fn missing_owner_falls_back_to_active_session() {
+        let sessions = vec![
+            session("idle", "idle", &["intro"]),
+            session("active", "busy", &["fixer"]),
+        ];
+        assert_eq!(choose_owned_session(&sessions, Some("gone")), Some(1));
+    }
+
     #[test]
     fn config_size_defaults_and_presets_work() {
         assert_eq!(size_from_config("small"), 80.0);

+ 88 - 4
src/companion/manager.test.ts

@@ -1,5 +1,6 @@
 import { afterEach, beforeEach, describe, expect, it } from 'bun:test';
 import {
+  chmodSync,
   existsSync,
   mkdirSync,
   readFileSync,
@@ -60,6 +61,10 @@ function attachFakeChild(manager: CompanionManager): { killed: () => boolean } {
   return { killed: () => killed };
 }
 
+function companionPidFile(): string {
+  return path.join(path.dirname(stateFilePath()), 'companion.pid');
+}
+
 describe('CompanionManager', () => {
   it('writes an intro entry on load', () => {
     const m = make();
@@ -254,6 +259,10 @@ describe('CompanionManager', () => {
     const first = make('reload-session');
     first.onLoad();
     const firstChild = attachFakeChild(first);
+    writeFileSync(companionPidFile(), String(process.pid));
+    (first as unknown as { wasSpawner: boolean }).wasSpawner = true;
+    (first as unknown as { spawnedCompanionPid: number }).spawnedCompanionPid =
+      process.pid;
 
     const second = make('reload-session');
     second.onLoad();
@@ -269,6 +278,11 @@ describe('CompanionManager', () => {
     const enabled = make('disable-session');
     enabled.onLoad();
     const child = attachFakeChild(enabled);
+    writeFileSync(companionPidFile(), String(process.pid));
+    (enabled as unknown as { wasSpawner: boolean }).wasSpawner = true;
+    (
+      enabled as unknown as { spawnedCompanionPid: number }
+    ).spawnedCompanionPid = process.pid;
 
     const disabled = new CompanionManager('disable-session', '/path', {
       enabled: false,
@@ -439,7 +453,7 @@ describe('CompanionManager', () => {
   it('skips spawn when PID file points to a live process', () => {
     // Write a PID file with our own PID (which is alive)
     mkdirSync(path.dirname(stateFilePath()), { recursive: true });
-    const pidFile = path.join(path.dirname(stateFilePath()), 'companion.pid');
+    const pidFile = companionPidFile();
     writeFileSync(pidFile, String(process.pid));
 
     const m = make('test-pid-guard');
@@ -453,7 +467,7 @@ describe('CompanionManager', () => {
 
   it('spawns when PID file contains a dead process', () => {
     mkdirSync(path.dirname(stateFilePath()), { recursive: true });
-    const pidFile = path.join(path.dirname(stateFilePath()), 'companion.pid');
+    const pidFile = companionPidFile();
     // Use an impossibly high PID that no kernel will ever assign
     writeFileSync(pidFile, '999999999');
 
@@ -466,6 +480,42 @@ describe('CompanionManager', () => {
     expect(state.sessions[0].session_id).toBe('test-stale-pid');
   });
 
+  it('skips spawn while another process holds the PID file lock', () => {
+    mkdirSync(path.dirname(stateFilePath()), { recursive: true });
+    const pidFile = companionPidFile();
+    const lock = `${pidFile}.lock`;
+    mkdirSync(lock);
+    writeFileSync(path.join(lock, 'owner'), String(process.pid));
+
+    const m = make('test-pending-pid');
+    m.onLoad();
+    m.onExit();
+
+    expect(existsSync(lock)).toBe(true);
+    expect(existsSync(pidFile)).toBe(false);
+  });
+
+  it('stores the spawned child PID in the PID file', () => {
+    const bin = path.join(TEST_DIR, 'fake-companion');
+    writeFileSync(bin, '#!/bin/sh\nexec sleep 30\n');
+    chmodSync(bin, 0o755);
+
+    const m = make('test-child-pid', '/path', {
+      enabled: true,
+      position: 'bottom-right',
+      size: 'medium',
+      binaryPath: bin,
+    });
+    m.onLoad();
+
+    const pid = Number(readFileSync(companionPidFile(), 'utf8'));
+    expect(Number.isInteger(pid)).toBe(true);
+    expect(pid).not.toBe(process.pid);
+    expect(pid).toBe(
+      (m as unknown as { spawnedCompanionPid: number }).spawnedCompanionPid,
+    );
+  });
+
   it('spawns when no PID file exists', () => {
     const m = make('test-no-pid');
     m.onLoad();
@@ -475,7 +525,7 @@ describe('CompanionManager', () => {
 
   it('cleans up PID file on exit when this manager was the spawner', () => {
     mkdirSync(path.dirname(stateFilePath()), { recursive: true });
-    const pidFile = path.join(path.dirname(stateFilePath()), 'companion.pid');
+    const pidFile = companionPidFile();
     writeFileSync(pidFile, '999999999'); // stale PID so spawn proceeds
 
     const m = make('test-pid-cleanup');
@@ -484,6 +534,8 @@ describe('CompanionManager', () => {
     // but the cleanup logic only fires when wasSpawner is true.
     writeFileSync(pidFile, String(process.pid));
     (m as unknown as { wasSpawner: boolean }).wasSpawner = true;
+    (m as unknown as { spawnedCompanionPid: number }).spawnedCompanionPid =
+      process.pid;
     m.onLoad();
     m.onExit();
 
@@ -492,7 +544,7 @@ describe('CompanionManager', () => {
 
   it('does not delete PID file on exit when this manager was not the spawner', () => {
     mkdirSync(path.dirname(stateFilePath()), { recursive: true });
-    const pidFile = path.join(path.dirname(stateFilePath()), 'companion.pid');
+    const pidFile = companionPidFile();
     writeFileSync(pidFile, String(process.pid));
 
     const m = make('test-pid-no-cleanup');
@@ -503,6 +555,38 @@ describe('CompanionManager', () => {
     expect(existsSync(pidFile)).toBe(true);
   });
 
+  it('does not delete a PID file owned by a different spawned child', () => {
+    mkdirSync(path.dirname(stateFilePath()), { recursive: true });
+    const pidFile = companionPidFile();
+    writeFileSync(pidFile, '222222222');
+
+    const m = make('test-pid-different-child');
+    (m as unknown as { wasSpawner: boolean }).wasSpawner = true;
+    (m as unknown as { spawnedCompanionPid: number }).spawnedCompanionPid =
+      111111111;
+    m.onExit();
+
+    expect(readFileSync(pidFile, 'utf8')).toBe('222222222');
+  });
+
+  it('does not kill the singleton when another session remains in state', () => {
+    const first = make('first-session');
+    const second = make('second-session');
+    first.onLoad();
+    second.onLoad();
+    const child = attachFakeChild(first);
+    const pidFile = companionPidFile();
+    writeFileSync(pidFile, String(process.pid));
+    (first as unknown as { wasSpawner: boolean }).wasSpawner = true;
+    (first as unknown as { spawnedCompanionPid: number }).spawnedCompanionPid =
+      process.pid;
+
+    first.onExit();
+
+    expect(child.killed()).toBe(false);
+    expect(readFileSync(pidFile, 'utf8')).toBe(String(process.pid));
+  });
+
   it('removes disabled session entries on load', () => {
     mkdirSync(path.dirname(stateFilePath()), { recursive: true });
     writeFileSync(

+ 128 - 37
src/companion/manager.ts

@@ -5,6 +5,7 @@ import {
   readFileSync,
   renameSync,
   rmSync,
+  statSync,
   writeFileSync,
 } from 'node:fs';
 import * as os from 'node:os';
@@ -77,12 +78,66 @@ function pidFilePath(): string {
 }
 
 function isProcessAlive(pid: number): boolean {
+  if (!Number.isInteger(pid) || pid <= 0) return false;
   try {
     process.kill(pid, 0);
     return true;
+  } catch (err) {
+    return (err as NodeJS.ErrnoException).code === 'EPERM';
+  }
+}
+
+function parsePidFile(raw: string): number | null {
+  const pid = Number(raw.trim());
+  if (!Number.isInteger(pid) || pid <= 0) return null;
+  return pid;
+}
+
+function acquirePidFileLock(file: string): (() => void) | null {
+  const lock = `${file}.lock`;
+  mkdirSync(path.dirname(lock), { recursive: true });
+  for (let attempt = 0; attempt < 2; attempt++) {
+    try {
+      mkdirSync(lock);
+      writeFileSync(path.join(lock, 'owner'), String(process.pid));
+      return () => {
+        try {
+          rmSync(lock, { recursive: true, force: true });
+        } catch {}
+      };
+    } catch (err) {
+      const code = (err as NodeJS.ErrnoException).code;
+      if (code !== 'EEXIST') throw err;
+      if (pidFileLockHasLiveOwner(lock)) return null;
+      log('[companion] removing stale PID file lock for dead process');
+      rmSync(lock, { recursive: true, force: true });
+    }
+  }
+  return null;
+}
+
+function acquirePidFileLockWithRetry(
+  file: string,
+  attempts: number,
+): (() => void) | null {
+  for (let attempt = 0; attempt < attempts; attempt++) {
+    const release = acquirePidFileLock(file);
+    if (release) return release;
+    Atomics.wait(new Int32Array(new SharedArrayBuffer(4)), 0, 0, 25);
+  }
+  return null;
+}
+
+function pidFileLockHasLiveOwner(lock: string): boolean {
+  try {
+    const owner = parsePidFile(readFileSync(path.join(lock, 'owner'), 'utf8'));
+    if (owner !== null) return isProcessAlive(owner);
   } catch {
-    return false;
+    try {
+      return Date.now() - statSync(lock).mtimeMs < 5000;
+    } catch {}
   }
+  return false;
 }
 
 function defaultBinaryPath(): string {
@@ -179,6 +234,7 @@ export class CompanionManager {
   private readonly config?: CompanionConfig;
   private companionProcess: ChildProcess | null = null;
   private wasSpawner = false;
+  private spawnedCompanionPid: number | null = null;
 
   constructor(sessionId: string, cwd: string, config?: CompanionConfig) {
     this.id = sessionId;
@@ -286,18 +342,6 @@ export class CompanionManager {
 
   onExit(): void {
     activeManagers.delete(this);
-    if (this.wasSpawner) {
-      try {
-        const pf = pidFilePath();
-        if (existsSync(pf)) rmSync(pf, { force: true });
-      } catch {}
-    }
-    if (this.companionProcess) {
-      try {
-        this.companionProcess.kill();
-      } catch {}
-      this.companionProcess = null;
-    }
     if (activeManagers.size === 0 && activeExitListener) {
       try {
         process.removeListener('exit', activeExitListener);
@@ -308,6 +352,37 @@ export class CompanionManager {
     writeState((state) => {
       state.sessions = state.sessions.filter((s) => s.session_id !== this.id);
     });
+    if (this.wasSpawner && this.removeOwnedPidFileIfNoSessionsRemain()) {
+      if (this.companionProcess) {
+        try {
+          this.companionProcess.kill();
+        } catch {}
+      }
+    }
+    this.companionProcess = null;
+  }
+
+  private removeOwnedPidFileIfNoSessionsRemain(): boolean {
+    if (this.spawnedCompanionPid == null) return true;
+    const file = pidFilePath();
+    const release = acquirePidFileLockWithRetry(file, 80);
+    if (!release) {
+      log('[companion] PID file lock busy during exit; leaving guard intact');
+      return false;
+    }
+    try {
+      if (readState().sessions.length > 0) return false;
+      if (!existsSync(file)) return true;
+      const parsed = parsePidFile(readFileSync(file, 'utf8'));
+      if (parsed === this.spawnedCompanionPid) {
+        rmSync(file, { force: true });
+      }
+      return true;
+    } catch {
+      return false;
+    } finally {
+      release();
+    }
   }
 
   /** One entry per running agent instance (two fixers → two cells). */
@@ -366,32 +441,37 @@ export class CompanionManager {
 
   private spawnIfAvailable(): void {
     if (this.config?.enabled !== true) return;
-    // Check if another process already spawned the companion
     const pidFile = pidFilePath();
+    let releasePidFileLock: (() => void) | null = null;
+    try {
+      releasePidFileLock = acquirePidFileLockWithRetry(pidFile, 80);
+      if (releasePidFileLock === null) {
+        log('[companion] another instance already running, skipping spawn');
+        return;
+      }
+    } catch (err) {
+      log('[companion] PID file lock failed', String(err));
+      return;
+    }
+    let spawnedChild: ChildProcess | null = null;
     try {
       if (existsSync(pidFile)) {
-        const existingPid = Number(readFileSync(pidFile, 'utf8').trim());
-        if (isProcessAlive(existingPid)) {
+        const existingPid = parsePidFile(readFileSync(pidFile, 'utf8'));
+        if (existingPid !== null && isProcessAlive(existingPid)) {
           log('[companion] another instance already running, skipping spawn');
           return;
         }
-        // Stale PID file — process died without cleanup
+        log('[companion] removing stale PID file for dead process');
+        rmSync(pidFile, { force: true });
+      }
+      const bin = resolveCompanionBinaryPath(this.config);
+      if (!bin) {
+        const expected = this.config.binaryPath?.trim() || defaultBinaryPath();
         log(
-          '[companion] removing stale PID file for dead process',
-          String(existingPid),
+          `[companion] enabled but companion binary not found at expected path: ${expected}. Please install/download the companion binary separately.`,
         );
-        rmSync(pidFile, { force: true });
+        return;
       }
-    } catch {}
-    const bin = resolveCompanionBinaryPath(this.config);
-    if (!bin) {
-      const expected = this.config.binaryPath?.trim() || defaultBinaryPath();
-      log(
-        `[companion] enabled but companion binary not found at expected path: ${expected}. Please install/download the companion binary separately.`,
-      );
-      return;
-    }
-    try {
       const child = spawn(bin, [], {
         detached: true,
         env: {
@@ -403,9 +483,19 @@ export class CompanionManager {
         },
         stdio: 'ignore',
       });
+      spawnedChild = child;
+      child.once('error', (err) => {
+        log('[companion] spawn failed', String(err));
+      });
       this.companionProcess = child;
       child.unref();
+      if (child.pid == null) {
+        log('[companion] spawn returned without a child PID, skipping guard');
+        return;
+      }
+      writeFileSync(pidFile, String(child.pid));
       this.wasSpawner = true;
+      this.spawnedCompanionPid = child.pid;
       log(
         '[companion] spawned',
         JSON.stringify({
@@ -414,14 +504,15 @@ export class CompanionManager {
           debug: this.config.debug === true,
         }),
       );
-      try {
-        if (child.pid != null) {
-          mkdirSync(path.dirname(pidFile), { recursive: true });
-          writeFileSync(pidFile, String(child.pid));
-        }
-      } catch {}
     } catch (err) {
-      log('[companion] spawn failed', String(err));
+      if (spawnedChild && !this.wasSpawner) {
+        try {
+          spawnedChild.kill();
+        } catch {}
+      }
+      log('[companion] spawn guard failed', String(err));
+    } finally {
+      releasePidFileLock?.();
     }
   }
 }