mirror of
https://github.com/hanzoai/python-sdk.git
synced 2026-10-23 20:23:00 +00:00
feat: add CloudClient and rewrite wire protocol for Rust compatibility
- Add CloudClient for luxfi/zap binary protocol (TCP + auto-TLS) - Rewrite wire.py Builder/ObjectBuilder to match Rust single-buffer architecture - Fix response field offsets (CLOUD_RESP_BODY=4, CLOUD_RESP_ERROR=12) - Use signed i32 relative offsets from absolute positions - 28 unit tests + cross-language integration test passing
This commit is contained in:
@@ -30,6 +30,7 @@ from .types import (
|
||||
)
|
||||
from .client import ZapClient
|
||||
from .server import ZapServer
|
||||
from .cloud import CloudClient
|
||||
from .playground import PlaygroundClient
|
||||
|
||||
__version__ = "0.7.0"
|
||||
@@ -37,6 +38,8 @@ __all__ = [
|
||||
# Wire protocol
|
||||
"ZapClient",
|
||||
"ZapServer",
|
||||
# Cloud (luxfi/zap binary protocol — compatible with Rust hanzo-zap)
|
||||
"CloudClient",
|
||||
# Playground
|
||||
"PlaygroundClient",
|
||||
# Core types
|
||||
|
||||
@@ -0,0 +1,175 @@
|
||||
"""
|
||||
ZAP Cloud Client — speaks the luxfi/zap binary wire protocol.
|
||||
|
||||
Connects to Hanzo Node (port 3692) or Engine via native binary transport.
|
||||
Compatible with the Rust `hanzo-zap` crate server.
|
||||
|
||||
Example:
|
||||
>>> from hanzo_zap import CloudClient
|
||||
>>> async with CloudClient.connect("localhost:3692") as client:
|
||||
... status, body, error = await client.call("chat.completions", "", body_bytes)
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
import ssl
|
||||
import struct
|
||||
from typing import Any
|
||||
|
||||
from .wire import (
|
||||
REQ_FLAG_REQ,
|
||||
REQ_FLAG_RESP,
|
||||
Message,
|
||||
build_cloud_request,
|
||||
build_handshake,
|
||||
parse_cloud_response,
|
||||
parse_handshake,
|
||||
read_frame,
|
||||
write_frame,
|
||||
)
|
||||
|
||||
DEFAULT_ENDPOINT = "localhost:3692"
|
||||
CLIENT_NODE_ID = "python-sdk"
|
||||
|
||||
|
||||
class CloudClient:
|
||||
"""
|
||||
A client that speaks the luxfi/zap binary wire protocol.
|
||||
|
||||
Supports both plain TCP (localhost) and TLS (remote endpoints).
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
reader: asyncio.StreamReader,
|
||||
writer: asyncio.StreamWriter,
|
||||
peer_id: str,
|
||||
) -> None:
|
||||
self._reader = reader
|
||||
self._writer = writer
|
||||
self._peer_id = peer_id
|
||||
self._req_id = 0
|
||||
|
||||
@classmethod
|
||||
async def connect(
|
||||
cls,
|
||||
endpoint: str | None = None,
|
||||
*,
|
||||
use_tls: bool | None = None,
|
||||
node_id: str = CLIENT_NODE_ID,
|
||||
) -> "CloudClient":
|
||||
"""
|
||||
Connect to a ZAP endpoint, perform handshake.
|
||||
|
||||
Args:
|
||||
endpoint: "host:port" (default: localhost:3692)
|
||||
use_tls: Force TLS on/off. None = auto-detect (TLS for non-localhost).
|
||||
node_id: Client node ID for handshake.
|
||||
"""
|
||||
addr = endpoint or DEFAULT_ENDPOINT
|
||||
host, _, port_str = addr.rpartition(":")
|
||||
if not host:
|
||||
host = addr
|
||||
port_str = "3692"
|
||||
port = int(port_str)
|
||||
|
||||
# Auto-detect TLS: skip for localhost
|
||||
if use_tls is None:
|
||||
use_tls = host not in ("localhost", "127.0.0.1", "::1")
|
||||
|
||||
ssl_ctx = None
|
||||
if use_tls:
|
||||
ssl_ctx = ssl.create_default_context()
|
||||
|
||||
reader, writer = await asyncio.open_connection(host, port, ssl=ssl_ctx)
|
||||
|
||||
# Send handshake
|
||||
hs_bytes = build_handshake(node_id)
|
||||
await write_frame(writer, hs_bytes)
|
||||
|
||||
# Read handshake response
|
||||
resp_data = await read_frame(reader)
|
||||
resp_msg = Message.parse(resp_data)
|
||||
peer_id = parse_handshake(resp_msg)
|
||||
|
||||
return cls(reader, writer, peer_id)
|
||||
|
||||
@property
|
||||
def peer_id(self) -> str:
|
||||
return self._peer_id
|
||||
|
||||
async def call(
|
||||
self,
|
||||
method: str,
|
||||
auth: str,
|
||||
body: bytes,
|
||||
) -> tuple[int, bytes, str]:
|
||||
"""
|
||||
Send a MsgType 100 cloud service request and return (status, body, error).
|
||||
"""
|
||||
self._req_id += 1
|
||||
req_id = self._req_id
|
||||
|
||||
# Build ZAP message
|
||||
msg_bytes = build_cloud_request(method, auth, body)
|
||||
|
||||
# Wrap with 8-byte Call correlation header
|
||||
wrapped = struct.pack("<II", req_id, REQ_FLAG_REQ) + msg_bytes
|
||||
|
||||
# Send
|
||||
await write_frame(self._writer, wrapped)
|
||||
|
||||
# Read response — loop until we match our req_id
|
||||
while True:
|
||||
data = await read_frame(self._reader)
|
||||
if len(data) < 8:
|
||||
continue
|
||||
|
||||
resp_id, resp_flag = struct.unpack_from("<II", data, 0)
|
||||
if resp_flag != REQ_FLAG_RESP:
|
||||
continue
|
||||
if resp_id != req_id:
|
||||
continue
|
||||
|
||||
# Parse ZAP message (skip 8-byte Call header)
|
||||
msg = Message.parse(data[8:])
|
||||
return parse_cloud_response(msg)
|
||||
|
||||
async def chat_completion(
|
||||
self,
|
||||
model: str,
|
||||
messages: list[dict[str, Any]],
|
||||
auth_token: str = "",
|
||||
**kwargs: Any,
|
||||
) -> dict[str, Any]:
|
||||
"""
|
||||
High-level: send an OpenAI-compatible chat completion request via ZAP.
|
||||
|
||||
Returns the parsed JSON response dict.
|
||||
"""
|
||||
request_body: dict[str, Any] = {"model": model, "messages": messages}
|
||||
request_body.update(kwargs)
|
||||
|
||||
body_bytes = json.dumps(request_body).encode("utf-8")
|
||||
auth = f"Bearer {auth_token}" if auth_token and not auth_token.startswith("Bearer ") else auth_token
|
||||
|
||||
status, resp_body, error = await self.call("chat.completions", auth, body_bytes)
|
||||
|
||||
if status != 200:
|
||||
err_msg = error or resp_body.decode("utf-8", errors="replace") or f"ZAP status {status}"
|
||||
raise RuntimeError(f"ZAP cloud error: {err_msg}")
|
||||
|
||||
return json.loads(resp_body)
|
||||
|
||||
async def close(self) -> None:
|
||||
"""Close the connection."""
|
||||
self._writer.close()
|
||||
await self._writer.wait_closed()
|
||||
|
||||
async def __aenter__(self) -> "CloudClient":
|
||||
return self
|
||||
|
||||
async def __aexit__(self, *args: Any) -> None:
|
||||
await self.close()
|
||||
@@ -0,0 +1,307 @@
|
||||
"""
|
||||
luxfi/zap binary wire protocol — Python implementation.
|
||||
|
||||
Compatible with the Rust `hanzo-zap` crate and Go `github.com/luxfi/zap` v0.2.0.
|
||||
|
||||
Wire format:
|
||||
Frame: [4-byte LE length][message bytes]
|
||||
Message header (16 bytes): magic(4) + version(2) + flags(2) + root_offset(4) + size(4)
|
||||
Object fields: inline primitives, (relOffset:i32 + length:u32) for text/bytes
|
||||
relOffset is relative to the field's absolute position in the buffer.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import struct
|
||||
|
||||
# ── Constants ────────────────────────────────────────────────────────────
|
||||
|
||||
ZAP_MAGIC = b"ZAP\x00"
|
||||
HEADER_SIZE = 16
|
||||
VERSION = 1
|
||||
ALIGNMENT = 8
|
||||
MAX_MESSAGE_SIZE = 10 * 1024 * 1024 # 10 MB
|
||||
|
||||
MSG_TYPE_CLOUD = 100
|
||||
|
||||
# Cloud request field byte offsets (each Text/Bytes = 8 bytes: relOffset + length)
|
||||
CLOUD_REQ_METHOD = 0
|
||||
CLOUD_REQ_AUTH = 8
|
||||
CLOUD_REQ_BODY = 16
|
||||
|
||||
# Cloud response field byte offsets
|
||||
# Layout: status(0:u32, 4 bytes) + body(4:Bytes, 8 bytes) + error(12:Text, 8 bytes)
|
||||
CLOUD_RESP_STATUS = 0 # u32 inline (4 bytes)
|
||||
CLOUD_RESP_BODY = 4 # (relOffset:i32 + length:u32)
|
||||
CLOUD_RESP_ERROR = 12 # (relOffset:i32 + length:u32)
|
||||
|
||||
# Call correlation flags
|
||||
REQ_FLAG_REQ = 1
|
||||
REQ_FLAG_RESP = 2
|
||||
|
||||
# Handshake constants
|
||||
HANDSHAKE_OBJ_SIZE = 64
|
||||
HANDSHAKE_ID_MAX = 60
|
||||
HANDSHAKE_ID_LEN_OFFSET = 60
|
||||
|
||||
|
||||
# ── Message ──────────────────────────────────────────────────────────────
|
||||
|
||||
class Message:
|
||||
"""Parsed ZAP binary message (owns the full buffer including header)."""
|
||||
|
||||
__slots__ = ("_data",)
|
||||
|
||||
def __init__(self, data: bytes | bytearray) -> None:
|
||||
self._data = bytes(data)
|
||||
|
||||
@classmethod
|
||||
def parse(cls, data: bytes | bytearray) -> "Message":
|
||||
if len(data) < HEADER_SIZE:
|
||||
raise ValueError(f"ZAP message too short: {len(data)} < {HEADER_SIZE}")
|
||||
if data[:4] != ZAP_MAGIC:
|
||||
raise ValueError(f"Bad ZAP magic: {data[:4]!r}")
|
||||
return cls(data)
|
||||
|
||||
@property
|
||||
def bytes(self) -> bytes:
|
||||
return self._data
|
||||
|
||||
@property
|
||||
def version(self) -> int:
|
||||
return struct.unpack_from("<H", self._data, 4)[0]
|
||||
|
||||
@property
|
||||
def flags(self) -> int:
|
||||
return struct.unpack_from("<H", self._data, 6)[0]
|
||||
|
||||
@property
|
||||
def msg_type(self) -> int:
|
||||
return self.flags >> 8
|
||||
|
||||
@property
|
||||
def root_offset(self) -> int:
|
||||
"""Absolute offset of the root object in the buffer."""
|
||||
return struct.unpack_from("<I", self._data, 8)[0]
|
||||
|
||||
@property
|
||||
def total_size(self) -> int:
|
||||
return struct.unpack_from("<I", self._data, 12)[0]
|
||||
|
||||
|
||||
# ── Object reader (works on full message buffer with absolute offsets) ──
|
||||
|
||||
def obj_uint32(data: bytes, obj_offset: int, field_offset: int) -> int:
|
||||
"""Read a u32 inline field from object at obj_offset."""
|
||||
pos = obj_offset + field_offset
|
||||
if pos + 4 > len(data):
|
||||
return 0
|
||||
return struct.unpack_from("<I", data, pos)[0]
|
||||
|
||||
|
||||
def obj_bytes(data: bytes, obj_offset: int, field_offset: int) -> bytes:
|
||||
"""Read a Bytes field (relOffset:i32 + length:u32) from object."""
|
||||
pos = obj_offset + field_offset
|
||||
if pos + 8 > len(data):
|
||||
return b""
|
||||
rel_off = struct.unpack_from("<i", data, pos)[0] # signed i32
|
||||
length = struct.unpack_from("<I", data, pos + 4)[0]
|
||||
if length == 0 or rel_off == 0:
|
||||
return b""
|
||||
abs_off = pos + rel_off
|
||||
if abs_off + length > len(data):
|
||||
return b""
|
||||
return data[abs_off:abs_off + length]
|
||||
|
||||
|
||||
def obj_text(data: bytes, obj_offset: int, field_offset: int) -> str:
|
||||
"""Read a Text field from object."""
|
||||
b = obj_bytes(data, obj_offset, field_offset)
|
||||
return b.decode("utf-8", errors="replace") if b else ""
|
||||
|
||||
|
||||
# ── Builder (matches Rust Builder exactly) ───────────────────────────────
|
||||
|
||||
def _align(pos: int) -> int:
|
||||
return (pos + ALIGNMENT - 1) & ~(ALIGNMENT - 1)
|
||||
|
||||
|
||||
class Builder:
|
||||
"""Builds a ZAP message in a single buffer, matching Rust Builder."""
|
||||
|
||||
def __init__(self, capacity: int = 256) -> None:
|
||||
cap = max(capacity, 256)
|
||||
self._buf = bytearray(cap)
|
||||
self._buf[:4] = ZAP_MAGIC
|
||||
struct.pack_into("<H", self._buf, 4, VERSION)
|
||||
self._pos = HEADER_SIZE
|
||||
self._root_offset = 0
|
||||
|
||||
def _grow(self, n: int) -> None:
|
||||
needed = self._pos + n
|
||||
if needed <= len(self._buf):
|
||||
return
|
||||
new_cap = max(len(self._buf) * 2, needed)
|
||||
self._buf.extend(b"\x00" * (new_cap - len(self._buf)))
|
||||
|
||||
def _align_pos(self) -> None:
|
||||
padding = (ALIGNMENT - (self._pos % ALIGNMENT)) % ALIGNMENT
|
||||
self._grow(padding)
|
||||
for _ in range(padding):
|
||||
self._buf[self._pos] = 0
|
||||
self._pos += 1
|
||||
|
||||
def start_object(self, data_size: int) -> "ObjectBuilder":
|
||||
self._align_pos()
|
||||
return ObjectBuilder(self, self._pos, data_size)
|
||||
|
||||
def finish(self, flags: int = 0) -> bytes:
|
||||
struct.pack_into("<H", self._buf, 6, flags)
|
||||
struct.pack_into("<I", self._buf, 8, self._root_offset)
|
||||
struct.pack_into("<I", self._buf, 12, self._pos)
|
||||
return bytes(self._buf[:self._pos])
|
||||
|
||||
|
||||
class ObjectBuilder:
|
||||
"""Builds a single object within a Builder, matching Rust ObjectBuilder."""
|
||||
|
||||
def __init__(self, builder: Builder, start_pos: int, data_size: int) -> None:
|
||||
self._builder = builder
|
||||
self._start = start_pos
|
||||
self._data_size = data_size
|
||||
self._deferred: list[tuple[int, bytes]] = [] # (field_offset, data)
|
||||
|
||||
def _ensure_field(self, end_offset: int) -> None:
|
||||
needed = self._start + end_offset
|
||||
if needed > self._builder._pos:
|
||||
self._builder._grow(needed - self._builder._pos)
|
||||
for i in range(self._builder._pos, needed):
|
||||
self._builder._buf[i] = 0
|
||||
self._builder._pos = needed
|
||||
|
||||
def set_u32(self, field_offset: int, v: int) -> None:
|
||||
self._ensure_field(field_offset + 4)
|
||||
struct.pack_into("<I", self._builder._buf, self._start + field_offset, v)
|
||||
|
||||
def set_u8(self, field_offset: int, v: int) -> None:
|
||||
self._ensure_field(field_offset + 1)
|
||||
self._builder._buf[self._start + field_offset] = v & 0xFF
|
||||
|
||||
def set_bytes(self, field_offset: int, data: bytes) -> None:
|
||||
self._ensure_field(field_offset + 8)
|
||||
pos = self._start + field_offset
|
||||
if not data:
|
||||
struct.pack_into("<I", self._builder._buf, pos, 0)
|
||||
struct.pack_into("<I", self._builder._buf, pos + 4, 0)
|
||||
return
|
||||
struct.pack_into("<I", self._builder._buf, pos + 4, len(data))
|
||||
self._deferred.append((field_offset, data))
|
||||
|
||||
def set_text(self, field_offset: int, text: str) -> None:
|
||||
self.set_bytes(field_offset, text.encode("utf-8"))
|
||||
|
||||
def finish_as_root(self) -> None:
|
||||
"""Finalize and set as root object."""
|
||||
self._ensure_field(self._data_size)
|
||||
for field_offset, data in self._deferred:
|
||||
data_pos = self._builder._pos
|
||||
self._builder._grow(len(data))
|
||||
start = self._builder._pos
|
||||
self._builder._buf[start:start + len(data)] = data
|
||||
self._builder._pos += len(data)
|
||||
field_abs = self._start + field_offset
|
||||
rel_offset = data_pos - field_abs
|
||||
struct.pack_into("<i", self._builder._buf, field_abs, rel_offset) # signed i32
|
||||
self._builder._root_offset = self._start
|
||||
|
||||
|
||||
# ── Cloud service helpers ────────────────────────────────────────────────
|
||||
|
||||
def build_cloud_request(method: str, auth: str, body: bytes) -> bytes:
|
||||
"""Build a MsgType 100 cloud service request message."""
|
||||
b = Builder(len(body) + len(method) + len(auth) + 128)
|
||||
obj = b.start_object(24) # 3 * 8 bytes
|
||||
obj.set_text(CLOUD_REQ_METHOD, method)
|
||||
obj.set_text(CLOUD_REQ_AUTH, auth)
|
||||
obj.set_bytes(CLOUD_REQ_BODY, body)
|
||||
obj.finish_as_root()
|
||||
return b.finish(flags=MSG_TYPE_CLOUD << 8)
|
||||
|
||||
|
||||
def build_cloud_response(status: int, body: bytes, error: str) -> bytes:
|
||||
"""Build a MsgType 100 cloud service response message."""
|
||||
b = Builder(len(body) + len(error) + 128)
|
||||
obj = b.start_object(20) # u32(4) + Bytes(8) + Text(8) = 20
|
||||
obj.set_u32(CLOUD_RESP_STATUS, status)
|
||||
obj.set_bytes(CLOUD_RESP_BODY, body)
|
||||
obj.set_text(CLOUD_RESP_ERROR, error)
|
||||
obj.finish_as_root()
|
||||
return b.finish(flags=MSG_TYPE_CLOUD << 8)
|
||||
|
||||
|
||||
def parse_cloud_request(msg: Message) -> tuple[str, str, bytes]:
|
||||
"""Parse a cloud request → (method, auth, body)."""
|
||||
data = msg.bytes
|
||||
off = msg.root_offset
|
||||
method = obj_text(data, off, CLOUD_REQ_METHOD)
|
||||
auth = obj_text(data, off, CLOUD_REQ_AUTH)
|
||||
body = obj_bytes(data, off, CLOUD_REQ_BODY)
|
||||
return method, auth, body
|
||||
|
||||
|
||||
def parse_cloud_response(msg: Message) -> tuple[int, bytes, str]:
|
||||
"""Parse a cloud response → (status, body, error)."""
|
||||
data = msg.bytes
|
||||
off = msg.root_offset
|
||||
status = obj_uint32(data, off, CLOUD_RESP_STATUS)
|
||||
body = obj_bytes(data, off, CLOUD_RESP_BODY)
|
||||
error = obj_text(data, off, CLOUD_RESP_ERROR)
|
||||
return status, body, error
|
||||
|
||||
|
||||
# ── Handshake ────────────────────────────────────────────────────────────
|
||||
|
||||
def build_handshake(node_id: str) -> bytes:
|
||||
"""Build a handshake message (msg_type=0, 64-byte fixed object)."""
|
||||
b = Builder(128)
|
||||
obj = b.start_object(HANDSHAKE_OBJ_SIZE)
|
||||
id_bytes = node_id.encode("utf-8")[:HANDSHAKE_ID_MAX]
|
||||
for i, byte in enumerate(id_bytes):
|
||||
obj.set_u8(i, byte)
|
||||
obj.set_u32(HANDSHAKE_ID_LEN_OFFSET, len(id_bytes))
|
||||
obj.finish_as_root()
|
||||
return b.finish()
|
||||
|
||||
|
||||
def parse_handshake(msg: Message) -> str:
|
||||
"""Parse a handshake message → peer node ID."""
|
||||
data = msg.bytes
|
||||
off = msg.root_offset
|
||||
id_len = obj_uint32(data, off, HANDSHAKE_ID_LEN_OFFSET)
|
||||
if id_len == 0:
|
||||
return ""
|
||||
start = off
|
||||
end = start + min(id_len, HANDSHAKE_ID_MAX)
|
||||
if end > len(data):
|
||||
return ""
|
||||
return data[start:end].decode("utf-8", errors="replace")
|
||||
|
||||
|
||||
# ── Frame I/O ────────────────────────────────────────────────────────────
|
||||
|
||||
async def read_frame(reader) -> bytes:
|
||||
"""Read a length-prefixed frame: [4-byte LE length][data]."""
|
||||
len_buf = await reader.readexactly(4)
|
||||
length = struct.unpack("<I", len_buf)[0]
|
||||
if length > MAX_MESSAGE_SIZE:
|
||||
raise ValueError(f"ZAP frame too large: {length}")
|
||||
if length == 0:
|
||||
return b""
|
||||
return await reader.readexactly(length)
|
||||
|
||||
|
||||
async def write_frame(writer, data: bytes) -> None:
|
||||
"""Write a length-prefixed frame."""
|
||||
writer.write(struct.pack("<I", len(data)))
|
||||
writer.write(data)
|
||||
await writer.drain()
|
||||
@@ -0,0 +1,340 @@
|
||||
"""Tests for luxfi/zap binary wire protocol compatibility."""
|
||||
import asyncio
|
||||
import struct
|
||||
|
||||
import pytest
|
||||
|
||||
from hanzo_zap.wire import (
|
||||
ZAP_MAGIC,
|
||||
HEADER_SIZE,
|
||||
VERSION,
|
||||
MSG_TYPE_CLOUD,
|
||||
REQ_FLAG_REQ,
|
||||
REQ_FLAG_RESP,
|
||||
CLOUD_REQ_METHOD,
|
||||
CLOUD_REQ_AUTH,
|
||||
CLOUD_REQ_BODY,
|
||||
CLOUD_RESP_STATUS,
|
||||
CLOUD_RESP_BODY,
|
||||
CLOUD_RESP_ERROR,
|
||||
MAX_MESSAGE_SIZE,
|
||||
Message,
|
||||
Builder,
|
||||
ObjectBuilder,
|
||||
build_cloud_request,
|
||||
build_cloud_response,
|
||||
parse_cloud_request,
|
||||
parse_cloud_response,
|
||||
build_handshake,
|
||||
parse_handshake,
|
||||
obj_uint32,
|
||||
obj_bytes,
|
||||
obj_text,
|
||||
read_frame,
|
||||
write_frame,
|
||||
)
|
||||
|
||||
|
||||
# ── constants ──────────────────────────────────────────────────────────────
|
||||
|
||||
def test_zap_magic():
|
||||
assert ZAP_MAGIC == b"ZAP\x00"
|
||||
assert HEADER_SIZE == 16
|
||||
|
||||
|
||||
def test_version():
|
||||
assert VERSION == 1
|
||||
|
||||
|
||||
def test_msg_type_cloud():
|
||||
assert MSG_TYPE_CLOUD == 100
|
||||
|
||||
|
||||
def test_req_flags():
|
||||
assert REQ_FLAG_REQ == 1
|
||||
assert REQ_FLAG_RESP == 2
|
||||
|
||||
|
||||
def test_cloud_resp_offsets():
|
||||
"""Cloud response layout: status(0:u32,4B) + body(4:Bytes,8B) + error(12:Text,8B)."""
|
||||
assert CLOUD_RESP_STATUS == 0
|
||||
assert CLOUD_RESP_BODY == 4
|
||||
assert CLOUD_RESP_ERROR == 12
|
||||
|
||||
|
||||
# ── Message parsing ────────────────────────────────────────────────────────
|
||||
|
||||
def test_message_parse_rejects_short():
|
||||
with pytest.raises(ValueError, match="too short"):
|
||||
Message.parse(b"ZAP")
|
||||
|
||||
|
||||
def test_message_parse_rejects_bad_magic():
|
||||
bad = b"BAD\x00" + b"\x00" * 12
|
||||
with pytest.raises(ValueError, match="Bad ZAP magic"):
|
||||
Message.parse(bad)
|
||||
|
||||
|
||||
def test_message_parse_valid_header():
|
||||
header = bytearray(HEADER_SIZE)
|
||||
header[:4] = ZAP_MAGIC
|
||||
struct.pack_into("<H", header, 4, VERSION)
|
||||
msg = Message.parse(bytes(header))
|
||||
assert msg.version == VERSION
|
||||
assert msg.bytes == bytes(header)
|
||||
|
||||
|
||||
def test_message_properties():
|
||||
flags = MSG_TYPE_CLOUD << 8 | 0x01
|
||||
header = bytearray(HEADER_SIZE)
|
||||
header[:4] = ZAP_MAGIC
|
||||
struct.pack_into("<H", header, 4, VERSION)
|
||||
struct.pack_into("<H", header, 6, flags)
|
||||
struct.pack_into("<I", header, 8, 0)
|
||||
struct.pack_into("<I", header, 12, 42)
|
||||
msg = Message.parse(bytes(header))
|
||||
assert msg.version == VERSION
|
||||
assert msg.msg_type == MSG_TYPE_CLOUD
|
||||
assert msg.flags == flags
|
||||
assert msg.root_offset == 0
|
||||
assert msg.total_size == 42
|
||||
|
||||
|
||||
# ── Builder ───────────────────────────────────────────────────────────────
|
||||
|
||||
def test_builder_header():
|
||||
"""Builder produces correct ZAP header."""
|
||||
b = Builder()
|
||||
result = b.finish(flags=MSG_TYPE_CLOUD << 8)
|
||||
assert result[:4] == ZAP_MAGIC
|
||||
version = struct.unpack_from("<H", result, 4)[0]
|
||||
assert version == VERSION
|
||||
flags = struct.unpack_from("<H", result, 6)[0]
|
||||
assert flags >> 8 == MSG_TYPE_CLOUD
|
||||
|
||||
|
||||
def test_builder_with_object():
|
||||
"""Builder + ObjectBuilder produce parseable message."""
|
||||
b = Builder()
|
||||
obj = b.start_object(4)
|
||||
obj.set_u32(0, 0xDEADBEEF)
|
||||
obj.finish_as_root()
|
||||
result = b.finish()
|
||||
msg = Message.parse(result)
|
||||
assert msg.root_offset == HEADER_SIZE # object starts right after header
|
||||
val = obj_uint32(msg.bytes, msg.root_offset, 0)
|
||||
assert val == 0xDEADBEEF
|
||||
|
||||
|
||||
def test_builder_text_field():
|
||||
"""Text field with relative offset is readable."""
|
||||
b = Builder()
|
||||
obj = b.start_object(8)
|
||||
obj.set_text(0, "hello")
|
||||
obj.finish_as_root()
|
||||
result = b.finish()
|
||||
msg = Message.parse(result)
|
||||
text = obj_text(msg.bytes, msg.root_offset, 0)
|
||||
assert text == "hello"
|
||||
|
||||
|
||||
def test_builder_bytes_field():
|
||||
"""Bytes field with relative offset is readable."""
|
||||
b = Builder()
|
||||
obj = b.start_object(8)
|
||||
payload = b"\x01\x02\x03\x04"
|
||||
obj.set_bytes(0, payload)
|
||||
obj.finish_as_root()
|
||||
result = b.finish()
|
||||
msg = Message.parse(result)
|
||||
data = obj_bytes(msg.bytes, msg.root_offset, 0)
|
||||
assert data == payload
|
||||
|
||||
|
||||
def test_builder_mixed_fields():
|
||||
"""Object with inline u32 + variable-length fields."""
|
||||
b = Builder()
|
||||
obj = b.start_object(20) # u32(4) + Bytes(8) + Text(8)
|
||||
obj.set_u32(0, 200)
|
||||
obj.set_bytes(4, b"body-data")
|
||||
obj.set_text(12, "error-msg")
|
||||
obj.finish_as_root()
|
||||
result = b.finish()
|
||||
msg = Message.parse(result)
|
||||
off = msg.root_offset
|
||||
assert obj_uint32(msg.bytes, off, 0) == 200
|
||||
assert obj_bytes(msg.bytes, off, 4) == b"body-data"
|
||||
assert obj_text(msg.bytes, off, 12) == "error-msg"
|
||||
|
||||
|
||||
def test_builder_empty_bytes():
|
||||
"""Empty bytes field roundtrips as empty."""
|
||||
b = Builder()
|
||||
obj = b.start_object(8)
|
||||
obj.set_bytes(0, b"")
|
||||
obj.finish_as_root()
|
||||
result = b.finish()
|
||||
msg = Message.parse(result)
|
||||
assert obj_bytes(msg.bytes, msg.root_offset, 0) == b""
|
||||
|
||||
|
||||
# ── Handshake ─────────────────────────────────────────────────────────────
|
||||
|
||||
def test_build_and_parse_handshake():
|
||||
hs_bytes = build_handshake("test-node")
|
||||
msg = Message.parse(hs_bytes)
|
||||
assert msg.msg_type == 0
|
||||
peer = parse_handshake(msg)
|
||||
assert peer == "test-node"
|
||||
|
||||
|
||||
def test_handshake_long_id_truncated():
|
||||
long_id = "x" * 100
|
||||
hs_bytes = build_handshake(long_id)
|
||||
msg = Message.parse(hs_bytes)
|
||||
peer = parse_handshake(msg)
|
||||
assert peer == "x" * 60
|
||||
|
||||
|
||||
def test_handshake_empty_id():
|
||||
hs_bytes = build_handshake("")
|
||||
msg = Message.parse(hs_bytes)
|
||||
peer = parse_handshake(msg)
|
||||
assert peer == ""
|
||||
|
||||
|
||||
# ── Cloud request/response ────────────────────────────────────────────────
|
||||
|
||||
def test_build_cloud_request():
|
||||
body = b'{"model":"test","messages":[]}'
|
||||
msg_bytes = build_cloud_request("chat.completions", "Bearer token123", body)
|
||||
msg = Message.parse(msg_bytes)
|
||||
assert msg.msg_type == MSG_TYPE_CLOUD
|
||||
|
||||
|
||||
def test_cloud_request_roundtrip():
|
||||
"""Build and parse a cloud request."""
|
||||
body = b'{"test":true}'
|
||||
msg_bytes = build_cloud_request("completions", "Bearer abc", body)
|
||||
msg = Message.parse(msg_bytes)
|
||||
method, auth, req_body = parse_cloud_request(msg)
|
||||
assert method == "completions"
|
||||
assert auth == "Bearer abc"
|
||||
assert req_body == body
|
||||
|
||||
|
||||
def test_cloud_response_roundtrip():
|
||||
"""Build and parse a cloud response."""
|
||||
resp_body = b'{"id":"test","choices":[]}'
|
||||
msg_bytes = build_cloud_response(200, resp_body, "")
|
||||
msg = Message.parse(msg_bytes)
|
||||
status, body, error = parse_cloud_response(msg)
|
||||
assert status == 200
|
||||
assert body == resp_body
|
||||
assert error == ""
|
||||
|
||||
|
||||
def test_cloud_response_error():
|
||||
msg_bytes = build_cloud_response(500, b"", "internal server error")
|
||||
msg = Message.parse(msg_bytes)
|
||||
status, body, error = parse_cloud_response(msg)
|
||||
assert status == 500
|
||||
assert error == "internal server error"
|
||||
|
||||
|
||||
def test_cloud_response_with_body_and_error():
|
||||
"""Response with both body and error text."""
|
||||
resp_body = b'{"detail":"something"}'
|
||||
msg_bytes = build_cloud_response(400, resp_body, "bad request")
|
||||
msg = Message.parse(msg_bytes)
|
||||
status, body, error = parse_cloud_response(msg)
|
||||
assert status == 400
|
||||
assert body == resp_body
|
||||
assert error == "bad request"
|
||||
|
||||
|
||||
# ── Frame I/O helpers ─────────────────────────────────────────────────────
|
||||
|
||||
async def _make_stream_pair():
|
||||
connected: asyncio.Future[tuple[asyncio.StreamReader, asyncio.StreamWriter]] = (
|
||||
asyncio.get_event_loop().create_future()
|
||||
)
|
||||
|
||||
async def on_connect(r: asyncio.StreamReader, w: asyncio.StreamWriter) -> None:
|
||||
connected.set_result((r, w))
|
||||
|
||||
server = await asyncio.start_server(on_connect, "127.0.0.1", 0)
|
||||
port = server.sockets[0].getsockname()[1]
|
||||
c_reader, c_writer = await asyncio.open_connection("127.0.0.1", port)
|
||||
s_reader, s_writer = await connected
|
||||
|
||||
async def cleanup() -> None:
|
||||
for w in (c_writer, s_writer):
|
||||
w.close()
|
||||
await w.wait_closed()
|
||||
server.close()
|
||||
await server.wait_closed()
|
||||
|
||||
return s_reader, c_writer, cleanup
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_frame_io():
|
||||
reader, writer, cleanup = await _make_stream_pair()
|
||||
try:
|
||||
test_data = b"hello ZAP"
|
||||
await write_frame(writer, test_data)
|
||||
result = await read_frame(reader)
|
||||
assert result == test_data
|
||||
finally:
|
||||
await cleanup()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_frame_io_empty():
|
||||
reader, writer, cleanup = await _make_stream_pair()
|
||||
try:
|
||||
await write_frame(writer, b"")
|
||||
result = await read_frame(reader)
|
||||
assert result == b""
|
||||
finally:
|
||||
await cleanup()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_frame_io_large():
|
||||
reader, writer, cleanup = await _make_stream_pair()
|
||||
try:
|
||||
big_data = b"X" * 65536
|
||||
await write_frame(writer, big_data)
|
||||
result = await read_frame(reader)
|
||||
assert result == big_data
|
||||
finally:
|
||||
await cleanup()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_frame_rejects_oversized():
|
||||
reader, writer, cleanup = await _make_stream_pair()
|
||||
try:
|
||||
writer.write(struct.pack("<I", MAX_MESSAGE_SIZE + 1))
|
||||
await writer.drain()
|
||||
with pytest.raises(ValueError, match="too large"):
|
||||
await read_frame(reader)
|
||||
finally:
|
||||
await cleanup()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_full_message_through_frames():
|
||||
"""Build a ZAP message, send through frame I/O, parse on the other side."""
|
||||
reader, writer, cleanup = await _make_stream_pair()
|
||||
try:
|
||||
hs = build_handshake("frame-test-node")
|
||||
await write_frame(writer, hs)
|
||||
received = await read_frame(reader)
|
||||
msg = Message.parse(received)
|
||||
peer = parse_handshake(msg)
|
||||
assert peer == "frame-test-node"
|
||||
finally:
|
||||
await cleanup()
|
||||
Reference in New Issue
Block a user