app/backend/fluksio/mcp/server.py

"""The tools an agent can call, each one a request to the flow API.

Tools do not reach into the engine: they call the same REST endpoints the
dashboard calls, over an in-process ASGI transport. That keeps one description
of what a flow is and how it may be changed — validation, the draft/publish
split, the version check that stops two clients overwriting each other — and it
means an agent cannot do anything a person could not do in the browser.

The caller's token rides along on every hop, so the API sees the agent's own
identity rather than some service account.

A tool takes a flow by its id or its key, resolved by `fluksio.sdk.resolve`
from the flows the caller can list, and calls the API with the id.
"""

from __future__ import annotations

import logging
from contextvars import ContextVar
from typing import Annotated, Any

import httpx
from mcp.server.context import CallNext, HandlerResult, ServerRequestContext
from mcp.server.mcpserver import MCPServer
from pydantic import Field

from fluksio.sdk.client import seg as _seg
from fluksio.sdk.resolve import Resolved, Unresolved, resolve

logger = logging.getLogger(__name__)

#: The ``Authorization`` value tool calls carry onto the REST API. The endpoint's
#: middleware sets it from each MCP request; the assistant sets it to the token of
#: the person chatting, for the length of their turn, since no MCP request is
#: involved there.
caller_authorization: ContextVar[str | None] = ContextVar(
    "caller_authorization", default=None
)


async def _carry_authorization(
    ctx: ServerRequestContext[Any, Any], call_next: CallNext
) -> HandlerResult:
    """Make the bearer token of an MCP request the caller's for its tool calls.

    The transport attaches the originating HTTP request to every JSON-RPC
    message, so the token is read per message rather than per connection.
    """
    headers = getattr(ctx.request, "headers", None)
    authorization = headers.get("authorization") if headers else None
    if not authorization:
        return await call_next(ctx)
    token = caller_authorization.set(authorization)
    try:
        return await call_next(ctx)
    finally:
        caller_authorization.reset(token)


mcp = MCPServer("fluksio", middleware=[_carry_authorization])

#: The tool calls never leave the process, so the host is a label, not a route.
_INTERNAL_BASE = "http://fluksio-mcp.internal"

#: The ASGI client that carries tool calls into the REST API without a network
#: hop. Built once by :func:`connect`, for the endpoint and the assistant alike.
_client: httpx.AsyncClient | None = None


def set_client(client: httpx.AsyncClient | None) -> None:
    global _client
    _client = client


def connect(app: object) -> None:
    """Wire the tools to *app*'s API, once."""
    if _client is not None:
        return
    set_client(
        httpx.AsyncClient(
            transport=httpx.ASGITransport(app=app),  # type: ignore[arg-type]
            base_url=_INTERNAL_BASE,
            # Running a flow is synchronous work on a threadpool, so a slow
            # flow must not look like a dead endpoint.
            timeout=120.0,
            event_hooks={"request": [forward_caller_auth]},
        )
    )


async def aclose() -> None:
    if _client is not None:
        await _client.aclose()
        set_client(None)


async def forward_caller_auth(request: httpx.Request) -> None:
    """Carry the caller's bearer token onto the REST call.

    A tool handler runs inside the context :data:`caller_authorization` was set
    in, by the endpoint's middleware or by the assistant.
    """
    authorization = caller_authorization.get()
    if authorization:
        request.headers["authorization"] = authorization


async def _call(method: str, path: str, **kwargs: Any) -> Any:
    """One REST call, with failures handed back as data rather than raised.

    An agent can read an error and try something else; a transport fault just
    ends the conversation.
    """
    if _client is None:
        return {"error": "The MCP endpoint is not running."}
    try:
        response = await _client.request(method, f"/api/v1{path}", **kwargs)
    except httpx.HTTPError as exc:
        logger.exception("MCP call to %s failed", path)
        return {"error": f"Could not reach the flow API: {exc}"}

    if response.status_code >= 400:
        detail: Any = response.text
        try:
            detail = response.json().get("detail", detail)
        except ValueError:
            pass
        return {"error": detail, "status": response.status_code}
    if not response.content:
        return {"ok": True}
    return response.json()


