dropp baes64 use binary

packages/node/meshnet_node/relay_bridge.py

@@ -19,7 +19,7 @@ log = logging.getLogger(__name__)
 
 DEFAULT_MAX_CONCURRENCY = 8
 
-# Activation tensors ride the relay as base64 inside one JSON frame, so the
+# Activation tensors ride the relay as one WebSocket frame per hop, so the
 # websockets default of 1 MiB rejects any real prefill (close code 1009).
 DEFAULT_WS_MAX_BYTES = 256 * 1024 * 1024
 
@@ -36,6 +36,26 @@ def ws_max_size() -> int | None:
     return None if value <= 0 else value
 
 
+# Binary relay frame: JSON header + raw body in one WebSocket binary message,
+# so activation bodies travel as bytes instead of base64 inside JSON. Same wire
+# format as meshnet_relay.server — duplicated because node and relay ship as
+# independent distributions.
+BINARY_FRAME_MAGIC = b"MRF1"
+
+
+def encode_binary_frame(header: dict, body: bytes) -> bytes:
+    header_bytes = json.dumps(header, separators=(",", ":")).encode()
+    return BINARY_FRAME_MAGIC + len(header_bytes).to_bytes(4, "big") + header_bytes + body
+
+
+def decode_binary_frame(frame: bytes) -> tuple[dict, bytes]:
+    if len(frame) < 8 or frame[:4] != BINARY_FRAME_MAGIC:
+        raise ValueError("not a meshnet binary relay frame")
+    header_len = int.from_bytes(frame[4:8], "big")
+    header = json.loads(frame[8:8 + header_len].decode())
+    return header, bytes(frame[8 + header_len:])
+
+
 @dataclass(frozen=True)
 class RelayBridgeInfo:
     peer_id: str
@@ -125,7 +145,9 @@ class RelayHttpBridge:
 
         while self._running:
             try:
-                with wsc.connect(self.relay_url, open_timeout=5, max_size=ws_max_size()) as ws:
+                with wsc.connect(
+                    self.relay_url, open_timeout=5, max_size=ws_max_size(), compression=None,
+                ) as ws:
                     self._ws = ws
                     self._connected.set()
                     ws.send(json.dumps(_make_envelope(
@@ -138,6 +160,17 @@ class RelayHttpBridge:
                             raw = ws.recv(timeout=1)
                         except TimeoutError:
                             continue
+                        if isinstance(raw, (bytes, bytearray)):
+                            try:
+                                payload, body = decode_binary_frame(bytes(raw))
+                            except (ValueError, json.JSONDecodeError):
+                                continue
+                            if payload.get("target_peer") not in {None, self.peer_id}:
+                                continue
+                            if self._executor is None:
+                                break
+                            self._executor.submit(self._process_request, payload, body)
+                            continue
                         try:
                             envelope = json.loads(raw)
                         except (TypeError, json.JSONDecodeError):
@@ -176,23 +209,41 @@ class RelayHttpBridge:
             log.debug("relay bridge send failed (request orphaned): %s", exc)
             return False
 
-    def _process_request(self, payload: dict) -> None:
+    def _send_binary_response_frame(self, header: dict, body: bytes) -> bool:
+        """Send one binary response frame; False if the socket is gone."""
+        ws = self._ws
+        if ws is None:
+            return False
+        frame = encode_binary_frame(header, body)
+        try:
+            with self._send_lock:
+                ws.send(frame)
+            return True
+        except Exception as exc:
+            log.debug("relay bridge binary send failed (request orphaned): %s", exc)
+            return False
+
+    def _process_request(self, payload: dict, binary_body: bytes | None = None) -> None:
         request_id = str(payload.get("request_id") or "")
         method = str(payload.get("method") or "POST").upper()
         path = str(payload.get("path") or "/")
         headers = payload.get("headers") if isinstance(payload.get("headers"), dict) else {}
+        binary_mode = binary_body is not None
 
         req_suffix = f" request_id={request_id}" if request_id else ""
         print(f"  [node] relay {method} {path}{req_suffix}", flush=True)
 
-        # body_base64 carries binary data (e.g. bfloat16 activation tensors) safely.
-        # Fallback to text "body" for backward-compat with non-binary requests.
-        body_b64 = payload.get("body_base64")
-        if body_b64:
-            data = base64.b64decode(body_b64)
+        if binary_mode:
+            data = binary_body
         else:
-            body_text = payload.get("body") or ""
-            data = body_text.encode() if isinstance(body_text, str) else bytes(body_text)
+            # Legacy JSON request: body_base64 carries binary data, text "body"
+            # covers non-binary requests.
+            body_b64 = payload.get("body_base64")
+            if body_b64:
+                data = base64.b64decode(body_b64)
+            else:
+                body_text = payload.get("body") or ""
+                data = body_text.encode() if isinstance(body_text, str) else bytes(body_text)
 
         url = f"{self.local_base_url}{path}"
         req = urllib.request.Request(url, data=data, headers=headers, method=method)
@@ -205,6 +256,13 @@ class RelayHttpBridge:
                     return
                 resp_bytes = resp.read()
                 # Forward all X-Meshnet-* headers so the caller can reconstruct the activation.
+                if binary_mode:
+                    self._send_binary_response_frame({
+                        "request_id": request_id,
+                        "status": resp.status,
+                        "headers": resp_headers,
+                    }, resp_bytes)
+                    return
                 is_binary = "octet-stream" in content_type
                 result: dict = {
                     "request_id": request_id,

packages/node/meshnet_node/torch_server.py

@@ -109,25 +109,33 @@ def _relay_hop(
 ) -> tuple[int, dict[str, str], bytes]:
     """Send a single HTTP-shaped request through a relay RPC WebSocket.
 
-    relay_addr is the wss://relay.../rpc/{peer_id} URL.
+    relay_addr is the wss://relay.../rpc/{peer_id} URL. The request and any
+    binary response travel as binary frames (JSON header + raw body); relay
+    error responses and legacy peers still answer with JSON text frames.
     Returns (status, response_headers_lower, response_body).
     Raises on connection failure so callers can fall back to direct.
     """
     import websockets.sync.client as wsc  # type: ignore[import]
 
-    from .relay_bridge import ws_max_size
+    from .relay_bridge import decode_binary_frame, encode_binary_frame, ws_max_size
 
     request_id = f"{time.time_ns():x}"
-    payload = json.dumps({
+    frame = encode_binary_frame({
         "request_id": request_id,
         "method": "POST",
         "path": path,
         "headers": headers,
-        "body_base64": base64.b64encode(body).decode(),
-    })
-    with wsc.connect(relay_addr, open_timeout=timeout, max_size=ws_max_size()) as ws:
-        ws.send(payload)
+    }, body)
+    with wsc.connect(
+        relay_addr, open_timeout=timeout, max_size=ws_max_size(), compression=None,
+    ) as ws:
+        ws.send(frame)
         raw = ws.recv(timeout=timeout)
+    if isinstance(raw, (bytes, bytearray)):
+        resp_header, resp_body = decode_binary_frame(bytes(raw))
+        status = int(resp_header.get("status", 503))
+        resp_headers = {k.lower(): v for k, v in (resp_header.get("headers") or {}).items()}
+        return status, resp_headers, resp_body
     resp = json.loads(raw)
     status = int(resp.get("status", 503))
     resp_headers = {k.lower(): v for k, v in (resp.get("headers") or {}).items()}
@@ -136,6 +144,20 @@ def _relay_hop(
     return status, resp_headers, resp_body
 
 
+# Below this, zstd overhead outweighs the win (per-token decode bodies are ~KBs).
+_COMPRESS_MIN_BYTES = 64 * 1024
+
+
+def _maybe_compress_activation(body: bytes) -> tuple[bytes, str | None]:
+    """zstd-compress large activation bodies; returns (wire_body, encoding)."""
+    if len(body) < _COMPRESS_MIN_BYTES:
+        return body, None
+    try:
+        return _compress_body(body, "zstd"), "zstd"
+    except Exception:
+        return body, None
+
+
 def _is_cache_miss_body(body: bytes) -> bool:
     try:
         return json.loads(body).get("error") == "cache_miss"
@@ -892,6 +914,7 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler):
                     + (f" relay={relay_addr}" if relay_addr else ""),
                     flush=True,
                 )
+            wire_body, wire_encoding = _maybe_compress_activation(current_body)
             headers: dict[str, str] = {
                 "Content-Type": "application/octet-stream",
                 "X-Meshnet-Wire": _WIRE_VERSION,
@@ -903,6 +926,8 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler):
                 "X-Meshnet-Hop-Index": str(hop_index),
                 "X-Meshnet-Start-Layer": str(start_layer),
             }
+            if wire_encoding:
+                headers["X-Meshnet-Encoding"] = wire_encoding
             if cache_mode:
                 headers["X-Meshnet-Cache"] = cache_mode
                 past_len = getattr(payload, "past_len", None)
@@ -915,7 +940,7 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler):
             if relay_addr:
                 try:
                     status, resp_headers, resp_body = _relay_hop(
-                        relay_addr, "/forward", current_body, headers, timeout=120.0,
+                        relay_addr, "/forward", wire_body, headers, timeout=120.0,
                     )
                     if status == 409 and _is_cache_miss_body(resp_body):
                         raise _PipelineCacheMiss(node_url)
@@ -938,7 +963,7 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler):
             if not relay_addr:
                 req = urllib.request.Request(
                     f"{node_url}/forward",
-                    data=current_body,
+                    data=wire_body,
                     headers=headers,
                     method="POST",
                 )
