From fa5624ceff7072fe2c3cad2d6e3b5bf318347329 Mon Sep 17 00:00:00 2001
From: faresobeid <111092724+faresobeid@users.noreply.github.com>
Date: Wed, 9 Sep 2026 00:00:10 +0000
Subject: [PATCH 1/2] feat(client): preserve paired routing-weight payloads

Keep compact expert IDs and FP32 coefficient data paired while parsing token
responses, without dropping usage metadata or changing the JSON fallback.

Assisted-by: Prime Agent
---
 renderers/client.py          |  79 +++++++++---
 tests/test_client.py         |  48 +++++++
 tests/test_client_routing.py | 239 +++++++++++++++++++++++++++++++++++
 3 files changed, 351 insertions(+), 15 deletions(-)
 create mode 100644 tests/test_client_routing.py

diff --git a/renderers/client.py b/renderers/client.py
index 33e27074..ba039b17 100644
--- a/renderers/client.py
+++ b/renderers/client.py
@@ -10,6 +10,7 @@
 import json
 import logging
 import math
+import re
 from collections.abc import Mapping
 from typing import Any, cast
 
@@ -28,6 +29,8 @@
 
 _request_logger = logging.getLogger("renderers.client")
 ROUTED_EXPERTS_DATA_PREFIX = b'"routed_experts":{"data":"'
+ROUTED_EXPERT_WEIGHTS_DATA_PREFIX = b'"weights":{"data":"'
+_BASE64_CHARACTERS = re.compile(rb"[A-Za-z0-9+/=]*")
 # vLLM uses this value both when sampled-token evidence is missing and as a
 # lower-bound clamp, so receiving it cannot prove the real logprob was returned.
 VLLM_LOGPROB_SENTINEL = -9999.0
@@ -109,28 +112,74 @@ async def _resolve_max_prompt_len(client: AsyncOpenAI, model: str) -> int | None
         return value
 
 
-def _strip_base64_field(raw: bytes, prefix: bytes) -> tuple[bytes, memoryview | None]:
-    """Splice a large base64 string field out of raw JSON bytes.
-
-    Avoids json-decoding megabytes of base64; the returned memoryview
-    references ``raw`` and is re-inserted into the parsed payload.
-    """
+def _base64_field_span(raw: bytes, prefix: bytes) -> tuple[int, int] | None:
+    """Find a compact JSON base64 value, excluding its surrounding quotes."""
     data_start = raw.find(prefix)
     if data_start < 0:
-        return raw, None
+        return None
 
     data_start += len(prefix)
-    data_end = raw.index(b'"', data_start)
-    data = memoryview(raw)[data_start:data_end]
-    stripped = raw[:data_start] + raw[data_end:]
-    return stripped, data
+    data_end = raw.find(b'"', data_start)
+    if data_end < 0 or _BASE64_CHARACTERS.fullmatch(raw, data_start, data_end) is None:
+        # Do not hide invalid JSON, or strip escaped strings before JSON has
+        # unescaped them. Actual base64/shape/dtype validation is downstream.
+        return None
+    return data_start, data_end
 
 
 def parse_generate_response(raw: bytes) -> dict[str, Any]:
-    stripped, routed_data = _strip_base64_field(raw, ROUTED_EXPERTS_DATA_PREFIX)
-    payload: dict[str, Any] = json.loads(stripped)
-    if routed_data is not None:
-        payload["choices"][0]["routed_experts"]["data"] = routed_data
+    """Parse a generate response without JSON-decoding large routing strings.
+
+    The compact server encoding uses raw-backed memoryviews for ID data and
+    optional FP32 weight data. All metadata is passed through unchanged.
+    Other JSON layouts use the normal decoder and retain base64 strings.
+    """
+    ids_span = _base64_field_span(raw, ROUTED_EXPERTS_DATA_PREFIX)
+    if ids_span is None:
+        return json.loads(raw)
+    weights_span = _base64_field_span(raw, ROUTED_EXPERT_WEIGHTS_DATA_PREFIX)
+
+    # Numeric placeholders bind each view to its exact JSON value through
+    # parse_float, rather than blindly reattaching a substring to choice[0].
+    # Reserve spellings absent from the input, including unrelated metadata.
+    # JSON numbers cannot use escaped spellings, unlike string sentinels.
+    markers = (b"0.0e+0", b"0.0e+1")
+    if any(marker in raw for marker in markers):
+        return json.loads(raw)
+    spans = [(ids_span, markers[0])]
+    if weights_span is not None:
+        spans.append((weights_span, markers[1]))
+
+    raw_view = memoryview(raw)
+    buffers: dict[str, memoryview] = {}
+    chunks: list[bytes | memoryview] = []
+    cursor = 0
+    for (start, end), marker in sorted(spans):
+        buffers[marker.decode("ascii")] = raw_view[start:end]
+        chunks.extend((raw_view[cursor : start - 1], marker))
+        cursor = end + 1
+    chunks.append(raw_view[cursor:])
+
+    def parse_float(value: str) -> float | memoryview:
+        data = buffers.get(value)
+        return float(value) if data is None else data
+
+    payload = json.loads(b"".join(chunks), parse_float=parse_float)
+    choices = payload.get("choices") if isinstance(payload, dict) else None
+    choice = choices[0] if isinstance(choices, list) and choices else None
+    routed = choice.get("routed_experts") if isinstance(choice, dict) else None
+    if not isinstance(routed, dict) or routed.get("data") is not buffers["0.0e+0"]:
+        # A prefix can occur in another choice, an unrelated object, or a
+        # duplicate key that the JSON decoder overwrites. Do not misattribute
+        # its buffer, or change any non-routing metadata in these cases.
+        return json.loads(raw)
+    if weights_span is not None:
+        weights = routed.get("weights")
+        if (
+            not isinstance(weights, dict)
+            or weights.get("data") is not buffers["0.0e+1"]
+        ):
+            return json.loads(raw)
     return payload
 
 
