lemming-cli 0.1.0__py3-none-any.whl

lemming/cli_test.py

@@ -0,0 +1,22 @@
+import unittest
+
+import click.testing
+
+from lemming import cli
+
+
+class TestCLIProxy(unittest.TestCase):
+    def setUp(self):
+        self.cli_runner = click.testing.CliRunner()
+
+    def test_cli_proxy_help(self):
+        # Verify that the proxy re-export works correctly
+        result = self.cli_runner.invoke(cli.cli, ["--help"])
+        self.assertEqual(result.exit_code, 0)
+        self.assertIn(
+            "Lemming: An autonomous, iterative task runner", result.output
+        )
+
+
+if __name__ == "__main__":
+    unittest.main()

lemming/conftest.py

@@ -0,0 +1,13 @@
+"""Shared pytest fixtures for the lemming test suite."""
+
+import os
+
+import pytest
+
+
+@pytest.fixture(autouse=True, scope="session")
+def setup_lemming_home(tmp_path_factory):
+    """Set up a temporary LEMMING_HOME for the entire test session."""
+    tmp_home = tmp_path_factory.mktemp("lemming_home")
+    os.environ["LEMMING_HOME"] = str(tmp_home)
+    return tmp_home

lemming/hooks.py

@@ -0,0 +1,145 @@
+"""Discovery and execution of orchestrator hooks for finished tasks."""
+
+import pathlib
+import traceback
+
+import click
+
+from . import prompts, runner, tasks
+
+
+def run_hooks(
+    tasks_file: pathlib.Path,
+    task_id: str,
+    runner_name: str,
+    yolo: bool,
+    runner_args: tuple,
+    no_defaults: bool,
+    verbose: bool,
+    hooks: list[str] | None = None,
+    working_dir: pathlib.Path | None = None,
+    final_status: tasks.TaskStatus | None = None,
+    time_limit: int = 0,
+) -> None:
+    """Discovers and executes orchestrator hooks for a finished task.
+
+    Args:
+        tasks_file: Path to the tasks YAML file.
+        task_id: ID of the task the hooks run against.
+        runner_name: Name of the runner CLI used to execute the hooks.
+        yolo: Whether to run the runner in unattended (yolo) mode.
+        runner_args: Extra arguments forwarded to the runner CLI.
+        no_defaults: Whether to skip the runner's default arguments.
+        verbose: Whether to echo prompts and hook diagnostics.
+        hooks: Explicit list of hooks to run. If None, uses config.hooks.
+        working_dir: Working directory for the hook runner processes.
+        final_status: If provided, mark the task with this status after hooks.
+        time_limit: Time limit in minutes for each hook run (0 disables it).
+    """
+    data = tasks.load_tasks(tasks_file)
+    task = next((t for t in data.tasks if t.id == task_id), None)
+    if not task:
+        return
+
+    # Use provided hooks or fall back to configuration
+    active_hooks = hooks if hooks is not None else data.config.hooks
+    if active_hooks is None:
+        active_hooks = prompts.list_hooks(tasks_file)
+
+    if final_status == tasks.TaskStatus.FAILED:
+        if "roadmap" in active_hooks:
+            active_hooks = ["roadmap"]
+        else:
+            active_hooks = []
+
+    if not active_hooks:
+        if final_status:
+            tasks.update_task(
+                tasks_file, task_id, status=final_status, force=True
+            )
+        return
+
+    for hook_name in active_hooks:
+        # Reload tasks every time to ensure each hook sees progress from
+        # previous hooks
+        data = tasks.load_tasks(tasks_file)
+        task = next((t for t in data.tasks if t.id == task_id), None)
+        if not task:
+            if verbose:
+                click.echo(
+                    f"Task {task_id} not found during hook '{hook_name}' run."
+                )
+            continue
+
+        try:
+            prompt = prompts.prepare_hook_prompt(
+                hook_name, data, task, tasks_file
+            )
+        except FileNotFoundError:
+            if verbose:
+                click.echo(f"Hook '{hook_name}' prompt not found, skipping.")
+            continue
+
+        if verbose:
+            click.secho(
+                f"\n=== Hook: {hook_name} Prompt ===", fg="magenta", bold=True
+            )
+            click.echo(prompt)
+            click.secho("========================\n", fg="magenta", bold=True)
+
+        cmd = runner.build_runner_command(
+            runner_name,
+            prompt,
+            yolo,
+            runner_args,
+            no_defaults,
+            verbose=verbose,
+            time_limit=time_limit,
+        )
+
+        try:
+            # Append hooks to the main runner log for a unified execution trace.
+            returncode, stdout, stderr = runner.run_with_heartbeat(
+                cmd,
+                tasks_file,
+                task_id,
+                verbose,
+                echo_fn=lambda line: click.echo(line, nl=False),
+                header=f"Hook: {hook_name}",
+                cwd=working_dir,
+                time_limit=time_limit,
+            )
+            if verbose:
+                if returncode != 0:
+                    click.echo(
+                        f"Hook '{hook_name}' exited with code {returncode}."
+                    )
+        except Exception as e:
+            click.echo(f"Hook '{hook_name}' error: {e}")
+
+    # Finally mark the task as completed or failed if requested.
+    # But first check whether a hook already changed the task (e.g. the
+    # roadmap hook reset a failed task for a new approach).  If the task
+    # is no longer IN_PROGRESS, a hook already intervened — skip
+    # finalization so we don't overwrite the recovery.
+    if final_status:
+        try:
+            data = tasks.load_tasks(tasks_file)
+            current = next((t for t in data.tasks if t.id == task_id), None)
+            if current and current.status != tasks.TaskStatus.IN_PROGRESS:
+                return
+            tasks.update_task(
+                tasks_file, task_id, status=final_status, force=True
+            )
+        except tasks.TaskNotFoundError:
+            # Task may have been deleted by the orchestrator (e.g. after
+            # a failure it decided to take a different approach).  This
+            # is expected and not an error worth a traceback.
+            click.echo(
+                f"Task {task_id} was removed before it could be "
+                "finalized — the orchestrator likely restructured "
+                "the plan."
+            )
+        except Exception as e:
+            click.echo(f"Error finalizing task {task_id}: {e}")
+            traceback.print_exc()

