Agentic-OS/backend/app/mcp_manager/commands.py

"""Redis queue for MCP maintenance commands (worker executes reload/upgrade/develop)."""
from __future__ import annotations

import asyncio
import json
import uuid
from collections.abc import Awaitable, Callable
from typing import Any

from app.services.events import get_redis

MCP_COMMAND_QUEUE = "agentic:mcp:commands"


def _result_key(request_id: str) -> str:
    return f"agentic:mcp:cmd_result:{request_id}"


async def _wait_result(request_id: str, timeout_seconds: float) -> dict:
    redis = get_redis()
    deadline = asyncio.get_event_loop().time() + timeout_seconds
    while asyncio.get_event_loop().time() < deadline:
        raw = await redis.get(_result_key(request_id))
        if raw:
            await redis.delete(_result_key(request_id))
            return json.loads(raw)
        await asyncio.sleep(0.5)
    return {"ok": False, "error": "Worker did not respond in time — is the worker running?"}


async def enqueue_mcp_command(
    action: str,
    *,
    server: str | None = None,
    timeout_seconds: float = 90,
    **extra: Any,
) -> dict:
    request_id = str(uuid.uuid4())
    cmd = {"request_id": request_id, "action": action, "server": server, **extra}
    await get_redis().rpush(MCP_COMMAND_QUEUE, json.dumps(cmd))
    return await _wait_result(request_id, timeout_seconds)


async def enqueue_mcp_command_async(
    action: str,
    *,
    server: str | None = None,
    **extra: Any,
) -> str:
    """Queue a command for the worker without waiting for the result."""
    request_id = str(uuid.uuid4())
    cmd = {"request_id": request_id, "action": action, "server": server, **extra}
    await get_redis().rpush(MCP_COMMAND_QUEUE, json.dumps(cmd))
    return request_id


async def store_mcp_command_result(request_id: str, result: dict) -> None:
    redis = get_redis()
    await redis.set(_result_key(request_id), json.dumps(result), ex=300)


async def process_mcp_command(cmd: dict) -> dict:
    from app.mcp_manager.manager import get_manager

    action = cmd.get("action")
    server = cmd.get("server")
    manager = get_manager()
    try:
        if action == "reload":
            if not server:
                raise ValueError("server name required")
            await manager.reload_server(server)
        elif action == "upgrade":
            if not server:
                raise ValueError("server name required")
            await manager.upgrade_server(server)
        elif action == "reload_all":
            await manager.reload_all()
        elif action == "develop":
            if not server:
                raise ValueError("server name required")
            return await _execute_develop_command(cmd, manager)
        else:
            raise ValueError(f"Unknown action: {action}")
        return {"ok": True, "action": action, "server": server}
    except Exception as exc:  # noqa: BLE001
        return {"ok": False, "action": action, "server": server, "error": str(exc)}


async def _execute_develop_command(cmd: dict, manager) -> dict:
    server = cmd["server"]
    result = await _develop_on_worker(cmd, manager)
    task_id = int(cmd.get("task_id") or 0)
    if task_id:
        await _finalize_develop_task(task_id, result, server)
    return result


async def run_develop_command_background(cmd: dict, manager) -> None:
    """Run MCP develop without blocking the worker task queue loop."""
    request_id = cmd.get("request_id", "")
    try:
        result = await _execute_develop_command(cmd, manager)
        await store_mcp_command_result(request_id, result)
    except Exception as exc:  # noqa: BLE001
        await store_mcp_command_result(
            request_id,
            {"ok": False, "action": "develop", "server": cmd.get("server"), "error": str(exc)},
        )


async def _finalize_develop_task(task_id: int, result: dict, server: str) -> None:
    from sqlalchemy import select

    from app.database import SessionLocal
    from app.models.enums import ApprovalStatus, TaskStatus
    from app.models.task import ApprovalRequest, Task
    from app.services.events import publish_event

    async with SessionLocal() as db:
        res = await db.execute(select(Task).where(Task.id == task_id))
        task = res.scalar_one_or_none()
        if not task:
            return

        pend = await db.execute(
            select(ApprovalRequest).where(
                ApprovalRequest.task_id == task_id,
                ApprovalRequest.status == ApprovalStatus.pending,
            )
        )
        has_pending = pend.scalars().first() is not None

        if result.get("ok"):
            task.report = {
                "kind": "mcp_develop",
                "mcp_server": server,
                "summary": result.get("summary"),
                "tools": result.get("tools"),
                "push": result.get("push"),
                "resolved": not has_pending,
            }
            task.error = None
            task.status = TaskStatus.waiting_approval if has_pending else TaskStatus.succeeded
        else:
            task.status = TaskStatus.failed
            task.error = result.get("error") or result.get("reason") or "MCP development failed"
        await db.commit()

    if result.get("ok"):
        msg = result.get("summary") or f"MCP develop finished for {server}"
        if has_pending:
            msg += " — approve git push on this task."
        await publish_event(task_id, "report", msg, {"kind": "mcp_develop", **result}, persist=True)
    else:
        err = result.get("error") or result.get("reason") or "MCP development failed"
        await publish_event(task_id, "error", err, result, persist=True)


async def _develop_on_worker(cmd: dict, manager) -> dict:
    from app.agent.mcp_development import attempt_mcp_repair
    from app.services.events import publish_event

    server = cmd["server"]
    task_id = int(cmd.get("task_id") or 0)

    async def emit(tid: int, kind: str, message: str, payload: dict | None = None):
        if tid:
            await publish_event(tid, kind, message, payload)

    if task_id:
        await emit(
            task_id,
            "log",
            f"MCP develop started for {server}",
            {"server": server, "kind": "mcp_develop"},
        )

    return await attempt_mcp_repair(
        task_id=task_id,
        server=server,
        tool=cmd.get("tool") or "unknown",
        args={},
        issue=cmd.get("issue") or "",
        error=cmd.get("error"),
        output=cmd.get("output"),
        manager=manager,
        emit=emit,
        task_title=cmd.get("title") or f"MCP develop: {server}",
        push=cmd.get("push"),
    )