diff --git a/tests/test_client.py b/tests/test_client.py
index 3a9e0cd9..a1843bd2 100644
--- a/tests/test_client.py
+++ b/tests/test_client.py
@@ -203,6 +203,54 @@ def test_generate_builds_request_body_and_parses_response(usage):
     assert tc.status == ToolCallParseStatus.OK
 
 
+@pytest.mark.parametrize("dtype", ["uint8", "uint16"])
+@pytest.mark.parametrize(
+    "usage",
+    [
+        None,
+        {
+            "prompt_tokens": 3,
+            "completion_tokens": 2,
+            "total_tokens": 5,
+            "prompt_tokens_details": {"cached_tokens": 2},
+            "completion_tokens_details": {"reasoning_tokens": 1},
+        },
+    ],
+)
+def test_generate_forwards_paired_routing_and_metadata(dtype, usage):
+    client = _FakeClient()
+    client.usage = usage
+    ids = np.array([[[5, 2]], [[1, 4]]], dtype=dtype)
+    weights = np.array([[[0.75, 0.25]], [[0.625, 0.375]]], dtype="<f4")
+    routed = {
+        "data": base64.b64encode(ids.tobytes()).decode("ascii"),
+        "shape": list(ids.shape),
+        "dtype": dtype,
+        "start": 3,
+        "format_version": 1,
+        "weights": {
+            "data": base64.b64encode(weights.tobytes()).decode("ascii"),
+            "dtype": "float32",
+        },
+    }
+    client.choice["routed_experts"] = routed
+
+    result = _run_generate(client)
+    actual = result["routed_experts"]
+
+    assert result["usage"] == usage
+    assert isinstance(actual["data"], memoryview)
+    assert isinstance(actual["weights"]["data"], memoryview)
+    assert base64.b64decode(actual["data"], validate=True) == ids.tobytes()
+    assert (
+        base64.b64decode(actual["weights"]["data"], validate=True) == weights.tobytes()
+    )
+    assert {
+        key: value for key, value in actual.items() if key not in {"data", "weights"}
+    } == {key: value for key, value in routed.items() if key not in {"data", "weights"}}
+    assert actual["weights"]["dtype"] == "float32"
+
+
 def test_generate_process_multimodal_false_sends_content_parts():
     class DeferredMultimodalRenderer(_FakeRenderer):
         supports_process_multimodal = True
