headroomlabs-ai
diff --git a/‎headroom/cli/proxy.py‎
Lines changed: 21 additions & 0 deletions b/‎headroom/cli/proxy.py‎
Lines changed: 21 additions & 0 deletions
diff --git a/‎tests/test_cli_proxy_malloc_reexec_guard.py‎
Lines changed: 103 additions & 0 deletions b/‎tests/test_cli_proxy_malloc_reexec_guard.py‎
Lines changed: 103 additions & 0 deletions
diff --git a/‎tests/test_gemini_ccr_continuation_usage.py‎
Lines changed: 4 additions & 0 deletions b/‎tests/test_gemini_ccr_continuation_usage.py‎
Lines changed: 4 additions & 0 deletions
diff --git a/‎tests/test_malloc_tuning.py‎
Lines changed: 4 additions & 0 deletions b/‎tests/test_malloc_tuning.py‎
Lines changed: 4 additions & 0 deletions
@@ -5,6 +5,7 @@
 import sys
 import warnings
 from importlib import import_module
+from pathlib import Path
 from typing import Any, Literal, cast
 
 import click
@@ -125,13 +126,33 @@ def _get_env_bool_optional(name: str) -> bool | None:
 }
 
 
+def _process_is_headroom_cli_entrypoint() -> bool:
+    """Is this process the Headroom CLI itself, rather than an embedder?
+
+    ``_reexec_with_malloc_tuning`` rebuilds the command line as
+    ``python -m headroom.cli <argv[1:]>``. That is only a faithful
+    reconstruction when the process really was started as the Headroom CLI. If
+    something else invoked the ``proxy`` command in-process — pytest's
+    ``CliRunner``, an embedding application, ``runpy`` — then ``argv[1:]``
+    belongs to *that* program, and ``os.execv`` would replace it with a Headroom
+    process parsing arguments that were never meant for us.
+    """
+    argv0 = Path(sys.argv[0] or "")
+    if argv0.name in {"headroom", "headroom.exe"}:
+        return True
+    # `python -m headroom.cli` sets argv[0] to .../headroom/cli/__main__.py.
+    return argv0.parts[-3:] == ("headroom", "cli", "__main__.py")
+
+
 def _reexec_with_malloc_tuning() -> None:
     if sys.platform != "darwin":
         return
     if not _get_env_bool("HEADROOM_MALLOC_TUNING", True):
         return
     if os.environ.get("_HEADROOM_MALLOC_TUNED") == "1":
         return
+    if not _process_is_headroom_cli_entrypoint():
+        return
     missing = {k: v for k, v in _MALLOC_TUNING.items() if k not in os.environ}
     # Set the loop guard before the re-exec so the replacement process (which
     # inherits this environment) skips this path instead of re-execing forever.
 
