Skip to content

Commit a754423

Browse files
committed
tests: add --mcp-mode to run the IDA test suite end-to-end over MCP
Adds a second test layer that routes every @tool call through a real MCP client/server HTTP round-trip and validates each response against its advertised outputSchema. Covers the runtime-conformance gap the synthetic fixtures can't: does a production tool actually return data matching its schema? Usage: ida-mcp-test <bin> --mcp-mode Implementation: - New tests/mcp_mode.py: boots McpServer on a random port, connects an mcp SDK client over streamable HTTP, monkeypatches every registered tool in its defining module with a proxy that calls through the client, validates structuredContent against outputSchema, and unwraps the {"result": ...} envelope for primitive returns. - Background asyncio loop lets sync tests drive the async client. - Patches are reverted on exit; server is stopped. Results on crackme03.elf: 299 passed / 5 failed / 0 skipped. Failures are useful findings: tests using partial module-level mocks return schema-incomplete responses (e.g. idalib_health mock session), or rely on session scoping that MCP round-trips don't preserve. Documenting rather than papering over, to surface the gap. New dev dep: mcp>=1.0 (official MCP Python SDK).
1 parent 307c8be commit a754423

4 files changed

Lines changed: 736 additions & 7 deletions

File tree

pyproject.toml

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -27,6 +27,7 @@ dependencies = [
2727
dev = [
2828
"coverage>=7.13.4",
2929
"jsonschema>=4.0",
30+
"mcp>=1.0",
3031
]
3132

3233
[project.urls]
Lines changed: 206 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,206 @@
1+
"""End-to-end MCP mode for the test framework.
2+
3+
When enabled, every @tool-decorated function is monkeypatched to route
4+
through a real MCP client/server HTTP round-trip and the response is
5+
validated against its advertised outputSchema. Lets the existing test
6+
suite double as a runtime-conformance check without rewriting anything.
7+
"""
8+
9+
import asyncio
10+
import inspect
11+
import socket
12+
import sys
13+
import threading
14+
from contextlib import contextmanager
15+
from functools import wraps
16+
from typing import Any, Callable, Iterator
17+
18+
from jsonschema import Draft202012Validator
19+
20+
from ..rpc import MCP_SERVER
21+
22+
23+
class _AsyncHarness:
24+
"""Event loop on a background thread; sync callers schedule coros."""
25+
26+
def __init__(self):
27+
self.loop: asyncio.AbstractEventLoop | None = None
28+
self.thread: threading.Thread | None = None
29+
self._ready = threading.Event()
30+
31+
def start(self):
32+
def loop_target():
33+
self.loop = asyncio.new_event_loop()
34+
asyncio.set_event_loop(self.loop)
35+
self._ready.set()
36+
self.loop.run_forever()
37+
38+
self.thread = threading.Thread(target=loop_target, daemon=True)
39+
self.thread.start()
40+
self._ready.wait()
41+
42+
def stop(self):
43+
if self.loop is not None:
44+
self.loop.call_soon_threadsafe(self.loop.stop)
45+
if self.thread is not None:
46+
self.thread.join(timeout=2)
47+
48+
def run(self, coro, timeout: float = 30.0):
49+
fut = asyncio.run_coroutine_threadsafe(coro, self.loop)
50+
return fut.result(timeout=timeout)
51+
52+
53+
class _McpMode:
54+
def __init__(self):
55+
self.harness = _AsyncHarness()
56+
self.host = "127.0.0.1"
57+
self.port: int | None = None
58+
self.session = None
59+
self._http_ctx = None
60+
self._session_ctx = None
61+
self.output_schemas: dict[str, dict] = {}
62+
self._patched: list[tuple[Any, str, Callable]] = []
63+
64+
def enable(self):
65+
self.harness.start()
66+
self.port = _pick_free_port(self.host)
67+
MCP_SERVER.serve(self.host, self.port, background=True)
68+
_wait_until_ready(f"http://{self.host}:{self.port}/mcp")
69+
self.harness.run(self._connect())
70+
self._patch_tools()
71+
72+
def disable(self):
73+
self._unpatch_tools()
74+
try:
75+
self.harness.run(self._disconnect())
76+
except Exception:
77+
pass
78+
try:
79+
MCP_SERVER.stop()
80+
except Exception:
81+
pass
82+
self.harness.stop()
83+
84+
async def _connect(self):
85+
from mcp import ClientSession
86+
from mcp.client.streamable_http import streamablehttp_client
87+
88+
self._http_ctx = streamablehttp_client(f"http://{self.host}:{self.port}/mcp")
89+
read, write, _ = await self._http_ctx.__aenter__()
90+
self._session_ctx = ClientSession(read, write)
91+
self.session = await self._session_ctx.__aenter__()
92+
await self.session.initialize()
93+
tools = await self.session.list_tools()
94+
for t in tools.tools:
95+
if getattr(t, "outputSchema", None):
96+
self.output_schemas[t.name] = t.outputSchema
97+
98+
async def _disconnect(self):
99+
if self._session_ctx is not None:
100+
await self._session_ctx.__aexit__(None, None, None)
101+
if self._http_ctx is not None:
102+
await self._http_ctx.__aexit__(None, None, None)
103+
104+
def _patch_tools(self):
105+
for name, original in MCP_SERVER.tools.methods.items():
106+
mod = sys.modules.get(original.__module__)
107+
if mod is None:
108+
continue
109+
if getattr(mod, name, None) is not original:
110+
continue
111+
proxy = self._make_proxy(name, original)
112+
self._patched.append((mod, name, original))
113+
setattr(mod, name, proxy)
114+
115+
def _unpatch_tools(self):
116+
for mod, name, original in self._patched:
117+
setattr(mod, name, original)
118+
self._patched.clear()
119+
120+
def _make_proxy(self, name: str, original: Callable) -> Callable:
121+
sig = inspect.signature(original)
122+
schema = self.output_schemas.get(name)
123+
harness = self.harness
124+
state = self
125+
126+
@wraps(original)
127+
def proxy(*args, **kwargs):
128+
bound = sig.bind(*args, **kwargs)
129+
bound.apply_defaults()
130+
arguments = dict(bound.arguments)
131+
132+
async def _call():
133+
return await state.session.call_tool(name, arguments=arguments)
134+
135+
result = harness.run(_call())
136+
137+
if getattr(result, "isError", False):
138+
msg = ""
139+
if result.content:
140+
msg = getattr(result.content[0], "text", str(result.content[0]))
141+
raise RuntimeError(f"MCP tool {name!r} returned isError: {msg}")
142+
143+
structured = getattr(result, "structuredContent", None)
144+
if schema is not None and structured is not None:
145+
Draft202012Validator(schema).validate(structured)
146+
147+
if (
148+
isinstance(structured, dict)
149+
and set(structured.keys()) == {"result"}
150+
):
151+
return structured["result"]
152+
return structured
153+
154+
return proxy
155+
156+
157+
_instance: _McpMode | None = None
158+
159+
160+
def enable_mcp_mode() -> None:
161+
global _instance
162+
assert _instance is None, "MCP mode already enabled"
163+
_instance = _McpMode()
164+
_instance.enable()
165+
166+
167+
def disable_mcp_mode() -> None:
168+
global _instance
169+
if _instance is not None:
170+
_instance.disable()
171+
_instance = None
172+
173+
174+
@contextmanager
175+
def mcp_mode() -> Iterator[None]:
176+
enable_mcp_mode()
177+
try:
178+
yield
179+
finally:
180+
disable_mcp_mode()
181+
182+
183+
def _pick_free_port(host: str) -> int:
184+
s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
185+
s.bind((host, 0))
186+
port = s.getsockname()[1]
187+
s.close()
188+
return port
189+
190+
191+
def _wait_until_ready(url: str, timeout: float = 2.0) -> None:
192+
import time
193+
import urllib.error
194+
import urllib.request
195+
196+
deadline = time.time() + timeout
197+
while time.time() < deadline:
198+
try:
199+
req = urllib.request.Request(url, method="OPTIONS")
200+
with urllib.request.urlopen(req, timeout=0.2):
201+
return
202+
except urllib.error.HTTPError:
203+
return
204+
except (urllib.error.URLError, ConnectionRefusedError, socket.timeout):
205+
time.sleep(0.02)
206+
raise RuntimeError(f"server at {url} not ready within {timeout}s")

src/ida_pro_mcp/test.py

Lines changed: 24 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -76,6 +76,12 @@ def main() -> int:
7676
action="store_true",
7777
help="Show IDA console messages",
7878
)
79+
parser.add_argument(
80+
"--mcp-mode",
81+
action="store_true",
82+
help="Route every @tool call through a real MCP client/server "
83+
"round-trip and validate responses against outputSchema",
84+
)
7985
args = parser.parse_args()
8086