@@ -970,7 +995,13 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler):
             # Binary activation — update and forward to next node
             shape_header = resp_headers.get("x-meshnet-shape", ",".join(str(d) for d in current_shape))
             current_shape = _parse_shape(shape_header)
-            current_body = resp_body
+            try:
+                current_body = _decompress_body(
+                    resp_body, resp_headers.get("x-meshnet-encoding")
+                )
+            except ValueError as exc:
+                print(f"  [node] pipeline hop {hop_index} bad response encoding: {exc}", flush=True)
+                return f"pipeline error at {node_url}: {exc}", None
             current_attn = resp_headers.get("x-meshnet-attn-mask")
             current_pos = resp_headers.get("x-meshnet-position-ids")
         return "", None

packages/relay/meshnet_relay/server.py

@@ -41,6 +41,26 @@ def ws_max_size() -> int | None:
     return None if value <= 0 else value
 
 
+# Binary relay frame: JSON header (request/response metadata) + raw body in one
+# WebSocket binary message. Activation bodies stay bytes end to end — no base64
+# inflation, no JSON re-encode of megabytes per hop. Text JSON frames remain the
+# control plane (gossip, peer-register, streamed SSE responses).
+BINARY_FRAME_MAGIC = b"MRF1"
+
+
+def encode_binary_frame(header: dict, body: bytes) -> bytes:
+    header_bytes = json.dumps(header, separators=(",", ":")).encode()
+    return BINARY_FRAME_MAGIC + len(header_bytes).to_bytes(4, "big") + header_bytes + body
+
+
+def decode_binary_frame(frame: bytes) -> tuple[dict, bytes]:
+    if len(frame) < 8 or frame[:4] != BINARY_FRAME_MAGIC:
+        raise ValueError("not a meshnet binary relay frame")
+    header_len = int.from_bytes(frame[4:8], "big")
+    header = json.loads(frame[8:8 + header_len].decode())
+    return header, bytes(frame[8 + header_len:])
+
+
 class RelayServer:
     """Async WebSocket relay server that runs in a background thread.
 
@@ -118,6 +138,9 @@ class RelayServer:
             self.port,
             ssl=ssl_ctx,
             max_size=ws_max_size(),
+            # Bulk payloads are zstd-compressed at the pipeline layer already;
+            # per-message deflate would recompress them on every hop for nothing.
+            compression=None,
         )
         # Record actual port after bind
         for sock in server.sockets or []:
@@ -162,6 +185,17 @@ class RelayServer:
 
         try:
             async for raw in ws:
+                # Binary frames are relay-http-response bodies from a bridge —
+                # route them to the waiting rpc requester as-is, never fan out.
+                if isinstance(raw, (bytes, bytearray)):
+                    try:
+                        header, _ = decode_binary_frame(bytes(raw))
+                    except (ValueError, json.JSONDecodeError):
+                        continue
+                    queue = self._pending_rpc.get(header.get("request_id"))
+                    if queue is not None:
+                        queue.put_nowait(bytes(raw))
+                    continue
                 try:
                     envelope = json.loads(raw)
                 except (json.JSONDecodeError, TypeError):
@@ -252,14 +286,27 @@ class RelayServer:
 
         try:
             raw = await asyncio.wait_for(ws_requester.recv(), timeout=30.0)
-            payload = json.loads(raw)
+            if isinstance(raw, (bytes, bytearray)):
+                header, body = decode_binary_frame(bytes(raw))
+                request_id = str(header.get("request_id") or uuid.uuid4())
+                header["request_id"] = request_id
+                header["target_peer"] = target_peer_id
+                outbound: str | bytes = encode_binary_frame(header, body)
+            else:
+                payload = json.loads(raw)
+                request_id = str(payload.get("request_id") or uuid.uuid4())
+                payload["request_id"] = request_id
+                payload["target_peer"] = target_peer_id
+                outbound = json.dumps({
+                    "topic": "relay-http-request",
+                    "version": 1,
+                    "from_peer": "relay",
+                    "payload": payload,
+                })
         except Exception:
             await ws_requester.close(1003, "invalid relay rpc request")
             return
 
-        request_id = str(payload.get("request_id") or uuid.uuid4())
-        payload["request_id"] = request_id
-        payload["target_peer"] = target_peer_id
         queue: asyncio.Queue = asyncio.Queue()
         self._pending_rpc[request_id] = queue
         overall_timeout = 310.0
@@ -267,15 +314,11 @@ class RelayServer:
         loop = asyncio.get_running_loop()
         deadline = loop.time() + overall_timeout
         try:
-            await target.ws.send(json.dumps({
-                "topic": "relay-http-request",
-                "version": 1,
-                "from_peer": "relay",
-                "payload": payload,
-            }))
+            await target.ws.send(outbound)
             # Forward frames until a terminal one: streamed responses (US-036)
             # end with {"stream": true, "done": true}; a frame without "stream"
-            # is a complete legacy single response.
+            # is a complete legacy single response. A binary frame is always a
+            # complete single response.
             while True:
                 remaining = deadline - loop.time()
                 if remaining <= 0:
@@ -283,6 +326,9 @@ class RelayServer:
                 frame = await asyncio.wait_for(
                     queue.get(), timeout=min(idle_timeout, remaining)
                 )
+                if isinstance(frame, (bytes, bytearray)):
+                    await ws_requester.send(frame)
+                    break
                 await ws_requester.send(json.dumps(frame))
                 if not frame.get("stream") or frame.get("done"):
                     break

tests/test_gossip_and_relay.py

@@ -372,12 +372,53 @@ def test_relay_rpc_round_trips_http_request_to_peer():
     assert json.loads(response["body"]) == {"ok": True, "path": "/v1/chat/completions"}
 
 
+def test_binary_relay_frame_codecs_interoperate():
+    """Node and relay ship the same binary frame format as separate copies."""
+    import os
+
+    from meshnet_node import relay_bridge
+    from meshnet_relay import server as relay_server
+
+    header = {"request_id": "r-1", "status": 200, "headers": {"X": "y"}}
+    body = os.urandom(4096)
+    for encoder, decoder in (
+        (relay_bridge.encode_binary_frame, relay_server.decode_binary_frame),
+        (relay_server.encode_binary_frame, relay_bridge.decode_binary_frame),
+    ):
+        assert decoder(encoder(header, body)) == (header, body)
+
+    try:
+        relay_server.decode_binary_frame(b"not a frame")
+    except ValueError:
+        pass
+    else:
+        raise AssertionError("garbage bytes must not decode as a binary frame")
+
+
+def test_activation_compression_round_trips_and_skips_small_bodies():
+    """Pipeline hops zstd-compress large activations; tiny decode bodies pass raw."""
+    import os
+
+    from meshnet_node.server import _decompress_body
+    from meshnet_node.torch_server import _maybe_compress_activation
+
+    small = os.urandom(1024)
+    assert _maybe_compress_activation(small) == (small, None)
+
+    large = os.urandom(96 * 1024) * 2  # repetition so zstd actually shrinks it
+    wire, encoding = _maybe_compress_activation(large)
+    assert encoding == "zstd"
+    assert len(wire) < len(large)
+    assert _decompress_body(wire, encoding) == large
+
+
 def test_relay_rpc_carries_activation_sized_frames():
     """A >1 MiB activation body must survive the full relay round trip.
 
     Regression: the websockets library caps frames at 1 MiB by default, so
     prefill activations forwarded via /rpc/<peer> died with close code 1009
     at every hop (requester → relay, relay → bridge, bridge → relay → requester).