async def _resolve(
    flow: str, *, mutation: bool = False, history: bool = False
) -> Resolved | dict[str, Any]:
    """The flow *flow* names, or the error to hand back.

    Read from what the caller may list, with the caller's token, so nothing
    is resolved that the API would not show them. ``history`` counts deleted
    flows too, under the key their runs last had.
    """
    listing = await _call("GET", "/flows/")
    if "error" in listing:
        return dict(listing)
    flows = [(str(row["id"]), str(row["name"])) for row in listing.get("data") or []]
    if history:
        overview = await _call("GET", "/runs/overview")
        if isinstance(overview, dict):
            return overview
        flows += [
            (str(row["flow"]), str(row.get("flow_name") or row["flow"]))
            for row in overview
        ]
    try:
        return resolve(flow, flows, mutation=mutation)
    except Unresolved as exc:
        return {"error": str(exc)}


def _warned(answer: Any, found: Resolved) -> Any:
    """*answer*, saying which flow was taken when the name could mean two."""
    if not found.warning:
        return answer
    if isinstance(answer, dict):
        return {**answer, "warning": found.warning}
    return {"result": answer, "warning": found.warning}


async def _on_flow(
    flow: str, method: str, tail: str = "", *, mutation: bool, **kwargs: Any
) -> Any:
    """One call to ``/flows/<id><tail>`` for the flow *flow* names.

    ``mutation`` is whether the call changes anything, which a name two flows
    answer to is refused for rather than guessed.
    """
    found = await _resolve(flow, mutation=mutation)
    if not isinstance(found, Resolved):
        return found
    answer = await _call(method, f"/flows/{_seg(found.id)}{tail}", **kwargs)
    return _warned(answer, found)


async def _narrowed(flow: str | None, path: str, params: dict[str, Any]) -> Any:
    """A history read at *path*, narrowed to the flow *flow* names when it
    names one, a deleted one included."""
    if not flow:
        return await _call("GET", path, params=params)
    found = await _resolve(flow, history=True)
    if not isinstance(found, Resolved):
        return found
    params["flow"] = found.id
    return _warned(await _call("GET", path, params=params), found)


#: Every tool's *flow* argument, described where an agent reads the schema.
FlowArg = Annotated[
    str,
    Field(
        description=(
            "The flow's id or its key. Write id:<x> or key:<x> when one flow's "
            "id is another's key."
        )
    ),
]

# -----------------------------------------------------------------------------
# Reading
# -----------------------------------------------------------------------------


@mcp.tool()
async def list_flows() -> Any:
    """Every flow, with its node count, errors, and whether it is running."""
    return await _call("GET", "/flows/")


@mcp.tool()
async def get_flow(flow: FlowArg) -> Any:
    """One flow: its definition, the state of its nodes, and its problems.

    The definition is the working copy — unpublished edits included — which is
    what to base a change on. It carries the flow's ``id`` and its key,
    ``name``.
    """
    return await _on_flow(flow, "GET", mutation=False)


@mcp.tool()
async def list_node_types() -> Any:
    """The node types that can be placed, with their parameter schemas."""
    return await _call("GET", "/flows/node-types")


@mcp.tool()
async def get_node_source(flow: FlowArg, node_id: str) -> Any:
    """The Python source of one node."""
    return await _on_flow(flow, "GET", f"/nodes/{_seg(node_id)}/source", mutation=False)


@mcp.tool()
async def get_flow_state(flow: FlowArg) -> Any:
    """The last value seen on every message of a flow."""
    return await _on_flow(flow, "GET", "/state", mutation=False)


@mcp.tool()
async def get_message_history(flow: FlowArg, message: str) -> Any:
    """Recent numeric values of one message, oldest first."""
    return await _on_flow(flow, "GET", f"/history/{_seg(message)}", mutation=False)


