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"