feat: add FastAPI manager API compatibility baseline

2026-07-24 08:03:53 +08:00 · 2026-07-20 17:00:13 +08:00
parent 7c58fa37b2
commit 804ddb51f2
140 changed files with 47169 additions and 1 deletions
@@ -0,0 +1,106 @@
+from __future__ import annotations
+
+import asyncio
+import base64
+import hashlib
+import json
+from typing import Any
+from urllib.parse import quote_plus, urlsplit, urlunsplit
+
+from cryptography.hazmat.primitives import padding
+from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes
+from websockets.asyncio.client import connect
+
+
+def _java_aes_key(value: str) -> bytes:
+    raw = value.encode("utf-8")
+    if len(raw) in {16, 24, 32}:
+        return raw
+    return raw[:32].ljust(32, b"\x00")
+
+
+def encrypt_agent_token(agent_id: str, key: str) -> str:
+    digest = hashlib.md5(agent_id.encode("utf-8"), usedforsecurity=False).hexdigest()
+    plain_text = f'{{"agentId": "{digest}"}}'.encode()
+    padder = padding.PKCS7(128).padder()
+    padded = padder.update(plain_text) + padder.finalize()
+    # ECB is required for byte-for-byte compatibility with Java AES/ECB/PKCS5Padding.
+    encryptor = Cipher(algorithms.AES(_java_aes_key(key)), modes.ECB()).encryptor()  # noqa: S305
+    encrypted = encryptor.update(padded) + encryptor.finalize()
+    return base64.b64encode(encrypted).decode("ascii")
+
+
+def build_agent_mcp_address(endpoint: str | None, agent_id: str) -> str | None:
+    if endpoint is None or not endpoint.strip() or endpoint == "null":
+        return None
+    parsed = urlsplit(endpoint)
+    if not parsed.scheme or not parsed.netloc:
+        raise ValueError("mcp的地址存在错误，请进入参数管理修改mcp接入点地址")
+    marker = "key="
+    marker_index = parsed.query.find(marker)
+    # Java takes everything following the first key= marker, including subsequent query text.
+    key = parsed.query[marker_index + len(marker) :] if marker_index >= 0 else parsed.query[3:]
+    ws_scheme = "wss" if parsed.scheme == "https" else "ws"
+    path = parsed.path
+    parent = path[: path.rfind("/")] if "/" in path else ""
+    base = urlunsplit((ws_scheme, parsed.netloc, parent, "", "")).rstrip("/")
+    token = quote_plus(encrypt_agent_token(agent_id, key), safe="")
+    return f"{base}/mcp/?token={token}"
+
+
+INITIALIZE_REQUEST = {
+    "jsonrpc": "2.0",
+    "method": "initialize",
+    "params": {
+        "protocolVersion": "2024-11-05",
+        "capabilities": {"roots": {"listChanged": False}, "sampling": {}},
+        "clientInfo": {"name": "xz-mcp-broker", "version": "0.0.1"},
+    },
+    "id": 1,
+}
+INITIALIZED_NOTIFICATION = {"jsonrpc": "2.0", "method": "notifications/initialized"}
+TOOLS_REQUEST = {"jsonrpc": "2.0", "method": "tools/list", "params": None, "id": 2}
+
+
+async def _receive_matching(websocket: Any, request_id: int, timeout: float) -> dict[str, Any] | None:
+    async def receive() -> dict[str, Any] | None:
+        async for message in websocket:
+            try:
+                value = json.loads(message)
+            except (TypeError, json.JSONDecodeError):
+                continue
+            if isinstance(value, dict) and value.get("id") == request_id:
+                return value
+        return None
+
+    return await asyncio.wait_for(receive(), timeout=timeout)
+
+
+async def list_mcp_tools(address: str, *, connect_timeout: float = 8.0, session_timeout: float = 10.0) -> list[str]:
+    call_address = address.replace("/mcp/", "/call/")
+    try:
+        async with connect(
+            call_address,
+            open_timeout=connect_timeout,
+            max_size=1024 * 1024,
+            close_timeout=1,
+        ) as websocket:
+            await websocket.send(json.dumps(INITIALIZE_REQUEST, ensure_ascii=False, separators=(",", ":")))
+            initialized = await _receive_matching(websocket, 1, session_timeout)
+            if not initialized or "result" not in initialized or "error" in initialized:
+                return []
+            await websocket.send(json.dumps(INITIALIZED_NOTIFICATION, separators=(",", ":")))
+            await websocket.send(json.dumps(TOOLS_REQUEST, separators=(",", ":")))
+            response = await _receive_matching(websocket, 2, session_timeout)
+            if not response or "error" in response:
+                return []
+            result = response.get("result")
+            tools = result.get("tools") if isinstance(result, dict) else None
+            if not isinstance(tools, list):
+                return []
+            return sorted(
+                item["name"] for item in tools if isinstance(item, dict) and isinstance(item.get("name"), str)
+            )
+    # Java treats every connect/protocol/parse failure as an empty tool list.
+    except Exception:
+        return []