@mcp.tool()
async def list_secrets() -> Any:
    """The names of stored secrets, for pointing a node parameter at one.

    Values are never returned. Reference one from a node parameter as
    ``{"$secret": "<name>"}``.
    """
    return await _call("GET", "/secrets/")


@mcp.tool()
async def list_shared_nodes() -> Any:
    """Node sources shared across flows, and which nodes use each."""
    return await _call("GET", "/flows/library")


@mcp.tool()
async def list_devices() -> Any:
    """Every paired screen or board, the dashboards it shows, whether it is
    connected and what it reports having."""
    return await _call("GET", "/devices/")


@mcp.tool()
async def get_device(device_id: str) -> Any:
    """One device: its dashboards in rail order, touch, and its own I/O ports."""
    return await _call("GET", f"/devices/{_seg(device_id)}")


@mcp.tool()
async def get_graph() -> Any:
    """Every flow as one graph, with nodes talking to the same thing merged.

    Nodes on the same broker topic, URL or bucket come back as a single entry,
    which is how wiring that runs between flows shows up.
    """
    return await _call("GET", "/flows/graph")


# -----------------------------------------------------------------------------
# Building
# -----------------------------------------------------------------------------


@mcp.tool()
async def create_flow(definition: dict[str, Any]) -> Any:
    """Create a new flow, as an unpublished draft.

    ``definition`` is a flow document like the one `get_flow` returns under
    ``definition``. Its ``name`` is the new flow's key, which no other flow
    may hold. The server gives the flow its id, so an ``id`` or ``version`` in
    the document is ignored; the answer carries both. Change an existing flow
    with `save_flow` instead.
    """
    document = {k: v for k, v in definition.items() if k not in ("id", "version")}
    return await _call("POST", "/flows/", json=document)


@mcp.tool()
async def save_flow(flow: FlowArg, definition: dict[str, Any]) -> Any:
    """Save changes to an existing flow as an unpublished draft.

    ``definition`` is a whole flow document — the shape `get_flow` returns
    under ``definition``. It must carry the ``version`` that was read, so an
    edit someone else made in between is refused rather than overwritten; on a
    conflict, read the flow again and reapply the change. Its ``name``, the
    key, cannot change here. A new flow is `create_flow`.

    Nothing here reaches the engine until `publish_flow`.
    """
    return await _on_flow(flow, "PUT", mutation=True, json=definition)


@mcp.tool()
async def save_node_source(flow: FlowArg, node_id: str, code: str) -> Any:
    """Save a node's Python source and report whether it compiles.

    A node defines ``process(...)``, taking one argument per input port and
    one per setting, and returns a dict keyed by output port. A parameter added
    or dropped there adds or drops the input port or setting in the draft,
    which moves its version on: read the flow again before saving it.
    """
    return await _on_flow(
        flow,
        "PUT",
        f"/nodes/{_seg(node_id)}/source",
        mutation=True,
        json={"code": code},
    )


@mcp.tool()
async def publish_flow(flow: FlowArg, version: int) -> Any:
    """Deploy a flow's unpublished changes: this is what puts them live."""
    return await _on_flow(
        flow, "POST", "/publish", mutation=True, json={"version": version}
    )


@mcp.tool()
async def discard_draft(flow: FlowArg) -> Any:
    """Throw unpublished changes away and go back to what is running."""
    return await _on_flow(flow, "POST", "/discard-draft", mutation=True)


@mcp.tool()
async def delete_flow(flow: FlowArg) -> Any:
    """Delete a flow and the code of its nodes."""
    return await _on_flow(flow, "DELETE", mutation=True)


@mcp.tool()
async def validate_flow(flow: FlowArg) -> Any:
    """What would keep this flow from running: loops, unconnected inputs."""
    return await _on_flow(flow, "POST", "/validate", mutation=False)


# -----------------------------------------------------------------------------
# Running
# -----------------------------------------------------------------------------