lemming/hooks_test.py

@@ -0,0 +1,180 @@
+import pathlib
+import shutil
+import tempfile
+import unittest
+import unittest.mock
+
+from lemming import tasks
+from lemming.hooks import run_hooks
+
+
+class TestHooks(unittest.TestCase):
+    def setUp(self):
+        self.test_dir = tempfile.mkdtemp()
+        self.test_tasks_file = pathlib.Path(self.test_dir) / "tasks_test.yml"
+
+        # Scaffold a valid file
+        data = tasks.Roadmap(
+            context="Initial context",
+            tasks=[
+                tasks.Task(
+                    id="12345678",
+                    description="Initial Task",
+                    status=tasks.TaskStatus.PENDING,
+                    attempts=0,
+                    progress=[],
+                )
+            ],
+            config=tasks.RoadmapConfig(
+                retries=3, runner="agy", hooks=["roadmap"]
+            ),
+        )
+        tasks.save_tasks(self.test_tasks_file, data)
+
+    def tearDown(self):
+        shutil.rmtree(self.test_dir)
+
+    @unittest.mock.patch("lemming.runner.run_with_heartbeat")
+    @unittest.mock.patch("lemming.prompts.list_hooks")
+    @unittest.mock.patch("lemming.prompts.prepare_hook_prompt")
+    def test_run_hooks_success(self, mock_prepare, mock_list, mock_run):
+        mock_list.return_value = ["roadmap"]
+        mock_prepare.return_value = "Hook Prompt"
+        mock_run.return_value = (0, "stdout", "")
+
+        # Task must be IN_PROGRESS for finalization to apply
+        tasks.update_task(
+            self.test_tasks_file,
+            "12345678",
+            status=tasks.TaskStatus.IN_PROGRESS,
+        )
+
+        run_hooks(
+            self.test_tasks_file,
+            "12345678",
+            "agy",
+            yolo=True,
+            runner_args=(),
+            no_defaults=False,
+            verbose=True,
+            final_status=tasks.TaskStatus.COMPLETED,
+        )
+
+        self.assertTrue(mock_run.called)
+        data = tasks.load_tasks(self.test_tasks_file)
+        self.assertEqual(data.tasks[0].status, tasks.TaskStatus.COMPLETED)
+
+    @unittest.mock.patch("lemming.runner.run_with_heartbeat")
+    def test_run_hooks_no_hooks(self, mock_run):
+        run_hooks(
+            self.test_tasks_file,
+            "12345678",
+            "agy",
+            yolo=True,
+            runner_args=(),
+            no_defaults=False,
+            verbose=True,
+            hooks=[],
+            final_status=tasks.TaskStatus.COMPLETED,
+        )
+
+        self.assertFalse(mock_run.called)
+        data = tasks.load_tasks(self.test_tasks_file)
+        self.assertEqual(data.tasks[0].status, tasks.TaskStatus.COMPLETED)
+
+    @unittest.mock.patch("lemming.hooks.prompts.prepare_hook_prompt")
+    @unittest.mock.patch("lemming.runner.run_with_heartbeat")
+    def test_run_hooks_failure_filters_hooks(self, mock_run, mock_prepare):
+        mock_run.return_value = (0, "stdout", "")
+        mock_prepare.return_value = "Mock Prompt"
+
+        # Task must be IN_PROGRESS for finalization to apply
+        tasks.update_task(
+            self.test_tasks_file,
+            "12345678",
+            status=tasks.TaskStatus.IN_PROGRESS,
+        )
+
+        run_hooks(
+            self.test_tasks_file,
+            "12345678",
+            "agy",
+            yolo=True,
+            runner_args=(),
+            no_defaults=False,
+            verbose=True,
+            hooks=["readability", "roadmap", "testing"],
+            final_status=tasks.TaskStatus.FAILED,
+        )
+
+        # It should only run 'roadmap' hook
+        self.assertEqual(mock_prepare.call_count, 1)
+        self.assertEqual(mock_prepare.call_args[0][0], "roadmap")
+
+    @unittest.mock.patch("lemming.hooks.prompts.prepare_hook_prompt")
+    @unittest.mock.patch("lemming.runner.run_with_heartbeat")
+    def test_run_hooks_skips_finalization_when_healed(
+        self, mock_run, mock_prepare
+    ):
+        """If a hook resets the task (heals it), skip finalization."""
+        mock_prepare.return_value = "Mock Prompt"
+
+        # Task must be IN_PROGRESS for hooks to run
+        tasks.update_task(
+            self.test_tasks_file,
+            "12345678",
+            status=tasks.TaskStatus.IN_PROGRESS,
+        )
+
+        # Simulate the hook resetting the task during execution
+        def hook_resets_task(*args, **kwargs):
+            tasks.reset_task(self.test_tasks_file, "12345678")
+            return (0, "stdout", "")
+
+        mock_run.side_effect = hook_resets_task
+
+        run_hooks(
+            self.test_tasks_file,
+            "12345678",
+            "agy",
+            yolo=True,
+            runner_args=(),
+            no_defaults=False,
+            verbose=True,
+            hooks=["roadmap"],
+            final_status=tasks.TaskStatus.FAILED,
+        )
+
+        # Task should remain PENDING (healed), not FAILED
+        data = tasks.load_tasks(self.test_tasks_file)
+        self.assertEqual(data.tasks[0].status, tasks.TaskStatus.PENDING)
+        self.assertEqual(data.tasks[0].attempts, 0)
+
+    @unittest.mock.patch("lemming.hooks.prompts.prepare_hook_prompt")
+    @unittest.mock.patch("lemming.runner.run_with_heartbeat")
+    def test_run_hooks_reloads_tasks(self, mock_run, mock_prepare):
+        # Create a real Roadmap object to return
+        real_data = tasks.load_tasks(self.test_tasks_file)
+        mock_run.return_value = (0, "stdout", "")
+        mock_prepare.return_value = "Mock Prompt"
+
+        with unittest.mock.patch(
+            "lemming.hooks.tasks.load_tasks", return_value=real_data
+        ) as mock_load:
+            run_hooks(
+                self.test_tasks_file,
+                "12345678",
+                "agy",
+                yolo=True,
+                runner_args=(),
+                no_defaults=False,
+                verbose=True,
+                hooks=["h1", "h2"],
+            )
+
+            # 1 initial load + 2 hook loads = 3
+            self.assertEqual(mock_load.call_count, 3)
+
+
+if __name__ == "__main__":
+    unittest.main()

