creator/backend/agents.py

"""Provider layer: runs agent calls via the Claude CLI or OpenCode (MiniMax).

Both runners are independent. If a binary/key is missing, only the
respective provider fails — the other keeps running unchanged.
"""

import asyncio
import heapq
import logging
import os
import re
import shutil
import signal
import tempfile
import time
import urllib.request
from contextlib import asynccontextmanager
from pathlib import Path

from config import (PROVIDERS, DEFAULT_PROVIDER, MAX_CONCURRENT_AGENTS,
                    MAX_CONCURRENT_AGENTS_PER_TOPIC, MAX_CONCURRENT_INTERACTIVE)

log = logging.getLogger("creator.agents")

_active_processes: dict[str, asyncio.subprocess.Process] = {}
_active_started: dict[str, float] = {}   # agent_key → wall-clock start (for the live runtime display)


def active_agents(scope_prefix: str | None = None) -> list[dict]:
    """Currently running agents and how long they've been running. Filter by key prefix
    (e.g. f"blocks-{topic}-") for one topic. → [{key, runtime}] sorted longest-first."""
    now = time.time()
    out = [{"key": k, "runtime": round(now - t, 1)}
           for k, t in list(_active_started.items())
           if k in _active_processes and (not scope_prefix or k.startswith(scope_prefix))]
    return sorted(out, key=lambda a: -a["runtime"])

# Cancelled scopes (key prefixes, symmetric to kill_process). An agent whose
# key starts with one of these prefixes aborts BEFORE the spawn — so agents WAITING
# in the semaphore queue are also stopped immediately on abort instead of still starting.
_cancelled_prefixes: set[str] = set()


def cancel_scope(prefix: str) -> None:
    _cancelled_prefixes.add(prefix)


def clear_scope(prefix: str) -> None:
    _cancelled_prefixes.discard(prefix)


def _scope_cancelled(agent_key: str) -> bool:
    return any(agent_key.startswith(p) for p in _cancelled_prefixes)

# Caps the real CLI processes — independent of the pipeline semaphore in
# generator.py. The acquire happens BEFORE the spawn so that queue wait time
# does not count against the agent timeout.
class _PrioritySemaphore:
    """asyncio.Semaphore variant: when slots are scarce, the LOWEST priority number is served first
    (FIFO within the same priority). Lets earlier pipeline columns grab agents before later ones."""
    def __init__(self, value: int):
        self._value = value
        self._waiters: list = []  # heap of [priority, seq, future]
        self._seq = 0

    async def acquire(self, priority: int = 100):
        if self._value > 0:
            self._value -= 1
            return
        fut = asyncio.get_event_loop().create_future()
        entry = [priority, self._seq, fut]
        self._seq += 1
        heapq.heappush(self._waiters, entry)
        try:
            await fut  # release() hands us the slot directly (no value change)
        except BaseException:
            entry[2] = None  # tombstone so release() skips this dead waiter
            if fut.done() and not fut.cancelled():
                self.release()  # granted just before we were cancelled → pass it on
            raise

    def release(self):
        while self._waiters:
            entry = heapq.heappop(self._waiters)
            if entry[2] is not None and not entry[2].done():
                entry[2].set_result(None)  # hand the slot straight to the highest-priority waiter
                return
        self._value += 1


_batch_sem = _PrioritySemaphore(MAX_CONCURRENT_AGENTS)
_interactive_sem = asyncio.Semaphore(MAX_CONCURRENT_INTERACTIVE)
# Per-topic caps (lazily created): each topic gets its own priority semaphore of size
# MAX_CONCURRENT_AGENTS_PER_TOPIC, nested INSIDE the global _batch_sem. Priority-based too, so the
# per-topic queue can't undo the global priority when one topic is the only load.
_topic_sems: dict[str, _PrioritySemaphore] = {}

# Earlier kanban columns get the scarce global slot first (smaller = higher priority).
_STAGE_PRIORITY = ("research", "verify", "naming", "small", "dep")


def _agent_priority(key: str) -> int:
    for i, tag in enumerate(_STAGE_PRIORITY):
        if f"-{tag}-" in key or key.endswith(f"-{tag}"):
            return i
    return len(_STAGE_PRIORITY)  # downstream agents (subblocks/facts/…) after the inventory columns


