lemming-cli 0.1.0__py3-none-any.whl

lemming/runner.py

@@ -0,0 +1,327 @@
+"""Runner command construction and subprocess execution with heartbeats."""
+
+import logging
+import os
+import pathlib
+import shlex
+import signal
+import subprocess
+import threading
+import time
+from typing import Callable
+
+from . import paths, tasks
+
+logger = logging.getLogger(__name__)
+
+# Sentinel return codes for non-standard process termination.
+RETURNCODE_TIMEOUT = -14
+
+
+def _pretty_quote(s: str) -> str:
+    """Quotes a string for shell execution, preferring readable quoting.
+
+    Strings containing single quotes are wrapped in double quotes for
+    readability. This function is idempotent for single shell words: if s
+    is already a valid shell-quoted single word, it is unquoted before
+    being re-quoted prettily.
+    """
+    if not s:
+        return "''"
+
+    # If it's already a single shell word, try to unquote it first to avoid
+    # compounding. We only unquote if shlex.split succeeds and returns
+    # exactly one token.
+    try:
+        parts = shlex.split(s)
+        if len(parts) == 1:
+            # Check if it was actually quoted or escaped.
+            # shlex.split("foo") is "foo", but shlex.split("'foo'") is also
+            # "foo". We want to canonicalize any quoted string back to its
+            # literal form.
+            if parts[0] != s or (s.startswith("'") or s.startswith('"')):
+                s = parts[0]
+    except Exception:
+        pass
+
+    # If shlex.quote says it doesn't need quotes, return as-is
+    if shlex.quote(s) == s:
+        return s
+
+    # If it contains single quotes, try to use double quotes for better
+    # readability
+    if "'" in s:
+        # If it has !, double quotes might trigger history expansion in
+        # interactive bash
+        if "!" in s:
+            return shlex.quote(s)
+
+        # We need to escape \, ", $, ` inside double quotes
+        escaped = (
+            s.replace("\\", "\\\\")
+            .replace('"', '\\"')
+            .replace("$", "\\$")
+            .replace("`", "\\`")
+        )
+        return f'"{escaped}"'
+
+    # Default to standard shlex.quote
+    return shlex.quote(s)
+
+
+def _shlex_join_pretty(cmd: list[str], max_len: int = -1) -> str:
+    """Joins command arguments into a single string with pretty quoting.
+
+    Args:
+        cmd: List of command arguments.
+        max_len: If > 0, truncate each quoted argument to this length.
+    """
+    parts = []
+    for arg in cmd:
+        quoted = _pretty_quote(arg)
+        if max_len > 0 and len(quoted) > max_len:
+            parts.append(quoted[:max_len] + "... [truncated]")
+        else:
+            parts.append(quoted)
+    return " ".join(parts)
+
+
+def build_runner_command(
+    runner_name: str,
+    prompt: str,
+    yolo: bool,
+    runner_args: tuple | None = None,
+    no_defaults: bool = False,
+    verbose: bool = False,
+    time_limit: int = 0,
+) -> list[str]:
+    """Constructs the CLI command for the specified runner.
+
+    Args:
+        runner_name: Name or path of the runner executable. May contain a
+            ``{{prompt}}`` placeholder; when present, the template is
+            shlex-split and the placeholder token is replaced with the
+            prompt text.  Default flag injection is skipped in template
+            mode.
+        prompt: The full prompt text to pass to the runner.
+        yolo: Whether to enable auto-approval/YOLO mode.
+        runner_args: Extra arguments to pass to the runner.
+        no_defaults: If True, do not inject default flags for known runners.
+        verbose: If True, enable verbose output for supported runners.
+        time_limit: Task time limit in minutes; used to size runner-side
+            timeouts for runners that enforce their own (0 means no limit).
+
+    Returns:
+        A list of command-line arguments.
+    """
+    # Template mode: {{prompt}} in runner_name means the user controls the
+    # full command layout.  Split, substitute, and return early.
+    if "{{prompt}}" in runner_name:
+        parts = shlex.split(runner_name)
+        cmd = [p.replace("{{prompt}}", prompt) for p in parts]
+        if runner_args:
+            cmd.extend(runner_args)
+        return cmd
+
+    parts = shlex.split(runner_name)
+    cmd = [parts[0]]
+    extra_parts = parts[1:]
+    prompt_arg = None
+
+    runner_base = os.path.basename(parts[0])
+
+    if not no_defaults:
+        if runner_base.startswith("agy"):
+            if yolo:
+                cmd.append("--dangerously-skip-permissions")
+
+            # Print mode buffers stdout until the run completes, so stream the
+            # internal log to stdout for live visibility in the task log.
+            cmd.extend(["--log-file", "/dev/stdout"])
+
+            # agy caps print mode at 5 minutes by default; extend it to the
+            # task time limit (or effectively unlimited when there is none).
+            print_timeout = f"{time_limit}m" if time_limit > 0 else "24h"
+            cmd.extend(["--print-timeout", print_timeout])
+
+            prompt_arg = "--prompt"
+        elif runner_base.startswith("aider"):
+            if yolo:
+                cmd.append("--yes")
+            if not verbose:
+                cmd.append("--quiet")
+            prompt_arg = "--message"
+        elif runner_base.startswith("claude"):
+            if yolo:
+                cmd.append("--dangerously-skip-permissions")
+            cmd.extend(["--output-format=stream-json", "--verbose"])
+            prompt_arg = "--print"
+        elif runner_base.startswith("codex"):
+            if yolo:
+                cmd.append("--yolo")
+            prompt_arg = "--instructions"
+
+    if extra_parts:
+        cmd.extend(extra_parts)
+    if runner_args:
+        cmd.extend(runner_args)
+
+    if prompt_arg:
+        cmd.extend([prompt_arg, prompt])
+    else:
+        cmd.append(prompt)
+
+    return cmd
+
+
+def _kill_process_tree(process: subprocess.Popen) -> None:
+    """Kills a process and its entire process group.
+
+    Tries SIGTERM on the process group first, falls back to killing the
+    process directly if the group signal fails.
+
+    Args:
+        process: The subprocess to kill.
+    """
+    try:
+        os.killpg(os.getpgid(process.pid), signal.SIGTERM)
+    except OSError:
+        try:
+            process.kill()
+        except OSError:
+            pass
+
+
+def run_with_heartbeat(
+    cmd: list[str],
+    tasks_file: pathlib.Path,
+    task_id: str,
+    verbose: bool,
+    echo_fn: Callable[[str], None] = print,
+    cwd: pathlib.Path | None = None,
+    header: str | None = None,
+    time_limit: int = 0,
+) -> tuple[int, str, str]:
+    """Runs the runner process and updates the task heartbeat periodically.
+
+    Args:
+        cmd: The command to execute as a list of strings.
+        tasks_file: Path to the tasks YAML file.
+        task_id: ID of the task being executed (used for heartbeat and
+            parent env var).
+        verbose: If True, echo runner output to the console.
+        echo_fn: Function to use for echoing output (defaults to print).
+        cwd: Optional working directory for the subprocess.
+        header: Optional header to write to the log (e.g. "Orchestrator
+            Hook: roadmap").
+        time_limit: Maximum execution time in minutes. 0 means no limit.
+
+    Returns:
+        A tuple of (returncode, stdout_log, stderr_log). Note: stderr is
+        currently merged into stdout_log.
+    """
+    log_file = paths.get_log_file(tasks_file, task_id)
+
+    # Use a separator for new attempts or hooks
+    # For the log file, we truncate long arguments (like the prompt) to
+    # avoid unreadable logs and exponential escaping growth when prompts
+    # are re-quoted.
+    log_command_str = _shlex_join_pretty(cmd, max_len=200)
+    command_str = _shlex_join_pretty(cmd)
+
+    with open(log_file, "a", encoding="utf-8") as f:
+        timestamp = time.strftime("%Y-%m-%d %H:%M:%S %Z")
+        f.write(f"\n--- Attempt started at {timestamp} ---\n")
+        if header:
+            f.write(f"{'=' * 80}\n")
+            f.write(f"{header.upper()} started at {timestamp}\n")
+            f.write(f"{'=' * 80}\n")
+        f.write(f"Command: {log_command_str}\n")
+        f.flush()
+
+    if verbose:
+        echo_fn(f"Executing: {command_str}\n\n")
+
+    # Start the process in a new session so we can kill its entire process
+    # tree if needed.
+    env = os.environ.copy()
+    env["LEMMING_PARENT_TASK_ID"] = task_id
+    env["LEMMING_PARENT_TASKS_FILE"] = str(tasks_file.resolve())
+
+    process = subprocess.Popen(
+        cmd,
+        env=env,
+        stdout=subprocess.PIPE,
+        stderr=subprocess.STDOUT,
+        text=True,
+        bufsize=1,
+        errors="replace",
+        start_new_session=True,
+        cwd=cwd,
+    )
+
+    full_log: list[str] = []
+
+    timed_out = False
+
+    try:
+        # Heartbeat and cancellation management
+        is_claimed = tasks.update_heartbeat(
+            tasks_file, task_id, pid=process.pid
+        )
+        start_time = time.monotonic()
+
+        def heartbeat_loop():
+            """Updates the task heartbeat while the process is running."""
+            nonlocal timed_out
+            while process.poll() is None:
+                if not tasks.update_heartbeat(tasks_file, task_id):
+                    # Task was cancelled or finished.
+                    _kill_process_tree(process)
+                    return
+
+                # Check time limit
+                if time_limit > 0:
+                    elapsed = time.monotonic() - start_time
+                    if elapsed >= time_limit * 60:
+                        timed_out = True
+                        tasks.add_progress(
+                            tasks_file,
+                            task_id,
+                            f"Task killed: time limit of {time_limit}"
+                            " minutes reached.",
+                        )
+                        _kill_process_tree(process)
+                        return
+
+                time.sleep(tasks.STALE_THRESHOLD // 2)
+
+        if is_claimed:
+            heartbeat_thread = threading.Thread(
+                target=heartbeat_loop, daemon=True
+            )
+            heartbeat_thread.start()
+
+        # Stream output to log file and optionally to console
+        try:
+            if process.stdout:
+                with open(log_file, "a", encoding="utf-8") as f:
+                    for line in process.stdout:
+                        full_log.append(line)
+                        f.write(line)
+                        f.flush()
+                        if verbose:
+                            echo_fn(line)
+
+            process.wait()
+        except BaseException:
+            _kill_process_tree(process)
+            raise
+
+        returncode = process.returncode
+        if timed_out:
+            returncode = RETURNCODE_TIMEOUT
+        return returncode, "".join(full_log), ""
+    finally:
+        if process.stdout and hasattr(process.stdout, "close"):
+            process.stdout.close()

lemming/runner_test.py

@@ -0,0 +1,378 @@
+import signal
+import subprocess
+import time
+import unittest.mock
+
+from lemming import paths, runner, tasks
+
+
+def test_build_runner_command_agy():
+    cmd = runner.build_runner_command("agy", "my prompt", yolo=True)
+    assert "--dangerously-skip-permissions" in cmd
+    assert "--prompt" in cmd
+    assert "my prompt" in cmd
+
+
+def test_build_runner_command_agy_streams_internal_log():
+    # agy print mode buffers stdout until the end, so its internal log is
+    # redirected to stdout to give live visibility in the task log.
+    cmd = runner.build_runner_command("agy", "my prompt", yolo=True)
+    assert cmd[cmd.index("--log-file") + 1] == "/dev/stdout"
+
+
+def test_build_runner_command_agy_print_timeout_matches_time_limit():
+    cmd = runner.build_runner_command(
+        "agy", "my prompt", yolo=True, time_limit=45
+    )
+    assert cmd[cmd.index("--print-timeout") + 1] == "45m"
+
+
+def test_build_runner_command_agy_print_timeout_without_time_limit():
+    # agy defaults --print-timeout to 5m, so an explicit large value is
+    # required even when the task has no time limit.
+    cmd = runner.build_runner_command(
+        "agy", "my prompt", yolo=True, time_limit=0
+    )
+    assert cmd[cmd.index("--print-timeout") + 1] == "24h"
+
+
+def test_build_runner_command_time_limit_ignored_by_other_runners():
+    cmd = runner.build_runner_command(
+        "claude", "my prompt", yolo=True, time_limit=45
+    )
+    assert "--print-timeout" not in cmd
+    assert "--log-file" not in cmd
+
+
+def test_build_runner_command_aider():
+    cmd = runner.build_runner_command("aider", "my prompt", yolo=True)
+    assert "--yes" in cmd
+    assert "--message" in cmd
+
+
+def test_build_runner_command_with_flags_in_name():
+    cmd = runner.build_runner_command(
+        "claude-corp -- --output-format=stream-json", "my prompt", yolo=True
+    )
+    assert cmd[0] == "claude-corp"
+    assert "--" in cmd
+    assert "--output-format=stream-json" in cmd
+    assert "--dangerously-skip-permissions" in cmd
+    assert "--print" in cmd
+    assert "my prompt" in cmd
+
+
+def test_build_runner_command_with_quoted_flags_in_name():
+    cmd = runner.build_runner_command(
+        'my-runner --model "gpt 4"', "my prompt", yolo=True, no_defaults=True
+    )
+    assert cmd[0] == "my-runner"
+    assert "--model" in cmd
+    assert "gpt 4" in cmd
+
+
+def test_build_runner_command_template_basic():
+    cmd = runner.build_runner_command(
+        "my-tool --input={{prompt}} --json", "hello world", yolo=True
+    )
+    assert cmd == ["my-tool", "--input=hello world", "--json"]
+
+
+def test_build_runner_command_template_standalone_placeholder():
+    cmd = runner.build_runner_command(
+        "my-tool --flag {{prompt}}", "hello world", yolo=True
+    )
+    assert cmd == ["my-tool", "--flag", "hello world"]
+
+
+def test_build_runner_command_template_with_runner_args():
+    cmd = runner.build_runner_command(
+        "my-tool {{prompt}}", "hello", yolo=True, runner_args=("--extra",)
+    )
+    assert cmd == ["my-tool", "hello", "--extra"]
+
+
+def test_build_runner_command_template_ignores_defaults():
+    # Even though runner starts with "agy", template mode should not
+    # inject --dangerously-skip-permissions etc.
+    cmd = runner.build_runner_command(
+        "agy --custom {{prompt}}", "do stuff", yolo=True
+    )
+    assert "--dangerously-skip-permissions" not in cmd
+    assert cmd == ["agy", "--custom", "do stuff"]
+
+
+def test_build_runner_command_template_prompt_in_flag_value():
+    cmd = runner.build_runner_command(
+        "my-tool --msg={{prompt}} --verbose", "hi there", yolo=True
+    )
+    assert cmd == ["my-tool", "--msg=hi there", "--verbose"]
+
+
+def test_pretty_quote():
+    # Test fallback to shlex
+    assert runner._pretty_quote("simple") == "simple"
+    assert runner._pretty_quote("has space") == "'has space'"
+
+    # Test readable double quotes for single quotes
+    assert (
+        runner._pretty_quote("has 'single' quotes") == "\"has 'single' quotes\""
+    )
+    assert runner._pretty_quote("You are 'Lemming'") == "\"You are 'Lemming'\""
+
+    # Test string with double quotes (should fall back to single quotes)
+    assert (
+        runner._pretty_quote('has "double" quotes') == "'has \"double\" quotes'"
+    )
+
+    # Test escaping specials inside double quotes
+    assert (
+        runner._pretty_quote("has 'single' and \"double\" quotes")
+        == '"has \'single\' and \\"double\\" quotes"'
+    )
+
+    # Test exclamation mark fallback
+    assert runner._pretty_quote("Hello!") == "'Hello!'"
+
+    assert runner._pretty_quote("has 'single' and !") == (
+        "'has '\"'\"'single'\"'\"' and !'"
+    )
+
+    # Test idempotency (should NOT compound quotes)
+    q = runner._pretty_quote("it's!")
+    assert q == "'it'\"'\"'s!'"
+    qq = runner._pretty_quote(q)
+    assert qq == q
+    qqq = runner._pretty_quote(qq)
+    assert qqq == q
+
+    # Test idempotent single quotes (should NOT compound to multiple escapes)
+    q_s = runner._pretty_quote("it's")
+    assert q_s == '"it\'s"'
+    qq_s = runner._pretty_quote(q_s)
+    assert qq_s == q_s
+
+    # Test complex shell-quoted strings
+    already_quoted = "'path with space' and \"double'quotes\""
+    # This is NOT a single shell word, so it won't be unquoted.
+    # But it will be double-quoted correctly.
+    q_complex = runner._pretty_quote(already_quoted)
+    assert q_complex.startswith('"')
+    assert q_complex.endswith('"')
+    assert '\\"double\'quotes\\"' in q_complex
+
+
+def test_shlex_join_pretty():
+    cmd = [
+        "example-cli",
+        "--dangerously-skip-permissions",
+        "--print",
+        "You are 'Lemming'",
+    ]
+    joined = runner._shlex_join_pretty(cmd)
+    assert (
+        joined == "example-cli --dangerously-skip-permissions --print"
+        " \"You are 'Lemming'\""
+    )
+
+    # Test truncation
+    long_arg = "a" * 300
+    joined_truncated = runner._shlex_join_pretty(["cli", long_arg], max_len=100)
+    assert "a" * 100 in joined_truncated
+    assert "... [truncated]" in joined_truncated
+    assert len(joined_truncated) < 150
+
+
+def test_run_with_heartbeat_truncation_only_affects_log(tmp_path):
+    tasks_file = tmp_path / "tasks.yml"
+    task_id = "test_task"
+    log_file = paths.get_log_file(tasks_file, task_id)
+    log_file.parent.mkdir(parents=True, exist_ok=True)
+
+    # Use a long command that would be truncated in logs
+    long_arg = "a" * 300
+    cmd = ["echo", long_arg]
+
+    mock_process = unittest.mock.MagicMock()
+    mock_process.returncode = 0
+    mock_process.stdout = None
+    mock_process.poll.return_value = 0
+
+    with unittest.mock.patch(
+        "subprocess.Popen", return_value=mock_process
+    ) as mock_popen:
+        runner.run_with_heartbeat(cmd, tasks_file, task_id, verbose=False)
+
+        # Verify Popen received the original untruncated cmd
+        mock_popen.assert_called_once()
+        called_cmd = mock_popen.call_args[0][0]
+        assert called_cmd == cmd
+        assert len(called_cmd[1]) == 300
+
+    # Verify log file contains the truncated command
+    content = log_file.read_text()
+    assert "Command: echo " in content
+    assert "a" * 200 in content
+    assert "... [truncated]" in content
+    assert "a" * 201 not in content
+
+
+def test_run_with_heartbeat_log_header(tmp_path):
+    tasks_file = tmp_path / "tasks.yml"
+    task_id = "test_task"
+    log_file = paths.get_log_file(tasks_file, task_id)
+    log_file.parent.mkdir(parents=True, exist_ok=True)
+
+    # Use a command that exits quickly
+    cmd = ["true"]
+
+    # 1. Run with a header
+    runner.run_with_heartbeat(
+        cmd, tasks_file, task_id, verbose=False, header="Hook: roadmap"
+    )
+
+    content = log_file.read_text()
+    assert "--- Attempt started at" in content
+    assert "HOOK: ROADMAP started at" in content
+    assert "=" * 80 in content
+
+    # 2. Run without a header (it should still have the attempt marker)
+    task_id_2 = "test_task_2"
+    log_file_2 = paths.get_log_file(tasks_file, task_id_2)
+    runner.run_with_heartbeat(
+        cmd, tasks_file, task_id_2, verbose=False, header=None
+    )
+
+    content_2 = log_file_2.read_text()
+    assert "--- Attempt started at" in content_2
+    assert "started at" not in content_2.replace("Attempt started at", "")
+    assert "=" * 80 not in content_2
+
+
+def test_run_with_heartbeat_interruption_cleanup(tmp_path):
+    tasks_file = tmp_path / "tasks.yml"
+    task_id = "test_task"
+
+    # 1. Setup a dummy Roadmap
+    roadmap = tasks.Roadmap(tasks=[tasks.Task(id=task_id, description="test")])
+    tasks.save_tasks(tasks_file, roadmap)
+
+    # 2. Mock subprocess.Popen and related functions
+    mock_process = unittest.mock.MagicMock()
+    mock_process.pid = 12345
+    mock_process.stdout = None
+    # We want process.wait() to raise a KeyboardInterrupt (BaseException)
+    mock_process.wait.side_effect = KeyboardInterrupt()
+
+    with (
+        unittest.mock.patch("subprocess.Popen", return_value=mock_process),
+        unittest.mock.patch("os.killpg") as mock_killpg,
+        unittest.mock.patch("os.getpgid", return_value=54321),
+    ):
+        # 3. Call run_with_heartbeat and expect it to re-raise KeyboardInterrupt
+        try:
+            runner.run_with_heartbeat(
+                ["long-running-cmd"], tasks_file, task_id, verbose=False
+            )
+        except KeyboardInterrupt:
+            pass
+        else:
+            assert False, "KeyboardInterrupt was not raised"
+
+        # 4. Verify cleanup was attempted
+        mock_killpg.assert_called_once_with(54321, signal.SIGTERM)
+
+
+def test_returncode_timeout_constant():
+    assert runner.RETURNCODE_TIMEOUT == -14
+
+
+def test_kill_process_tree_killpg():
+    """Verifies _kill_process_tree uses killpg first."""
+    process = unittest.mock.MagicMock(spec=subprocess.Popen)
+    process.pid = 12345
+
+    with (
+        unittest.mock.patch("os.killpg") as mock_killpg,
+        unittest.mock.patch("os.getpgid", return_value=54321),
+    ):
+        runner._kill_process_tree(process)
+        mock_killpg.assert_called_once_with(54321, signal.SIGTERM)
+        process.kill.assert_not_called()
+
+
+def test_kill_process_tree_fallback():
+    """Verifies _kill_process_tree falls back to process.kill on OSError."""
+    process = unittest.mock.MagicMock(spec=subprocess.Popen)
+    process.pid = 12345
+
+    with (
+        unittest.mock.patch("os.killpg", side_effect=OSError),
+        unittest.mock.patch("os.getpgid", return_value=54321),
+    ):
+        runner._kill_process_tree(process)
+        process.kill.assert_called_once()
+
+
+def test_run_with_heartbeat_timeout(tmp_path):
+    """Verifies run_with_heartbeat kills and records progress on timeout."""
+    tasks_file = tmp_path / "tasks.yml"
+    task_id = "timeout_task"
+
+    # Setup a task so heartbeat updates work
+    roadmap = tasks.Roadmap(
+        tasks=[tasks.Task(id=task_id, description="test timeout")]
+    )
+    tasks.save_tasks(tasks_file, roadmap)
+    tasks.mark_task_in_progress(tasks_file, task_id)
+
+    # Use a 1-minute time limit. Mock time.monotonic to simulate elapsed
+    # time so the heartbeat loop detects the timeout immediately without
+    # waiting 60 real seconds.
+    real_monotonic = time.monotonic
+    call_count = 0
+
+    def fast_monotonic():
+        nonlocal call_count
+        call_count += 1
+        # After the first call (start_time), jump 2 minutes ahead
+        if call_count > 1:
+            return real_monotonic() + 120
+        return real_monotonic()
+
+    with unittest.mock.patch("time.monotonic", side_effect=fast_monotonic):
+        returncode, stdout, stderr = runner.run_with_heartbeat(
+            ["sleep", "60"],
+            tasks_file,
+            task_id,
+            verbose=False,
+            time_limit=1,
+        )
+
+    assert returncode == runner.RETURNCODE_TIMEOUT
+
+    # Verify the timeout progress was recorded
+    data = tasks.load_tasks(tasks_file)
+    task = next(t for t in data.tasks if t.id == task_id)
+    assert any("time limit" in o for o in task.progress)
+
+
+def test_run_with_heartbeat_no_timeout(tmp_path):
+    """Verifies that time_limit=0 does not enforce any timeout."""
+    tasks_file = tmp_path / "tasks.yml"
+    task_id = "no_timeout"
+
+    roadmap = tasks.Roadmap(
+        tasks=[tasks.Task(id=task_id, description="test no timeout")]
+    )
+    tasks.save_tasks(tasks_file, roadmap)
+
+    returncode, _, _ = runner.run_with_heartbeat(
+        ["true"],
+        tasks_file,
+        task_id,
+        verbose=False,
+        time_limit=0,
+    )
+
+    assert returncode == 0