From 8c68f48903b354f8ff19c74df7652bbcc3693185 Mon Sep 17 00:00:00 2001 From: JD Davis Date: Sun, 12 Jul 2026 16:19:46 +0000 Subject: [PATCH] refactor(proxy): extract memory golden replay policy (#2007) ## Description Extracts memory-tool golden byte replay and canonicalization from `headroom.proxy.helpers.apply_session_sticky_memory_tools` into a focused policy module. The session tracker, skip/deduplication decisions, and logging remain in the existing helper; the byte-level replay policy now has direct coverage. Closes # ## Type of Change - [ ] 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 - [x] Code refactoring (no functional changes) ## Changes Made - Added `headroom.proxy.memory_golden_policy` for replaying stored memory golden bytes and canonicalizing fresh memory tool definitions. - Updated `apply_session_sticky_memory_tools` to delegate golden-byte decode/canonicalization while preserving tracker and logging behavior. - Added direct tests for golden replay, invalid/corrupt bytes, non-UTF-8 bytes, and serializer parity with the existing helper. ## Testing - [x] Unit tests pass (`pytest`) - [x] Linting passes (`ruff check .`) - [x] Type checking passes (`mypy headroom`) - [x] New tests added for new functionality - [ ] Manual testing performed ### Test Output ```text python -m pytest tests/test_memory_golden_policy.py tests/test_memory_tool_session_sticky.py tests/test_corrupt_golden_bytes_recovery.py tests/test_issue_728_empty_tools_injection.py 50 passed in 0.87s python -m ruff check . All checks passed! python -m ruff format --check . 1069 files already formatted python -m mypy headroom --ignore-missing-imports Success: no issues found in 410 source files gitleaks protect --staged --no-banner --redact no leaks found ``` ## Real Behavior Proof - Environment: Windows, Python 3.13.13, clean worktree from `headroomlabs/main` at `d2170b19`. - Exact command / steps: Ran targeted memory golden replay/sticky injection/corrupt-byte regression tests plus ruff, ruff-format, mypy, and staged gitleaks scan. - Observed result: All targeted tests and local gates passed; staged secret scan found no leaks. - Not tested: Full Docker/native wrapper CI locally; covered by repository CI. ## 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. ## Additional Notes Documentation and changelog updates are not applicable for this internal refactor. The push reported existing default-branch Dependabot vulnerabilities; this PR's staged gitleaks scan passed and CI security checks are expected to validate the branch. Co-authored-by: Tejas Chopra --- headroom/proxy/helpers.py | 21 +++++++---- headroom/proxy/memory_golden_policy.py | 41 +++++++++++++++++++++ tests/test_memory_golden_policy.py | 50 ++++++++++++++++++++++++++ 3 files changed, 105 insertions(+), 7 deletions(-) create mode 100644 headroom/proxy/memory_golden_policy.py create mode 100644 tests/test_memory_golden_policy.py diff --git a/headroom/proxy/helpers.py b/headroom/proxy/helpers.py index 0ccd3cfda..33f3ed073 100644 --- a/headroom/proxy/helpers.py +++ b/headroom/proxy/helpers.py @@ -74,6 +74,10 @@ from headroom.proxy.internal_header_policy import ( resolve_strip_internal_headers_mode, strip_internal_headers, ) +from headroom.proxy.memory_golden_policy import ( + replay_golden_memory_tool_definition, + serialize_memory_tool_definition_canonical, +) from headroom.proxy.tool_injection_config import ( ToolInjectionStickyMode, ) @@ -2017,7 +2021,7 @@ def apply_session_sticky_memory_tools( continue tools_out.append(tool_def) existing_names.add(tn) - added_bytes += len(serialize_tool_definition_canonical(tool_def)) + added_bytes += len(serialize_memory_tool_definition_canonical(tool_def)) log_tool_injection_decision( provider=provider, session_id=session_id, @@ -2047,7 +2051,7 @@ def apply_session_sticky_memory_tools( continue tools_out.append(tool_def) existing_names.add(tn) - added_bytes += len(serialize_tool_definition_canonical(tool_def)) + added_bytes += len(serialize_memory_tool_definition_canonical(tool_def)) log_tool_injection_decision( provider=provider, session_id=None, @@ -2072,7 +2076,10 @@ def apply_session_sticky_memory_tools( # Their bytes win (the client's choice, not ours to gate). continue try: - tool_def = json.loads(golden_bytes.decode("utf-8")) + replay = replay_golden_memory_tool_definition( + tool_name=tool_name, + golden_tool_bytes=golden_bytes, + ) except (UnicodeDecodeError, json.JSONDecodeError) as exc: logger.error( "corrupt golden tool bytes for session %s tool %s: %s — skipping tool injection", @@ -2082,9 +2089,9 @@ def apply_session_sticky_memory_tools( exc_info=True, ) continue - tools_out.append(tool_def) - existing_names.add(tool_name) - replay_bytes += len(golden_bytes) + tools_out.append(replay.tool_definition) + existing_names.add(replay.tool_name) + replay_bytes += len(replay.canonical_bytes) log_tool_injection_decision( provider=provider, session_id=session_id, @@ -2111,7 +2118,7 @@ def apply_session_sticky_memory_tools( tn = _extract_tool_name(tool_def) if tn is None or tn in existing_names: continue - golden_bytes = serialize_tool_definition_canonical(tool_def) + golden_bytes = serialize_memory_tool_definition_canonical(tool_def) tracker.record_injection( provider=provider, session_id=session_id, diff --git a/headroom/proxy/memory_golden_policy.py b/headroom/proxy/memory_golden_policy.py new file mode 100644 index 000000000..8ed2d1f92 --- /dev/null +++ b/headroom/proxy/memory_golden_policy.py @@ -0,0 +1,41 @@ +"""Policy helpers for replaying memory-tool golden definitions.""" + +from __future__ import annotations + +import json +from dataclasses import dataclass +from typing import Any, cast + + +@dataclass(frozen=True) +class MemoryToolDefinitionReplay: + """Memory tool definition selected for sticky replay.""" + + tool_name: str + tool_definition: dict[str, Any] + canonical_bytes: bytes + + +def serialize_memory_tool_definition_canonical(tool_definition: dict[str, Any]) -> bytes: + """Return stable canonical bytes for a memory tool definition.""" + + return json.dumps( + tool_definition, + ensure_ascii=False, + separators=(",", ":"), + ).encode("utf-8") + + +def replay_golden_memory_tool_definition( + *, + tool_name: str, + golden_tool_bytes: bytes, +) -> MemoryToolDefinitionReplay: + """Decode a stored memory tool definition and preserve its original bytes.""" + + tool_definition = json.loads(golden_tool_bytes.decode("utf-8")) + return MemoryToolDefinitionReplay( + tool_name=tool_name, + tool_definition=cast(dict[str, Any], tool_definition), + canonical_bytes=golden_tool_bytes, + ) diff --git a/tests/test_memory_golden_policy.py b/tests/test_memory_golden_policy.py new file mode 100644 index 000000000..35c71bdaa --- /dev/null +++ b/tests/test_memory_golden_policy.py @@ -0,0 +1,50 @@ +from __future__ import annotations + +import pytest + +from headroom.proxy.helpers import serialize_tool_definition_canonical +from headroom.proxy.memory_golden_policy import ( + replay_golden_memory_tool_definition, + serialize_memory_tool_definition_canonical, +) + + +def test_replays_golden_memory_tool_without_reserializing() -> None: + golden = b'{ "name" : "memory_save" , "description" : "client bytes" }' + + replay = replay_golden_memory_tool_definition( + tool_name="memory_save", + golden_tool_bytes=golden, + ) + + assert replay.tool_name == "memory_save" + assert replay.tool_definition["name"] == "memory_save" + assert replay.canonical_bytes == golden + + +def test_rejects_invalid_golden_json() -> None: + with pytest.raises(ValueError): + replay_golden_memory_tool_definition( + tool_name="memory_save", + golden_tool_bytes=b"not-json", + ) + + +def test_rejects_non_utf8_golden_bytes() -> None: + with pytest.raises(UnicodeDecodeError): + replay_golden_memory_tool_definition( + tool_name="memory_save", + golden_tool_bytes=b"\x80\x81", + ) + + +def test_memory_canonical_serializer_matches_existing_helper() -> None: + tool_definition = { + "name": "memory_search", + "description": "Find memory", + "input_schema": {"type": "object", "properties": {"query": {"type": "string"}}}, + } + + assert serialize_memory_tool_definition_canonical( + tool_definition + ) == serialize_tool_definition_canonical(tool_definition)