atlas-iac/services/hermes/scripts/cli_lane_provider.py

#!/usr/bin/env python3
"""Native Codex and Claude process lifecycle for Hermes CLI workers."""

from __future__ import annotations

import contextlib
import json
import os
import selectors
import signal
import subprocess
import time
import uuid
from dataclasses import asdict
from pathlib import Path
from typing import Any, Callable

import cli_lane_goal
from cli_lane_config import (
    CAPACITY_PATTERN,
    CLAUDE_BIN,
    CLAUDE_SESSION_COLLISION,
    CLAUDE_SETTINGS,
    CODEX_BIN,
    DATA_ROOT,
    HEARTBEAT_SECONDS,
    NO_CLAUDE_SESSION,
    NO_CODEX_THREAD,
    ProcessResult,
    RESULT_SCHEMA,
    RESULT_SCHEMA_PATH,
    Route,
    utc_now,
)
from cli_lane_files import _result_path, atomic_json, load_json
from cli_lane_prompt import _event_payload


def stream_process(
    command: list[str],
    *,
    provider: str,
    cwd: Path,
    env: dict[str, str],
    log_path: Path,
    state: dict[str, Any],
    state_file: Path,
    heartbeat: Callable[[str], bool],
    max_runtime: int,
) -> ProcessResult:
    """Stream JSONL to Kanban logs while maintaining the authoritative lease."""
    process = subprocess.Popen(
        command,
        cwd=cwd,
        env=env,
        stdout=subprocess.PIPE,
        stderr=subprocess.STDOUT,
        text=True,
        bufsize=1,
        start_new_session=True,
    )
    assert process.stdout is not None
    selector = selectors.DefaultSelector()
    selector.register(process.stdout, selectors.EVENT_READ)
    started = time.monotonic()
    last_heartbeat = 0.0
    known_descendants: dict[int, tuple[int, int]] = {}
    lines: list[str] = []
    structured: dict[str, Any] | None = None
    log_path.parent.mkdir(parents=True, exist_ok=True)
    forced_failure = ""
    with log_path.open("a", encoding="utf-8") as log:
        log.write(f"\n[{utc_now()}] starting {provider} worker\n")
        log.flush()
        while process.poll() is None:
            known_descendants.update(_descendant_processes(process.pid))
            now = time.monotonic()
            if now - started > max_runtime:
                _terminate_worker_process(process, known_descendants)
                forced_failure = "worker exceeded its maximum runtime"
                lines.append(forced_failure + "\n")
                break
            if now - last_heartbeat >= HEARTBEAT_SECONDS:
                if not heartbeat(f"{provider} worker active for {round(now - started)}s"):
                    _terminate_worker_process(process, known_descendants)
                    forced_failure = "Kanban lease was lost; provider process terminated"
                    lines.append(forced_failure + "\n")
                    break
                last_heartbeat = now
            for key, _ in selector.select(timeout=1.0):
                line = key.fileobj.readline()
                if not line:
                    continue
                lines.append(line)
                if len(lines) > 4000:
                    lines = lines[-4000:]
                log.write(line)
                log.flush()
                parsed = _event_payload(provider, line, state, state_file)
                structured = parsed or structured
        _terminate_worker_process(process, known_descendants)
        remainder = process.stdout.read()
        if remainder:
            lines.append(remainder)
            log.write(remainder)
            for line in remainder.splitlines():
                parsed = _event_payload(provider, line, state, state_file)
                structured = parsed or structured
        if forced_failure:
            log.write(forced_failure + "\n")
        log.write(f"\n[{utc_now()}] {provider} exit={process.returncode}\n")
    selector.close()
    output = "".join(lines)[-100000:]
    returncode = int(process.returncode if process.returncode is not None else 1)
    return ProcessResult(
        returncode=returncode,
        output=output,
        structured=structured,
        capacity_failure=returncode != 0 and bool(CAPACITY_PATTERN.search(output)),
    )

def _process_record(pid: int) -> tuple[int, int, int] | None:
    """Return one Linux process's parent, group, and start-time identity."""
    try:
        raw = Path(f"/proc/{pid}/stat").read_text(encoding="utf-8")
        fields = raw[raw.rfind(")") + 2 :].split()
        return int(fields[1]), int(fields[2]), int(fields[19])
    except (IndexError, OSError, ValueError):
        return None

