background.py

"""Background task execution — dispatch slow shell commands to background
threads and inject completion notifications into the conversation."""

import json
import threading
from typing import Callable

from config import get_logger
from wcode_errors import WcodeError

logger = get_logger(__name__)


# ── Heuristics (module-level, stateless) ────────────────────────────────


def is_slow_operation(tool_name: str, tool_input: dict) -> bool:
    """Return True if the tool call is likely to be slow (install, build, test…)."""
    if tool_name != "run_bash":
        return False
    cmd = tool_input.get("command", "").lower()
    slow_keywords = [
        "install", "build", "test", "deploy", "compile",
        "docker build", "pip install", "npm install",
        "cargo build", "pytest", "make",
    ]
    return any(kw in cmd for kw in slow_keywords)


def should_run_background(tool_name: str, tool_input: dict) -> bool:
    """Decide whether a tool call should be dispatched to a background thread."""
    if tool_input.get("run_in_background"):
        return True
    return is_slow_operation(tool_name, tool_input)


# ── BackgroundExecutor ──────────────────────────────────────────────────


class BackgroundExecutor:
    """Dispatches slow tool calls to background threads and collects results.

    Constructor receives a ``tool_registry`` for tool execution (injected
    to avoid circular imports).  The agent loop polls ``collect_results()``
    each turn.

    Usage::

        bg = BackgroundExecutor(tool_registry)
        if should_run_background(name, args):
            bg_id = bg.start_task(tool_call)
        ...
        for notification in bg.collect_results():
            messages.append({"role": "user", "content": notification})
    """

    def __init__(self, execute_tool_fn: Callable) -> None:
        self._execute = execute_tool_fn
        self._tasks: dict[str, dict] = {}
        self._results: dict[str, str] = {}
        self._lock = threading.Lock()
        self._counter = 0

    def execute_tool(self, block) -> str:
        """Execute a single tool call synchronously via the injected handler."""
        name = block.function.name
        args = json.loads(block.function.arguments)
        return self._execute(name, args)

    def start_task(self, block) -> str:
        """Dispatch *block* to a background thread, returning a task ID."""
        self._counter += 1
        args = json.loads(block.function.arguments)
        bg_id = f"bg_{self._counter:04d}"
        cmd = args.get("command")

        def worker():
            result = self.execute_tool(block)
            with self._lock:
                self._tasks[bg_id]["status"] = "completed"
                self._results[bg_id] = result

        with self._lock:
            self._tasks[bg_id] = {
                "tool_call_id": block.id,
                "status": "running",
                "command": cmd,
            }
        thread = threading.Thread(target=worker, daemon=True)
        thread.start()
        logger.info(f"[background] dispatched {bg_id}: {cmd[:40] if cmd else '?'}")
        return bg_id

    def collect_results(self) -> list[str]:
        """Collect and return XML notifications for completed background tasks."""
        notifications = []
        with self._lock:
            ready_ids = [
                bg_id
                for bg_id, task in self._tasks.items()
                if task["status"] == "completed"
            ]
            for bg_id in ready_ids:
                task = self._tasks.pop(bg_id)
                output = self._results.pop(bg_id, "")
                notifications.append(
                    "<task_notification>\n"
                    f"  <task_id>{bg_id}</task_id>\n"
                    "  <status>completed</status>\n"
                    f"  <command>{task['command']}</command>\n"
                    f"  <summary>{output}</summary>\n"
                    "</task_notification>"
                )
                logger.info(
                    f"[background done] {bg_id}: "
                    f"{task['command'][:40]} ({len(output)} chars)"
                )
        return notifications


# ── Module-level shims ──────────────────────────────────────────────────

_default_executor: BackgroundExecutor | None = None


def _get_default() -> BackgroundExecutor:
    global _default_executor
    if _default_executor is None:
        from tools import TOOL_HANDLER

        def _execute(name: str, args: dict) -> str:
            handler = TOOL_HANDLER.get(name)
            if handler is None:
                return f"Error: Unknown tool '{name}'"
            try:
                return handler(**args)
            except WcodeError as e:
                logger.warning(f"[tool] {name} error: {e}")
                return f"Error: {e}"

        _default_executor = BackgroundExecutor(_execute)
    return _default_executor


def execute_tool(block) -> str:
    return _get_default().execute_tool(block)


def start_background_task(block) -> str:
    return _get_default().start_task(block)


def collect_background_results() -> list[str]:
    return _get_default().collect_results()