@asynccontextmanager
async def _batch_gate(scope: str | None, priority: int):
    """Per-topic slot FIRST (fair), then the GLOBAL slot by priority (earlier columns win when
    agents are scarce). Order matters — a waiter holds only its per-topic slot while queueing globally."""
    topic_sem = _topic_sems.setdefault(scope, _PrioritySemaphore(MAX_CONCURRENT_AGENTS_PER_TOPIC)) if scope else None
    if topic_sem is not None:
        await topic_sem.acquire(priority)
    await _batch_sem.acquire(priority)
    try:
        yield
    finally:
        _batch_sem.release()
        if topic_sem is not None:
            topic_sem.release()

# Serialize OpenCode starts: processes starting simultaneously collide on the
# internal session DB ("database is locked", exit after <1s). The short
# stagger spreads out the starts; afterwards the processes run in parallel normally.
_opencode_start_lock = asyncio.Lock()
_OPENCODE_START_DELAY = 1.0

# Capability → Claude --allowedTools
_CLAUDE_TOOLS = {
    "full": "Write,Bash,Read,WebSearch,WebFetch",
    "files": "Read,Bash,Write",
    "read": "Read",
    "none": None,
}

# Capability → OpenCode agent (tool permissions defined in dev-ops/opencode.json)
_OPENCODE_AGENTS = {
    "full": "full",
    "files": "files",
    "read": "readonly",
    "none": "text",
}


def provider_available(provider: str) -> bool:
    cfg = PROVIDERS.get(provider)
    if not cfg:
        return False
    if shutil.which(cfg["cli"]) is None:
        return False
    env_key = cfg.get("env_key")
    if env_key and not os.environ.get(env_key):
        return False
    check_url = cfg.get("check_url")
    if check_url:
        try:
            urllib.request.urlopen(check_url, timeout=1)
        except Exception:
            return False
    return True


def _kill(process) -> None:
    """Kill the agent and its child processes via the process group (otherwise the
    children spawned by the CLI survive, keep the pipes open and block communicate())."""
    try:
        os.killpg(os.getpgid(process.pid), signal.SIGKILL)
    except (ProcessLookupError, PermissionError):
        try:
            process.kill()
        except ProcessLookupError:
            pass


def kill_process(agent_key_prefix: str) -> None:
    """Kill all active processes whose key starts with the prefix (covers -plan/-w1…)."""
    for key, process in list(_active_processes.items()):
        if process.returncode is not None:  # clean up dead entries while iterating
            _active_processes.pop(key, None)
            _active_started.pop(key, None)
            continue
        if key.startswith(agent_key_prefix):
            log.debug("kill agent %s", key)
            _kill(process)


async def run_agent(
    agent_key: str,
    prompt: str,
    timeout: int,
    provider: str = DEFAULT_PROVIDER,
    role: str = "fast",
    capabilities: str = "none",
    lane: str = "batch",
    scope: str | None = None,
    on_line=None,
) -> tuple[int, str, str]:
    if _scope_cancelled(agent_key):  # before queueing: don't even enter the queue
        return 1, "", "cancelled"
    if provider not in PROVIDERS:
        return 1, "", f"Unknown provider: {provider}"
    if shutil.which(PROVIDERS[provider]["cli"]) is None:
        return 1, "", f"CLI '{PROVIDERS[provider]['cli']}' not installed (provider: {provider})"
    gate = _interactive_sem if lane == "interactive" else _batch_gate(scope, _agent_priority(agent_key))
    async with gate:
        if _scope_cancelled(agent_key):  # after the acquire: cancelled in the queue → no spawn
            return 1, "", "cancelled"
        if PROVIDERS[provider]["cli"] == "opencode":
            return await _run_opencode(agent_key, prompt, timeout, provider, role, capabilities, on_line=on_line)
        return await _run_claude_cli(agent_key, prompt, timeout, role, capabilities)


