From 0f846e5a8fb58942431b1edc22dfda9ab7d6de70 Mon Sep 17 00:00:00 2001 From: JD Davis Date: Sat, 11 Jul 2026 05:03:15 +0000 Subject: [PATCH] refactor(proxy): extract tool injection config (#2010) ## Description Extracts memory tool-injection operator config parsing from `headroom.proxy.helpers` into a focused config policy module. Existing helper functions and imports remain available while the environment parsing is now directly 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.tool_injection_config` for `HEADROOM_TOOL_INJECTION_STICKY` and `HEADROOM_TOOL_TRACKER_MAX_SESSIONS` parsing. - Updated `helpers.get_tool_injection_sticky_mode` and `helpers.get_tool_tracker_max_sessions` to delegate to the config module while preserving existing import paths. - Added direct tests for defaults, valid values, invalid values, and helper wrapper compatibility. ## 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_tool_injection_config.py tests/test_memory_tool_session_sticky.py tests/test_issue_728_empty_tools_injection.py 46 passed in 0.53s python -m ruff check . All checks passed! python -m ruff format --check . 1078 files already formatted python -m mypy headroom --ignore-missing-imports Success: no issues found in 415 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 `cb38f793`. - Exact command / steps: Ran targeted tool-injection config, memory session sticky, and empty-tool 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/helpers.py | 38 +++++--------- headroom/proxy/tool_injection_config.py | 43 ++++++++++++++++ tests/test_tool_injection_config.py | 66 +++++++++++++++++++++++++ 3 files changed, 120 insertions(+), 27 deletions(-) create mode 100644 headroom/proxy/tool_injection_config.py create mode 100644 tests/test_tool_injection_config.py diff --git a/headroom/proxy/helpers.py b/headroom/proxy/helpers.py index ef6f1e105..68d1ae7f4 100644 --- a/headroom/proxy/helpers.py +++ b/headroom/proxy/helpers.py @@ -38,6 +38,15 @@ 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.tool_injection_config import ( + ToolInjectionStickyMode, +) +from headroom.proxy.tool_injection_config import ( + get_tool_injection_sticky_mode as _get_tool_injection_sticky_mode, +) +from headroom.proxy.tool_injection_config import ( + get_tool_tracker_max_sessions as _get_tool_tracker_max_sessions, +) if TYPE_CHECKING: import httpx @@ -1923,13 +1932,6 @@ def log_beta_header_merge( # silent fallback. It exists for diagnostic shadow tracing / emergency # rollback only. -_TOOL_INJECTION_STICKY_ENV = "HEADROOM_TOOL_INJECTION_STICKY" -ToolInjectionStickyMode = Literal["enabled", "disabled"] -_TOOL_INJECTION_STICKY_DEFAULT: ToolInjectionStickyMode = "enabled" - -_TOOL_TRACKER_MAX_SESSIONS_ENV = "HEADROOM_TOOL_TRACKER_MAX_SESSIONS" -_TOOL_TRACKER_MAX_SESSIONS_DEFAULT = 1000 - def get_tool_injection_sticky_mode() -> ToolInjectionStickyMode: """Return the active memory-tool stickiness mode. @@ -1938,30 +1940,12 @@ def get_tool_injection_sticky_mode() -> ToolInjectionStickyMode: restart. Unknown values raise loudly per the no-silent-fallback build constraint. """ - raw = os.environ.get(_TOOL_INJECTION_STICKY_ENV, "").strip().lower() - if not raw: - return _TOOL_INJECTION_STICKY_DEFAULT - if raw in ("enabled", "disabled"): - return cast(ToolInjectionStickyMode, raw) - raise ValueError( - f"Invalid {_TOOL_INJECTION_STICKY_ENV}={raw!r}; expected 'enabled' or 'disabled'" - ) + return _get_tool_injection_sticky_mode() def get_tool_tracker_max_sessions() -> int: """Return the LRU bound for `SessionToolTracker` (sessions cap).""" - raw = os.environ.get(_TOOL_TRACKER_MAX_SESSIONS_ENV, "").strip() - if not raw: - return _TOOL_TRACKER_MAX_SESSIONS_DEFAULT - try: - value = int(raw) - except ValueError as exc: - raise ValueError( - f"Invalid {_TOOL_TRACKER_MAX_SESSIONS_ENV}={raw!r}; expected positive int" - ) from exc - if value <= 0: - raise ValueError(f"Invalid {_TOOL_TRACKER_MAX_SESSIONS_ENV}={raw!r}; expected positive int") - return value + return _get_tool_tracker_max_sessions() def serialize_tool_definition_canonical(tool_definition: dict[str, Any]) -> bytes: diff --git a/headroom/proxy/tool_injection_config.py b/headroom/proxy/tool_injection_config.py new file mode 100644 index 000000000..01ebdf0bf --- /dev/null +++ b/headroom/proxy/tool_injection_config.py @@ -0,0 +1,43 @@ +"""Operator configuration policy for proxy tool injection.""" + +from __future__ import annotations + +import os +from typing import Literal, cast + +TOOL_INJECTION_STICKY_ENV = "HEADROOM_TOOL_INJECTION_STICKY" +ToolInjectionStickyMode = Literal["enabled", "disabled"] +TOOL_INJECTION_STICKY_DEFAULT: ToolInjectionStickyMode = "enabled" + +TOOL_TRACKER_MAX_SESSIONS_ENV = "HEADROOM_TOOL_TRACKER_MAX_SESSIONS" +TOOL_TRACKER_MAX_SESSIONS_DEFAULT = 1000 + + +def get_tool_injection_sticky_mode() -> ToolInjectionStickyMode: + """Return the active memory-tool stickiness mode.""" + + raw = os.environ.get(TOOL_INJECTION_STICKY_ENV, "").strip().lower() + if not raw: + return TOOL_INJECTION_STICKY_DEFAULT + if raw in ("enabled", "disabled"): + return cast(ToolInjectionStickyMode, raw) + raise ValueError( + f"Invalid {TOOL_INJECTION_STICKY_ENV}={raw!r}; expected 'enabled' or 'disabled'" + ) + + +def get_tool_tracker_max_sessions() -> int: + """Return the LRU bound for memory tool session tracking.""" + + raw = os.environ.get(TOOL_TRACKER_MAX_SESSIONS_ENV, "").strip() + if not raw: + return TOOL_TRACKER_MAX_SESSIONS_DEFAULT + try: + value = int(raw) + except ValueError as exc: + raise ValueError( + f"Invalid {TOOL_TRACKER_MAX_SESSIONS_ENV}={raw!r}; expected positive int" + ) from exc + if value <= 0: + raise ValueError(f"Invalid {TOOL_TRACKER_MAX_SESSIONS_ENV}={raw!r}; expected positive int") + return value diff --git a/tests/test_tool_injection_config.py b/tests/test_tool_injection_config.py new file mode 100644 index 000000000..73ecb98d6 --- /dev/null +++ b/tests/test_tool_injection_config.py @@ -0,0 +1,66 @@ +from __future__ import annotations + +import pytest + +from headroom.proxy.helpers import ( + get_tool_injection_sticky_mode as helper_get_tool_injection_sticky_mode, +) +from headroom.proxy.helpers import ( + get_tool_tracker_max_sessions as helper_get_tool_tracker_max_sessions, +) +from headroom.proxy.tool_injection_config import ( + get_tool_injection_sticky_mode, + get_tool_tracker_max_sessions, +) + + +def test_sticky_mode_defaults_enabled(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.delenv("HEADROOM_TOOL_INJECTION_STICKY", raising=False) + + assert get_tool_injection_sticky_mode() == "enabled" + + +def test_sticky_mode_accepts_enabled_and_disabled(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("HEADROOM_TOOL_INJECTION_STICKY", "enabled") + assert get_tool_injection_sticky_mode() == "enabled" + + monkeypatch.setenv("HEADROOM_TOOL_INJECTION_STICKY", " DISABLED ") + assert get_tool_injection_sticky_mode() == "disabled" + + +def test_sticky_mode_rejects_unknown_values(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("HEADROOM_TOOL_INJECTION_STICKY", "maybe") + + with pytest.raises(ValueError, match="HEADROOM_TOOL_INJECTION_STICKY"): + get_tool_injection_sticky_mode() + + +def test_tracker_max_sessions_defaults_to_1000(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.delenv("HEADROOM_TOOL_TRACKER_MAX_SESSIONS", raising=False) + + assert get_tool_tracker_max_sessions() == 1000 + + +def test_tracker_max_sessions_accepts_positive_int(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("HEADROOM_TOOL_TRACKER_MAX_SESSIONS", "42") + + assert get_tool_tracker_max_sessions() == 42 + + +@pytest.mark.parametrize("raw", ["0", "-1", "not-int"]) +def test_tracker_max_sessions_rejects_invalid_values( + monkeypatch: pytest.MonkeyPatch, + raw: str, +) -> None: + monkeypatch.setenv("HEADROOM_TOOL_TRACKER_MAX_SESSIONS", raw) + + with pytest.raises(ValueError, match="HEADROOM_TOOL_TRACKER_MAX_SESSIONS"): + get_tool_tracker_max_sessions() + + +def test_helpers_keep_existing_config_import_paths(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("HEADROOM_TOOL_INJECTION_STICKY", "disabled") + monkeypatch.setenv("HEADROOM_TOOL_TRACKER_MAX_SESSIONS", "12") + + assert helper_get_tool_injection_sticky_mode() == get_tool_injection_sticky_mode() + assert helper_get_tool_tracker_max_sessions() == get_tool_tracker_max_sessions()