diff --git a/tests/test_client_routing.py b/tests/test_client_routing.py
new file mode 100644
index 00000000..c345764d
--- /dev/null
+++ b/tests/test_client_routing.py
@@ -0,0 +1,239 @@
+"""CPU-only coverage for the routing transport JSON parser."""
+
+import base64
+import json
+import struct
+
+import pytest
+from renderers.client import parse_generate_response
+
+
+def _routing(dtype="uint8", *, with_weights=True):
+    # Non-sorted expert slots and distinct exact FP32 coefficients catch any
+    # stream swaps, slot changes, or accidental coefficient recomputation.
+    ids = [7, 2, 0, 3, 1, 5, 2, 6]
+    if dtype == "uint16":
+        ids = [expert + 256 for expert in ids]
+    id_bytes = struct.pack("<8B" if dtype == "uint8" else "<8H", *ids)
+    weight_bytes = struct.pack(
+        "<8f", 0.875, 0.125, 0.625, 0.375, 0.75, 0.25, 0.5625, 0.4375
+    )
+    routed = {
+        "data": base64.b64encode(id_bytes).decode("ascii"),
+        "shape": [2, 2, 2],
+        "dtype": dtype,
+        "start": 11,
+        "source": {"name": "capture", "route_scale": 1.0},
+    }
+    if with_weights:
+        routed.update(
+            format_version=1,
+            weights={
+                "data": base64.b64encode(weight_bytes).decode("ascii"),
+                "dtype": "float32",
+                "coefficient_convention": "normalized_topk",
+            },
+        )
+    return routed, id_bytes, weight_bytes
+
+
+def _payload(routed):
+    return {
+        "request_id": "test-routing",
+        "metadata": {"latency": 0.125},
+        "choices": [{"index": 0, "token_ids": [7, 8], "routed_experts": routed}],
+    }
+
+
+def _plain_values(value):
+    """Compare transport metadata without requiring a fast-path layout."""
+    if isinstance(value, memoryview):
+        return value.tobytes().decode("ascii")
+    if isinstance(value, dict):
+        return {key: _plain_values(item) for key, item in value.items()}
+    if isinstance(value, list):
+        return [_plain_values(item) for item in value]
+    return value
+
+
+@pytest.mark.parametrize("dtype", ["uint8", "uint16"])
+@pytest.mark.parametrize("with_weights", [False, True], ids=["legacy", "paired"])
+def test_compact_routing_keeps_exact_raw_backed_streams(
+    dtype, with_weights, monkeypatch
+):
+    routed, id_bytes, weight_bytes = _routing(dtype, with_weights=with_weights)
+    expected = _payload(routed)
+    raw = json.dumps(expected, separators=(",", ":")).encode()
+    original_loads = json.loads
+    decoded_inputs = []
+
+    def record_loads(value, **kwargs):
+        decoded_inputs.append(value)
+        return original_loads(value, **kwargs)
+
+    monkeypatch.setattr("renderers.client.json.loads", record_loads)
+    parsed = parse_generate_response(raw)
+    actual = parsed["choices"][0]["routed_experts"]
+    streams = [(actual["data"], id_bytes)]
+    if with_weights:
+        streams.append((actual["weights"]["data"], weight_bytes))
+    else:
+        assert "weights" not in actual
+        assert "format_version" not in actual
+    assert len(decoded_inputs) == 1
+    for encoded, expected_bytes in streams:
+        assert isinstance(encoded, memoryview)
+        assert encoded.obj is raw
+        assert encoded.readonly
+        assert encoded.tobytes() == base64.b64encode(expected_bytes)
+        assert base64.b64decode(encoded, validate=True) == expected_bytes
+        assert encoded.tobytes() not in decoded_inputs[0]
+    assert _plain_values(parsed) == expected
+
+
+@pytest.mark.parametrize("with_weights", [False, True], ids=["legacy", "paired"])
+def test_compact_parser_does_not_copy_raw_chunks_before_join(with_weights):
+    class UnslicedBytes(bytes):
+        def __getitem__(self, key):
+            if isinstance(key, slice):
+                raise AssertionError("Use memoryviews instead of copying raw chunks")
+            return super().__getitem__(key)
+
+    routed, _, _ = _routing(with_weights=with_weights)
+    expected = _payload(routed)
+    raw = UnslicedBytes(json.dumps(expected, separators=(",", ":")).encode())
+
+    parsed = parse_generate_response(raw)
+
+    assert _plain_values(parsed) == expected
+    actual = parsed["choices"][0]["routed_experts"]
+    assert actual["data"].obj is raw
+    if with_weights:
+        assert actual["weights"]["data"].obj is raw
+
+
+@pytest.mark.parametrize("dtype", ["uint8", "uint16"])
+@pytest.mark.parametrize("with_weights", [False, True], ids=["legacy", "paired"])
+@pytest.mark.parametrize("layout", ["whitespace", "reordered", "weights_reordered"])
+def test_json_layout_fallback_preserves_exact_streams(dtype, with_weights, layout):
+    routed, id_bytes, weight_bytes = _routing(dtype, with_weights=with_weights)
+    if layout == "reordered":
+        routed = dict(reversed(list(routed.items())))
+    elif layout == "weights_reordered" and with_weights:
+        routed["weights"] = dict(reversed(list(routed["weights"].items())))
+    expected = _payload(routed)
+    kwargs = {"indent": 2} if layout == "whitespace" else {"separators": (",", ":")}
+    raw = json.dumps(expected, **kwargs).encode()
+
+    parsed = parse_generate_response(raw)
+
+    actual = parsed["choices"][0]["routed_experts"]
+    assert base64.b64decode(actual["data"], validate=True) == id_bytes
+    if with_weights:
+        assert (
+            base64.b64decode(actual["weights"]["data"], validate=True) == weight_bytes
+        )
+    assert _plain_values(parsed) == expected
+
+
+@pytest.mark.parametrize(
+    "raw",
+    [
+        b'{"routed_experts":{"data":"T1RIRVI="},"choices":[{"routed_experts":{"data":"SUQ="}}]}',
+        b'{"choices":[{}, {"routed_experts":{"data":"SUQ=","weights":{"data":"Vw=="}}}]}',
+        b'{"weights":{"data":"T1RIRVI="},"choices":[{"routed_experts":{"data":"SUQ=","weights":{"data":"Vw=="}}}]}',
+        b'{"choices":[{"routed_experts":{"data":"SUQ=","metadata":{"weights":{"data":"T1RIRVI="}},"weights":{"data":"Vw=="}}}]}',
+        b'{"choices":[{"routed_experts":{"data":"SUQ=","weights":{"data":"Vw=="}}},{"routed_experts":{"data":"T1RIRVI=","weights":{"data":"V1JPTkc="}}}]}',
+        b'{"choices":[{"routed_experts":{"data":"SUQ="},"weights":{"data":"T1RIRVI="}}]}',
+        b'{"choices":[],"routed_experts":{"data":"SUQ="}}',
+        b'[{"routed_experts":{"data":"SUQ="}}]',
+    ],
+    ids=[
+        "unrelated-routing",
+        "other-choice-only",
+        "unrelated-weights-first",
+        "nested-unrelated-weights",
+        "multiple-choices",
+        "weights-outside-routing",
+        "empty-choices",
+        "non-object-payload",
+    ],
+)
+def test_does_not_attach_buffers_to_the_wrong_field(raw):
+    assert _plain_values(parse_generate_response(raw)) == json.loads(raw)
+
+
+@pytest.mark.parametrize(
+    "raw",
+    [
+        b'{"choices":[{"routed_experts":{"data":"T1RIRVI=","data":"SUQ="}}]}',
+        b'{"choices":[{"routed_experts":{"data":"SUQ=","weights":{"data":"T1RIRVI=","data":"Vw=="}}}]}',
+        b'{"choices":[{"routed_experts":{"data":"T1RIRVI="},"routed_experts":{"data":"SUQ="}}]}',
+        b'{"choices":[{"routed_experts":{"data":"SUQ=","weights":{"data":"T1RIRVI="},"weights":{"data":"Vw=="}}}]}',
+        b'{"choices":[{"routed_experts":{"data":"T1RIRVI="}}],"choices":[{"routed_experts":{"data":"SUQ="}}]}',
+        rb'{"choices":[{"routed_experts":{"data":"T1RIRVI=","d\u0061ta":"SUQ="}}]}',
+    ],
+    ids=["ids", "weights-data", "routing", "weights", "choices", "escaped-key"],
+)
+def test_duplicate_keys_keep_decoder_semantics_without_cross_wiring(raw):
+    assert _plain_values(parse_generate_response(raw)) == json.loads(raw)
+
+
+@pytest.mark.parametrize("stream", ["ids", "weights"])
+@pytest.mark.parametrize("bad_bytes", [b"\x00", b"\n", b"\t", b"\xff", b"\\q"])
+def test_fast_path_does_not_hide_invalid_json_inside_routing_data(stream, bad_bytes):
+    routed, _, _ = _routing()
+    raw = json.dumps(_payload(routed), separators=(",", ":")).encode()
+    encoded = routed["data"] if stream == "ids" else routed["weights"]["data"]
+    raw = raw.replace(encoded.encode(), bad_bytes)
+
+    with pytest.raises((json.JSONDecodeError, UnicodeDecodeError)):
+        parse_generate_response(raw)
+
+
+@pytest.mark.parametrize(
+    "raw",
+    [
+        b'{"choices":[{"routed_experts":{"data":"AQ==',
+        b'{"choices":[{"routed_experts":{"data":"AQ==","weights":{"data":"AA==',
+        b'{"choices":[{"routed_experts":{"data":"AQ==",}}]}',
+        b'{"choices":[{"routed_experts":{"data":"AQ=="}}]} trailing',
+    ],
+)
+def test_truncated_or_malformed_response_fails_as_json(raw):
+    with pytest.raises(json.JSONDecodeError):
+        parse_generate_response(raw)
+
+
+@pytest.mark.parametrize("stream", ["ids", "weights"])
+def test_escaped_base64_is_unescaped_by_json_before_forwarding(stream):
+    routed, id_bytes, weight_bytes = _routing()
+    raw = json.dumps(_payload(routed), separators=(",", ":")).encode()
+    encoded = routed["data"] if stream == "ids" else routed["weights"]["data"]
+    escaped = f"\\u{ord(encoded[0]):04x}".encode() + encoded[1:].encode()
+    raw = raw.replace(encoded.encode(), escaped)
+
+    parsed = parse_generate_response(raw)
+
+    actual = parsed["choices"][0]["routed_experts"]
+    assert base64.b64decode(actual["data"], validate=True) == id_bytes
+    assert base64.b64decode(actual["weights"]["data"], validate=True) == weight_bytes
+    assert _plain_values(parsed) == _payload(routed)
+
+
+def test_empty_streams_are_not_treated_as_missing():
+    raw = b'{"choices":[{"routed_experts":{"data":"","weights":{"data":""}}}]}'
+    routed = parse_generate_response(raw)["choices"][0]["routed_experts"]
+    assert isinstance(routed["data"], memoryview)
+    assert isinstance(routed["weights"]["data"], memoryview)
+    assert routed["data"].tobytes() == routed["weights"]["data"].tobytes() == b""
+
+
+@pytest.mark.parametrize("literal", [b"0.0e+0", b"0.0e+1", b'"0.0e+0"'])
+def test_numeric_metadata_cannot_be_mistaken_for_a_stream(literal):
+    raw = (
+        b'{"metadata":'
+        + literal
+        + b',"choices":[{"routed_experts":{"data":"SUQ=","weights":{"data":"Vw=="}}}]}'
+    )
+    assert _plain_values(parse_generate_response(raw)) == json.loads(raw)