lemming/integration_test.py

@@ -0,0 +1,299 @@
+import pathlib
+import shutil
+import tempfile
+import threading
+import time
+import unittest
+import unittest.mock
+
+from lemming import models, orchestrator, runner, tasks
+
+
+class TestIntegration(unittest.TestCase):
+    def setUp(self):
+        self.test_dir = tempfile.mkdtemp()
+        self.test_tasks_file = pathlib.Path(self.test_dir) / "tasks_test.yml"
+
+        # Scaffold a valid file with one task
+        self.initial_data = models.Roadmap(
+            context="Initial context",
+            tasks=[
+                models.Task(
+                    id="task1",
+                    description="Task 1",
+                    status=models.TaskStatus.PENDING,
+                    attempts=0,
+                    progress=[],
+                )
+            ],
+            config=models.RoadmapConfig(retries=3, runner="true"),
+        )
+        tasks.save_tasks(self.test_tasks_file, self.initial_data)
+
+    def tearDown(self):
+        shutil.rmtree(self.test_dir)
+
+    def test_heartbeat_updates_during_execution(self):
+        """Verify a long-running task updates its heartbeat periodically."""
+        # We need to speed up the threshold for testing
+        with (
+            unittest.mock.patch("lemming.tasks.STALE_THRESHOLD", 2),
+            unittest.mock.patch("lemming.persistence.STALE_THRESHOLD", 2),
+            unittest.mock.patch("lemming.tasks.lifecycle.STALE_THRESHOLD", 2),
+        ):
+            # 0. Mark task as in progress first, as
+            # runner.run_with_heartbeat expects it
+            tasks.mark_task_in_progress(self.test_tasks_file, "task1")
+
+            # Start a task that runs for 3 seconds (longer than threshold)
+            cmd = ["sleep", "3"]
+
+            # Run in a separate thread so we can check the file
+            def run_task():
+                runner.run_with_heartbeat(
+                    cmd, self.test_tasks_file, "task1", verbose=False
+                )
+
+            t = threading.Thread(target=run_task)
+            t.start()
+
+            # Wait a bit for it to start and set initial heartbeat
+            time.sleep(0.5)
+            data = tasks.load_tasks(self.test_tasks_file)
+            h1 = data.tasks[0].last_heartbeat
+            self.assertIsNotNone(h1, "Initial heartbeat should be set")
+
+            # Wait for next heartbeat (interval is threshold // 2 = 1s)
+            time.sleep(1.5)
+            data = tasks.load_tasks(self.test_tasks_file)
+            h2 = data.tasks[0].last_heartbeat
+            self.assertIsNotNone(h2, "Second heartbeat should be set")
+            self.assertGreater(h2, h1, "Heartbeat should increase over time")
+
+            t.join()
+
+    def test_task_reclaimed_if_heartbeat_stops(self):
+        """Verify a task can be reclaimed if its runner stops heartbeats."""
+        with (
+            unittest.mock.patch("lemming.tasks.STALE_THRESHOLD", 1),
+            unittest.mock.patch("lemming.persistence.STALE_THRESHOLD", 1),
+            unittest.mock.patch("lemming.tasks.lifecycle.STALE_THRESHOLD", 1),
+        ):
+            # 1. Claim the task manually with a PID that doesn't exist
+            # Note: is_pid_alive will return False for 999999 (likely)
+            tasks.claim_task(self.test_tasks_file, "task1", pid=999999)
+
+            data = tasks.load_tasks(self.test_tasks_file)
+            self.assertEqual(
+                data.tasks[0].status, models.TaskStatus.IN_PROGRESS
+            )
+
+            # 2. Wait for it to become stale
+            time.sleep(1.1)
+
+            # 3. Try to claim it again (as if another orchestrator is running)
+            # claim_task should allow reclaiming if stale
+            task = tasks.claim_task(self.test_tasks_file, "task1", pid=88888)
+            self.assertIsNotNone(
+                task, "Task should be reclaimable after heartbeat timeout"
+            )
+            self.assertEqual(task.pid, 88888)
+
+    def test_orchestrator_retries_on_runner_failure(self):
+        """Verify the orchestrator retries if the runner does not succeed."""
+        # Mocking time.sleep to speed up tests
+        with unittest.mock.patch("time.sleep", return_value=None):
+            # Configure roadmap with 2 retries
+            data = tasks.load_tasks(self.test_tasks_file)
+            data.config.retries = 2
+            data.config.runner = "true"  # 'true' command just exits 0
+            tasks.save_tasks(self.test_tasks_file, data)
+
+            orchestrator.run_loop(
+                self.test_tasks_file,
+                verbose=False,
+                retry_delay=0,
+                yolo=True,
+                no_defaults=False,
+                runner_args=(),
+            )
+
+            # After 2 attempts it should be FAILED
+            data = tasks.load_tasks(self.test_tasks_file)
+            self.assertEqual(data.tasks[0].attempts, 2)
+            self.assertEqual(data.tasks[0].status, models.TaskStatus.FAILED)
+
+    def test_runner_terminates_if_reclaimed(self):
+        """Verify the original runner terminates if its task is reclaimed."""
+        with (
+            unittest.mock.patch("lemming.tasks.STALE_THRESHOLD", 2),
+            unittest.mock.patch("lemming.persistence.STALE_THRESHOLD", 2),
+            unittest.mock.patch("lemming.tasks.lifecycle.STALE_THRESHOLD", 2),
+        ):
+            # 0. Mark task as in progress
+            tasks.mark_task_in_progress(self.test_tasks_file, "task1")
+
+            # Start a long-running task
+            cmd = ["sleep", "10"]
+
+            def run_task():
+                runner.run_with_heartbeat(
+                    cmd, self.test_tasks_file, "task1", verbose=True
+                )
+
+            t = threading.Thread(target=run_task)
+            t.start()
+
+            # Wait for it to start
+            time.sleep(0.5)
+
+            # 2. Reclaim the task manually by changing the PID in the file
+            # runner.py check:
+            #   if not tasks.update_heartbeat(tasks_file, task_id):
+            # update_heartbeat returns True ONLY if task is IN_PROGRESS.
+            # Wait, update_heartbeat in lifecycle.py DOES NOT check PID if
+            # it matches! It just updates it.
+
+            # Let's re-read update_heartbeat:
+            # def update_heartbeat(tasks_file, task_id, pid=None):
+            #     for task in data.tasks:
+            #         if task.id == task_id:
+            #             if task.status != models.TaskStatus.IN_PROGRESS:
+            #                 return False
+            #             task.last_heartbeat = time.time()
+            #             if pid is not None:
+            #                 task.pid = pid
+            #             break
+
+            # Ah! It doesn't check if the PID matches. So if another process
+            # just calls update_heartbeat, it will succeed.
+
+            # BUT, mark_task_in_progress/claim_task will succeed if it's stale.
+            # And once it's COMPLETED or FAILED, update_heartbeat returns False.
+
+            # Let's simulate cancellation/completion by another process
+            with tasks.lock_tasks(self.test_tasks_file):
+                data = tasks.load_tasks(self.test_tasks_file)
+                data.tasks[0].status = models.TaskStatus.COMPLETED
+                tasks.save_tasks(self.test_tasks_file, data)
+
+            # Wait for the next heartbeat (threshold // 2 = 1s)
+            t.join(timeout=5)
+            self.assertFalse(
+                t.is_alive(),
+                "Runner thread should have terminated after task was "
+                "marked COMPLETED",
+            )
+
+    def test_orchestrator_retries_on_runner_crash(self):
+        """Verify the orchestrator retries if the runner exits non-zero."""
+        with unittest.mock.patch("time.sleep", return_value=None):
+            data = tasks.load_tasks(self.test_tasks_file)
+            data.config.retries = 3
+            data.config.runner = "false"  # 'false' command exits with 1
+            tasks.save_tasks(self.test_tasks_file, data)
+
+            orchestrator.run_loop(
+                self.test_tasks_file,
+                verbose=False,
+                retry_delay=0,
+                yolo=True,
+                no_defaults=False,
+                runner_args=(),
+            )
+
+            data = tasks.load_tasks(self.test_tasks_file)
+            # It should have attempted 3 times and then failed
+            self.assertEqual(data.tasks[0].attempts, 3)
+            self.assertEqual(data.tasks[0].status, models.TaskStatus.FAILED)
+
+    def test_orchestrator_stops_on_cancellation(self):
+        """Verify the orchestrator stops if the runner exits with SIGTERM."""
+        with (
+            unittest.mock.patch("time.sleep", return_value=None),
+            unittest.mock.patch(
+                "lemming.runner.run_with_heartbeat", return_value=(-15, "", "")
+            ),
+        ):
+            data = tasks.load_tasks(self.test_tasks_file)
+            data.config.retries = 3
+            tasks.save_tasks(self.test_tasks_file, data)
+
+            # This should NOT loop 3 times. It should break immediately.
+            orchestrator.run_loop(
+                self.test_tasks_file,
+                verbose=False,
+                retry_delay=0,
+                yolo=True,
+                no_defaults=False,
+                runner_args=(),
+            )
+
+            data = tasks.load_tasks(self.test_tasks_file)
+            # It should have only 1 attempt because it broke
+            self.assertEqual(data.tasks[0].attempts, 1)
+            # Status should be PENDING (as finish_task_attempt sets it if
+            # no requested_status)
+            self.assertEqual(data.tasks[0].status, models.TaskStatus.PENDING)
+
+    def test_claim_task_race_condition(self):
+        """Verify that claim_task is atomic and prevents double claiming."""
+        # Mock is_pid_alive to return True so that even if we use fake PIDs,
+        # they are not considered dead and thus not reclaimable immediately.
+        with unittest.mock.patch(
+            "lemming.tasks.lifecycle.is_pid_alive", return_value=True
+        ):
+            # We'll use multiple threads to try to claim the same task
+            results = []
+
+            def try_claim(pid):
+                res = tasks.claim_task(self.test_tasks_file, "task1", pid=pid)
+                if res:
+                    results.append(pid)
+
+            threads = []
+            for i in range(1, 11):  # Use 1 to 10
+                t = threading.Thread(target=try_claim, args=(i,))
+                threads.append(t)
+
+            for t in threads:
+                t.start()
+            for t in threads:
+                t.join()
+
+            # Only one should have succeeded
+            self.assertEqual(len(results), 1)
+            data = tasks.load_tasks(self.test_tasks_file)
+            self.assertEqual(
+                data.tasks[0].status, models.TaskStatus.IN_PROGRESS
+            )
+            self.assertIn(data.tasks[0].pid, results)
+
+    def test_reclaim_if_pid_dead(self):
+        """Verify a task with a dead PID is immediately reclaimable.
+
+        This holds even if the heartbeat is fresh.
+        """
+        with unittest.mock.patch(
+            "lemming.tasks.lifecycle.is_pid_alive", return_value=False
+        ):
+            # 1. Claim it with some PID
+            tasks.claim_task(self.test_tasks_file, "task1", pid=12345)
+
+            data = tasks.load_tasks(self.test_tasks_file)
+            self.assertEqual(
+                data.tasks[0].status, models.TaskStatus.IN_PROGRESS
+            )
+            self.assertEqual(data.tasks[0].pid, 12345)
+
+            # 2. Heartbeat is fresh, but is_pid_alive is mocked to False
+            # Try to claim it again with another PID
+            task = tasks.claim_task(self.test_tasks_file, "task1", pid=67890)
+            self.assertIsNotNone(
+                task, "Task should be reclaimable if PID is dead"
+            )
+            self.assertEqual(task.pid, 67890)
+
+
+if __name__ == "__main__":
+    unittest.main()

lemming/main.py

@@ -0,0 +1,6 @@
+"""Entry point for running the Lemming CLI as a module."""
+
+from .cli import cli
+
+if __name__ == "__main__":
+    cli()