Skip to content

Commit 0f846e5

Browse files
authored
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.
1 parent b3a559b commit 0f846e5

3 files changed

Lines changed: 120 additions & 27 deletions

File tree

headroom/proxy/helpers.py

Lines changed: 11 additions & 27 deletions
Original file line numberDiff line numberDiff line change
@@ -38,6 +38,15 @@
3838
prepare_outbound_body_bytes as prepare_outbound_body_bytes, # noqa: F401 - compatibility export
3939
)
4040
from headroom.proxy.body_forwarding import serialize_body_canonical
41+
from headroom.proxy.tool_injection_config import (
42+
ToolInjectionStickyMode,
43+
)
44+
from headroom.proxy.tool_injection_config import (
45+
get_tool_injection_sticky_mode as _get_tool_injection_sticky_mode,
46+
)
47+
from headroom.proxy.tool_injection_config import (
48+
get_tool_tracker_max_sessions as _get_tool_tracker_max_sessions,
49+
)
4150

4251
if TYPE_CHECKING:
4352
import httpx
@@ -1923,13 +1932,6 @@ def log_beta_header_merge(
19231932
# silent fallback. It exists for diagnostic shadow tracing / emergency
19241933
# rollback only.
19251934

1926-
_TOOL_INJECTION_STICKY_ENV = "HEADROOM_TOOL_INJECTION_STICKY"
1927-
ToolInjectionStickyMode = Literal["enabled", "disabled"]
1928-
_TOOL_INJECTION_STICKY_DEFAULT: ToolInjectionStickyMode = "enabled"
1929-
1930-
_TOOL_TRACKER_MAX_SESSIONS_ENV = "HEADROOM_TOOL_TRACKER_MAX_SESSIONS"
1931-
_TOOL_TRACKER_MAX_SESSIONS_DEFAULT = 1000
1932-
19331935

19341936
def get_tool_injection_sticky_mode() -> ToolInjectionStickyMode:
19351937
"""Return the active memory-tool stickiness mode.
@@ -1938,30 +1940,12 @@ def get_tool_injection_sticky_mode() -> ToolInjectionStickyMode:
19381940
restart. Unknown values raise loudly per the no-silent-fallback
19391941
build constraint.
19401942
"""
1941-
raw = os.environ.get(_TOOL_INJECTION_STICKY_ENV, "").strip().lower()
1942-
if not raw:
1943-
return _TOOL_INJECTION_STICKY_DEFAULT
1944-
if raw in ("enabled", "disabled"):
1945-
return cast(ToolInjectionStickyMode, raw)
1946-
raise ValueError(
1947-
f"Invalid {_TOOL_INJECTION_STICKY_ENV}={raw!r}; expected 'enabled' or 'disabled'"
1948-
)
1943+
return _get_tool_injection_sticky_mode()
19491944

19501945

19511946
def get_tool_tracker_max_sessions() -> int:
19521947
"""Return the LRU bound for `SessionToolTracker` (sessions cap)."""
1953-
raw = os.environ.get(_TOOL_TRACKER_MAX_SESSIONS_ENV, "").strip()
1954-
if not raw:
1955-
return _TOOL_TRACKER_MAX_SESSIONS_DEFAULT
1956-
try:
1957-
value = int(raw)
1958-
except ValueError as exc:
1959-
raise ValueError(
1960-
f"Invalid {_TOOL_TRACKER_MAX_SESSIONS_ENV}={raw!r}; expected positive int"
1961-
) from exc
1962-
if value <= 0:
1963-
raise ValueError(f"Invalid {_TOOL_TRACKER_MAX_SESSIONS_ENV}={raw!r}; expected positive int")
1964-
return value
1948+
return _get_tool_tracker_max_sessions()
19651949

19661950

19671951
def serialize_tool_definition_canonical(tool_definition: dict[str, Any]) -> bytes:
Lines changed: 43 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,43 @@
1+
"""Operator configuration policy for proxy tool injection."""
2+
3+
from __future__ import annotations
4+
5+
import os
6+
from typing import Literal, cast
7+
8+
TOOL_INJECTION_STICKY_ENV = "HEADROOM_TOOL_INJECTION_STICKY"
9+
ToolInjectionStickyMode = Literal["enabled", "disabled"]
10+
TOOL_INJECTION_STICKY_DEFAULT: ToolInjectionStickyMode = "enabled"
11+
12+
TOOL_TRACKER_MAX_SESSIONS_ENV = "HEADROOM_TOOL_TRACKER_MAX_SESSIONS"
13+
TOOL_TRACKER_MAX_SESSIONS_DEFAULT = 1000
14+
15+
16+
def get_tool_injection_sticky_mode() -> ToolInjectionStickyMode:
17+
"""Return the active memory-tool stickiness mode."""
18+
19+
raw = os.environ.get(TOOL_INJECTION_STICKY_ENV, "").strip().lower()
20+
if not raw:
21+
return TOOL_INJECTION_STICKY_DEFAULT
22+
if raw in ("enabled", "disabled"):
23+
return cast(ToolInjectionStickyMode, raw)
24+
raise ValueError(
25+
f"Invalid {TOOL_INJECTION_STICKY_ENV}={raw!r}; expected 'enabled' or 'disabled'"
26+
)
27+
28+
29+
def get_tool_tracker_max_sessions() -> int:
30+
"""Return the LRU bound for memory tool session tracking."""
31+
32+
raw = os.environ.get(TOOL_TRACKER_MAX_SESSIONS_ENV, "").strip()
33+
if not raw:
34+
return TOOL_TRACKER_MAX_SESSIONS_DEFAULT
35+
try:
36+
value = int(raw)
37+
except ValueError as exc:
38+
raise ValueError(
39+
f"Invalid {TOOL_TRACKER_MAX_SESSIONS_ENV}={raw!r}; expected positive int"
40+
) from exc
41+
if value <= 0:
42+
raise ValueError(f"Invalid {TOOL_TRACKER_MAX_SESSIONS_ENV}={raw!r}; expected positive int")
43+
return value
Lines changed: 66 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,66 @@
1+
from __future__ import annotations
2+
3+
import pytest
4+
5+
from headroom.proxy.helpers import (
6+
get_tool_injection_sticky_mode as helper_get_tool_injection_sticky_mode,
7+
)
8+
from headroom.proxy.helpers import (
9+
get_tool_tracker_max_sessions as helper_get_tool_tracker_max_sessions,
10+
)
11+
from headroom.proxy.tool_injection_config import (
12+
get_tool_injection_sticky_mode,
13+
get_tool_tracker_max_sessions,
14+
)
15+
16+
17+
def test_sticky_mode_defaults_enabled(monkeypatch: pytest.MonkeyPatch) -> None:
18+
monkeypatch.delenv("HEADROOM_TOOL_INJECTION_STICKY", raising=False)
19+
20+
assert get_tool_injection_sticky_mode() == "enabled"
21+
22+
23+
def test_sticky_mode_accepts_enabled_and_disabled(monkeypatch: pytest.MonkeyPatch) -> None:
24+
monkeypatch.setenv("HEADROOM_TOOL_INJECTION_STICKY", "enabled")
25+
assert get_tool_injection_sticky_mode() == "enabled"
26+
27+
monkeypatch.setenv("HEADROOM_TOOL_INJECTION_STICKY", " DISABLED ")
28+
assert get_tool_injection_sticky_mode() == "disabled"
29+
30+
31+
def test_sticky_mode_rejects_unknown_values(monkeypatch: pytest.MonkeyPatch) -> None:
32+
monkeypatch.setenv("HEADROOM_TOOL_INJECTION_STICKY", "maybe")
33+
34+
with pytest.raises(ValueError, match="HEADROOM_TOOL_INJECTION_STICKY"):
35+
get_tool_injection_sticky_mode()
36+
37+
38+
def test_tracker_max_sessions_defaults_to_1000(monkeypatch: pytest.MonkeyPatch) -> None:
39+
monkeypatch.delenv("HEADROOM_TOOL_TRACKER_MAX_SESSIONS", raising=False)
40+
41+
assert get_tool_tracker_max_sessions() == 1000
42+
43+
44+
def test_tracker_max_sessions_accepts_positive_int(monkeypatch: pytest.MonkeyPatch) -> None:
45+
monkeypatch.setenv("HEADROOM_TOOL_TRACKER_MAX_SESSIONS", "42")
46+
47+
assert get_tool_tracker_max_sessions() == 42
48+
49+
50+
@pytest.mark.parametrize("raw", ["0", "-1", "not-int"])
51+
def test_tracker_max_sessions_rejects_invalid_values(
52+
monkeypatch: pytest.MonkeyPatch,
53+
raw: str,
54+
) -> None:
55+
monkeypatch.setenv("HEADROOM_TOOL_TRACKER_MAX_SESSIONS", raw)
56+
57+
with pytest.raises(ValueError, match="HEADROOM_TOOL_TRACKER_MAX_SESSIONS"):
58+
get_tool_tracker_max_sessions()
59+
60+
61+
def test_helpers_keep_existing_config_import_paths(monkeypatch: pytest.MonkeyPatch) -> None:
62+
monkeypatch.setenv("HEADROOM_TOOL_INJECTION_STICKY", "disabled")
63+
monkeypatch.setenv("HEADROOM_TOOL_TRACKER_MAX_SESSIONS", "12")
64+
65+
assert helper_get_tool_injection_sticky_mode() == get_tool_injection_sticky_mode()
66+
assert helper_get_tool_tracker_max_sessions() == get_tool_tracker_max_sessions()

0 commit comments

Comments
 (0)