From 0e7f0f3486cf7f35bbd114724343fe11c6977198 Mon Sep 17 00:00:00 2001
From: faresobeid <111092724+faresobeid@users.noreply.github.com>
Date: Sat, 24 Oct 2026 18:45:27 +0000
Subject: [PATCH 2/2] fix(client): routing fast path falls back when a stream
 object has no data

---
 renderers/client.py          | 3 +--
 tests/test_client_routing.py | 4 ++++
 2 files changed, 5 insertions(+), 2 deletions(-)

diff --git a/renderers/client.py b/renderers/client.py
index d0106c66..c0a3aac4 100644
--- a/renderers/client.py
+++ b/renderers/client.py
@@ -194,14 +194,13 @@ def parse_float(value: str) -> float | memoryview:
         return float(value) if data is None else data
 
     payload = json.loads(b"".join(chunks), parse_float=parse_float)
-    placed = {id(obj[key]) for obj, key in _stream_slots(payload)}
+    placed = {id(obj.get(key)) for obj, key in _stream_slots(payload)}
     if any(id(data) not in placed for data in buffers.values()):
         # A prefix can occur in an unrelated object, or in a duplicate key that
         # the JSON decoder overwrites. Do not misattribute its buffer, or change
         # any non-stream metadata in these cases.
         return json.loads(raw)
     return payload
