hanzoai/python-sdk
2
0
Fork 0
mirror of https://github.com/hanzoai/python-sdk.git synced 2026-10-23 20:23:00 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix(zap): replace websockets with native ZAP TCP transport

Browse Source 
- Rewrite zap_server.py to use asyncio TCP (not websockets library)
- ZAP server now speaks native binary protocol over raw TCP
- Wire format: [magic:4][type:1][length:4 BE][JSON payload]
- Remove websockets from ide optional dependency
- Fix type annotations: any -> Any, dict -> dict[str, Any]
- Add proper frame reading with asyncio.StreamReader.readexactly()
- Support MAX_MESSAGE_SIZE (16MB) limit per ZAP spec
This commit is contained in:
zeekay
2026-03-12 16:23:41 -07:00
parent 367b604e40
commit b4c4e5b796
3 changed files with 118 additions and 103 deletions
Download Patch File Download Diff File Expand all files Collapse all files
pkg/hanzo-mcp/hanzo_mcp/server.py
+4 -4
View File
@@ -13,7 +13,7 @@ import secrets
import signal
import threading
import warnings
from typing import TYPE_CHECKING, Literal, cast, final
from typing import TYPE_CHECKING, Any, Literal, cast, final
# Type-only imports - don't execute at runtime
if TYPE_CHECKING:
@@ -355,7 +355,7 @@ class HanzoMCPServer:
return
# Collect tool manifest from registered MCP tools
tool_list: list[dict] = []
tool_list: list[dict[str, Any]] = []
try:
# FastMCP _tool_manager._tools is a dict[str, Tool]
tm = getattr(self.mcp, "_tool_manager", None)
@@ -387,7 +387,7 @@ class HanzoMCPServer:
except Exception:
pass
async def call_tool(name: str, args: dict) -> any:
async def call_tool(name: str, args: dict) -> Any: # type: ignore[type-arg]
"""Route ZAP tool calls to the MCP server."""
try:
result = await self.mcp.call_tool(name, args)
@@ -395,7 +395,7 @@ class HanzoMCPServer:
except Exception as e:
return {"error": str(e)}
async def handle_method(method: str, params: any) -> any:
async def handle_method(method: str, params: Any) -> Any:
"""Pass-through for ALL MCP methods — full protocol parity over ZAP.
Routes resources/*, prompts/*, and any other MCP method to the
pkg/hanzo-mcp/hanzo_mcp/zap_server.py
+114 -98
View File
@@ -1,13 +1,15 @@
"""ZAP (Zero-latency Agent Protocol) Server for hanzo-mcp.
"""ZAP (Zero-copy Agent Protocol) Server for hanzo-mcp.
Allows Hanzo browser extensions to discover this MCP server
and call tools directly over a binary WebSocket protocol.
Allows Hanzo browser extensions and agents to discover this MCP server
and call tools directly over the native ZAP binary protocol.
Protocol: 9-byte header + JSON payload
Wire format: 9-byte header + payload
[0x5A 0x41 0x50 0x01] magic "ZAP\x01"
[type] 1 byte message type
[length] 4 bytes big-endian payload length
[payload] UTF-8 JSON
[payload] UTF-8 JSON (Cap'n Proto in future)
Transport: Raw TCP on ports 9999-9995 (first available).
"""
from __future__ import annotations
@@ -27,11 +29,13 @@ MSG_HANDSHAKE = 0x01
MSG_HANDSHAKE_OK = 0x02
MSG_REQUEST = 0x10
MSG_RESPONSE = 0x11
MSG_STREAM = 0x20
MSG_PING = 0xFE
MSG_PONG = 0xFF
ZAP_PORTS = [9999, 9998, 9997, 9996, 9995]
HEADER_SIZE = 9 # 4 (magic) + 1 (type) + 4 (length)
MAX_MESSAGE_SIZE = 16 * 1024 * 1024 # 16 MB
# ── Encode / Decode ─────────────────────────────────────────────────────
@@ -45,47 +49,23 @@ def encode(msg_type: int, payload: Any) -> bytes:
def decode(buf: bytes) -> tuple[int, Any] | None:
"""Decode a ZAP message. Supports both MCP ZAP and hanzo/dev wire formats.
"""Decode a ZAP message from a complete frame.
Format 1 (MCP ZAP): [magic:4 "ZAP\\x01"][type:1][length:4 BE][JSON]
Format 2 (hanzo/dev): [length:4 LE][type:1][payload]
Format 3: Plain JSON fallback
Expects exactly one ZAP frame: [magic:4][type:1][length:4 BE][JSON].
"""
if len(buf) < 5:
if len(buf) < HEADER_SIZE:
return None
# Format 1: MCP ZAP — magic header + big-endian length
if buf[:4] == ZAP_MAGIC:
if len(buf) < HEADER_SIZE:
return None
msg_type = buf[4]
length = struct.unpack("!L", buf[5:9])[0]
if len(buf) < HEADER_SIZE + length:
return None
try:
payload = json.loads(
buf[HEADER_SIZE : HEADER_SIZE + length].decode("utf-8")
)
return (msg_type, payload)
except (json.JSONDecodeError, UnicodeDecodeError):
return None
if buf[:4] != ZAP_MAGIC:
return None
# Format 2: hanzo/dev ZAP — little-endian length, no magic
le_length = struct.unpack("<L", buf[:4])[0]
if 0 < le_length <= 16 * 1024 * 1024 and len(buf) >= 5 + le_length:
msg_type = buf[4]
if msg_type <= 0x45 or msg_type >= 0xFE:
payload_buf = buf[5 : 5 + le_length]
try:
payload = json.loads(payload_buf.decode("utf-8"))
return (msg_type, payload)
except (json.JSONDecodeError, UnicodeDecodeError):
return (msg_type, {"raw": payload_buf})
# Format 3: Plain JSON fallback
msg_type = buf[4]
length = struct.unpack("!L", buf[5:9])[0]
if len(buf) < HEADER_SIZE + length:
return None
try:
payload = json.loads(buf.decode("utf-8"))
return (MSG_REQUEST, payload)
payload = json.loads(buf[HEADER_SIZE : HEADER_SIZE + length].decode("utf-8"))
return (msg_type, payload)
except (json.JSONDecodeError, UnicodeDecodeError):
return None
@@ -102,7 +82,7 @@ class ZapClient:
client_id: str = "unknown",
browser: str = "unknown",
version: str = "0",
):
) -> None:
self.writer = writer
self.client_id = client_id
self.browser = browser
@@ -110,15 +90,15 @@ class ZapClient:
self.connected_at = time.time()
# ── WebSocket-like framing over raw TCP ─────────────────────────────────
# We use asyncio websockets for the actual WS server.
# ── ZAP TCP Server ──────────────────────────────────────────────────────
class ZapServer:
"""ZAP WebSocket server for browser extension discovery.
"""ZAP TCP server for agent and browser extension discovery.
Provides full MCP protocol parity — any MCP method (tools/*, resources/*, prompts/*)
can be called over the ZAP binary transport via the handle_method pass-through.
Provides full MCP protocol parity — any MCP method (tools/*, resources/*,
prompts/*) can be called over the ZAP binary transport via the
handle_method pass-through.
"""
def __init__(
@@ -127,14 +107,14 @@ class ZapServer:
call_tool: Callable[[str, dict[str, Any]], Awaitable[Any]],
handle_method: Callable[[str, Any], Awaitable[Any]] | None = None,
name: str = "hanzo-mcp",
):
) -> None:
self.tools = tools
self.call_tool = call_tool
self.handle_method = handle_method
self.name = name
self.server_id = f"mcp-{int(time.time()):x}"
self.clients: dict[Any, ZapClient] = {}
self._server: Any = None
self.clients: dict[str, ZapClient] = {}
self._server: asyncio.Server | None = None
self.port: int | None = None
self.tool_manifest = [
@@ -146,17 +126,56 @@ class ZapServer:
for t in self.tools
]
async def _handle_connection(self, websocket: Any) -> None:
"""Handle a single WebSocket connection."""
client: ZapClient | None = None
async def _read_frame(
self, reader: asyncio.StreamReader
) -> tuple[int, Any] | None:
"""Read a single ZAP frame from the stream."""
try:
async for raw in websocket:
if isinstance(raw, str):
raw = raw.encode("utf-8")
header = await reader.readexactly(HEADER_SIZE)
except (asyncio.IncompleteReadError, ConnectionError):
return None
result = decode(raw)
if header[:4] != ZAP_MAGIC:
return None
msg_type = header[4]
length = struct.unpack("!L", header[5:9])[0]
if length > MAX_MESSAGE_SIZE:
return None
try:
payload_bytes = await reader.readexactly(length)
except (asyncio.IncompleteReadError, ConnectionError):
return None
try:
payload = json.loads(payload_bytes.decode("utf-8"))
return (msg_type, payload)
except (json.JSONDecodeError, UnicodeDecodeError):
return None
async def _send_frame(
self, writer: asyncio.StreamWriter, msg_type: int, payload: Any
) -> None:
"""Send a ZAP frame to the stream."""
frame = encode(msg_type, payload)
writer.write(frame)
await writer.drain()
async def _handle_connection(
self, reader: asyncio.StreamReader, writer: asyncio.StreamWriter
) -> None:
"""Handle a single TCP connection."""
peer = writer.get_extra_info("peername")
client: ZapClient | None = None
client_key = f"{peer[0]}:{peer[1]}" if peer else f"unknown-{time.time()}"
try:
while True:
result = await self._read_frame(reader)
if result is None:
continue
break
msg_type, payload = result
@@ -165,44 +184,48 @@ class ZapServer:
browser = payload.get("browser", "unknown")
version = payload.get("version", "0")
client = ZapClient(
writer=websocket,
writer=writer,
client_id=client_id,
browser=browser,
version=version,
)
self.clients[websocket] = client
self.clients[client_key] = client
logger.info(
f"[ZAP] Client connected: {client_id} ({browser} v{version})"
)
await websocket.send(
encode(
MSG_HANDSHAKE_OK,
{
"serverId": self.server_id,
"name": self.name,
"tools": self.tool_manifest,
},
)
await self._send_frame(
writer,
MSG_HANDSHAKE_OK,
{
"serverId": self.server_id,
"name": self.name,
"tools": self.tool_manifest,
},
)
elif msg_type == MSG_REQUEST:
req_id = payload.get("id", "")
method = payload.get("method", "")
params = payload.get("params", {})
await self._handle_request(websocket, req_id, method, params)
await self._handle_request(writer, req_id, method, params)
elif msg_type == MSG_PING:
await websocket.send(encode(MSG_PONG, {}))
await self._send_frame(writer, MSG_PONG, {})
except Exception as e:
logger.debug(f"[ZAP] Connection error: {e}")
finally:
if websocket in self.clients:
client = self.clients.pop(websocket)
if client_key in self.clients:
client = self.clients.pop(client_key)
logger.info(f"[ZAP] Client disconnected: {client.client_id}")
writer.close()
try:
await writer.wait_closed()
except Exception:
pass
async def _handle_request(
self, websocket: Any, req_id: str, method: str, params: Any
self, writer: asyncio.StreamWriter, req_id: str, method: str, params: Any
) -> None:
"""Handle a ZAP RPC request."""
try:
@@ -217,7 +240,7 @@ class ZapServer:
result = await self.call_tool(tool_name, args)
elif method.startswith("notifications/"):
# Client→server notifications (fire-and-forget)
# Client -> server notifications (fire-and-forget)
result = {"acknowledged": True}
else:
@@ -228,30 +251,22 @@ class ZapServer:
else:
raise ValueError(f"Unsupported method: {method}")
await websocket.send(encode(MSG_RESPONSE, {"id": req_id, "result": result}))
await self._send_frame(
writer, MSG_RESPONSE, {"id": req_id, "result": result}
)
except Exception as e:
await websocket.send(
encode(
MSG_RESPONSE,
{
"id": req_id,
"error": {"code": -1, "message": str(e)},
},
)
await self._send_frame(
writer,
MSG_RESPONSE,
{
"id": req_id,
"error": {"code": -1, "message": str(e)},
},
)
async def start(self, preferred_port: int | None = None) -> bool:
"""Start the ZAP server on the first available port."""
try:
import websockets
except ImportError:
logger.warning(
"[ZAP] websockets package not installed. "
"Install with: pip install websockets"
)
return False
"""Start the ZAP server on the first available TCP port."""
ports = (
[preferred_port, *[p for p in ZAP_PORTS if p != preferred_port]]
if preferred_port
@@ -260,14 +275,14 @@ class ZapServer:
for port in ports:
try:
self._server = await websockets.serve(
self._server = await asyncio.start_server(
self._handle_connection,
"127.0.0.1",
port,
)
self.port = port
logger.info(
f"[ZAP] Server listening on ws://127.0.0.1:{port} "
f"[ZAP] Server listening on zap://127.0.0.1:{port} "
f"({len(self.tool_manifest)} tools)"
)
return True
@@ -285,9 +300,9 @@ class ZapServer:
self._server.close()
self._server = None
self.port = None
for ws in list(self.clients.keys()):
for client in list(self.clients.values()):
try:
asyncio.ensure_future(ws.close())
client.writer.close()
except Exception:
pass
self.clients.clear()
@@ -305,7 +320,8 @@ async def start_zap_server(
Args:
tools: List of tool dicts with name, description, inputSchema.
call_tool: Async callable (name, args) -> result to route tool calls.
handle_method: Optional async callable (method, params) -> result for full MCP parity.
handle_method: Optional async callable (method, params) -> result for
full MCP parity.
name: Server name for handshake.
preferred_port: Preferred port to try first.
pkg/hanzo-mcp/pyproject.toml
-1
View File
@@ -90,7 +90,6 @@ repl = [
# IDE integration (VS Code, Cursor, JetBrains)
ide = [
"hanzo-tools-ide>=0.1.0",
"websockets>=12.0",
]
# Full interactive suite (REPL + IDE + Browser)