@@ -0,0 +1,103 @@
+"""The macOS malloc re-exec must never replace an embedder's process.
+
+``headroom proxy`` re-execs itself once on Darwin to apply two libmalloc knobs
+that libmalloc only reads before ``main()`` (#2820). The re-exec rebuilds the
+command as ``python -m headroom.cli <argv[1:]>``, which is only a faithful
+reconstruction when this process really is the Headroom CLI.
+
+When the ``proxy`` command is invoked *in-process* — pytest's ``CliRunner``, an
+embedding application — ``os.execv`` replaces that process instead. The whole
+pytest run is destroyed mid-suite with no traceback, and the replacement
+Headroom process is handed pytest's own argv.
+
+CI cannot catch this: the tuning is Darwin-only and no CI runner is macOS, so
+these tests assert the guard's *logic* on every platform rather than relying on
+the re-exec being reachable.
+"""
+
+from __future__ import annotations
+
+import sys
+
+import pytest
+
+from headroom.cli import proxy as proxy_cli
+
+
+@pytest.mark.parametrize(
+    "argv0",
+    [
+        "/usr/local/bin/headroom",
+        "/opt/homebrew/bin/headroom",
+    ],
+)
+def test_console_script_is_recognised_as_the_entrypoint(
+    argv0: str, monkeypatch: pytest.MonkeyPatch
+) -> None:
+    monkeypatch.setattr(sys, "argv", [argv0, "proxy"])
+    assert proxy_cli._process_is_headroom_cli_entrypoint() is True
+
+
+def test_module_invocation_is_recognised_as_the_entrypoint(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    monkeypatch.setattr(
+        sys, "argv", ["/venv/lib/python3.12/site-packages/headroom/cli/__main__.py", "proxy"]
+    )
+    assert proxy_cli._process_is_headroom_cli_entrypoint() is True
+
+
+@pytest.mark.parametrize(
+    "argv0",
+    [
+        "/venv/bin/pytest",
+        # `python -m pytest` — same basename as a module run, different package.
+        "/venv/lib/python3.12/site-packages/pytest/__main__.py",
+        "/usr/bin/uvicorn",
+        "",
+    ],
+)
+def test_embedders_are_not_mistaken_for_the_entrypoint(
+    argv0: str, monkeypatch: pytest.MonkeyPatch
+) -> None:
+    monkeypatch.setattr(sys, "argv", [argv0, "proxy"])
+    assert proxy_cli._process_is_headroom_cli_entrypoint() is False
+
+
+def test_reexec_does_not_exec_when_embedded(monkeypatch: pytest.MonkeyPatch) -> None:
+    """The end-to-end guard: no execv when another program owns the process."""
+    monkeypatch.setattr(sys, "platform", "darwin")
+    monkeypatch.setattr(sys, "argv", ["/venv/bin/pytest", "tests/"])
+    monkeypatch.delenv("_HEADROOM_MALLOC_TUNED", raising=False)
+    for key in proxy_cli._MALLOC_TUNING:
+        monkeypatch.delenv(key, raising=False)
+
+    calls: list[object] = []
+    monkeypatch.setattr(proxy_cli.os, "execv", lambda *a, **k: calls.append(a))
+
+    proxy_cli._reexec_with_malloc_tuning()
+
+    assert calls == []
+    # The loop guard must not be set either: this process never applied the
+    # tuning, so a genuine CLI child inheriting the env must still be free to.
+    assert "_HEADROOM_MALLOC_TUNED" not in proxy_cli.os.environ
+
+
+def test_reexec_still_execs_for_a_real_cli_launch(monkeypatch: pytest.MonkeyPatch) -> None:
+    """The fix must not disable the feature it is guarding."""
+    monkeypatch.setattr(sys, "platform", "darwin")
+    monkeypatch.setattr(sys, "argv", ["/usr/local/bin/headroom", "proxy", "--port", "8787"])
+    monkeypatch.delenv("_HEADROOM_MALLOC_TUNED", raising=False)
+    for key in proxy_cli._MALLOC_TUNING:
+        monkeypatch.delenv(key, raising=False)
+
+    calls: list[tuple] = []
+    monkeypatch.setattr(proxy_cli.os, "execv", lambda *a, **k: calls.append(a))
+
+    proxy_cli._reexec_with_malloc_tuning()
+
+    assert len(calls) == 1
+    _executable, argv = calls[0]
+    assert argv[1:] == ["-m", "headroom.cli", "proxy", "--port", "8787"]
+    for key, value in proxy_cli._MALLOC_TUNING.items():
+        assert proxy_cli.os.environ[key] == value
@@ -25,6 +25,10 @@ def __init__(self) -> None:
         self.headers: dict[str, str] = {}
         self.query_params: dict[str, str] = {}
         self.url = SimpleNamespace(path="/v1beta/models/gemini-pro:generateContent", query="")
+        # Every real Starlette Request carries a scope, and the Gemini handler
+        # binds the savings-attribution ledger to it (#3051). Without this the
+        # double is a shape that cannot occur in production.
+        self.scope: dict = {"type": "http", "method": "POST"}
 
 
 class _CcrToolCallResponse:
 
@@ -65,6 +65,10 @@ def test_reexec_guard_prevents_loop(monkeypatch):
 
 def test_reexec_skips_when_operator_already_set_vars(monkeypatch):
     monkeypatch.setattr(proxy_cli.sys, "platform", "darwin")
+    # A real CLI launch, like the sibling exec test below: the tuning path is
+    # only reachable when this process is the Headroom CLI entrypoint, and
+    # under pytest argv[0] is pytest's own.
+    monkeypatch.setattr(proxy_cli.sys, "argv", ["headroom", "proxy"])
     monkeypatch.setenv("MallocAggressiveMadvise", "1")
     monkeypatch.setenv("MallocLargeCache", "0")
     rec: dict = {}