-    return payload
 
 
 def _parse_packed_completion_logprobs(
diff --git a/tests/test_client_routing.py b/tests/test_client_routing.py
index c345764d..7eeaa2e1 100644
--- a/tests/test_client_routing.py
+++ b/tests/test_client_routing.py
@@ -147,6 +147,8 @@ def test_json_layout_fallback_preserves_exact_streams(dtype, with_weights, layou
         b'{"choices":[{"routed_experts":{"data":"SUQ="},"weights":{"data":"T1RIRVI="}}]}',
         b'{"choices":[],"routed_experts":{"data":"SUQ="}}',
         b'[{"routed_experts":{"data":"SUQ="}}]',
+        b'{"choices":[{"routed_experts":{"shape":[1]},"compact_logprobs":{"data":"SUQ="}}]}',
+        b'{"choices":[{"routed_experts":{"data":"SUQ=","weights":{"dtype":"float32"}}}]}',
     ],
     ids=[
         "unrelated-routing",
@@ -157,6 +159,8 @@ def test_json_layout_fallback_preserves_exact_streams(dtype, with_weights, layou
         "weights-outside-routing",
         "empty-choices",
         "non-object-payload",
+        "routing-without-data",
+        "weights-without-data",
     ],
 )
 def test_does_not_attach_buffers_to_the_wrong_field(raw):