path: root/debug/risc-v-gsp-rm-vm/tests/test_mcp.py

"""Tests for the green-eye MCP server plugin (roadmap R3).

The protocol and tool tests drive an in-memory Session over a stub RSP-shaped
client, so they need no QEMU and no simulator.  A standalone subprocess test
launches the native backend end-to-end and skips when that backend is absent.
"""
from __future__ import annotations

import io
import json
import os
import subprocess
import sys
import tempfile
import unittest

HERE = os.path.dirname(os.path.abspath(__file__))
TOOL_ROOT = os.path.abspath(os.path.join(HERE, ".."))
PLUGINS_DIR = os.path.join(TOOL_ROOT, "plugins")
MCP_SERVER = os.path.join(PLUGINS_DIR, "mcp_server.py")
DEMO_ELF = os.path.join(TOOL_ROOT, "firmware", "uproc_demo", "build", "uproc_demo.elf")
for _path in (TOOL_ROOT, PLUGINS_DIR):
    if _path not in sys.path:
        sys.path.insert(0, _path)

import mcp_server as mcp  # noqa: E402
from greeneye.rsp import StopReply  # noqa: E402
from greeneye.session import Session  # noqa: E402
from greeneye.symbols import Symbol, SymbolDB  # noqa: E402

BASE = 0x80000000


class StubClient:
    """RSP-shaped in-memory client: enough for Session without an emulator."""

    def __init__(self, size: int = 0x20000):
        self.base = BASE
        self.mem = bytearray(size)
        self.regs = {"pc": BASE, "ra": 0, "sp": BASE + size, "a0": 0, "a1": 0}
        self.sw_bps: set[int] = set()
        self.hw_bps: set[int] = set()
        self.watches: dict[tuple[int, int, str], None] = {}
        self.closed = False
        self.features = "swbreak+;hwbreak+"

    def _span(self, addr: int, size: int) -> int:
        start = addr - self.base
        if start < 0 or start + size > len(self.mem):
            raise ValueError(f"out of range: {addr:#x}+{size}")
        return start

    def read_mem(self, addr: int, size: int, chunk: int = 0) -> bytes:
        start = self._span(addr, size)
        return bytes(self.mem[start : start + size])

    def write_mem(self, addr: int, data: bytes, chunk: int = 0) -> None:
        start = self._span(addr, len(data))
        self.mem[start : start + len(data)] = data

    def read_registers(self) -> dict[str, int]:
        return dict(self.regs)

    def read_reg(self, name: str) -> int:
        return int(self.regs.get(str(name), 0))

    def write_reg(self, name: str, value: int) -> None:
        self.regs[str(name)] = int(value)

    @property
    def pc(self) -> int:
        return self.regs["pc"]

    def set_sw_break(self, addr: int, kind: int = 4) -> None:
        self.sw_bps.add(addr)

    def clear_sw_break(self, addr: int, kind: int = 4) -> None:
        self.sw_bps.discard(addr)

    def set_hw_break(self, addr: int, kind: int = 4) -> None:
        self.hw_bps.add(addr)

    def clear_hw_break(self, addr: int, kind: int = 4) -> None:
        self.hw_bps.discard(addr)

    def set_watch(self, addr: int, size: int = 8, kind: str = "w") -> None:
        self.watches[(addr, size, kind)] = None

    def clear_watch(self, addr: int, size: int = 8, kind: str = "w") -> None:
        self.watches.pop((addr, size, kind), None)

    def continue_(self, timeout: float | None = None) -> StopReply:
        return StopReply(raw="T05", signal=5)

    def step(self, timeout: float | None = None) -> StopReply:
        return StopReply(raw="T05", signal=5)

    def interrupt(self, timeout: float | None = None) -> StopReply:
        return StopReply(raw="T02", signal=2)

    def close(self) -> None:
        self.closed = True


class McpTestCase(unittest.TestCase):
    def setUp(self):
        self.client = StubClient()
        symbols = SymbolDB()
        symbols.add(Symbol(addr=BASE + 0x100, name="rpc_ping", size=0x20, source="test"))
        symbols.add(Symbol(addr=BASE + 0x200, name="gsp_ctrl", size=0x40, source="test"))
        self.session = Session(client=self.client, symbols=symbols, record=False, name="mcp-test")
        self.addCleanup(self.session.close)

    def make_server(self, **kw):
        self.inp = io.StringIO()
        self.out = io.StringIO()
        self.log = io.StringIO()
        options = dict(input_stream=self.inp, output_stream=self.out, log_stream=self.log)
        options.update(kw)
        return mcp.GreenEyeMCPServer(self.session, **options)

    def request(self, server, method, params=None, request_id=1):
        message = {"jsonrpc": "2.0", "id": request_id, "method": method}
        if params is not None:
            message["params"] = params
        return server.handle_message(message)

    def call_tool(self, server, name, **arguments):
        response = self.request(server, "tools/call", {"name": name, "arguments": arguments}, request_id=9)
        self.assertIn("result", response, response)
        result = response["result"]
        payload = json.loads(result["content"][0]["text"])
        return payload, result.get("isError", False)