def _descendant_processes(root_pid: int) -> dict[int, tuple[int, int]]:
    """Snapshot descendants as pid -> (process group, start-time identity)."""
    records: dict[int, tuple[int, int, int]] = {}
    for entry in Path("/proc").iterdir():
        if not entry.name.isdigit():
            continue
        pid = int(entry.name)
        record = _process_record(pid)
        if record is not None:
            records[pid] = record
    family = {root_pid}
    changed = True
    while changed:
        changed = False
        for pid, (parent, _, _) in records.items():
            if pid not in family and parent in family:
                family.add(pid)
                changed = True
    return {
        pid: (records[pid][1], records[pid][2])
        for pid in family
        if pid != root_pid and pid in records
    }

def _process_identity_matches(pid: int, start_time: int) -> bool:
    record = _process_record(pid)
    return record is not None and record[2] == start_time

def _signal_worker_tree(
    root_pid: int,
    descendants: dict[int, tuple[int, int]],
    sig: signal.Signals,
) -> None:
    """Signal captured descendants across terminal-created process groups."""
    groups = {root_pid}
    for pid, (process_group, start_time) in descendants.items():
        if _process_identity_matches(pid, start_time):
            groups.add(process_group)
    for process_group in groups:
        with contextlib.suppress(ProcessLookupError):
            os.killpg(process_group, sig)
    for pid, (_, start_time) in descendants.items():
        if not _process_identity_matches(pid, start_time):
            continue
        with contextlib.suppress(ProcessLookupError):
            os.kill(pid, sig)

def _terminate_worker_process(
    process: subprocess.Popen[str],
    known_descendants: dict[int, tuple[int, int]] | None = None,
) -> None:
    """Reap a worker and terminal descendants that created their own groups."""
    descendants = dict(known_descendants or {})
    descendants.update(_descendant_processes(process.pid))
    if process.poll() is None:
        _signal_worker_tree(process.pid, descendants, signal.SIGTERM)
        with contextlib.suppress(subprocess.TimeoutExpired):
            process.wait(timeout=10)

    # Provider terminal tools create their own sessions, so killing only the
    # native CLI's process group leaves test/build subprocesses orphaned. The
    # identity check prevents a reused PID from targeting an unrelated worker.
    _signal_worker_tree(process.pid, descendants, signal.SIGKILL)
    if process.poll() is None:
        process.wait(timeout=10)

def _base_env() -> dict[str, str]:
    """Build a worker environment while preserving Vault-backed CLI homes."""
    env = os.environ.copy()
    # Subscription CLIs only: a vendor API key in the environment could
    # silently flip a worker onto a metered path, so drop every key here.
    for secret in ("ANTHROPIC_API_KEY", "CLAUDE_API_KEY", "OPENAI_API_KEY", "API_SERVER_KEY"):
        env.pop(secret, None)
    env.update(
        {
            "HOME": str(DATA_ROOT / "home"),
            "CODEX_HOME": os.environ.get("CODEX_HOME", "/runtime-access/codex"),
            "CLAUDE_CONFIG_DIR": os.environ.get(
                "CLAUDE_CONFIG_DIR", "/runtime-access/claude"
            ),
            "GIT_TERMINAL_PROMPT": "0",
        }
    )
    env["PATH"] = f"{DATA_ROOT / 'tools/bin'}:/opt/coordinator:" + env.get("PATH", "")
    return env

def _codex_command(route: Route, prompt: str, workspace: Path, state: dict[str, Any], result_file: Path) -> list[str]:
    base = [
        str(CODEX_BIN),
        "--dangerously-bypass-approvals-and-sandbox",
        "--dangerously-bypass-hook-trust",
        "exec",
    ]
    thread_id = str(state.get("codex_thread_id") or "")
    common = [
        "--json",
        "-m",
        route.model,
        "-c",
        f'model_reasoning_effort="{route.effort}"',
        "--output-schema",
        str(RESULT_SCHEMA_PATH),
        "-o",
        str(result_file),
    ]
    if thread_id:
        return [*base, "resume", *common, thread_id, prompt]
    return [
        *base,
        *common,
        "-C",
        str(workspace),
        prompt,
    ]

