Expose the docs assistant as an MCP server

"""Test the MCP server in memory: the real MCP protocol, with a counting model instead of Gemini.

Run: python -m unittest -v test_server.py
"""
import asyncio
import json
import logging
import os
import unittest
import warnings
from unittest.mock import AsyncMock, MagicMock, patch
from datetime import timedelta
from pathlib import Path

warnings.filterwarnings("ignore")
logging.disable(logging.CRITICAL)

from google.adk.models.base_llm import BaseLlm  # noqa: E402
from google.adk.models.llm_response import LlmResponse  # noqa: E402
from google.adk.tools.mcp_tool import StdioConnectionParams, to_mcp_server  # noqa: E402
from google.genai import types  # noqa: E402
from mcp import McpError, StdioServerParameters  # noqa: E402
from mcp.shared.memory import create_connected_server_and_client_session as connect  # noqa: E402

import server  # noqa: E402
import release_bot  # noqa: E402

RECORD = json.loads(Path(__file__).with_name("recorded-runs.json").read_text(encoding="utf-8"))


class CountingLlm(BaseLlm):
    """Answers with the number of developer messages in the request, after an optional delay."""
    delay: float = 0.0

    async def generate_content_async(self, llm_request, stream=False):
        await asyncio.sleep(self.delay)
        seen = sum(1 for c in llm_request.contents if c.role == "user" and any(p.text for p in c.parts))
        yield LlmResponse(content=types.Content(role="model", parts=[types.Part(text=f"messages seen: {seen}")]))


def serve(model) -> object:
    return to_mcp_server(server.root_agent.clone(update={"model": model}))


def run(coro):
    return asyncio.run(coro)


class AgentAsMcpServer(unittest.TestCase):
    def test_the_server_offers_one_tool_that_takes_a_request(self):
        async def go():
            async with connect(serve(CountingLlm(model="count"))) as client:
                return (await client.list_tools()).tools
        tools = run(go())
        self.assertEqual([t.name for t in tools], ["docuctl_docs"])
        self.assertEqual(tools[0].inputSchema["required"], ["request"])
        self.assertEqual(tools[0].description, server.root_agent.description)

    def test_calls_on_one_connection_share_a_session(self):
        mcp_server = serve(CountingLlm(model="count"))

        async def go():
            async with connect(mcp_server) as first:
                a = await first.call_tool("docuctl_docs", {"request": "one"})
                b = await first.call_tool("docuctl_docs", {"request": "two"})
            async with connect(mcp_server) as second:
                c = await second.call_tool("docuctl_docs", {"request": "three"})
            return [r.content[0].text for r in (a, b, c)]
        self.assertEqual(run(go()), ["messages seen: 1", "messages seen: 2", "messages seen: 1"])

    def test_an_agent_error_comes_back_as_an_error_result(self):
        saved = {k: os.environ.pop(k) for k in ("GOOGLE_API_KEY", "GEMINI_API_KEY") if k in os.environ}
        try:
            async def go():
                async with connect(serve("gemini-3.8-flash")) as client:
                    return await client.call_tool("docuctl_docs", {"request": "How do I roll back?"})
            result = run(go())
        finally:
            os.environ.update(saved)
        self.assertTrue(result.isError)
        self.assertIn("No API key was provided", result.content[0].text)

    def test_a_slow_agent_exceeds_the_client_read_timeout(self):
        async def go():
            async with connect(serve(CountingLlm(model="count", delay=1.0)),
                               read_timeout_seconds=timedelta(seconds=0.2)) as client:
                try:
                    await client.call_tool("docuctl_docs", {"request": "slow"})
                except McpError as error:
                    return str(error)
        self.assertIn("Timed out while waiting for response", run(go()))

    def test_stdio_connections_default_to_a_five_second_timeout(self):
        params = StdioConnectionParams(server_params=StdioServerParameters(command="python"))
        self.assertEqual(params.timeout, 5.0)
        errors = [r for r in RECORD["failures"]["default_timeout"]["results"] if "error" in r]
        self.assertIn("Waited 5.0 seconds", errors[0]["error"])


class ClientCleanup(unittest.TestCase):
    def test_a_failed_run_closes_the_runner_and_toolset(self):
        async def failing_run(**kwargs):
            raise RuntimeError("transport failed")
            yield  # keep the Runner.run_async interface

        runner = MagicMock()
        runner.run_async = failing_run
        runner.close = AsyncMock()
        toolset = MagicMock()
        toolset.close = AsyncMock()
        with patch.object(release_bot, "docs_toolset", return_value=toolset), \
             patch.object(release_bot, "make_bot"), patch.object(release_bot, "Runner", return_value=runner):
            with self.assertRaisesRegex(RuntimeError, "transport failed"):
                run(release_bot.ask("Read the rollback page."))
        runner.close.assert_awaited_once()
        toolset.close.assert_awaited_once()


