From 1c50b1454a0ecc323aa879015ad1307d1b5c54b4 Mon Sep 17 00:00:00 2001 From: AIOSAI Date: Sun, 26 Apr 2026 17:30:24 -0700 Subject: [PATCH] feat(ai_mail): test(ai_mail): fix lock ordering test for DPLAN-0155 Phase 5 Co-Authored-By: @ai_mail --- .../ai_mail/apps/handlers/dispatch/daemon.py | 18 ++++++++++--- .../ai_mail/apps/handlers/dispatch/wake.py | 25 +++++++++++++------ src/aipass/ai_mail/tests/test_wake.py | 16 ++++-------- 3 files changed, 37 insertions(+), 22 deletions(-) diff --git a/src/aipass/ai_mail/apps/handlers/dispatch/daemon.py b/src/aipass/ai_mail/apps/handlers/dispatch/daemon.py index 70cf9b97..ae890e36 100644 --- a/src/aipass/ai_mail/apps/handlers/dispatch/daemon.py +++ b/src/aipass/ai_mail/apps/handlers/dispatch/daemon.py @@ -384,6 +384,13 @@ def spawn_agent( if key.startswith("CLAUDE") or key == "AIPASS_BOT_ID": spawn_env.pop(key) + # Acquire lock BEFORE spawn to prevent TOCTOU race (DPLAN-0155 Phase 5). + # Use current PID as placeholder; overwrite with monitor PID after spawn. + acquired, lock_msg = _acquire_lock(branch_path, os.getpid()) + if not acquired: + logger.info(f"Lock acquisition failed for {branch_email}: {lock_msg}") + return False + try: process = subprocess.Popen( monitor_cmd, @@ -396,10 +403,10 @@ def spawn_agent( monitor_pid = process.pid - # Lock PID = monitor PID (stays alive as long as claude does) - acquired, lock_msg = _acquire_lock(branch_path, monitor_pid) - if not acquired: - logger.info(f"Lock acquisition failed after spawn for {branch_email}: {lock_msg}") + # Update lock with real monitor PID + lock_file = branch_path / ".ai_mail.local" / ".dispatch.lock" + lock_data = {"pid": monitor_pid, "timestamp": datetime.now().isoformat(), "branch": str(branch_path)} + _write_json(lock_file, lock_data) # Track session cycles for rotation cycles = state.get("session_cycles", {}) @@ -428,6 +435,9 @@ def spawn_agent( return True except Exception as e: + # Release lock on spawn failure so branch isn't stuck locked + lock_file = branch_path / ".ai_mail.local" / ".dispatch.lock" + lock_file.unlink(missing_ok=True) logger.info(f"SPAWN FAILED {branch_email}: {e}") log_dispatch(branch_email, None, "failed", error_msg=str(e)) _notify_telegram(f"[Dispatch FAILED] {branch_email}\n{type(e).__name__}: {e}") diff --git a/src/aipass/ai_mail/apps/handlers/dispatch/wake.py b/src/aipass/ai_mail/apps/handlers/dispatch/wake.py index 6536bfbe..d126478d 100644 --- a/src/aipass/ai_mail/apps/handlers/dispatch/wake.py +++ b/src/aipass/ai_mail/apps/handlers/dispatch/wake.py @@ -480,6 +480,13 @@ def wake_branch( if key.startswith("CLAUDE") or key == "AIPASS_BOT_ID": spawn_env.pop(key) + # Acquire lock BEFORE spawn to prevent TOCTOU race (DPLAN-0155 Phase 5). + acquired, lock_msg = _acquire_lock(branch_path, os.getpid()) + if not acquired: + status.fail("lock-acquire", f"Lock failed: {lock_msg}") + return status, False + status.ok("lock-acquire", "Dispatch lock acquired") + try: process = subprocess.Popen( monitor_cmd, @@ -491,24 +498,28 @@ def wake_branch( ) monitor_pid = process.pid + + # Update lock with real monitor PID + lock_file = branch_path / ".ai_mail.local" / ".dispatch.lock" + lock_data = {"pid": monitor_pid, "timestamp": time.strftime("%Y-%m-%dT%H:%M:%S"), "branch": str(branch_path)} + with open(lock_file, "w", encoding="utf-8") as f: + json.dump(lock_data, f, indent=2) + status.ok("spawn", f"Monitor started (PID {monitor_pid})") except FileNotFoundError as e: + lock_file = branch_path / ".ai_mail.local" / ".dispatch.lock" + lock_file.unlink(missing_ok=True) logger.warning("[wake] Spawn failed — script not found: %s", e) status.fail("spawn", "Python or monitor script not found") return status, False except Exception as e: + lock_file = branch_path / ".ai_mail.local" / ".dispatch.lock" + lock_file.unlink(missing_ok=True) logger.warning("[wake] Spawn failed for %s: %s", branch_email, e) status.fail("spawn", f"{type(e).__name__}: {e}") return status, False - # Step 8: Acquire lock (with monitor PID — stays alive as long as agent) - acquired, lock_msg = _acquire_lock(branch_path, monitor_pid) - if not acquired: - status.warn("lock-acquire", f"Lock failed: {lock_msg}") - else: - status.ok("lock-acquire", "Dispatch lock acquired") - # Step 9: Liveness check (brief wait then verify) time.sleep(2) if _check_pid_alive(monitor_pid): diff --git a/src/aipass/ai_mail/tests/test_wake.py b/src/aipass/ai_mail/tests/test_wake.py index 5cdc38b3..2ce16c40 100644 --- a/src/aipass/ai_mail/tests/test_wake.py +++ b/src/aipass/ai_mail/tests/test_wake.py @@ -1024,21 +1024,15 @@ class TestWakeBranch: assert ok is False assert any(s[0] == "fail" and "RuntimeError" in s[2] for s in status.steps) - # --- post-spawn: lock acquisition failure --- + # --- pre-spawn: lock acquisition failure (DPLAN-0155) --- - def test_lock_acquisition_fails_after_spawn_warns(self, tmp_path, monkeypatch): - """Lock fails after spawn -> warn step but still succeeds.""" + def test_lock_acquisition_fails_before_spawn(self, tmp_path, monkeypatch): + """Lock fails before spawn -> fail step, returns False (DPLAN-0155 lock-before-spawn).""" _make_wake_fixtures(tmp_path, monkeypatch) _patch_wake_deps(monkeypatch, _acquire_lock=lambda p, pid: (False, "Lock file already exists")) - monkeypatch.setattr("subprocess.Popen", lambda *a, **kw: _FakeProc()) - monkeypatch.setattr( - "aipass.ai_mail.apps.handlers.notify.send_notification", - lambda *a, **kw: None, - raising=False, - ) status, ok = wake_branch("@testbranch") - assert ok is True - assert any(s[0] == "warn" and "lock-acquire" in s[1] for s in status.steps) + assert ok is False + assert any(s[0] == "fail" and "lock-acquire" in s[1] for s in status.steps) # --- alive check fails ---