From f359f21424a54e9d9ef34ca7de49a9d11aa50589 Mon Sep 17 00:00:00 2001 From: JD Davis Date: Mon, 13 Jul 2026 04:34:25 +0000 Subject: [PATCH] refactor(proxy): extract beta header merge policy (#1993) ## Description Extracts deterministic beta-header token parsing and merge rules from `headroom.proxy.helpers` into a focused module. Existing helper names remain available for Anthropic/OpenAI handlers and the session beta tracker. 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.beta_header_merge` for beta token splitting and deterministic merge behavior. - Re-exported the existing `merge_anthropic_beta` and `merge_openai_beta` helper names from `helpers.py` for compatibility. - Added direct unit tests for token splitting, ordering, case-insensitive dedupe, empty required tokens, and provider wrappers. ## 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_beta_header_merge.py tests/test_anthropic_beta_session_sticky.py tests/test_openai_beta_session_sticky.py 48 passed in 0.38s 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 - Exact command / steps: Ran focused beta merge tests, existing Anthropic/OpenAI beta sticky suites, full ruff, format check, mypy, and staged gitleaks scan. - Observed result: Existing beta merge and tracker behavior remains green while extracted merge rules are covered directly. - Not tested: Full repository pytest suite locally; CI covers the broader matrix. ## 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 default-branch Dependabot alerts reported during push are pre-existing and unrelated to this PR. --- headroom/proxy/beta_header_merge.py | 53 ++++++++++++++++++ headroom/proxy/helpers.py | 83 +++++------------------------ tests/test_beta_header_merge.py | 33 ++++++++++++ 3 files changed, 98 insertions(+), 71 deletions(-) create mode 100644 headroom/proxy/beta_header_merge.py create mode 100644 tests/test_beta_header_merge.py diff --git a/headroom/proxy/beta_header_merge.py b/headroom/proxy/beta_header_merge.py new file mode 100644 index 000000000..572a3be89 --- /dev/null +++ b/headroom/proxy/beta_header_merge.py @@ -0,0 +1,53 @@ +"""Deterministic merge helpers for provider beta request headers.""" + +from __future__ import annotations + + +def split_beta_tokens(value: str | None) -> list[str]: + """Split a comma-separated beta-header value into trimmed tokens.""" + + if not value: + return [] + out: list[str] = [] + for raw in value.split(","): + token = raw.strip() + if token: + out.append(token) + return out + + +def merge_beta_tokens(client_value: str | None, headroom_required: list[str]) -> str: + """Merge client beta tokens with Headroom-required tokens deterministically.""" + + seen_lower: set[str] = set() + out: list[str] = [] + for token in split_beta_tokens(client_value): + lower = token.lower() + if lower in seen_lower: + continue + seen_lower.add(lower) + out.append(token) + for token in headroom_required: + if not token: + continue + token = token.strip() + if not token: + continue + lower = token.lower() + if lower in seen_lower: + continue + seen_lower.add(lower) + out.append(token) + return ",".join(out) + + +def merge_anthropic_beta(client_value: str | None, headroom_required: list[str]) -> str: + """Merge client anthropic-beta value with Headroom-required tokens.""" + + return merge_beta_tokens(client_value, headroom_required) + + +def merge_openai_beta(client_value: str | None, headroom_required: list[str]) -> str: + """Merge client OpenAI-Beta value with Headroom-required tokens.""" + + return merge_beta_tokens(client_value, headroom_required) diff --git a/headroom/proxy/helpers.py b/headroom/proxy/helpers.py index 33f3ed073..1eb41c124 100644 --- a/headroom/proxy/helpers.py +++ b/headroom/proxy/helpers.py @@ -33,6 +33,16 @@ from headroom.proxy import ( wire_debug_format_policy, wire_debug_redaction_policy, ) +from headroom.proxy.beta_header_merge import ( + merge_anthropic_beta as merge_anthropic_beta, +) +from headroom.proxy.beta_header_merge import ( + merge_beta_tokens, + split_beta_tokens, +) +from headroom.proxy.beta_header_merge import ( + merge_openai_beta as merge_openai_beta, +) from headroom.proxy.beta_header_policy import ( BETA_HEADER_STICKY_DEFAULT, BETA_HEADER_STICKY_ENV, @@ -1565,79 +1575,10 @@ def get_beta_tracker_max_sessions() -> int: return resolve_beta_tracker_max_sessions(os.environ.get(_BETA_TRACKER_MAX_SESSIONS_ENV)) -def _split_beta_tokens(value: str | None) -> list[str]: - """Split a comma-separated beta-header value into trimmed tokens. - - Empty/whitespace-only entries are dropped. Pure function, no regex. - """ - if not value: - return [] - out: list[str] = [] - for raw in value.split(","): - token = raw.strip() - if token: - out.append(token) - return out +_split_beta_tokens = split_beta_tokens -def _merge_beta_tokens(client_value: str | None, headroom_required: list[str]) -> str: - """Shared deterministic merge for `anthropic-beta` / `OpenAI-Beta` tokens. - - Rules (per Anthropic guide §6.3 #6 "sticky-on; add but never reorder"): - - * Client tokens come first, in their original order. - * Headroom-required tokens append in the order given, skipping any - token already present (case-insensitive). - * Dedupe is case-insensitive but the FIRST occurrence's casing wins - (prevents drift when client uses one casing across turns). - * Returns ``""`` when both inputs are empty. - - Pure function. No regex. No global state. - """ - seen_lower: set[str] = set() - out: list[str] = [] - for token in _split_beta_tokens(client_value): - lower = token.lower() - if lower in seen_lower: - continue - seen_lower.add(lower) - out.append(token) - for token in headroom_required: - if not token: - continue - token = token.strip() - if not token: - continue - lower = token.lower() - if lower in seen_lower: - continue - seen_lower.add(lower) - out.append(token) - return ",".join(out) - - -def merge_anthropic_beta(client_value: str | None, headroom_required: list[str]) -> str: - """Merge client `anthropic-beta` value with Headroom-required tokens. - - See `_merge_beta_tokens` for full semantics. Order is deterministic: - client tokens first (in their original order), then headroom tokens - (in the order passed). No sorting — sticky-on per Anthropic guide - §6.3 #6 means we add but never reorder. Dedupe is case-insensitive - but preserves the original casing of the first occurrence. - - Returns ``""`` when both inputs are empty. - """ - return _merge_beta_tokens(client_value, headroom_required) - - -def merge_openai_beta(client_value: str | None, headroom_required: list[str]) -> str: - """Merge client `OpenAI-Beta` value with Headroom-required tokens. - - Mirror of `merge_anthropic_beta`. Same semantics — the OpenAI header - follows the same comma-separated convention and the same cache-stable - rules apply. - """ - return _merge_beta_tokens(client_value, headroom_required) +_merge_beta_tokens = merge_beta_tokens class SessionBetaTracker: diff --git a/tests/test_beta_header_merge.py b/tests/test_beta_header_merge.py new file mode 100644 index 000000000..4dba8937f --- /dev/null +++ b/tests/test_beta_header_merge.py @@ -0,0 +1,33 @@ +from __future__ import annotations + +from headroom.proxy.beta_header_merge import ( + merge_anthropic_beta, + merge_beta_tokens, + merge_openai_beta, + split_beta_tokens, +) + + +def test_split_beta_tokens_drops_empty_entries() -> None: + assert split_beta_tokens(None) == [] + assert split_beta_tokens("") == [] + assert split_beta_tokens(" alpha, , beta ,, ") == ["alpha", "beta"] + + +def test_merge_beta_tokens_preserves_client_order_then_appends_required() -> None: + assert merge_beta_tokens("client-1,client-2", ["required-1", "required-2"]) == ( + "client-1,client-2,required-1,required-2" + ) + + +def test_merge_beta_tokens_dedupes_case_insensitively_with_first_casing() -> None: + assert merge_beta_tokens("Foo,foo", ["FOO", "bar"]) == "Foo,bar" + + +def test_merge_beta_tokens_skips_empty_required_values() -> None: + assert merge_beta_tokens("alpha", ["", " beta ", " "]) == "alpha,beta" + + +def test_provider_wrappers_share_merge_semantics() -> None: + assert merge_anthropic_beta("a", ["b"]) == "a,b" + assert merge_openai_beta("a", ["b"]) == "a,b"