from __future__ import annotations
import argparse
import hashlib
import json
import pathlib
import sys
from typing import Any
from mcp_safety_check import McpClient, package_version
def canonicalize(value: Any) -> Any:
if isinstance(value, dict):
return {key: canonicalize(value[key]) for key in sorted(value)}
if isinstance(value, list):
return [canonicalize(item) for item in value]
return value
def catalog_entry(tool: dict[str, Any]) -> dict[str, Any]:
return {
"name": tool["name"],
"description": tool.get("description") or "",
"inputSchema": tool.get("inputSchema") or {"type": "object"},
"outputSchema": tool.get("outputSchema"),
"annotations": tool.get("annotations") or {},
}
def runtime_entry(tool: dict[str, Any]) -> dict[str, Any]:
return {
"name": tool["name"],
"description": tool.get("description") or "",
"inputSchema": tool.get("inputSchema") or {"type": "object"},
"annotations": tool.get("annotations") or {},
}
def catalog_digest(catalog: list[dict[str, Any]]) -> str:
return hashlib.sha256(
json.dumps(
canonicalize(catalog),
ensure_ascii=False,
separators=(",", ":"),
).encode("utf-8")
).hexdigest()
def render_catalog(
version: str,
tools: list[dict[str, Any]],
optional_tools: list[dict[str, Any]],
) -> str:
catalog = sorted((catalog_entry(tool) for tool in tools), key=lambda tool: tool["name"])
optional_catalog = sorted(
(catalog_entry(tool) for tool in optional_tools), key=lambda tool: tool["name"]
)
canonical_catalog = canonicalize(catalog)
canonical_optional_catalog = canonicalize(optional_catalog)
digest = catalog_digest(canonical_catalog)
shell_digest = catalog_digest(
sorted(canonical_catalog + canonical_optional_catalog, key=lambda tool: tool["name"])
)
runtime_catalog = canonicalize(
sorted((runtime_entry(tool) for tool in tools), key=lambda tool: tool["name"])
)
runtime_optional_catalog = canonicalize(
sorted(
(runtime_entry(tool) for tool in optional_tools),
key=lambda tool: tool["name"],
)
)
catalog_json = json.dumps(
runtime_catalog,
ensure_ascii=False,
indent=2,
sort_keys=True,
)
optional_catalog_json = json.dumps(
runtime_optional_catalog,
ensure_ascii=False,
indent=2,
sort_keys=True,
)
return f"""\
// Generated by scripts/generate_pi_tool_catalog.py. Do not edit by hand.
export interface GeneratedMcpToolDefinition {{
\tname: string;
\tdescription: string;
\tinputSchema: Record<string, unknown>;
\tannotations: Record<string, unknown>;
}}
export const GENERATED_SERVER_VERSION = {json.dumps(version)};
export const GENERATED_TOOL_CATALOG_HASH = {json.dumps(digest)};
export const GENERATED_SHELL_TOOL_CATALOG_HASH = {json.dumps(shell_digest)};
export const GENERATED_MCP_TOOLS =
{catalog_json} as const satisfies readonly GeneratedMcpToolDefinition[];
export const GENERATED_OPTIONAL_MCP_TOOLS =
{optional_catalog_json} as const satisfies readonly GeneratedMcpToolDefinition[];
"""
def read_tools(
binary: pathlib.Path, extra_env: dict[str, str] | None = None
) -> list[dict[str, Any]]:
client = McpClient(binary, extra_env)
try:
client.request(
"initialize",
{
"protocolVersion": "2024-11-05",
"capabilities": {},
"clientInfo": {
"name": "computer-use-linux-pi-catalog",
"version": "0",
},
},
)
client.notify("notifications/initialized", {})
result = client.request("tools/list", {})["result"]
tools = result.get("tools") or []
if not isinstance(tools, list):
raise AssertionError("tools/list did not return a tools array")
return tools
finally:
client.close()
def main() -> int:
parser = argparse.ArgumentParser()
parser.add_argument("--binary", default="target/debug/computer-use-linux")
parser.add_argument("--repo", default=".")
parser.add_argument("--output", default="pi/extension/generated-tools.ts")
parser.add_argument("--check", action="store_true")
args = parser.parse_args()
repo = pathlib.Path(args.repo).resolve()
binary = pathlib.Path(args.binary).resolve()
output = pathlib.Path(args.output)
if not output.is_absolute():
output = repo / output
if not binary.exists():
raise AssertionError(f"binary does not exist: {binary}")
default_tools = read_tools(binary)
shell_tools = read_tools(binary, {"COMPUTER_USE_LINUX_ENABLE_SHELL": "1"})
default_by_name = {tool["name"]: catalog_entry(tool) for tool in default_tools}
shell_by_name = {tool["name"]: catalog_entry(tool) for tool in shell_tools}
changed_common = {
name
for name, tool in default_by_name.items()
if shell_by_name.get(name) != tool
}
if changed_common:
raise AssertionError(
f"shell opt-in changed existing tool definitions: {sorted(changed_common)}"
)
optional_names = sorted(set(shell_by_name) - set(default_by_name))
if optional_names != ["run_shell"]:
raise AssertionError(
f"unexpected shell-only tools: {optional_names}"
)
optional_tools = [
tool for tool in shell_tools if tool.get("name") in optional_names
]
rendered = render_catalog(
package_version(repo),
default_tools,
optional_tools,
)
if args.check:
if not output.exists():
raise AssertionError(f"generated catalog does not exist: {output}")
current = output.read_text(encoding="utf-8")
if current != rendered:
raise AssertionError(
"Pi tool catalog is stale; run "
f"{pathlib.Path(__file__).name} --binary {binary} --output {output}"
)
print(f"Pi tool catalog is current: {output}")
return 0
output.parent.mkdir(parents=True, exist_ok=True)
output.write_text(rendered, encoding="utf-8")
print(f"Wrote Pi tool catalog: {output}")
return 0
if __name__ == "__main__":
try:
raise SystemExit(main())
except Exception as exc:
print(f"generate_pi_tool_catalog.py: {exc}", file=sys.stderr)
raise SystemExit(1)