diff --git a/headroom/transforms/content_router.py b/headroom/transforms/content_router.py index f317bfbd3..75ee193c4 100644 --- a/headroom/transforms/content_router.py +++ b/headroom/transforms/content_router.py @@ -2807,6 +2807,7 @@ class ContentRouter(Transform): language: str | None = None, question: str | None = None, bias: float = 1.0, + _allow_embedded: bool = True, ) -> tuple[str, int, list[str]]: """Apply a compression strategy to content. @@ -2827,6 +2828,37 @@ class ContentRouter(Transform): log]``). Log readers use this to see *how* we got to the final compressor without parsing decision_reason strings. """ + # ── STRUCTURAL (embedded) JSON routing ─────────────────────────────── + # Before anything else: if this block is not a single JSON value but + # CONTAINS balanced JSON span(s), route each span through this very + # dispatch and splice the result back (surrounding bytes kept exact). + # This is how nested/embedded JSON reaches the JSON compressors at all — + # today's linear splitter never sees it. Each span goes through the + # UNCHANGED path, so SmartCrusher/CodeCompressor register their + # `<>` markers exactly as for a whole-block JSON (CCR is hash- + # keyed → location-agnostic). `_allow_embedded=False` on the recursive + # call is a one-shot re-entrancy guard (NOT a depth cap). Deterministic + + # benefit-gated (no size/min thresholds) → prefix-cache- and CCR-store- + # stable, and a strict no-op when the block has no embedded JSON. + if _allow_embedded: + from headroom.transforms.recursive_json import route_embedded_json + + def _dispatch_span(span: str) -> str | None: + strat = self._strategy_from_detection_type(_detect_content(span).content_type) + text, _t, _c = self._apply_strategy_to_content( + span, + strat, + context, + question=question, + bias=bias, + _allow_embedded=False, + ) + return text if text != span else None + + routed = route_embedded_json(content, _dispatch_span, tok=_estimate_tokens) + if routed is not None: + return routed, _estimate_tokens(routed), ["embedded_json"] + # Track original tokens for TOIN recording original_tokens = _estimate_tokens(content) compressed: str | None = None diff --git a/headroom/transforms/recursive_json.py b/headroom/transforms/recursive_json.py new file mode 100644 index 000000000..e6d7b1665 --- /dev/null +++ b/headroom/transforms/recursive_json.py @@ -0,0 +1,162 @@ +"""Structural (recursive) JSON routing for the ContentRouter. + +Today the router is *linear*: it splits a block into textual sections and picks +one strategy per section. It never looks *inside* a structure, so JSON embedded +in a larger payload (a ``gh api`` dump, an MCP result, a ``curl | jq`` tail) is +invisible to the JSON compressors — even though, in practice, that embedded shape +is the overwhelming majority of JSON the agent ever sees. + +This module adds the missing structural step: find balanced JSON spans at any +offset in a block and route each one through the router's *existing* dispatch, +splicing the result back in place with the surrounding bytes kept exact. + +Why this is CCR-safe by construction +------------------------------------- +Each span is handed to the router's own ``_apply_strategy_to_content`` — the same +code path a whole-block JSON already takes — so SmartCrusher / CodeCompressor +register their ``<>`` retrieval markers exactly as they do today. CCR +is hash-keyed and therefore location-agnostic: a marker resolves whether it sits +at the top of a block or nested inside one. This module never touches the CCR +store; it only relocates where the dispatch is invoked. + +Safety invariants (no thresholds — outcome-gated only): + * A span that already contains a ``< int | None: + """Index just past the balanced JSON container opening at ``start`` (honoring + string/escape rules), or ``None`` if it never balances.""" + stack: list[str] = [] + in_str = esc = False + for j in range(start, len(text)): + ch = text[j] + if in_str: + if esc: + esc = False + elif ch == "\\": + esc = True + elif ch == '"': + in_str = False + continue + if ch == '"': + in_str = True + elif ch in _OPEN: + stack.append(ch) + elif ch in _CLOSE: + if not stack or stack[-1] != _PAIR[ch]: + return None + stack.pop() + if not stack: + return j + 1 + return None + + +def _spans(text: str) -> list[tuple[int, int]]: + """Deterministic list of ``(start, end)`` for top-level balanced JSON spans. + Nested spans are not returned separately — the dispatch handles depth.""" + out: list[tuple[int, int]] = [] + i, n = 0, len(text) + while i < n: + if text[i] in _OPEN: + end = _match_span(text, i) + if end is not None: + out.append((i, end)) + i = end + continue + i += 1 + return out + + +def _has_routable_json(span: str) -> bool: + """True if ``span`` parses and contains an array of objects somewhere — the + shape the JSON compressors actually act on. Cheap structural check, no size + threshold.""" + try: + v = json.loads(span) + except (ValueError, TypeError): + return False + + found = False + + def walk(x: object) -> None: + nonlocal found + if found: + return + if isinstance(x, list): + if len(x) >= 2 and sum(isinstance(e, dict) for e in x) >= 0.8 * len(x): + found = True + return + for e in x: + walk(e) + elif isinstance(x, dict): + for e in x.values(): + walk(e) + + walk(v) + return found + + +def route_embedded_json( + content: str, + dispatch: Dispatch, + *, + tok: Callable[[str], int] | None = None, +) -> str | None: + """Route every embedded JSON span in ``content`` through ``dispatch`` and + splice the results back in place. Returns the rewritten block, or ``None`` + when nothing safe/smaller applied. + + ``content`` that is itself a single JSON value is intentionally skipped — the + caller already routes pure-JSON blocks; this exists for the *embedded* case. + """ + tok = tok or (lambda s: max(1, len(s) // 4)) + spans = _spans(content) + if not spans: + return None + # Whole-block JSON is the caller's job, not ours. + if len(spans) == 1 and spans[0] == (0, len(content.strip())): + return None + + repls: list[tuple[int, int, str]] = [] + for a, b in spans: + chunk = content[a:b] + if "< str | None: + """Fake compressor: returns a shorter deterministic stand-in for any span.""" + try: + v = json.loads(span) + except ValueError: + return None + return f"" if isinstance(v, list) else "" + + +def test_embedded_json_routed_and_surroundings_exact() -> None: + payload = json.dumps([{"id": i, "ok": True} for i in range(6)], separators=(",", ":")) + content = f"Fetched rows from API:\n{payload}\nDone (200 OK)." + out = route_embedded_json(content, _upper_dispatch) + assert out is not None + assert out.startswith("Fetched rows from API:\n") + assert out.endswith("\nDone (200 OK).") + assert "
" in out + + +def test_ccr_marker_span_passed_through() -> None: + # A span already carrying a CCR marker must never be re-routed (R1). + content = 'prefix [{"a":1,"b":2},{"a":3,"b":"<>"}] suffix' + out = route_embedded_json(content, _upper_dispatch) + assert out is None # only span contains a marker → skipped → nothing to do + + +def test_no_json_is_noop() -> None: + assert route_embedded_json("just prose, nothing structured here", _upper_dispatch) is None + + +def test_whole_block_json_is_callers_job() -> None: + # A block that IS a single JSON value is skipped (routed by the caller). + content = json.dumps([{"a": i} for i in range(5)], separators=(",", ":")) + assert route_embedded_json(content, _upper_dispatch) is None + + +def test_benefit_gate_declines_when_not_smaller() -> None: + payload = json.dumps([{"a": i} for i in range(5)], separators=(",", ":")) + content = f"x {payload} y" + # Dispatch that returns something LARGER → must be declined (outcome gate). + assert route_embedded_json(content, lambda s: s + " " * 999) is None + + +def test_deterministic() -> None: + payload = json.dumps([{"k": i} for i in range(8)], separators=(",", ":")) + content = f"a {payload} b {payload} c" + r1 = route_embedded_json(content, _upper_dispatch) + r2 = route_embedded_json(content, _upper_dispatch) + assert r1 == r2 and r1 is not None + assert r1.count("
") == 2 # both embedded spans routed + + +def test_scalar_array_not_routed() -> None: + # array of scalars is not a "routable" JSON shape (no dict rows) + content = "nums: [1,2,3,4,5,6,7,8] done" + assert route_embedded_json(content, _upper_dispatch) is None