class RecordedRun(unittest.TestCase):
    """Checks on the recorded Gemini run: these read recorded-runs.json and make no model call."""

    def test_the_follow_up_on_the_same_connection_answered_for_1x(self):
        same, fresh = RECORD["same_connection"]["calls"][1]["text"], RECORD["new_connection"]["calls"][0]["text"]
        self.assertIn("shipctl release list", same)
        self.assertNotIn("1.x", fresh)

    def test_the_server_without_the_key_failed_every_call(self):
        self.assertTrue(all(c["is_error"] for c in RECORD["failures"]["no_key"]["calls"]))


if __name__ == "__main__":
    unittest.main()

"""Record MCP clients talking to the docs assistant's MCP server over stdio, on Gemini.

Run from this folder with GOOGLE_API_KEY set: python record.py
Writes recorded-runs.json and run-output.txt. The key is passed to the server process, never recorded.
"""
import asyncio
import json
import os
import sys
import time
import warnings
from pathlib import Path

warnings.filterwarnings("ignore")

from mcp import ClientSession, StdioServerParameters  # noqa: E402
from mcp.client.stdio import stdio_client  # noqa: E402

import release_bot  # noqa: E402

HERE = Path(__file__).resolve().parent
FIRST = "I'm on docuctl 1.x. How do I roll back production to r41?"
FOLLOW_UP = "And how do I list the releases of production?"


def params(with_key: bool = True) -> StdioServerParameters:
    env = {"GOOGLE_API_KEY": os.environ["GOOGLE_API_KEY"]} if with_key else None
    return StdioServerParameters(command=sys.executable, args=[str(HERE / "server.py")], env=env, cwd=str(HERE))


def text_of(result) -> str:
    return "".join(getattr(block, "text", "") for block in result.content)


async def connection(questions: list[str], with_key: bool = True) -> dict:
    """One MCP connection: initialize, list tools, then call the agent tool once per question."""
    async with stdio_client(params(with_key)) as (read, write), ClientSession(read, write) as session:
        init = await session.initialize()
        tools = (await session.list_tools()).tools
        calls = []
        for question in questions:
            start = time.perf_counter()
            result = await session.call_tool(tools[0].name, {"request": question})
            calls.append({"request": question, "is_error": result.isError, "text": text_of(result),
                          "seconds": round(time.perf_counter() - start, 1)})
    return {"server": {"name": init.serverInfo.name, "instructions": init.instructions,
                       "capabilities": sorted(k for k, v in init.capabilities.model_dump().items() if v)},
            "tools": [t.model_dump(mode="json", exclude_none=True) for t in tools], "calls": calls}


async def main() -> None:
    first = await connection([FIRST, FOLLOW_UP])
    fresh = await connection([FOLLOW_UP])
    no_key = await connection([FIRST], with_key=False)
    bot = await release_bot.record(timeout=60)
    bot_default_timeout = await release_bot.record(timeout=None)
    record = {"model": release_bot.MODEL, "same_connection": first, "new_connection": fresh,
              "failures": {"no_key": no_key, "default_timeout": bot_default_timeout}, "release_bot": bot}
    text = json.dumps(record, indent=2)
    assert os.environ["GOOGLE_API_KEY"] not in text
    (HERE / "recorded-runs.json").write_text(text + "\n", encoding="utf-8")
    tool = first["tools"][0]
    lines = [f"server {first['server']['name']}: {first['server']['instructions']}",
             f"tool   {tool['name']}({', '.join(tool['inputSchema']['properties'])}): {tool['description']}", ""]
    for label, run in [("connection 1", first), ("connection 2", fresh), ("no GOOGLE_API_KEY in the server env", no_key)]:
        lines.append(f"# {label}")
        for c in run["calls"]:
            lines += [f"> {c['request']}", f"{'error ' if c['is_error'] else 'answer'} {c['text'][:260]}  ({c['seconds']}s)"]
        lines.append("")
    for label, run in [("release_bot through McpToolset, timeout=60", bot),
                       ("release_bot through McpToolset, default timeout", bot_default_timeout)]:
        lines += [f"# {label}", f"> {run['question']}"]
        lines += [f"call   {c['name']} {json.dumps(c['args'])[:160]}" for c in run["calls"]]
        lines += [f"result {json.dumps(r)[:220]}" for r in run["results"]]
        lines += ["answer " + run["answer"].replace("\n", " ")[:300], ""]
    (HERE / "run-output.txt").write_text("\n".join(lines) + "\n", encoding="utf-8")
    print("\n".join(lines))


if __name__ == "__main__":
    asyncio.run(main())