From d6f0f0f64269bfbdf36070cb304703c606c64b72 Mon Sep 17 00:00:00 2001 From: Ashish Date: Sat, 13 Jun 2026 21:24:12 -0700 Subject: [PATCH] fix(compression): correct JSON array item counting and entropy gate (#887) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Description Two bugs in `JSONStructureHandler` that jointly defeated the "keep first N array items fully" design. (1) Every comma under `array_depth > 0` was counted as an array item separator — including commas *between keys inside objects* — so for arrays of objects the first record's own keys exhausted `max_array_items_full` and dropped values belonging to item 0. (2) Fixing that unmasked a second bug: self-normalized Shannon entropy scores English prose at 0.90+, above the 0.85 "identifier" threshold, so every long description was preserved as a fake high-entropy identifier. Closes # ## Type of Change - [x] Bug fix (non-breaking change that fixes an issue) - [ ] New feature (non-breaking change that adds functionality) - [ ] Breaking change (fix or feature that would cause existing functionality to change) - [ ] Documentation update - [ ] Performance improvement - [ ] Code refactoring (no functional changes) ## Changes Made - `headroom/compression/handlers/json_handler.py`: replace depth-keyed comma counting with a container stack so only commas whose immediate enclosing container is an array advance that array's item index. - `headroom/compression/handlers/json_handler.py`: gate the entropy preservation check on a no-spaces identifier signal, so UUIDs/hashes still pass but prose compresses. - `tests/test_compression/test_json_handler.py`: regression tests for object-comma counting, items past the threshold, and prose-vs-identifier entropy. ## Testing - [x] Unit tests pass (`pytest`) - [x] Linting passes (`ruff check .`) - [x] Type checking passes (`mypy headroom`) - [x] New tests added for new functionality - [x] Manual testing performed ### Test Output ```text $ pytest tests/test_compression/test_json_handler.py -q 32 passed ``` ## Real Behavior Proof - Environment: local macOS, Python 3.11, branch `fix/json-array-item-count`. - Exact command / steps: `pytest tests/test_compression/test_json_handler.py -q` plus an empirical mask dump on `[{"a":1,"b":2,...}]`. - Observed result: Values inside array item 0 are now preserved; long prose values compress while UUIDs are retained (prose scored 0.906-0.929, UUID 0.956 — the threshold alone could not separate them). - Not tested: End-to-end through the live proxy pipeline (the handler is not yet wired into the proxy hot path). ## Review Readiness - [x] I have performed a self-review - [x] This PR is ready for human review ## Checklist - [x] My code follows the project's style guidelines - [x] I have performed a self-review of my code - [x] I have commented my code, particularly in hard-to-understand areas - [ ] I have made corresponding changes to the documentation - [x] My changes generate no new warnings - [x] I have added tests that prove my fix is effective or that my feature works - [x] New and existing unit tests pass locally with my changes - [ ] I have updated the CHANGELOG.md if applicable ## Screenshots (if applicable) N/A — library/compression change with no UI. See Test Output. ## Additional Notes First of a 7-PR compression-handler review series. 🤖 Generated with [Claude Code](https://claude.com/claude-code) --------- Co-authored-by: Claude Fable 5 --- headroom/compression/handlers/json_handler.py | 54 ++++++++++----- tests/test_compression/test_json_handler.py | 69 +++++++++++++++++++ 2 files changed, 105 insertions(+), 18 deletions(-) diff --git a/headroom/compression/handlers/json_handler.py b/headroom/compression/handlers/json_handler.py index 941ef2b87..67e41af37 100644 --- a/headroom/compression/handlers/json_handler.py +++ b/headroom/compression/handlers/json_handler.py @@ -138,30 +138,44 @@ class JSONStructureHandler(BaseStructureHandler): # Build character-level mask mask = [False] * len(content) - # Track array depth for selective preservation - array_depth = 0 - array_item_counts: dict[int, int] = {} # depth -> count + # Track containers so commas inside objects are not counted as + # array item separators — only commas whose immediate enclosing + # container is an array advance that array's item index. Values + # nested in an object that is itself an array item inherit the + # innermost enclosing array's index via array_item_stack[-1]. + container_stack: list[str] = [] # "[" or "{" per open container + array_item_stack: list[int] = [] # item count per open array for token in json_tokens: - # Track array items + # Track containers if token.token_type == JSONTokenType.BRACKET: - if token.text == "[": - array_depth += 1 - array_item_counts[array_depth] = 0 + if token.text in "{[": + container_stack.append(token.text) + if token.text == "[": + array_item_stack.append(0) + elif token.text == "}": + if container_stack and container_stack[-1] == "{": + container_stack.pop() elif token.text == "]": - if array_depth in array_item_counts: - del array_item_counts[array_depth] - array_depth = max(0, array_depth - 1) + if container_stack and container_stack[-1] == "[": + container_stack.pop() + if array_item_stack: + array_item_stack.pop() - # Count array items at commas - if token.token_type == JSONTokenType.COMMA and array_depth > 0: - array_item_counts[array_depth] = array_item_counts.get(array_depth, 0) + 1 + # Count array items only at the array's own commas + if ( + token.token_type == JSONTokenType.COMMA + and container_stack + and container_stack[-1] == "[" + and array_item_stack + ): + array_item_stack[-1] += 1 # Determine if this token should be preserved preserve = self._should_preserve_token( token, - array_depth, - array_item_counts.get(array_depth, 0), + len(array_item_stack), + array_item_stack[-1] if array_item_stack else 0, ) # Mark in mask @@ -225,9 +239,13 @@ class JSONStructureHandler(BaseStructureHandler): if self.preserve_high_entropy: # Strip quotes for entropy calculation value = token.text.strip('"') - score = EntropyScore.compute(value, self.entropy_threshold) - if score.should_preserve: - return True + # Entropy targets identifiers (UUIDs, hashes, API keys). + # Self-normalized entropy also scores English prose >0.85, + # so gate on the cheapest identifier signal: no spaces. + if " " not in value: + score = EntropyScore.compute(value, self.entropy_threshold) + if score.should_preserve: + return True return False diff --git a/tests/test_compression/test_json_handler.py b/tests/test_compression/test_json_handler.py index ed30f07aa..c86b7e20f 100644 --- a/tests/test_compression/test_json_handler.py +++ b/tests/test_compression/test_json_handler.py @@ -153,6 +153,75 @@ class TestJSONStructureHandler: assert result.mask.mask[first_id] is True assert result.mask.mask[second_id] is True + def test_object_commas_do_not_advance_array_item_index(self): + """Commas between keys inside an object must not count as array + item separators. + + Regression: depth-keyed comma counting treated every comma under + array_depth > 0 as an item boundary, so an array's FIRST object + exhausted max_array_items_full within its own keys and later + values were dropped despite being in item 0. + """ + handler = JSONStructureHandler( + preserve_short_values=True, + short_value_threshold=20, + max_array_items_full=3, + ) + # Single array item (index 0) with 5 short values — all must + # stay eligible for preservation. + content = '[{"a": "valuea", "b": "valueb", "c": "valuec", "d": "valued", "e": "valuee"}]' + result = handler.get_mask(content) + + for marker in ('"valuea"', '"valueb"', '"valuec"', '"valued"', '"valuee"'): + start = content.index(marker) + for i in range(start, start + len(marker)): + assert result.mask.mask[i] is True, f"{marker} in array item 0 should be preserved" + + def test_array_items_past_threshold_compressed(self): + """Items at index >= max_array_items_full are still compressed.""" + handler = JSONStructureHandler( + preserve_short_values=True, + short_value_threshold=20, + max_array_items_full=2, + ) + content = '[{"v": "itemzero"}, {"v": "itemone"}, {"v": "itemtwo"}]' + result = handler.get_mask(content) + + # Items 0 and 1 preserved + for marker in ('"itemzero"', '"itemone"'): + start = content.index(marker) + assert result.mask.mask[start + 1] is True, f"{marker} should be preserved" + + # Item 2 is past the threshold — value chars compressed + start = content.index('"itemtwo"') + inner = range(start + 1, start + len('"itemtwo"') - 1) + assert not any(result.mask.mask[i] for i in inner), ( + "values in array items past max_array_items_full should be compressible" + ) + + def test_prose_not_preserved_by_entropy(self): + """Long natural-language strings must not pass the entropy gate. + + Regression: self-normalized Shannon entropy scores English prose + above the 0.85 threshold, so every description survived + compression. Identifiers (no spaces) are still preserved. + """ + handler = JSONStructureHandler(short_value_threshold=20) + prose = "Context optimization layer for LLM applications with caching" + uuid = "8f14e45f-ceea-4123-8f14-e45fceea4123" + content = f'{{"description": "{prose}", "id": "{uuid}"}}' + result = handler.get_mask(content) + + prose_start = content.index(prose) + assert not any(result.mask.mask[i] for i in range(prose_start, prose_start + len(prose))), ( + "long prose value should be compressible" + ) + + uuid_start = content.index(uuid) + assert all(result.mask.mask[i] for i in range(uuid_start, uuid_start + len(uuid))), ( + "UUID value should be preserved via entropy" + ) + def test_metadata_contains_key_count(self, handler): """Test that metadata includes key count.""" content = '{"a": 1, "b": 2, "c": 3}'