+    The body now travels as binary frames — raw bytes, no base64.
     """
     import http.server
     import os

tests/test_kv_cache_distributed.py

@@ -278,6 +278,43 @@ def test_session_is_stable_and_decode_payloads_are_single_token():
     assert head_backend.released == [session_id]
 
 
+def test_large_prefill_activation_survives_zstd_compressed_hop():
+    """A prefill body above _COMPRESS_MIN_BYTES travels the hop zstd-compressed.
+
+    The head compresses and sets X-Meshnet-Encoding; the tail's /forward must
+    decompress before shape validation, so a passing generation proves the
+    compressed round trip (a mishandled encoding fails validation with 400).
+    """
+
+    class _BigHeadBackend(_CachedHeadBackend):
+        def encode_prompt(self, prompt, session_id=None):
+            self.prefills.append(session_id)
+            if session_id:
+                self._seq[session_id] = 2048
+            return TensorPayload(
+                body=b"\x00" * (1 * 2048 * 32 * 2),  # 128 KiB, above the zstd threshold
+                shape=[1, 2048, 32],
+                attention_mask_header=None,
+                position_ids_header=None,
+            )
+
+    head_backend = _BigHeadBackend()
+    tail_backend = _CachedTailBackend([(" a", 1), (" b", 2)])
+    head = TorchNodeServer(backend=head_backend, tracker_mode=True)
+    tail = TorchNodeServer(backend=tail_backend)
+    head_port = head.start()
+    tail_port = tail.start()
+    try:
+        content = _chat_once(head_port, tail_port, max_tokens=2)
+    finally:
+        head.stop()
+        tail.stop()
+
+    assert content == " a b"
+    assert tail_backend.calls[0]["mode"] == "prefill"
+    assert tail_backend.calls[0]["shape"] == [1, 2048, 32]
+
+
 def test_eos_token_id_stops_generation():
     head_backend = _CachedHeadBackend()
     tail_backend = _CachedTailBackend([(" a", 1), ("", 99)])