Merge pull request #654 from mhenke/fix/companion-singleton-pid-guard

fix: PID file singleton guard for companion spawn (#650)

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);

src/companion/manager.test.ts

@@ -1,5 +1,12 @@
 import { afterEach, beforeEach, describe, expect, it } from 'bun:test';
-import { mkdirSync, readFileSync, rmSync, writeFileSync } from 'node:fs';
+import {
+  chmodSync,
+  existsSync,
+  mkdirSync,
+  readFileSync,
+  rmSync,
+  writeFileSync,
+} from 'node:fs';
 import * as os from 'node:os';
 import * as path from 'node:path';
 import {
@@ -54,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();
@@ -248,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();
@@ -263,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,
@@ -430,6 +450,143 @@ describe('CompanionManager', () => {
     expect(state.config.enabled).toBe(true);
   });
 
+  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 = companionPidFile();
+    writeFileSync(pidFile, String(process.pid));
+
+    const m = make('test-pid-guard');
+    m.onLoad();
+
+    // Should not have spawned — PID file guard prevented it
+    // The session should still be written to state
+    const state = readState();
+    expect(state.sessions[0].session_id).toBe('test-pid-guard');
+  });
+
+  it('spawns when PID file contains a dead process', () => {
+    mkdirSync(path.dirname(stateFilePath()), { recursive: true });
+    const pidFile = companionPidFile();
+    // Use an impossibly high PID that no kernel will ever assign
+    writeFileSync(pidFile, '999999999');
+
+    const m = make('test-stale-pid');
+    m.onLoad();
+
+    // Stale PID file should have been cleaned up
+    expect(existsSync(pidFile)).toBe(false);
+    const state = readState();
+    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();
+    const state = readState();
+    expect(state.sessions[0].session_id).toBe('test-no-pid');
+  });
+
+  it('cleans up PID file on exit when this manager was the spawner', () => {
+    mkdirSync(path.dirname(stateFilePath()), { recursive: true });
+    const pidFile = companionPidFile();
+    writeFileSync(pidFile, '999999999'); // stale PID so spawn proceeds
+
+    const m = make('test-pid-cleanup');
+    // Simulate a spawner by writing a PID file as if spawn succeeded.
+    // In reality the binary doesn't exist so spawn fails before writing,
+    // 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();
+
+    expect(existsSync(pidFile)).toBe(false);
+  });
+
+  it('does not delete PID file on exit when this manager was not the spawner', () => {
+    mkdirSync(path.dirname(stateFilePath()), { recursive: true });
+    const pidFile = companionPidFile();
+    writeFileSync(pidFile, String(process.pid));
+
+    const m = make('test-pid-no-cleanup');
+    m.onLoad(); // skips spawn because PID is alive, wasSpawner stays false
+    m.onExit();
+
+    // Non-spawner must not delete the guard file
+    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(

src/companion/manager.ts

@@ -5,6 +5,7 @@ import {
   readFileSync,
   renameSync,
   rmSync,
+  statSync,
   writeFileSync,
 } from 'node:fs';
 import * as os from 'node:os';
@@ -61,6 +62,84 @@ export function stateFilePath(): string {
   );
 }
 
+function pidFilePath(): string {
+  const xdg = process.env.XDG_DATA_HOME?.trim();
+  const base =
+    xdg && path.isAbsolute(xdg)
+      ? xdg
+      : path.join(os.homedir(), '.local', 'share');
+  return path.join(
+    base,
+    'opencode',
+    'storage',
+    'oh-my-opencode-slim',
+    'companion.pid',
+  );
+}
+
+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 {
+    try {
+      return Date.now() - statSync(lock).mtimeMs < 5000;
+    } catch {}
+  }
+  return false;
+}
+
 function defaultBinaryPath(): string {
   const xdg = process.env.XDG_DATA_HOME?.trim();
   const base =
@@ -154,6 +233,8 @@ export class CompanionManager {
   private readonly busyAgentSessions = new Map<string, string>();
   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;
@@ -261,12 +342,6 @@ export class CompanionManager {
 
   onExit(): void {
     activeManagers.delete(this);
-    if (this.companionProcess) {
-      try {
-        this.companionProcess.kill();
-      } catch {}
-      this.companionProcess = null;
-    }
     if (activeManagers.size === 0 && activeExitListener) {
       try {
         process.removeListener('exit', activeExitListener);
@@ -277,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). */
@@ -335,15 +441,37 @@ export class CompanionManager {
 
   private spawnIfAvailable(): void {
     if (this.config?.enabled !== true) return;
-    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.`,
-      );
+    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 = parsePidFile(readFileSync(pidFile, 'utf8'));
+        if (existingPid !== null && isProcessAlive(existingPid)) {
+          log('[companion] another instance already running, skipping spawn');
+          return;
+        }
+        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] enabled but companion binary not found at expected path: ${expected}. Please install/download the companion binary separately.`,
+        );
+        return;
+      }
       const child = spawn(bin, [], {
         detached: true,
         env: {
@@ -355,8 +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({
@@ -366,7 +505,14 @@ export class CompanionManager {
         }),
       );
     } 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?.();
     }
   }
 }