MCP servers as tools

"""MCP servers as tools on AgentOS: the model runs through the gateway, the MCP tools run inside the agent."""

import json
import sys
from contextlib import asynccontextmanager

from mcp import Client, StdioServerParameters

import agentos

ORDERS_SERVER = StdioServerParameters(command=sys.executable, args=["orders_server.py"])
NEEDS_APPROVAL = {"cancel_order"}   # tools a person must approve before they run
MAX_TURNS = 10


@asynccontextmanager
async def mcp_client(server):
    """An MCP client that lets errors through as they are. The client's task group wraps whatever is raised in
    its block (the SDK's ask, a Fail, a guardrail block) in an ExceptionGroup, which the SDK can't read."""
    try:
        async with Client(server) as client:
            yield client
    except BaseExceptionGroup as group:
        error = group
        while isinstance(error, BaseExceptionGroup) and len(error.exceptions) == 1:
            error = error.exceptions[0]
        raise error from None


def definitions(tools) -> list[dict]:
    """MCP tools as OpenAI function tools for the gateway."""
    return [{"type": "function", "function": {"name": t.name, "description": t.description or "",
                                              "parameters": t.input_schema}} for t in tools]


def as_text(result) -> str:
    if result.structured_content is not None:
        return json.dumps(result.structured_content)
    return "\n".join(getattr(part, "text", "") for part in result.content)


async def tool_loop(inv: agentos.Invocation, mcp: Client, saved: dict) -> str:
    """Model → MCP tool calls → model, until the model answers. `saved` (messages and the tool calls not run yet)
    is kept in run state when the loop stops to ask a person, so the resume continues where it stopped."""
    tools = definitions((await mcp.list_tools()).tools)
    messages, pending = saved["messages"], saved["pending"]
    decision = inv.resume                       # on a resume: the decision about the first pending call
    for _ in range(MAX_TURNS):
        while pending:
            call = pending[0]
            name, args = call["function"]["name"], json.loads(call["function"]["arguments"] or "{}")
            if name in NEEDS_APPROVAL:
                if decision is None:
                    await inv.state.run.aput("conversation", {"messages": messages, "pending": pending})
                    await inv.aask(f"Run {name}?", [{"id": call["id"], "title": name, "data": args}],
                                   actions=["approve", "edit", "reject"])   # the run waits for a person here
                item = decision.decided_items[0] if decision.decided_items else None
                decision = None
                if item is None or item.action == "reject":
                    result = "A person declined this tool call."
                else:
                    result = as_text(await mcp.call_tool(name, item.data))   # an edit's arguments win
            else:
                result = as_text(await mcp.call_tool(name, args))
            messages.append({"role": "tool", "tool_call_id": call["id"], "content": result})
            pending.pop(0)

        message = (await inv.llm.achat(messages, tools=tools))["choices"][0]["message"]
        messages.append({"role": "assistant", "content": message.get("content"),
                         **({"tool_calls": message["tool_calls"]} if message.get("tool_calls") else {})})
        pending = list(message.get("tool_calls") or [])
        if not pending:
            return message.get("content") or ""
    raise agentos.Fail("too_many_turns", f"no answer after {MAX_TURNS} model turns")


async def support(inv: agentos.Invocation):
    """A stdio MCP server bundled in the container."""
    if inv.resume is None:
        saved = {"messages": [{"role": "system", "content": "You help customers with their orders."},
                              {"role": "user", "content": inv.input["message"]}], "pending": []}
    else:
        saved = await inv.state.run.aget("conversation")
    async with mcp_client(ORDERS_SERVER) as mcp:
        return {"answer": await tool_loop(inv, mcp, saved)}


async def ask_docs(inv: agentos.Invocation):
    """A remote MCP server over Streamable HTTP. Its host must be in network.allow (agentos.yaml)."""
    saved = {"messages": [{"role": "system", "content": "Answer with the DeepWiki tools. Name the repository."},
                          {"role": "user", "content": inv.input["question"]}], "pending": []}
    async with mcp_client(inv.config["params"]["docs_mcp_url"]) as mcp:
        return {"answer": await tool_loop(inv, mcp, saved)}


if __name__ == "__main__":
    agentos.serve({"support": support, "ask_docs": ask_docs}, port=8080)

"""A small MCP server bundled with the agent. The agent starts it as a subprocess and talks to it over stdio."""

from mcp.server.mcpserver import MCPServer

server = MCPServer("orders")

ORDERS = {"A-1001": {"status": "shipped", "total": "84.90"}, "A-1002": {"status": "processing", "total": "250.00"}}


@server.tool()
def lookup_order(order_id: str) -> dict:
    """Look up an order's status and total."""
    return {"order_id": order_id, **ORDERS.get(order_id, {"status": "not found"})}


@server.tool()
def cancel_order(order_id: str, reason: str) -> dict:
    """Cancel an order. Irreversible."""
    if order_id not in ORDERS:
        return {"order_id": order_id, "cancelled": False, "error": "not found"}
    ORDERS[order_id]["status"] = "cancelled"
    return {"order_id": order_id, "cancelled": True, "reason": reason}


if __name__ == "__main__":
    server.run()  # stdio

name: mcp-support
title: MCP support agent
description: Uses MCP servers as tools - one bundled in the container, one remote - and asks a person before a risky tool call.
runtime: stateless
port: 8080
capabilities:
  support:
    description: Answers order questions with the bundled orders MCP server; cancelling an order needs a person's approval.
    input:
      type: object
      properties: {message: {type: string}}
      required: [message]
    output:
      type: object
      properties: {answer: {type: string}}
      required: [answer]
    timeout: 5m
  ask_docs:
    description: Answers questions about a GitHub repository with the remote DeepWiki MCP server.
    input:
      type: object
      properties: {question: {type: string}}
      required: [question]
    output:
      type: object
      properties: {answer: {type: string}}
      required: [answer]
    timeout: 5m
config:
  models: {main: openai/gpt-4o-mini}
  params: {docs_mcp_url: "https://mcp.deepwiki.com/mcp"}
network:
  allow: [mcp.deepwiki.com]

FROM python:3.12-slim
WORKDIR /app
ENV PYTHONUNBUFFERED=1
COPY requirements.txt .
COPY vendor/ vendor/
RUN pip install --no-cache-dir -r requirements.txt vendor/*.whl
COPY app.py orders_server.py ./
USER nobody
EXPOSE 8080
CMD ["python", "app.py"]

mcp==2.3.0