async def _communicate(agent_key: str, cmd: list[str], stdin_data: bytes | None, timeout: int, stagger: bool = False, on_line=None) -> tuple[int, str, str]:
    start = time.monotonic()

    async def spawn():
        return await asyncio.create_subprocess_exec(
            *cmd,
            stdin=asyncio.subprocess.PIPE if stdin_data is not None else asyncio.subprocess.DEVNULL,
            stdout=asyncio.subprocess.PIPE,
            stderr=asyncio.subprocess.PIPE,
            start_new_session=True,  # own process group → killpg also kills child processes
        )

    if stagger:
        async with _opencode_start_lock:
            process = await spawn()
            await asyncio.sleep(_OPENCODE_START_DELAY)
    else:
        process = await spawn()
    _active_processes[agent_key] = process
    _active_started[agent_key] = time.time()
    try:
        try:
            if on_line is not None:
                # Streaming path: read stdout line by line, hand each raw line to on_line LIVE.
                out_chunks: list[str] = []
                async def _pump():
                    async for raw in process.stdout:
                        s = raw.decode("utf-8", errors="replace")
                        out_chunks.append(s)
                        try:
                            on_line(s)
                        except Exception:
                            log.debug("on_line callback failed", exc_info=True)
                await asyncio.wait_for(_pump(), timeout=timeout)
                await process.wait()
                stderr_b = await process.stderr.read()
                stdout, stderr = "".join(out_chunks).encode("utf-8"), stderr_b
            else:
                stdout, stderr = await asyncio.wait_for(
                    process.communicate(input=stdin_data),
                    timeout=timeout,
                )
        except asyncio.TimeoutError:
            _kill(process)
            try:
                await asyncio.wait_for(process.wait(), timeout=5)
            except asyncio.TimeoutError:
                pass
            log.info("agent %s: timeout after %ds", agent_key, timeout)
            raise
        log.info(
            "agent %s: exit %s after %.1fs (%d bytes stdout)",
            agent_key, process.returncode, time.monotonic() - start, len(stdout),
        )
        return process.returncode, stdout.decode("utf-8", errors="replace"), stderr.decode("utf-8", errors="replace")
    finally:
        # Pop only on identity: a slot restart under the same key must not evict
        # the NEW process from tracking.
        if _active_processes.get(agent_key) is process:
            del _active_processes[agent_key]
            _active_started.pop(agent_key, None)


async def _run_claude_cli(agent_key: str, prompt: str, timeout: int, role: str, capabilities: str) -> tuple[int, str, str]:
    cfg = PROVIDERS["claude"]
    cmd = [cfg["cli"], "-p", "--model", cfg[role]]
    tools = _CLAUDE_TOOLS.get(capabilities)
    if tools:
        cmd += ["--allowedTools", tools]
    cmd += ["--dangerously-skip-permissions"]
    return await _communicate(agent_key, cmd, prompt.encode("utf-8"), timeout)


async def _run_opencode(agent_key: str, prompt: str, timeout: int, provider: str, role: str, capabilities: str, on_line=None) -> tuple[int, str, str]:
    cfg = PROVIDERS[provider]
    # Prompt via temp file instead of argv (ARG_MAX protection for large project prompts)
    with tempfile.NamedTemporaryFile("w", suffix=".md", delete=False, encoding="utf-8", dir=tempfile.gettempdir()) as f:
        f.write(prompt)
        prompt_path = Path(f.name)
    # The positional message MUST come before -f: -f is an array flag and
    # would otherwise eat the text as a second file name ("File not found").
    cmd = [
        cfg["cli"], "run",
        "Folge exakt den Anweisungen in der angehängten Datei. Sie sind der vollständige Auftrag.",
        "-m", cfg[role],
        "--agent", _OPENCODE_AGENTS.get(capabilities, "text"),
        "--dangerously-skip-permissions",
        "-f", str(prompt_path),
    ]
    if on_line is not None:
        cmd += ["--format", "json"]  # raw JSON events → parsed live by on_line
    try:
        rc, stdout, stderr = await _communicate(agent_key, cmd, None, timeout, stagger=True, on_line=on_line)
        return rc, (stdout if on_line is not None else _clean_opencode_output(stdout)), stderr
    finally:
        prompt_path.unlink(missing_ok=True)


_ANSI_RE = re.compile(r"\x1b\[[0-9;]*m")


def _clean_opencode_output(text: str) -> str:
    """Strip ANSI codes and the leading banner ("> agent · model")."""
    text = _ANSI_RE.sub("", text)
    lines = text.splitlines()
    while lines and (not lines[0].strip() or lines[0].lstrip().startswith(">")):
        lines.pop(0)
    return "\n".join(lines).strip()