def _claude_command(route: Route, prompt: str, state: dict[str, Any], resume: bool) -> list[str]:
    session_id = str(state["claude_session_id"])
    session = ["--resume", session_id] if resume else ["--session-id", session_id]
    denied = [
        "Bash(git push --force *)",
        "Bash(git reset --hard *)",
        "Bash(git clean -f *)",
    ]
    return [
        str(CLAUDE_BIN),
        "--dangerously-skip-permissions",
        "--autocompact",
        "auto",
        "--settings",
        str(CLAUDE_SETTINGS),
        "--disallowedTools",
        *denied,
        "--model",
        route.model,
        "--effort",
        route.effort,
        "--output-format",
        "stream-json",
        "--verbose",
        "--json-schema",
        json.dumps(RESULT_SCHEMA, separators=(",", ":")),
        *session,
        "-p",
        prompt,
    ]

def run_provider(
    route: Route,
    prompt: str,
    workspace: Path,
    state: dict[str, Any],
    state_file: Path,
    log_path: Path,
    heartbeat: Callable[[str], bool],
    max_runtime: int,
) -> ProcessResult:
    """Run or resume one provider session using its pinned structured CLI."""
    state.setdefault("attempts", []).append({"route": asdict(route), "started_at": utc_now()})
    if route.provider == "claude" and not state.get("claude_session_id"):
        # The reservation is durable before Claude starts, closing the crash gap.
        state["claude_session_id"] = str(uuid.uuid4())
    state["current_route"] = asdict(route)
    state["updated_at"] = utc_now()
    env = _base_env()
    if route.provider == "codex":
        result_sequence = max(0, int(state.get("result_sequence", 0) or 0)) + 1
        result_file = _result_path(state_file, state.get("run_id"), result_sequence)
        while result_file.exists():
            result_sequence += 1
            result_file = _result_path(
                state_file,
                state.get("run_id"),
                result_sequence,
            )
        result_file.parent.mkdir(parents=True, exist_ok=True)
        result_file.touch(mode=0o600, exist_ok=False)
        state["result_sequence"] = result_sequence
        state["current_result_file"] = str(result_file)
        atomic_json(state_file, state)
        command = _codex_command(route, prompt, workspace, state, result_file)
        result = stream_process(
            command,
            provider="codex",
            cwd=workspace,
            env=env,
            log_path=log_path,
            state=state,
            state_file=state_file,
            heartbeat=heartbeat,
            max_runtime=max_runtime,
        )
        if (
            state.get("codex_thread_id")
            and result.returncode != 0
            and NO_CODEX_THREAD in result.output.lower()
        ):
            state.pop("codex_thread_id", None)
            result_sequence += 1
            result_file = _result_path(
                state_file,
                state.get("run_id"),
                result_sequence,
            )
            while result_file.exists():
                result_sequence += 1
                result_file = _result_path(
                    state_file,
                    state.get("run_id"),
                    result_sequence,
                )
            result_file.touch(mode=0o600, exist_ok=False)
            state["result_sequence"] = result_sequence
            state["current_result_file"] = str(result_file)
            atomic_json(state_file, state)
            result = stream_process(
                _codex_command(route, prompt, workspace, state, result_file),
                provider="codex",
                cwd=workspace,
                env=env,
                log_path=log_path,
                state=state,
                state_file=state_file,
                heartbeat=heartbeat,
                max_runtime=max_runtime,
            )
        file_result = load_json(result_file)
        if file_result.get("status") in cli_lane_goal.RESULT_STATUSES:
            result.structured = file_result
            with contextlib.suppress(OSError):
                result_file.chmod(0o600)
        return result

    atomic_json(state_file, state)
    resume = bool(state.get("claude_started"))
    result = stream_process(
        _claude_command(route, prompt, state, resume),
        provider="claude",
        cwd=workspace,
        env=env,
        log_path=log_path,
        state=state,
        state_file=state_file,
        heartbeat=heartbeat,
        max_runtime=max_runtime,
    )
    if resume and result.returncode == 1 and NO_CLAUDE_SESSION in result.output:
        result = stream_process(
            _claude_command(route, prompt, state, False),
            provider="claude",
            cwd=workspace,
            env=env,
            log_path=log_path,
            state=state,
            state_file=state_file,
            heartbeat=heartbeat,
            max_runtime=max_runtime,
        )
    elif not resume and result.returncode == 1 and CLAUDE_SESSION_COLLISION in result.output:
        result = stream_process(
            _claude_command(route, prompt, state, True),
            provider="claude",
            cwd=workspace,
            env=env,
            log_path=log_path,
            state=state,
            state_file=state_file,
            heartbeat=heartbeat,
            max_runtime=max_runtime,
        )
    if result.returncode == 0 or NO_CLAUDE_SESSION not in result.output:
        state["claude_started"] = True
        atomic_json(state_file, state)
    return result