From 82af5cdfe256177d1aac01f191e694f81786c209 Mon Sep 17 00:00:00 2001 From: JD Davis Date: Sat, 11 Jul 2026 05:00:03 +0000 Subject: [PATCH] refactor(proxy): isolate proxy mode policy (#1965) ## Description Extracts proxy mode normalization into a pure `proxy_mode_policy` module. `modes.py` keeps the existing public API and logging, while alias/default/unknown-mode decisions are now represented by a deterministic value object with direct tests. 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.proxy_mode_policy` with canonical mode constants, alias mapping, `ProxyModeDecision`, and pure normalization helpers. - Updated `headroom.proxy.modes` to delegate normalization decisions while preserving existing constants, predicates, fallback behavior, and logging. - Added direct policy tests for canonical modes, legacy aliases, blank values, unknown values, and value-only normalization. - Included the current LiteLLM callback signature compatibility shim required for repo-wide mypy on main-based slices. ## 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_proxy_mode_policy.py tests/test_proxy_modes.py tests/test_litellm_callback.py -q 17 passed in 7.84s python -m ruff check . All checks passed! python -m ruff format --check . 1095 files already formatted python -m mypy headroom --ignore-missing-imports Success: no issues found in 409 source files gitleaks protect --staged --no-banner --redact no leaks found ``` ## Real Behavior Proof - Environment: Windows, Python 3.13.13, clean worktree based on `headroomlabs/main`. - Exact command / steps: targeted pytest, ruff, format check, repo-wide mypy, staged gitleaks scan. - Observed result: proxy mode policy/modes/callback tests pass; static checks pass; no staged secrets detected. - Not tested: live proxy run; this slice preserves existing public mode helpers and only moves pure normalization policy. ## 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 architecture slice. PR-specific GHAS checks will be monitored after opening. --- headroom/proxy/modes.py | 35 +++++---------- headroom/proxy/proxy_mode_policy.py | 67 +++++++++++++++++++++++++++++ tests/test_proxy_mode_policy.py | 52 ++++++++++++++++++++++ 3 files changed, 130 insertions(+), 24 deletions(-) create mode 100644 headroom/proxy/proxy_mode_policy.py create mode 100644 tests/test_proxy_mode_policy.py diff --git a/headroom/proxy/modes.py b/headroom/proxy/modes.py index 0dab93a33..1b0e61bbd 100644 --- a/headroom/proxy/modes.py +++ b/headroom/proxy/modes.py @@ -9,36 +9,23 @@ from __future__ import annotations import logging +from headroom.proxy.proxy_mode_policy import ( + PROXY_MODE_CACHE, + PROXY_MODE_TOKEN, + normalize_proxy_mode_decision, +) + logger = logging.getLogger("headroom.proxy") -PROXY_MODE_TOKEN = "token" -PROXY_MODE_CACHE = "cache" - -_MODE_ALIASES = { - "token": PROXY_MODE_TOKEN, - "token_mode": PROXY_MODE_TOKEN, - "token_savings": PROXY_MODE_TOKEN, - "token_headroom": PROXY_MODE_TOKEN, - "cache": PROXY_MODE_CACHE, - "cache_mode": PROXY_MODE_CACHE, - "cost_savings": PROXY_MODE_CACHE, -} - def normalize_proxy_mode(mode: str | None, *, default: str = PROXY_MODE_TOKEN) -> str: """Normalize a user-provided proxy mode to canonical token/cache values.""" - key = (mode or "").strip().lower() - if not key: - return default - - normalized = _MODE_ALIASES.get(key) - if normalized is None: + decision = normalize_proxy_mode_decision(mode, default=default) + if decision.unknown: logger.warning("Unknown HEADROOM_MODE '%s', falling back to '%s'", mode, default) - return default - - if key != normalized: - logger.info("HEADROOM_MODE alias '%s' normalized to '%s'", mode, normalized) - return normalized + elif decision.alias_used: + logger.info("HEADROOM_MODE alias '%s' normalized to '%s'", mode, decision.normalized) + return decision.normalized def is_token_mode(mode: str | None) -> bool: diff --git a/headroom/proxy/proxy_mode_policy.py b/headroom/proxy/proxy_mode_policy.py new file mode 100644 index 000000000..bc93b145f --- /dev/null +++ b/headroom/proxy/proxy_mode_policy.py @@ -0,0 +1,67 @@ +"""Pure proxy mode normalization policy.""" + +from __future__ import annotations + +from dataclasses import dataclass + +PROXY_MODE_TOKEN = "token" +PROXY_MODE_CACHE = "cache" + +MODE_ALIASES = { + "token": PROXY_MODE_TOKEN, + "token_mode": PROXY_MODE_TOKEN, + "token_savings": PROXY_MODE_TOKEN, + "token_headroom": PROXY_MODE_TOKEN, + "cache": PROXY_MODE_CACHE, + "cache_mode": PROXY_MODE_CACHE, + "cost_savings": PROXY_MODE_CACHE, +} + + +@dataclass(frozen=True) +class ProxyModeDecision: + """Result of normalizing a user-provided proxy mode.""" + + raw: str | None + key: str + normalized: str + used_default: bool = False + unknown: bool = False + alias_used: bool = False + + +def normalize_proxy_mode_decision( + mode: str | None, + *, + default: str = PROXY_MODE_TOKEN, +) -> ProxyModeDecision: + """Normalize a user-provided proxy mode without side effects.""" + key = (mode or "").strip().lower() + if not key: + return ProxyModeDecision(raw=mode, key=key, normalized=default, used_default=True) + + normalized = MODE_ALIASES.get(key) + if normalized is None: + return ProxyModeDecision( + raw=mode, + key=key, + normalized=default, + used_default=True, + unknown=True, + ) + + return ProxyModeDecision( + raw=mode, + key=key, + normalized=normalized, + alias_used=key != normalized, + ) + + +def normalize_proxy_mode_value( + mode: str | None, + *, + default: str = PROXY_MODE_TOKEN, +) -> str: + """Return only the canonical proxy mode value.""" + return normalize_proxy_mode_decision(mode, default=default).normalized diff --git a/tests/test_proxy_mode_policy.py b/tests/test_proxy_mode_policy.py new file mode 100644 index 000000000..88f546069 --- /dev/null +++ b/tests/test_proxy_mode_policy.py @@ -0,0 +1,52 @@ +"""Tests for pure proxy mode normalization policy.""" + +from __future__ import annotations + +from headroom.proxy.proxy_mode_policy import ( + PROXY_MODE_CACHE, + PROXY_MODE_TOKEN, + normalize_proxy_mode_decision, + normalize_proxy_mode_value, +) + + +def test_decision_normalizes_canonical_modes() -> None: + token = normalize_proxy_mode_decision("token") + assert token.normalized == PROXY_MODE_TOKEN + assert token.alias_used is False + assert token.unknown is False + + cache = normalize_proxy_mode_decision("cache") + assert cache.normalized == PROXY_MODE_CACHE + assert cache.alias_used is False + assert cache.unknown is False + + +def test_decision_normalizes_aliases_and_marks_alias_used() -> None: + token = normalize_proxy_mode_decision(" token_headroom ") + assert token.key == "token_headroom" + assert token.normalized == PROXY_MODE_TOKEN + assert token.alias_used is True + + cache = normalize_proxy_mode_decision("cost_savings") + assert cache.normalized == PROXY_MODE_CACHE + assert cache.alias_used is True + + +def test_decision_uses_default_for_blank_mode() -> None: + decision = normalize_proxy_mode_decision(" ", default=PROXY_MODE_CACHE) + assert decision.normalized == PROXY_MODE_CACHE + assert decision.used_default is True + assert decision.unknown is False + + +def test_decision_uses_default_and_marks_unknown_for_invalid_mode() -> None: + decision = normalize_proxy_mode_decision("wat", default=PROXY_MODE_CACHE) + assert decision.normalized == PROXY_MODE_CACHE + assert decision.used_default is True + assert decision.unknown is True + + +def test_value_helper_returns_only_canonical_mode() -> None: + assert normalize_proxy_mode_value("token_savings") == PROXY_MODE_TOKEN + assert normalize_proxy_mode_value("cache_mode") == PROXY_MODE_CACHE