@mcp.tool()
async def run_flow(flow: FlowArg, inputs: dict[str, Any] | None = None) -> Any:
    """Run every node of a flow once and return the resulting state.

    With unpublished changes this runs the draft, so a change can be tried
    before it is published.
    """
    return await _on_flow(
        flow, "POST", "/run", mutation=True, json={"inputs": inputs or {}}
    )


@mcp.tool()
async def trigger_node(
    flow: FlowArg, node_id: str, values: dict[str, Any] | None = None
) -> Any:
    """Feed values into one node and run everything downstream of it."""
    return await _on_flow(
        flow,
        "POST",
        f"/nodes/{_seg(node_id)}/trigger",
        mutation=True,
        json={"values": values or {}},
    )


@mcp.tool()
async def start_flow(flow: FlowArg) -> Any:
    """Let the engine run this flow: subscriptions, schedules and webhooks."""
    return await _on_flow(flow, "POST", "/start", mutation=True)


@mcp.tool()
async def stop_flow(flow: FlowArg) -> Any:
    """Take a flow off the engine. Nothing of it stays subscribed or scheduled."""
    return await _on_flow(flow, "POST", "/stop", mutation=True)


@mcp.tool()
async def pause_flow(flow: FlowArg) -> Any:
    """Hold a flow's nodes while its incoming values keep arriving."""
    return await _on_flow(flow, "POST", "/pause", mutation=True)


@mcp.tool()
async def resume_flow(flow: FlowArg) -> Any:
    """Let a paused flow carry on, running whatever was held back."""
    return await _on_flow(flow, "POST", "/resume", mutation=True)


@mcp.tool()
async def cancel_node(flow: FlowArg, node_id: str) -> Any:
    """Stop a node's code while it is running. Nothing to stop is not an error."""
    return await _on_flow(flow, "POST", f"/nodes/{_seg(node_id)}/cancel", mutation=True)


# -----------------------------------------------------------------------------
# Modules
# -----------------------------------------------------------------------------


@mcp.tool()
async def get_modules() -> Any:
    """The python packages node code can import, and the manifest asking for them."""
    return await _call("GET", "/modules/")


@mcp.tool()
async def apply_modules(requirements: str) -> Any:
    """Install exactly these requirements, one pip line each.

    This replaces the whole manifest: a package left out is uninstalled. A
    manifest that does not resolve changes nothing.
    """
    return await _call("POST", "/modules/apply", json={"requirements": requirements})


# -----------------------------------------------------------------------------
# Observability
# -----------------------------------------------------------------------------


@mcp.tool()
async def get_health() -> Any:
    """How the engine is doing right now: flows, nodes, queue and loop lag."""
    return await _call("GET", "/observability/summary")


#: The optional flow a history question is narrowed to.
FlowFilter = Annotated[
    str | None,
    Field(
        description=(
            "Only this flow: its id or its key, or the key it had before it was "
            "deleted. Write id:<x> or key:<x> when the name could mean two."
        )
    ),
]


@mcp.tool()
async def get_metrics(
    flow: FlowFilter = None, node: str | None = None, hours: int = 24
) -> Any:
    """Executions, errors and timings per minute over the last few hours.

    ``node`` is ``<flow key>.<node id>``.
    """
    params: dict[str, Any] = {"hours": hours}
    if node:
        params["node"] = node
    return await _narrowed(flow, "/observability/timeseries", params)


@mcp.tool()
async def list_failures(flow: FlowFilter = None, limit: int = 50) -> Any:
    """Recent failures with their tracebacks, newest first."""
    params: dict[str, Any] = {"kind": "failure", "limit": limit}
    return await _narrowed(flow, "/observability/events", params)


@mcp.tool()
async def list_runs(flow: FlowFilter = None, limit: int = 50) -> Any:
    """Recent cascades: what triggered them, how long they took, how they ended.

    A page of rows plus the total number matching, which says whether the limit
    cut anything off.
    """
    params: dict[str, Any] = {"limit": limit}
    return await _narrowed(flow, "/observability/runs", params)