class ProtocolTests(McpTestCase):
    def test_initialize(self):
        response = self.request(
            self.make_server(),
            "initialize",
            {"protocolVersion": "2024-11-05", "capabilities": {}, "clientInfo": {"name": "t", "version": "0"}},
        )
        result = response["result"]
        self.assertEqual(response["jsonrpc"], "2.0")
        self.assertEqual(result["protocolVersion"], "2024-11-05")
        self.assertEqual(result["serverInfo"]["name"], "greeneye-mcp")
        self.assertIn("tools", result["capabilities"])

    def test_initialize_negotiation_and_fallback(self):
        server = self.make_server()
        self.assertEqual(
            self.request(server, "initialize", {"protocolVersion": "2025-03-26"})["result"]["protocolVersion"],
            "2025-03-26",
        )
        self.assertEqual(
            self.request(server, "initialize", {"protocolVersion": "1999-01-01"})["result"]["protocolVersion"],
            mcp.DEFAULT_PROTOCOL,
        )

    def test_initialized_notification_has_no_response(self):
        server = self.make_server()
        self.assertIsNone(server.handle_message({"jsonrpc": "2.0", "method": "notifications/initialized"}))

    def test_ping(self):
        self.assertEqual(self.request(self.make_server(), "ping")["result"], {})

    def test_tools_list_catalogue(self):
        tools = self.request(self.make_server(), "tools/list")["result"]["tools"]
        names = [tool["name"] for tool in tools]
        self.assertEqual(len(names), 24)
        self.assertEqual(set(names), mcp._TOOL_NAMES)
        for tool in tools:
            self.assertEqual(tool["inputSchema"]["type"], "object")
            self.assertTrue(tool["description"])

    def test_unknown_method(self):
        response = self.request(self.make_server(), "no/such")
        self.assertEqual(response["error"]["code"], mcp.METHOD_NOT_FOUND)
        self.assertEqual(response["id"], 1)

    def test_unknown_tool(self):
        response = self.request(self.make_server(), "tools/call", {"name": "nope", "arguments": {}})
        self.assertEqual(response["error"]["code"], mcp.INVALID_PARAMS)

    def test_tools_call_before_initialize_is_allowed(self):
        payload, is_error = self.call_tool(self.make_server(), "session_info")
        self.assertFalse(is_error)
        self.assertEqual(payload["name"], "mcp-test")

    def test_parse_error_is_reported(self):
        server = self.make_server()
        server.process_line("{not json")
        line = self.out.getvalue().strip()
        self.assertEqual(json.loads(line)["error"]["code"], mcp.PARSE_ERROR)

    def test_serve_stream_end_to_end_and_no_stdout_pollution(self):
        server = self.make_server()
        lines = [
            json.dumps({"jsonrpc": "2.0", "id": 1, "method": "initialize", "params": {}}),
            json.dumps({"jsonrpc": "2.0", "method": "notifications/initialized"}),
            json.dumps({"jsonrpc": "2.0", "id": 2, "method": "tools/list"}),
            json.dumps(
                {"jsonrpc": "2.0", "id": 3, "method": "tools/call", "params": {"name": "session_info", "arguments": {}}}
            ),
        ]
        server.serve_stream(io.StringIO("\n".join(lines) + "\n"), self.out)
        raw = self.out.getvalue()
        responses = [json.loads(line) for line in raw.splitlines()]
        self.assertEqual([response["id"] for response in responses], [1, 2, 3])
        self.assertTrue(all(response["jsonrpc"] == "2.0" for response in responses))
        self.assertNotIn("[greeneye-mcp]", raw)
        self.assertNotIn("[greeneye-mcp]", self.inp.getvalue())


