Agentic-OS/backend/app/agent/tool_runner.py

"""Run MCP tools with live events (start → done) and auto-reload on failure."""
from __future__ import annotations

from collections.abc import Awaitable, Callable

from app.agent.mcp_helpers import tool_result_ok
from app.agent.tool_events import tool_event_payload
from app.mcp_manager.recovery import call_tool_resilient
from app.services.events import get_redis
from app.services.task_runtime import ensure_not_cancelled

EmitFn = Callable[[int, str, str, dict | None], Awaitable[None]]


def _repair_key(task_id: int, server: str, tool: str) -> str:
    return f"agentic:task:{task_id}:mcp_repair:{server}:{tool}"


async def _repair_attempted(task_id: int, server: str, tool: str) -> bool:
    return bool(await get_redis().get(_repair_key(task_id, server, tool)))


async def _mark_repair_attempted(task_id: int, server: str, tool: str) -> None:
    await get_redis().set(_repair_key(task_id, server, tool), "1", ex=86400)


async def invoke_tool(
    task_id: int,
    device: str,
    server: str,
    tool: str,
    args: dict,
    manager,
    emit: EmitFn,
    *,
    task_title: str = "",
    task_issue: str = "",
) -> tuple[dict, bool]:
    await ensure_not_cancelled(task_id)
    await emit(
        task_id,
        "tool_start",
        f"{device} › {tool}",
        tool_event_payload(device=device, tool=tool, args=args, status="running"),
    )
    try:
        res = await call_tool_resilient(
            manager, server, tool, args, task_id=task_id, emit=emit
        )
        ok = tool_result_ok(res)
        if not ok:
            retry = await _maybe_mcp_repair(
                task_id,
                device,
                server,
                tool,
                args,
                res.get("text"),
                manager,
                emit,
                task_title=task_title,
                task_issue=task_issue,
            )
            if retry is not None:
                return retry
        await emit(
            task_id,
            "tool_call",
            f"{device} › {tool}",
            tool_event_payload(
                device=device,
                tool=tool,
                args=args,
                status="done",
                ok=ok,
                output=res.get("text", ""),
            ),
        )
        return res, ok
    except Exception as exc:
        retry = await _maybe_mcp_repair(
            task_id,
            device,
            server,
            tool,
            args,
            str(exc),
            manager,
            emit,
            task_title=task_title,
            task_issue=task_issue,
            error=str(exc),
        )
        if retry is not None:
            return retry
        await emit(
            task_id,
            "tool_call",
            f"{device} › {tool} failed",
            tool_event_payload(
                device=device,
                tool=tool,
                args=args,
                status="error",
                ok=False,
                error=str(exc),
            ),
        )
        return {"text": str(exc), "isError": True}, False


async def _maybe_mcp_repair(
    task_id: int,
    device: str,
    server: str,
    tool: str,
    args: dict,
    output: str | None,
    manager,
    emit: EmitFn,
    *,
    task_title: str,
    task_issue: str,
    error: str | None = None,
) -> tuple[dict, bool] | None:
    if await _repair_attempted(task_id, server, tool):
        return None
    from app.agent.mcp_development import maybe_repair_after_tool_failure

    await _mark_repair_attempted(task_id, server, tool)
    return await maybe_repair_after_tool_failure(
        task_id=task_id,
        device=device,
        server=server,
        tool=tool,
        args=args,
        issue=task_issue or task_title,
        error=error,
        output=output,
        manager=manager,
        emit=emit,
        task_title=task_title,
    )


async def invoke_connect(
    task_id: int,
    device: str,
    server: str,
    tool: str,
    args: dict,
    manager,
    emit: EmitFn,
) -> bool:
    try:
        res = await call_tool_resilient(
            manager, server, tool, args, task_id=task_id, emit=emit
        )
        ok = tool_result_ok(res)
        msg = f"{device}: {tool} ok" if ok else f"{device}: {tool} failed — {res.get('text', '')[:200]}"
        await emit(task_id, "connect", msg, {"output": res.get("text", "")[:400], "ok": ok, "tool": tool})
        return ok
    except Exception as exc:
        await emit(task_id, "connect", f"{device}: {tool} failed — {exc}", {"ok": False, "tool": tool})
        return False