diff --git a/src/praisonai-bot/praisonai_bot/gateway/kanban_dispatcher.py b/src/praisonai-bot/praisonai_bot/gateway/kanban_dispatcher.py index c2a15b684..0b8d2a71b 100644 --- a/src/praisonai-bot/praisonai_bot/gateway/kanban_dispatcher.py +++ b/src/praisonai-bot/praisonai_bot/gateway/kanban_dispatcher.py @@ -199,6 +199,10 @@ class KanbanDispatcher: # Clean up completed processes first self._cleanup_completed_tasks(store) + # Prove liveness for our own children before the reclaim sweep runs, + # so a healthy long-running worker is never mistaken for a stale one. + self._heartbeat_children(store) + # Reclaim stale/crashed claims before dispatching new work so that # tasks stranded by dead workers return to 'ready' and get re-dispatched. self._reclaim_stale_claims(store) @@ -221,6 +225,14 @@ class KanbanDispatcher: spawned = 0 for task in ready_tasks[:available_slots]: + live = self.running_tasks.get(task.id) + if live is not None and live.poll() is None: + # Never overwrite the handle of a still-running worker. + logger.error( + "Task %s is 'ready' but its worker (pid %s) is still running; " + "skipping re-dispatch", task.id, live.pid, + ) + continue try: # Try to claim the task atomically with a lease + owner PID. claimed = self._claim_task(store, task) @@ -394,6 +406,69 @@ class KanbanDispatcher: (pid, datetime.now(timezone.utc).isoformat(), task_id, self.worker_id), ) + def _heartbeat_children(self, store: Any): + """Renew lease + heartbeat for every child worker that is still alive. + + The spawned worker cannot heartbeat for itself, so the supervisor that + holds its ``Popen`` handle does it. If the claim is gone (an operator + released or moved the card), the child is terminated rather than left + running unowned. + """ + beat = getattr(store, 'heartbeat', None) + if beat is None: + return + for task_id, process in list(self.running_tasks.items()): + if process.poll() is not None: + continue + try: + owned = beat(task_id, self.worker_id, ttl_seconds=self.claim_ttl_seconds) + except Exception as e: + logger.debug(f"Heartbeat for task {task_id} failed: {e}") + continue + if not owned: + logger.warning( + f"Lost claim on task {task_id} while worker {process.pid} " + f"was running; terminating it" + ) + self._terminate_unowned(store, task_id, process) + + def _terminate_unowned(self, store: Any, task_id: str, process: subprocess.Popen): + """Stop a child whose claim is no longer ours and forget its tracking.""" + try: + process.terminate() + process.wait(timeout=5) + except subprocess.TimeoutExpired: + process.kill() + except OSError as e: + logger.error(f"Failed to terminate worker for task {task_id}: {e}") + self.running_tasks.pop(task_id, None) + entry = getattr(self, '_worktrees', {}).get(task_id) + if entry: + # Lossless teardown: keep tracking (and say why) if work remains. + preserved = self._remove_worktree(*entry) + if preserved is None: + self._worktrees.pop(task_id, None) + else: + try: + store.add_comment( + task_id, self.worker_id, + f"worktree_preserved at {entry[0]}: {preserved}", + ) + except Exception: + pass + log_path = getattr(self, '_temp_logs', {}).pop(task_id, None) + if log_path: + try: + os.unlink(log_path) + except OSError: + pass + run_id = self._task_runs.get(task_id) + if run_id is None or self._close_run_safe( + store, run_id, 'failed', + error='Claim lost while worker was running; worker terminated', + ): + self._task_runs.pop(task_id, None) + def _reclaim_stale_claims(self, store: Any): """Sweep running tasks and reclaim those stranded by dead/stale workers.""" reclaim = getattr(store, 'reclaim_stale_claims', None) @@ -409,6 +484,11 @@ class KanbanDispatcher: for task_id in reclaimed or []: logger.warning(f"Reclaimed stale kanban task {task_id} back to 'ready'") + # The store closed the abandoned run as 'crashed' atomically with + # the reclaim; count the attempt so the circuit breaker sees it. + self._record_failure_safe( + store, task_id, error='Worker died or went stale; claim reclaimed' + ) # The running -> ready transition is announced as # KANBAN_TASK_MOVED (reclaimed=True) by the store itself; # 'KANBAN_TASK_RECLAIMED' was never a HookEvent member and so diff --git a/src/praisonai-bot/praisonai_bot/kanban/sqlite_store.py b/src/praisonai-bot/praisonai_bot/kanban/sqlite_store.py index 3567929d7..0c3210155 100644 --- a/src/praisonai-bot/praisonai_bot/kanban/sqlite_store.py +++ b/src/praisonai-bot/praisonai_bot/kanban/sqlite_store.py @@ -1094,7 +1094,8 @@ class SQLiteKanbanStore: with self._get_connection() as conn: cursor = conn.execute(""" - SELECT id, claim_lock, claim_expires, worker_pid, last_heartbeat_at, version + SELECT id, claim_lock, claim_expires, worker_pid, last_heartbeat_at, version, + current_run_id FROM tasks WHERE status = 'running' AND claim_lock IS NOT NULL AND claim_lock != '' """) @@ -1127,7 +1128,7 @@ class SQLiteKanbanStore: result = conn.execute(""" UPDATE tasks SET claim_lock = NULL, claim_expires = NULL, worker_pid = NULL, - last_heartbeat_at = NULL, + last_heartbeat_at = NULL, current_run_id = NULL, updated_at = ?, status = 'ready', version = version + 1 WHERE id = ? AND status = 'running' AND claim_lock = ? AND version = ? @@ -1136,6 +1137,20 @@ class SQLiteKanbanStore: if result.rowcount > 0: reclaimed.append(row['id']) reason = 'dead' if not worker_alive else 'stale_heartbeat' + # Close the abandoned attempt in the same transaction, so a + # dispatcher that claims the task next can never have its + # own fresh run mistaken for this one. + if row.get('current_run_id') is not None and conn.execute( + "UPDATE task_runs SET outcome = ?, error = ?, ended_at = ? " + "WHERE id = ? AND ended_at IS NULL", + (RunOutcome.CRASHED.value, + f'Worker {reason}; claim reclaimed', + now.isoformat(), row['current_run_id']), + ).rowcount: + self._log_event(conn, row['id'], 'run_closed', { + 'run_id': row['current_run_id'], + 'outcome': RunOutcome.CRASHED.value, + }) reclaim_details[row['id']] = { 'worker_id': row['claim_lock'], 'worker_pid': pid, diff --git a/src/praisonai-bot/tests/unit/gateway/test_kanban_dispatcher_heartbeat.py b/src/praisonai-bot/tests/unit/gateway/test_kanban_dispatcher_heartbeat.py new file mode 100644 index 000000000..92e5805d9 --- /dev/null +++ b/src/praisonai-bot/tests/unit/gateway/test_kanban_dispatcher_heartbeat.py @@ -0,0 +1,167 @@ +"""Regression tests for #6025: the dispatcher heartbeats its own children. + +A healthy worker that outlives the claim lease + stale timeout must not be +reclaimed and re-dispatched while it is still running. +""" + +import asyncio +import subprocess +import sys + +import pytest + +from praisonai_bot.gateway.kanban_dispatcher import KanbanDispatcher +from praisonai_bot.kanban.sqlite_store import SQLiteKanbanStore + +pytestmark = pytest.mark.skipif(sys.platform == "win32", reason="uses sleep stub") + + +class _StubDispatcher(KanbanDispatcher): + def __init__(self, store, seconds, **kwargs): + super().__init__(**kwargs) + self._store = store + self._seconds = seconds + + def _get_kanban_store(self): + return self._store + + def _build_execution_command(self, task): + return [sys.executable, "-c", f"import time; time.sleep({self._seconds})"] + + +@pytest.fixture +def store(tmp_path, monkeypatch): + monkeypatch.setenv("PRAISONAI_KANBAN_DB", str(tmp_path / "kanban.db")) + return SQLiteKanbanStore() + + +def _run(coro): + return asyncio.run(coro) + + +def test_long_running_worker_is_not_reclaimed(store): + store.create_task({"id": "t1", "title": "long job", "status": "ready", + "workspace_kind": "default"}) + d = _StubDispatcher(store, 5, max_concurrent=3, poll_interval=0.5, + claim_ttl_seconds=1, stale_timeout_seconds=2) + + async def drive(): + handles = [] + loop = asyncio.get_running_loop() + t0 = loop.time() + while loop.time() - t0 < 4.5: + await d.dispatch_once() + proc = d.running_tasks.get("t1") + if proc is not None and proc not in handles: + handles.append(proc) + await asyncio.sleep(0.25) + return handles + + handles = _run(drive()) + try: + assert len(handles) == 1 + runs = store.get_runs("t1") + assert len(runs) == 1 + assert runs[0].ended_at is None + assert not [e for e in store.list_events() + if e.task_id == "t1" and e.event_type == "reclaimed"] + + handles[0].wait(timeout=10) + _run(d.dispatch_once()) + runs = store.get_runs("t1") + assert len(runs) == 1 + assert runs[0].outcome == "completed" + assert store.get_task("t1").status == "done" + finally: + for p in handles: + if p.poll() is None: + p.kill() + + +def test_lost_claim_terminates_child(store): + store.create_task({"id": "t2", "title": "job", "status": "ready", + "workspace_kind": "default"}) + d = _StubDispatcher(store, 30, max_concurrent=1, poll_interval=0.5) + _run(d.dispatch_once()) + proc = d.running_tasks["t2"] + try: + store.release_claim("t2", d.worker_id) # operator takes the card back + d._heartbeat_children(store) + assert "t2" not in d.running_tasks + assert proc.wait(timeout=10) is not None + runs = store.get_runs("t2") + assert len(runs) == 1 and runs[0].outcome == "failed" + finally: + if proc.poll() is None: + proc.kill() + + +def test_reclaim_of_dead_worker_closes_run_and_counts_failure(store): + store.create_task({"id": "t3", "title": "job", "status": "ready", + "workspace_kind": "default"}) + # Simulate another dispatcher whose worker has died. + assert store.claim_task("t3", "other", ttl_seconds=0, worker_pid=0) + run_id = store.start_run("t3", profile="other") + d = KanbanDispatcher(claim_ttl_seconds=1, stale_timeout_seconds=1) + + d._reclaim_stale_claims(store) + + run = [r for r in store.get_runs("t3") if r.id == run_id][0] + assert run.outcome == "crashed" + assert run.ended_at + assert store.get_task("t3").consecutive_failures == 1 + + +def test_live_handle_is_never_overwritten(store): + store.create_task({"id": "t4", "title": "job", "status": "ready", + "workspace_kind": "default"}) + d = _StubDispatcher(store, 30, max_concurrent=2) + proc = subprocess.Popen([sys.executable, "-c", "import time; time.sleep(30)"]) + d.running_tasks["t4"] = proc + d._heartbeat_children = lambda _store: None # isolate the spawn guard + try: + assert _run(d.dispatch_once()) == 0 + assert d.running_tasks["t4"] is proc + finally: + proc.kill() + proc.wait() + + +def test_reclaim_closes_run_atomically_so_replacement_run_survives(store): + store.create_task({"id": "t5", "title": "job", "status": "ready", + "workspace_kind": "default"}) + assert store.claim_task("t5", "other", ttl_seconds=0, worker_pid=0) + old_run = store.start_run("t5", profile="other") + + assert store.reclaim_stale_claims(stale_timeout_seconds=1) == ["t5"] + assert store.get_task("t5").current_run_id is None + assert any(e.event_type == "run_closed" and e.task_id == "t5" + for e in store.list_events()) + + # Another dispatcher wins the reclaimed task before ours reacts. + assert store.claim_task("t5", "rival", ttl_seconds=60) + new_run = store.start_run("t5", profile="rival") + KanbanDispatcher()._reclaim_stale_claims(store) + + runs = {r.id: r for r in store.get_runs("t5")} + assert runs[old_run].outcome == "crashed" + assert runs[new_run].ended_at is None + + +def test_lost_claim_keeps_worktree_tracking_when_work_is_preserved(store): + store.create_task({"id": "t6", "title": "job", "status": "ready", + "workspace_kind": "default"}) + d = _StubDispatcher(store, 30, max_concurrent=1) + _run(d.dispatch_once()) + proc = d.running_tasks["t6"] + d._worktrees = {"t6": ("/tmp/wt-t6", "kanban/t6")} + d._remove_worktree = lambda path, branch: "worktree has uncommitted changes" + try: + store.release_claim("t6", d.worker_id) + d._heartbeat_children(store) + assert "t6" not in d.running_tasks + assert d._worktrees == {"t6": ("/tmp/wt-t6", "kanban/t6")} + assert any("worktree_preserved" in c.text for c in store.get_comments("t6")) + finally: + if proc.poll() is None: + proc.kill()