class ReadOnlyToolTests(McpTestCase):
    def test_symbols_resolve(self):
        payload, is_error = self.call_tool(self.make_server(), "symbols_resolve", target="rpc_ping")
        self.assertFalse(is_error)
        self.assertEqual(payload["address"], BASE + 0x100)
        self.assertEqual(payload["symbol"], "rpc_ping")

    def test_read_mem_and_hexdump(self):
        self.client.mem[0x10:0x14] = b"\xde\xad\xbe\xef"
        server = self.make_server()
        payload, is_error = self.call_tool(server, "read_mem", addr=BASE + 0x10, size=4)
        self.assertFalse(is_error)
        self.assertEqual(payload["hex"], "deadbeef")
        payload, is_error = self.call_tool(server, "hexdump", addr=BASE + 0x10, size=4)
        self.assertFalse(is_error)
        self.assertIn("de ad be ef", payload["text"])

    def test_bp_set_and_clear(self):
        server = self.make_server()
        payload, is_error = self.call_tool(server, "bp_set", target="rpc_ping")
        self.assertFalse(is_error)
        self.assertIn(BASE + 0x100, self.client.sw_bps)
        payload, is_error = self.call_tool(server, "bp_clear", target="rpc_ping")
        self.assertFalse(is_error)
        self.assertEqual(payload["cleared"], 1)
        self.assertNotIn(BASE + 0x100, self.client.sw_bps)

    def test_watch_set(self):
        payload, is_error = self.call_tool(self.make_server(), "watch_set", target="gsp_ctrl", size=8, kind="w")
        self.assertFalse(is_error)
        self.assertEqual(payload["addr"], BASE + 0x200)
        self.assertIn((BASE + 0x200, 8, "w"), self.client.watches)

    def test_snapshot_digest(self):
        payload, is_error = self.call_tool(self.make_server(), "snapshot", addr=BASE, size=16, label="unit")
        self.assertFalse(is_error)
        self.assertEqual(payload["size"], 16)
        self.assertEqual(len(payload["sha256"]), 64)

    def test_session_info_capabilities(self):
        payload, is_error = self.call_tool(
            self.make_server(allow_write=True, allow_inject=True), "session_info"
        )
        self.assertFalse(is_error)
        self.assertEqual(payload["name"], "mcp-test")
        self.assertTrue(payload["capabilities"]["write"])
        self.assertTrue(payload["capabilities"]["inject"])
        self.assertFalse(payload["capabilities"]["trace"])

    def test_plugins_list(self):
        payload, is_error = self.call_tool(self.make_server(), "plugins_list")
        self.assertFalse(is_error)
        self.assertIn("mcp_server.py", payload["bundled"])

    def test_analyze_current_reads_active_recording(self):
        with tempfile.TemporaryDirectory(prefix="mcp-rec-") as directory:
            session = Session(client=StubClient(), symbols=SymbolDB(), record=directory, name="active")
            try:
                server = mcp.GreenEyeMCPServer(
                    session,
                    input_stream=io.StringIO(),
                    output_stream=io.StringIO(),
                    log_stream=io.StringIO(),
                )
                payload, is_error = self.call_tool(server, "analyze_current")
                self.assertFalse(is_error)
                self.assertIn("# green-eye session summary", payload["report"])
            finally:
                session.close()


class GatingTests(McpTestCase):
    def test_write_mem_refused_without_allow_write(self):
        server = self.make_server()
        payload, is_error = self.call_tool(server, "write_mem", addr=BASE + 0x10, hex="0102")
        self.assertTrue(is_error)
        self.assertEqual(payload["error"], "write_disabled")
        self.assertEqual(bytes(self.client.mem[0x10:0x12]), b"\x00\x00")

    def test_write_mem_enabled(self):
        server = self.make_server(allow_write=True)
        payload, is_error = self.call_tool(server, "write_mem", addr=BASE + 0x10, hex="0x01 02")
        self.assertFalse(is_error)
        self.assertEqual(payload["written"], 2)
        self.assertEqual(bytes(self.client.mem[0x10:0x12]), b"\x01\x02")

    def test_call_function_refused_without_allow_write(self):
        payload, is_error = self.call_tool(self.make_server(), "call_function", target="rpc_ping", args=[1])
        self.assertTrue(is_error)
        self.assertEqual(payload["error"], "write_disabled")

    def test_inject_refused_without_allow_inject(self):
        payload, is_error = self.call_tool(self.make_server(), "inject_service", service="/nonexistent.elf")
        self.assertTrue(is_error)
        self.assertEqual(payload["error"], "inject_disabled")

    def test_inject_missing_file_after_gate(self):
        payload, is_error = self.call_tool(
            self.make_server(allow_inject=True), "inject_service", service="/nonexistent.elf"
        )
        self.assertTrue(is_error)
        self.assertEqual(payload["error"], "invalid_arguments")


