@hybridlabor-api/bdb-hardware-pcb 0.1.0

package/scripts/test_mcp_connection.py

@@ -0,0 +1,626 @@
+#!/usr/bin/env python3
+"""
+================================================================================
+AOS MCP Multi-Harness Verification & Connection Test Suite
+================================================================================
+Implements the Model Context Protocol (MCP) JSON-RPC 2.0 protocol over stdio.
+Validates server initialization, handshake, tool enumeration, schema integrity,
+and clean termination without requiring any third-party Python packages.
+
+Features:
+- Pure Python standard library (zero external dependencies)
+- Supports JSON config (Antigravity, Claude Code) and TOML config (Codex)
+- Resilient non-blocking stdio I/O handling newline-delimited & framed JSON-RPC
+- Pagination handling for tools/list via nextCursor
+- In-depth tool schema validation
+- Pretty tabular output and JSON report generation
+- Antigravity lazy-tool schema generator
+================================================================================
+"""
+
+import sys
+import os
+import json
+import time
+import argparse
+import subprocess
+import select
+from pathlib import Path
+from typing import Dict, Any, List, Optional, Tuple
+
+# Attempt to import tomllib (Python 3.11+) or fallback to simple TOML parser
+try:
+    import tomllib  # type: ignore
+except ImportError:
+    tomllib = None  # type: ignore
+
+
+PROTOCOL_VERSION = "2024-11-05"
+
+
+class ToolSchemaValidationError(Exception):
+    """Raised when a tool definition fails MCP schema validation."""
+    pass
+
+
+class MCPStdioClient:
+    """Zero-dependency stdio JSON-RPC client for Model Context Protocol servers."""
+
+    def __init__(
+        self,
+        command: str,
+        args: Optional[List[str]] = None,
+        env: Optional[Dict[str, str]] = None,
+        cwd: Optional[str] = None,
+        timeout: float = 15.0,
+    ):
+        self.command = command
+        self.args = args or []
+        self.env = {**os.environ, **(env or {})}
+        self.cwd = cwd or os.getcwd()
+        self.timeout = timeout
+        self.process: Optional[subprocess.Popen] = None
+        self.msg_id = 1
+        self.server_info: Dict[str, Any] = {}
+        self.server_instructions: str = ""
+        self.protocol_version: str = ""
+        self.capabilities: Dict[str, Any] = {}
+
+    def start(self):
+        """Spawns the MCP server subprocess."""
+        cmd = [self.command] + self.args
+        try:
+            self.process = subprocess.Popen(
+                cmd,
+                stdin=subprocess.PIPE,
+                stdout=subprocess.PIPE,
+                stderr=subprocess.PIPE,
+                text=True,
+                bufsize=0,
+                cwd=self.cwd,
+                env=self.env,
+            )
+        except Exception as exc:
+            raise RuntimeError(f"Failed to spawn MCP server command '{cmd}': {exc}") from exc
+
+    def send_message(self, message: Dict[str, Any]):
+        """Sends a JSON-RPC message over stdio."""
+        if not self.process or not self.process.stdin:
+            raise RuntimeError("MCP server process is not running")
+        payload = json.dumps(message)
+        # We send standard newline-delimited JSON payload
+        self.process.stdin.write(f"{payload}\n")
+        self.process.stdin.flush()
+
+    def read_message(self, timeout: Optional[float] = None) -> Dict[str, Any]:
+        """Reads a JSON-RPC response, supporting newline-delimited and Content-Length framing."""
+        if not self.process or not self.process.stdout:
+            raise RuntimeError("MCP server process is not running")
+
+        wait_time = timeout if timeout is not None else self.timeout
+        start_time = time.time()
+        content_length = None
+
+        while time.time() - start_time < wait_time:
+            # Check if process died
+            if self.process.poll() is not None:
+                stderr_output = ""
+                if self.process.stderr:
+                    try:
+                        stderr_output = self.process.stderr.read()
+                    except Exception:
+                        pass
+                raise RuntimeError(
+                    f"Server process terminated unexpectedly with exit code {self.process.returncode}. Stderr: {stderr_output}"
+                )
+
+            # Use select on UNIX systems for non-blocking read
+            rlist, _, _ = select.select([self.process.stdout], [], [], 0.1)
+            if not rlist:
+                continue
+
+            line = self.process.stdout.readline()
+            if not line:
+                continue
+
+            # Check for HTTP/LSP style Content-Length header
+            stripped = line.strip()
+            if line.lower().startswith("content-length:"):
+                try:
+                    content_length = int(line.split(":", 1)[1].strip())
+                except ValueError:
+                    pass
+                continue
+            elif stripped == "" and content_length is not None:
+                # Header separator, read exact body
+                body = self.process.stdout.read(content_length)
+                return json.loads(body)
+            elif stripped.startswith("{") and stripped.endswith("}"):
+                # Direct newline-delimited JSON
+                try:
+                    return json.loads(stripped)
+                except json.JSONDecodeError:
+                    continue
+
+        # Timeout reached: collect stderr if available
+        stderr_info = ""
+        if self.process.stderr:
+            try:
+                rlist, _, _ = select.select([self.process.stderr], [], [], 0.05)
+                if rlist:
+                    stderr_info = self.process.stderr.read()
+            except Exception:
+                pass
+        raise TimeoutError(
+            f"Timed out after {wait_time:.1f}s waiting for MCP server response. Stderr: {stderr_info.strip()}"
+        )
+
+    def initialize(self) -> Dict[str, Any]:
+        """Performs the MCP initialize handshake and initialized notification."""
+        init_req = {
+            "jsonrpc": "2.0",
+            "id": self.msg_id,
+            "method": "initialize",
+            "params": {
+                "protocolVersion": PROTOCOL_VERSION,
+                "capabilities": {
+                    "roots": {"listChanged": True},
+                    "sampling": {}
+                },
+                "clientInfo": {
+                    "name": "aos-verification-test-suite",
+                    "version": "1.0.0"
+                }
+            }
+        }
+        self.send_message(init_req)
+        self.msg_id += 1
+
+        resp = self.read_message()
+        if "error" in resp:
+            raise RuntimeError(f"MCP initialize error: {resp['error']}")
+
+        result = resp.get("result", {})
+        self.server_info = result.get("serverInfo", {})
+        self.protocol_version = result.get("protocolVersion", "")
+        self.server_instructions = result.get("instructions", "")
+        self.capabilities = result.get("capabilities", {})
+
+        # Send initialized notification
+        init_notif = {
+            "jsonrpc": "2.0",
+            "method": "notifications/initialized",
+            "params": {}
+        }
+        self.send_message(init_notif)
+        return result
+
+    def list_tools(self) -> List[Dict[str, Any]]:
+        """Enumerates all tools exposed by the MCP server, with pagination support."""
+        all_tools: List[Dict[str, Any]] = []
+        cursor = None
+
+        while True:
+            params: Dict[str, Any] = {}
+            if cursor:
+                params["cursor"] = cursor
+
+            list_req = {
+                "jsonrpc": "2.0",
+                "id": self.msg_id,
+                "method": "tools/list",
+                "params": params
+            }
+            self.send_message(list_req)
+            self.msg_id += 1
+
+            resp = self.read_message()
+            if "error" in resp:
+                raise RuntimeError(f"tools/list error: {resp['error']}")
+
+            result = resp.get("result", {})
+            tools = result.get("tools", [])
+            all_tools.extend(tools)
+
+            cursor = result.get("nextCursor")
+            if not cursor:
+                break
+
+        return all_tools
+
+    def close(self):
+        """Cleanly terminates the MCP server process."""
+        if not self.process:
+            return
+        try:
+            if self.process.stdin:
+                try:
+                    self.process.stdin.close()
+                except Exception:
+                    pass
+            self.process.terminate()
+            self.process.wait(timeout=2.0)
+        except Exception:
+            try:
+                self.process.kill()
+            except Exception:
+                pass
+
+
+def validate_tool_schema(tool: Dict[str, Any]) -> Tuple[bool, str, int, int]:
+    """
+    Validates a tool definition against the MCP tool specification.
+    Returns: (is_valid, message, properties_count, required_count)
+    """
+    name = tool.get("name")
+    if not name or not isinstance(name, str) or not name.strip():
+        return False, "Missing or invalid 'name' field", 0, 0
+
+    desc = tool.get("description")
+    if desc is not None and not isinstance(desc, str):
+        return False, "Description must be a string", 0, 0
+
+    schema = tool.get("inputSchema")
+    if schema is None:
+        return False, "Missing 'inputSchema' definition", 0, 0
+
+    if not isinstance(schema, dict):
+        return False, "'inputSchema' must be a JSON object", 0, 0
+
+    schema_type = schema.get("type")
+    if schema_type != "object":
+        return False, f"Expected inputSchema.type == 'object', got '{schema_type}'", 0, 0
+
+    props = schema.get("properties", {})
+    if not isinstance(props, dict):
+        return False, "inputSchema.properties must be a dict", 0, 0
+
+    required = schema.get("required", [])
+    if not isinstance(required, list):
+        return False, "inputSchema.required must be a list", len(props), 0
+
+    return True, "Valid MCP Tool Schema", len(props), len(required)
+
+
+def parse_codex_toml(content: str) -> Dict[str, Any]:
+    """Parse Codex TOML configuration into a standard dict."""
+    if tomllib is not None:
+        return tomllib.loads(content)
+
+    # Minimal fallback parser for standard codex TOML files
+    result: Dict[str, Any] = {"mcp_servers": {}}
+    current_server = None
+    current_env = False
+
+    for raw_line in content.splitlines():
+        line = raw_line.strip()
+        if not line or line.startswith("#"):
+            continue
+
+        if line.startswith("[mcp_servers."):
+            section = line.strip("[]")
+            parts = section.split(".")
+            if len(parts) == 2:
+                current_server = parts[1]
+                current_env = False
+                result["mcp_servers"].setdefault(current_server, {"command": "", "args": [], "env": {}})
+            elif len(parts) == 3 and parts[2] == "env":
+                current_server = parts[1]
+                current_env = True
+                result["mcp_servers"].setdefault(current_server, {"command": "", "args": [], "env": {}})
+            continue
+
+        if current_server and "=" in line:
+            key, val = [x.strip() for x in line.split("=", 1)]
+            val = val.strip(' \t\n\r"')
+            if current_env:
+                result["mcp_servers"][current_server]["env"][key] = val
+            else:
+                if key == "command":
+                    result["mcp_servers"][current_server]["command"] = val
+                elif key == "args":
+                    try:
+                        result["mcp_servers"][current_server]["args"] = json.loads(val.replace("'", '"'))
+                    except Exception:
+                        result["mcp_servers"][current_server]["args"] = []
+                elif key == "enabled":
+                    result["mcp_servers"][current_server]["enabled"] = val.lower() == "true"
+
+    return result
+
+
+def export_antigravity_lazy_schemas(server_name: str, tools: List[Dict[str, Any]], output_dir: str):
+    """
+    Exports tools as individual JSON schema files for Google Antigravity lazy loading:
+    ~/.gemini/antigravity-cli/mcp/<serverName>/<toolName>.json
+    """
+    dest_dir = Path(output_dir) / server_name
+    dest_dir.mkdir(parents=True, exist_ok=True)
+
+    for tool in tools:
+        tool_name = tool["name"]
+        schema_file = dest_dir / f"{tool_name}.json"
+        schema_payload = {
+            "name": tool_name,
+            "description": tool.get("description", ""),
+            "parameters": tool.get("inputSchema", {
+                "type": "object",
+                "properties": {},
+                "required": []
+            })
+        }
+        with open(schema_file, "w", encoding="utf-8") as f:
+            json.dump(schema_payload, f, indent=2, ensure_ascii=False)
+
+    # Also create an instructions.md if none exists
+    instructions_file = dest_dir / "instructions.md"
+    if not instructions_file.exists():
+        with open(instructions_file, "w", encoding="utf-8") as f:
+            f.write(f"# {server_name} MCP Server Best Practices\n\n")
+            f.write(f"This directory contains lazy tool definitions for `{server_name}`.\n")
+            f.write(f"Total registered tools: {len(tools)}.\n")
+
+
+def run_verification_test(
+    server_name: str,
+    command: str,
+    args: Optional[List[str]] = None,
+    env: Optional[Dict[str, str]] = None,
+    timeout: float = 15.0,
+    export_dir: Optional[str] = None,
+) -> Dict[str, Any]:
+    """
+    Executes a complete verification cycle on a single MCP server.
+    """
+    args = args or []
+    env = env or {}
+    print(f"\n{'=' * 85}")
+    print(f"🚀 VERIFYING MCP SERVER: {server_name}")
+    print(f"Command : {command} {' '.join(args)}")
+    if env:
+        print(f"Env     : {env}")
+    print(f"{'=' * 85}")
+
+    client = MCPStdioClient(command, args, env, timeout=timeout)
+    test_result: Dict[str, Any] = {
+        "server_name": server_name,
+        "command": command,
+        "args": args,
+        "status": "FAILED",
+        "error": None,
+        "server_info": {},
+        "protocol_version": None,
+        "tool_count": 0,
+        "tools": [],
+        "latency_ms": 0.0,
+    }
+
+    t0 = time.time()
+    try:
+        # Step 1: Subprocess Startup
+        print("[1/4] Spawning server process over stdio...")
+        client.start()
+        print("      ✓ Process launched successfully (PID: {})".format(client.process.pid if client.process else "unknown"))
+
+        # Step 2: Initialize Handshake
+        print("[2/4] Executing JSON-RPC 'initialize' handshake...")
+        init_result = client.initialize()
+        test_result["server_info"] = client.server_info
+        test_result["protocol_version"] = client.protocol_version
+        print(f"      ✓ Handshake acknowledged.")
+        print(f"        Server Name    : {client.server_info.get('name', 'unknown')}")
+        print(f"        Server Version : {client.server_info.get('version', 'unknown')}")
+        print(f"        Protocol Ver   : {client.protocol_version}")
+
+        # Step 3: Tool Enumeration
+        print("[3/4] Enumerating tool registry via 'tools/list'...")
+        tools = client.list_tools()
+        dt_ms = (time.time() - t0) * 1000.0
+        test_result["tool_count"] = len(tools)
+        test_result["latency_ms"] = round(dt_ms, 2)
+        print(f"      ✓ Received {len(tools)} tools in {dt_ms:.1f}ms")
+
+        # Step 4: Tool Schema Validation
+        print("[4/4] Validating tool schema definitions...")
+        all_schemas_valid = True
+        tool_records = []
+
+        print("\n" + "-" * 85)
+        print(f"{'#':<3} | {'TOOL NAME':<32} | {'PROPS':<5} | {'REQ':<4} | {'STATUS':<7} | {'DESCRIPTION':<24}")
+        print("-" * 85)
+
+        for idx, tool in enumerate(tools, 1):
+            is_valid, msg, prop_cnt, req_cnt = validate_tool_schema(tool)
+            if not is_valid:
+                all_schemas_valid = False
+                status_str = "FAIL"
+            else:
+                status_str = "PASS"
+
+            desc = (tool.get("description") or "").replace("\n", " ").strip()
+            if len(desc) > 23:
+                desc = desc[:21] + "..."
+
+            tool_name = tool.get("name", f"tool_{idx}")
+            if len(tool_name) > 31:
+                tool_name = tool_name[:29] + ".."
+
+            print(f"{idx:<3} | {tool_name:<32} | {prop_cnt:<5} | {req_cnt:<4} | {status_str:<7} | {desc:<24}")
+
+            tool_records.append({
+                "name": tool.get("name"),
+                "description": tool.get("description"),
+                "schema_valid": is_valid,
+                "validation_message": msg,
+                "property_count": prop_cnt,
+                "required_count": req_cnt,
+                "inputSchema": tool.get("inputSchema"),
+            })
+
+        print("-" * 85)
+
+        if not all_schemas_valid:
+            raise ToolSchemaValidationError("One or more tool schemas failed MCP validation.")
+
+        # Optional Antigravity lazy export
+        if export_dir:
+            print(f"\n[+] Exporting lazy tool schemas for Antigravity to: {export_dir}/{server_name}/")
+            export_antigravity_lazy_schemas(server_name, tools, export_dir)
+            print(f"    ✓ {len(tools)} schema JSON files written.")
+
+        test_result["status"] = "PASSED"
+        test_result["tools"] = tool_records
+        print(f"\n✅ MCP SERVER '{server_name}' VERIFICATION PASSED ({len(tools)} tools validated)\n")
+
+    except Exception as exc:
+        test_result["status"] = "FAILED"
+        test_result["error"] = str(exc)
+        print(f"\n❌ MCP SERVER '{server_name}' VERIFICATION FAILED: {exc}\n", file=sys.stderr)
+    finally:
+        client.close()
+
+    return test_result
+
+
+def load_servers_from_config(config_path: str) -> Dict[str, Dict[str, Any]]:
+    """Loads server definitions from either JSON or TOML config file."""
+    path = Path(config_path)
+    if not path.exists():
+        raise FileNotFoundError(f"Configuration file not found: {config_path}")
+
+    content = path.read_text(encoding="utf-8")
+    servers: Dict[str, Dict[str, Any]] = {}
+
+    if path.suffix.lower() == ".toml":
+        parsed = parse_codex_toml(content)
+        raw_servers = parsed.get("mcp_servers", {})
+        for name, cfg in raw_servers.items():
+            if cfg.get("enabled", True):
+                servers[name] = {
+                    "command": cfg.get("command", ""),
+                    "args": cfg.get("args", []),
+                    "env": cfg.get("env", {})
+                }
+    else:
+        parsed = json.loads(content)
+        raw_servers = parsed.get("mcpServers", {})
+        for name, cfg in raw_servers.items():
+            servers[name] = {
+                "command": cfg.get("command", ""),
+                "args": cfg.get("args", []),
+                "env": cfg.get("env", {})
+            }
+
+    return servers
+
+
+def main():
+    parser = argparse.ArgumentParser(
+        description="AOS MCP Multi-Harness Connection & Verification Test Suite"
+    )
+    parser.add_argument("--config", help="Path to config file (JSON or TOML)")
+    parser.add_argument("--server", help="Specific server name to test from config")
+    parser.add_argument("--command", help="Direct command to run over stdio")
+    parser.add_argument("--arg", dest="args", action="append", default=[], help="Command arguments")
+    parser.add_argument("--env", action="append", default=[], help="Environment variable (KEY=VALUE)")
+    parser.add_argument("--timeout", type=float, default=15.0, help="Response timeout in seconds (default: 15.0)")
+    parser.add_argument("--export-antigravity", help="Directory to export Antigravity lazy tool JSON schemas")
+    parser.add_argument("--json-report", help="Path to write JSON test report")
+    parsed = parser.parse_args()
+
+    custom_env = {}
+    for item in parsed.env:
+        if "=" in item:
+            k, v = item.split("=", 1)
+            custom_env[k] = v
+
+    results: List[Dict[str, Any]] = []
+
+    if parsed.config:
+        servers = load_servers_from_config(parsed.config)
+        if parsed.server:
+            if parsed.server not in servers:
+                print(f"Error: Server '{parsed.server}' not found in {parsed.config}", file=sys.stderr)
+                print(f"Available servers: {list(servers.keys())}", file=sys.stderr)
+                sys.exit(1)
+            servers = {parsed.server: servers[parsed.server]}
+
+        for sname, scfg in servers.items():
+            merged_env = {**scfg.get("env", {}), **custom_env}
+            res = run_verification_test(
+                server_name=sname,
+                command=scfg["command"],
+                args=scfg.get("args", []),
+                env=merged_env,
+                timeout=parsed.timeout,
+                export_dir=parsed.export_antigravity,
+            )
+            results.append(res)
+
+    elif parsed.command:
+        res = run_verification_test(
+            server_name="direct-cli",
+            command=parsed.command,
+            args=parsed.args,
+            env=custom_env,
+            timeout=parsed.timeout,
+            export_dir=parsed.export_antigravity,
+        )
+        results.append(res)
+    else:
+        # Default self-test using project's default runner scripts
+        repo_dir = Path(__file__).resolve().parent.parent
+        kicad_runner = str(repo_dir / "scripts" / "run_kicad_mcp.sh")
+        openscad_runner = str(repo_dir / "scripts" / "run_openscad_mcp.sh")
+
+        print("No --config or --command provided. Running default verification on KiCad and OpenSCAD runners...")
+        res_kicad = run_verification_test(
+            server_name="kicad",
+            command=kicad_runner,
+            timeout=parsed.timeout,
+            export_dir=parsed.export_antigravity,
+        )
+        results.append(res_kicad)
+
+        res_openscad = run_verification_test(
+            server_name="openscad",
+            command=openscad_runner,
+            timeout=parsed.timeout,
+            export_dir=parsed.export_antigravity,
+        )
+        results.append(res_openscad)
+
+    # Print Overall Summary
+    total = len(results)
+    passed = sum(1 for r in results if r["status"] == "PASSED")
+    failed = total - passed
+
+    print("\n" + "=" * 85)
+    print("📋 MCP VERIFICATION TEST SUITE SUMMARY")
+    print("=" * 85)
+    for r in results:
+        sym = "✅ PASS" if r["status"] == "PASSED" else "❌ FAIL"
+        tools_str = f"{r['tool_count']} tools" if r['status'] == "PASSED" else f"Error: {r['error']}"
+        print(f"[{sym}] {r['server_name']:<20} | {tools_str:<35} | {r['latency_ms']}ms")
+    print("-" * 85)
+    print(f"Total: {total} | Passed: {passed} | Failed: {failed}")
+    print("=" * 85 + "\n")
+
+    if parsed.json_report:
+        report_path = Path(parsed.json_report)
+        report_path.parent.mkdir(parents=True, exist_ok=True)
+        with open(report_path, "w", encoding="utf-8") as f:
+            json.dump({
+                "timestamp": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()),
+                "total": total,
+                "passed": passed,
+                "failed": failed,
+                "results": results
+            }, f, indent=2)
+        print(f"Wrote machine-readable test report to: {parsed.json_report}")
+
+    sys.exit(0 if failed == 0 else 1)
+
+
+if __name__ == "__main__":
+    main()