From 7c9a032f5044f24fdde71126f591f3eac0a93da1 Mon Sep 17 00:00:00 2001 From: JD Davis Date: Sun, 12 Jul 2026 15:49:12 +0000 Subject: [PATCH] refactor(proxy): extract ccr golden replay policy (#2006) ## Description Extracts CCR golden tool replay and fresh-definition canonicalization from `headroom.proxy.helpers.apply_session_sticky_ccr_tool` into a focused policy module. This keeps sticky CCR orchestration in helpers while making the byte replay/regeneration behavior independently testable. 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.ccr_golden_policy` for replaying stored CCR golden bytes and creating canonical fresh CCR tool definitions. - Updated `apply_session_sticky_ccr_tool` to delegate CCR golden replay/fresh definition policy while preserving tracker coordination and logging decisions. - Added direct tests for golden-byte replay, invalid/corrupt bytes, non-UTF-8 bytes, and fresh canonical definition generation. ## 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_ccr_golden_policy.py tests/test_ccr_tool_always_on.py tests/test_corrupt_golden_bytes_recovery.py tests/test_proxy/test_ccr_frozen_prefix_coupling.py 30 passed in 0.34s 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 CCR 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. --- headroom/proxy/ccr_golden_policy.py | 52 +++++++++++++++++++++++++++++ headroom/proxy/helpers.py | 37 ++++++++++---------- tests/test_ccr_golden_policy.py | 38 +++++++++++++++++++++ 3 files changed, 109 insertions(+), 18 deletions(-) create mode 100644 headroom/proxy/ccr_golden_policy.py create mode 100644 tests/test_ccr_golden_policy.py diff --git a/headroom/proxy/ccr_golden_policy.py b/headroom/proxy/ccr_golden_policy.py new file mode 100644 index 000000000..4d54fddb7 --- /dev/null +++ b/headroom/proxy/ccr_golden_policy.py @@ -0,0 +1,52 @@ +"""Policy helpers for replaying CCR golden tool definitions.""" + +from __future__ import annotations + +import json +from dataclasses import dataclass +from typing import Any, Literal, cast + +from headroom.ccr.tool_injection import create_ccr_tool_definition + + +@dataclass(frozen=True) +class CcrToolDefinitionReplay: + """CCR tool definition selected for sticky replay or fresh injection.""" + + tool_definition: dict[str, Any] + canonical_bytes: bytes + used_golden_bytes: bool + + +def serialize_ccr_tool_definition_canonical(tool_definition: dict[str, Any]) -> bytes: + """Return stable canonical bytes for a CCR tool definition.""" + + return json.dumps( + tool_definition, + ensure_ascii=False, + separators=(",", ":"), + ).encode("utf-8") + + +def replay_golden_ccr_tool_definition(golden_tool_bytes: bytes) -> CcrToolDefinitionReplay: + """Decode a stored CCR tool definition and preserve its original bytes.""" + + tool_definition = json.loads(golden_tool_bytes.decode("utf-8")) + return CcrToolDefinitionReplay( + tool_definition=cast(dict[str, Any], tool_definition), + canonical_bytes=golden_tool_bytes, + used_golden_bytes=True, + ) + + +def create_fresh_ccr_tool_definition( + provider: Literal["anthropic", "openai", "google"], +) -> CcrToolDefinitionReplay: + """Create and canonicalize a fresh CCR tool definition for ``provider``.""" + + tool_definition = create_ccr_tool_definition(provider) + return CcrToolDefinitionReplay( + tool_definition=tool_definition, + canonical_bytes=serialize_ccr_tool_definition_canonical(tool_definition), + used_golden_bytes=False, + ) diff --git a/headroom/proxy/helpers.py b/headroom/proxy/helpers.py index f3c95e2a8..1e299cbe1 100644 --- a/headroom/proxy/helpers.py +++ b/headroom/proxy/helpers.py @@ -38,6 +38,10 @@ from headroom.proxy.body_forwarding import ( prepare_outbound_body_bytes as prepare_outbound_body_bytes, # noqa: F401 - compatibility export ) from headroom.proxy.body_forwarding import serialize_body_canonical +from headroom.proxy.ccr_golden_policy import ( + create_fresh_ccr_tool_definition, + replay_golden_ccr_tool_definition, +) from headroom.proxy.ccr_session_tracker import SessionCcrTracker as _SessionCcrTracker from headroom.proxy.internal_header_policy import ( INTERNAL_HEADER_PREFIX, @@ -2303,7 +2307,7 @@ def apply_session_sticky_ccr_tool( Returns ``(updated_tools, was_injected)``. ``updated_tools`` is a fresh list (caller-safe). """ - from headroom.ccr.tool_injection import CCR_TOOL_NAME, create_ccr_tool_definition + from headroom.ccr.tool_injection import CCR_TOOL_NAME if provider not in ("anthropic", "openai", "google"): raise ValueError(f"unsupported provider: {provider!r}") @@ -2337,14 +2341,13 @@ def apply_session_sticky_ccr_tool( request_id=request_id, ) return tools_out, False - tool_def = create_ccr_tool_definition(provider) - canonical = serialize_tool_definition_canonical(tool_def) - tools_out.append(tool_def) + replay = create_fresh_ccr_tool_definition(provider) + tools_out.append(replay.tool_definition) log_tool_injection_decision( provider=provider, session_id=None, decision="inject_first_time", - tool_definition_bytes_count=len(canonical), + tool_definition_bytes_count=len(replay.canonical_bytes), request_id=request_id, ) return tools_out, True @@ -2361,13 +2364,13 @@ def apply_session_sticky_ccr_tool( golden = tracker.get_golden_tool_bytes(provider, session_id) if golden is not None: try: - tool_def = json.loads(golden.decode("utf-8")) - tools_out.append(tool_def) + replay = replay_golden_ccr_tool_definition(golden) + tools_out.append(replay.tool_definition) log_tool_injection_decision( provider=provider, session_id=session_id, decision="inject_sticky_replay", - tool_definition_bytes_count=len(golden), + tool_definition_bytes_count=len(replay.canonical_bytes), request_id=request_id, ) return tools_out, True @@ -2381,15 +2384,14 @@ def apply_session_sticky_ccr_tool( # Fall through to fresh creation below # Tracker says "done CCR" but has no golden bytes (or they were corrupt). Pin # them now so future turns are stable. - tool_def = create_ccr_tool_definition(provider) - canonical = serialize_tool_definition_canonical(tool_def) - tracker.record_ccr_done(provider, session_id, canonical) - tools_out.append(tool_def) + replay = create_fresh_ccr_tool_definition(provider) + tracker.record_ccr_done(provider, session_id, replay.canonical_bytes) + tools_out.append(replay.tool_definition) log_tool_injection_decision( provider=provider, session_id=session_id, decision="inject_sticky_replay", - tool_definition_bytes_count=len(canonical), + tool_definition_bytes_count=len(replay.canonical_bytes), request_id=request_id, ) return tools_out, True @@ -2405,15 +2407,14 @@ def apply_session_sticky_ccr_tool( ) return tools_out, False - tool_def = create_ccr_tool_definition(provider) - canonical = serialize_tool_definition_canonical(tool_def) - tracker.record_ccr_done(provider, session_id, canonical) - tools_out.append(tool_def) + replay = create_fresh_ccr_tool_definition(provider) + tracker.record_ccr_done(provider, session_id, replay.canonical_bytes) + tools_out.append(replay.tool_definition) log_tool_injection_decision( provider=provider, session_id=session_id, decision="inject_first_time", - tool_definition_bytes_count=len(canonical), + tool_definition_bytes_count=len(replay.canonical_bytes), request_id=request_id, ) return tools_out, True diff --git a/tests/test_ccr_golden_policy.py b/tests/test_ccr_golden_policy.py new file mode 100644 index 000000000..1604f7936 --- /dev/null +++ b/tests/test_ccr_golden_policy.py @@ -0,0 +1,38 @@ +from __future__ import annotations + +import pytest + +from headroom.ccr.tool_injection import CCR_TOOL_NAME, create_ccr_tool_definition +from headroom.proxy.ccr_golden_policy import ( + create_fresh_ccr_tool_definition, + replay_golden_ccr_tool_definition, + serialize_ccr_tool_definition_canonical, +) + + +def test_replays_golden_definition_without_reserializing() -> None: + golden = b'{ "name" : "headroom_retrieve" , "description" : "client bytes" }' + + replay = replay_golden_ccr_tool_definition(golden) + + assert replay.tool_definition["name"] == CCR_TOOL_NAME + assert replay.canonical_bytes == golden + assert replay.used_golden_bytes is True + + +def test_rejects_invalid_golden_json() -> None: + with pytest.raises(ValueError): + replay_golden_ccr_tool_definition(b"not-json") + + +def test_rejects_non_utf8_golden_bytes() -> None: + with pytest.raises(UnicodeDecodeError): + replay_golden_ccr_tool_definition(b"\x80\x81") + + +def test_fresh_definition_uses_canonical_bytes() -> None: + replay = create_fresh_ccr_tool_definition("anthropic") + + assert replay.tool_definition == create_ccr_tool_definition("anthropic") + assert replay.canonical_bytes == serialize_ccr_tool_definition_canonical(replay.tool_definition) + assert replay.used_golden_bytes is False