class AvailabilityTests(McpTestCase):
    def test_analyze_current_without_recording(self):
        payload, is_error = self.call_tool(self.make_server(), "analyze_current")
        self.assertTrue(is_error)
        self.assertEqual(payload["error"], "not_available")

    def test_trace_start_on_non_native_backend(self):
        payload, is_error = self.call_tool(self.make_server(), "trace_start")
        self.assertTrue(is_error)
        self.assertEqual(payload["error"], "not_available")

    def test_trace_stop_without_active_trace(self):
        payload, is_error = self.call_tool(self.make_server(), "trace_stop")
        self.assertTrue(is_error)
        self.assertEqual(payload["error"], "invalid_arguments")

    def test_mutation_missing_script_is_structured(self):
        server = self.make_server()
        original = mcp.MUTATION_SCRIPT
        mcp.MUTATION_SCRIPT = "/nonexistent/mutation_harness.py"
        try:
            payload, is_error = self.call_tool(server, "mutation_run")
        finally:
            mcp.MUTATION_SCRIPT = original
        self.assertTrue(is_error)
        self.assertEqual(payload["error"], "not_available")

    def test_portability_missing_script_is_structured(self):
        server = self.make_server()
        original = mcp.MATRIX_SCRIPT
        mcp.MATRIX_SCRIPT = "/nonexistent/portability_matrix.py"
        try:
            payload, is_error = self.call_tool(server, "portability_matrix_run", targets=["/x.elf"])
        finally:
            mcp.MATRIX_SCRIPT = original
        self.assertTrue(is_error)
        self.assertEqual(payload["error"], "not_available")

    def test_probe_missing_elf(self):
        payload, is_error = self.call_tool(self.make_server(), "probe_firmware", elf="/nonexistent/fw.elf")
        self.assertTrue(is_error)
        self.assertEqual(payload["error"], "invalid_arguments")


class CloseAndTransportTests(McpTestCase):
    def test_close_stops_serving_but_leaves_plugin_session(self):
        server = self.make_server(owns_session=False)
        payload, is_error = self.call_tool(server, "close")
        self.assertFalse(is_error)
        self.assertFalse(payload["session_closed"])
        self.assertTrue(server._stop)
        self.assertFalse(self.client.closed)

    def test_close_closes_owned_session(self):
        server = self.make_server(owns_session=True)
        payload, is_error = self.call_tool(server, "close")
        self.assertFalse(is_error)
        self.assertTrue(payload["session_closed"])
        self.assertTrue(self.client.closed)

    def test_parse_tcp_is_loopback_only(self):
        self.assertEqual(mcp.parse_tcp("127.0.0.1:8765"), ("127.0.0.1", 8765))
        self.assertEqual(mcp.parse_tcp("localhost:0")[0], "localhost")
        with self.assertRaises(Exception):
            mcp.parse_tcp("0.0.0.0:8765")


class StandaloneEntryTests(unittest.TestCase):
    @unittest.skipUnless(os.path.isfile(MCP_SERVER), "mcp_server.py missing")
    def test_help(self):
        result = subprocess.run(
            [sys.executable, MCP_SERVER, "--help"], capture_output=True, text=True, cwd=TOOL_ROOT
        )
        self.assertEqual(result.returncode, 0, result.stderr)
        self.assertIn("--stdio", result.stdout)
        self.assertIn("--allow-write", result.stdout)

    @unittest.skipUnless(os.path.isfile(MCP_SERVER), "mcp_server.py missing")
    def test_stdio_native_session(self):
        try:
            import greeneye.sim_backend  # noqa: F401
        except Exception as exc:
            self.skipTest(f"native backend unavailable: {exc}")
        with tempfile.TemporaryDirectory(prefix="mcp-native-") as dumps:
            env = dict(os.environ, GREENEYE_DUMPS=dumps)
            requests = [
                {"jsonrpc": "2.0", "id": 1, "method": "initialize", "params": {"protocolVersion": "2024-11-05"}},
                {"jsonrpc": "2.0", "method": "notifications/initialized"},
                {"jsonrpc": "2.0", "id": 2, "method": "tools/call", "params": {"name": "session_info", "arguments": {}}},
                {"jsonrpc": "2.0", "id": 3, "method": "tools/call", "params": {"name": "close", "arguments": {}}},
            ]
            payload = "".join(json.dumps(request) + "\n" for request in requests)
            result = subprocess.run(
                [sys.executable, MCP_SERVER, "--stdio", "--backend", "native", "--record"],
                input=payload,
                capture_output=True,
                text=True,
                cwd=TOOL_ROOT,
                env=env,
                timeout=180,
            )
            self.assertEqual(result.returncode, 0, result.stderr[-2000:])
            responses = [json.loads(line) for line in result.stdout.splitlines() if line.strip()]
            self.assertEqual(len(responses), 3)
            by_id = {response.get("id"): response for response in responses}
            info = json.loads(by_id[2]["result"]["content"][0]["text"])
            self.assertEqual(info["backend"], "native")
            self.assertFalse(by_id[3]["result"].get("isError", False))
            self.assertTrue(any(name.endswith("-mcp") for name in os.listdir(dumps)))


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