8187
# Check binary exists
@@ -144,13 +150,24 @@ def main() -> int:
144150
in_ci = os.environ.get("CI", "").lower() not in ("", "0", "false", "no")
145151
interactive_output = sys.stdout.isatty()
146152
show_all_test_output = (not args.quiet) and (interactive_output or in_ci)
147-
results = run_tests(
148-
pattern=args.pattern,
149-
category=args.category,
150-
verbose=show_all_test_output,
151-
stop_on_failure=args.stop_on_failure,
152-
failures_only=(not args.quiet) and not show_all_test_output,
153-
)
153+
154+
def _run():
155+
return run_tests(
156+
pattern=args.pattern,
157+
category=args.category,
158+
verbose=show_all_test_output,
159+
stop_on_failure=args.stop_on_failure,
160+
failures_only=(not args.quiet) and not show_all_test_output,
161+
)
162+
163+
if args.mcp_mode:
164+
from ida_pro_mcp.ida_mcp.tests.mcp_mode import mcp_mode
165+
166+
print("[MCP] Running tests in end-to-end MCP mode.")
167+
with mcp_mode():
168+
results = _run()
169+
else:
170+
results = _run()
154171

155172
# No matched tests is likely a configuration/test-selection mistake
156173
if not results.results:

0 commit comments

Comments
 (0)