From ec950f7ef131fb124b60a8e75bc6af7ab733cc7f Mon Sep 17 00:00:00 2001 From: Tejas Chopra Date: Wed, 8 Jul 2026 20:38:45 -0400 Subject: [PATCH] feat(proxy): add turn-hook extension point for buffered model turns (#1891) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Description Adds a small, neutral **extension point** to the proxy: a "turn hook" that lets an opt-in extension observe and optionally re-drive a single buffered model turn, without touching the core request/response flow for anyone who has no extension installed. A hook can: - `on_request(ctx)` — inspect or rewrite the outbound tools/messages before they go upstream (the extensible counterpart to the built-in tool-search deferral that already lives at that point). - `on_response(ctx, response, call_model)` — inspect the model's response and, if it wants, call the model again (via `call_model`) and return a **replacement** response — transparently to the client. This is the capability that can't be done from ASGI middleware: it reuses the proxy-internal re-call path (the same `api_call_fn` the CCR handler already drives). The module is **inert unless a hook is registered**: the runners return their input unchanged and are gated on the registry, so with no extension the proxy is byte-identical to today. A failing hook is logged and skipped — it can never take the proxy down. Closes # ## Type of Change - [ ] Bug fix (non-breaking change that fixes an issue) - [x] New feature (non-breaking change that adds functionality) - [ ] Breaking change (fix or feature that would cause existing functionality to change) - [ ] Documentation update - [ ] Performance improvement - [ ] Code refactoring (no functional changes) ## Changes Made - Add `headroom/proxy/turn_hooks.py`: `TurnContext`, the `TurnHook` protocol (`on_request` / `on_response`), a module registry (`register_turn_hook` / `registered_turn_hooks` / `clear_turn_hooks`), and the runners `run_request_hooks` / `run_response_hooks`. Inert when empty; never raises. - Wire it at four seams, each gated so an empty registry is a byte-identical no-op: - Anthropic — pre-send (right after the existing tool-search deferral) + the CCR response seam. - OpenAI — the Responses tool-shaping point (right after the existing tool-search deferral, copy-on-write-safe) + the CCR response seam. - Add `tests/test_turn_hooks.py`. ## Testing - [x] Unit tests pass (`pytest`) - [x] Linting passes (`ruff check .`) - [x] Type checking passes (`mypy headroom`) - [x] New tests added for new functionality - [x] Manual testing performed (no-op regression across CCR/handler suites) ### Test Output ```text $ ruff check headroom/proxy/turn_hooks.py tests/test_turn_hooks.py \ headroom/proxy/handlers/anthropic.py headroom/proxy/handlers/openai.py All checks passed! $ ruff format --check 4 files already formatted $ mypy headroom Success: no issues found in 408 source files $ pytest tests/test_turn_hooks.py -q 9 passed in 0.17s $ pytest tests/test_turn_hooks.py tests/test_ccr_response_handler.py \ tests/test_ccr_tool_injection.py tests/test_proxy_ccr.py \ tests/test_openai_tool_search_deferral.py \ tests/test_openai_responses_compression_units.py \ tests/test_handler_outcome_tag_invariant.py -q 135 passed (+ 1 pre-existing cross-file flake in test_proxy_ccr::test_health_endpoint, which passes in isolation and in its own file: `pytest tests/test_proxy_ccr.py` -> 19 passed) ``` ## Real Behavior Proof - **Environment:** local macOS, project `.venv` (Python 3.12.6); `ruff` pinned to CI's `0.15.17` via `uvx ruff@0.15.17`; `mypy` from the venv. - **Exact command / steps:** branched off `upstream/main`; added the hook module + wired the four handler seams; ran the ruff/format/mypy/pytest commands above. - **Observed result:** The unit tests exercise the whole contract — registry, `on_request` mutating `ctx.tools`, `on_response` returning a replacement, the `await call_model(...)` re-drive loop, replacement-chaining across hooks, and the never-raise guarantee. The existing CCR + handler suites pass unchanged, which is the point: with no hook registered the added code is a no-op (the runners short-circuit on an empty registry). - **Not tested:** the live interactive re-drive path with a *registered* hook against a real upstream — no hook ships in this repo, so that path is covered here only by the unit test's fake `call_model`. The `on_request` seam fires on the Anthropic pre-send and OpenAI Responses paths (where the existing tool-search deferral runs); other send paths (e.g. chat-completions, streaming) are not wired in this PR. ## 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 --- headroom/proxy/handlers/anthropic.py | 46 +++++++ headroom/proxy/handlers/openai.py | 45 +++++++ headroom/proxy/turn_hooks.py | 121 ++++++++++++++++++ tests/test_turn_hooks.py | 184 +++++++++++++++++++++++++++ 4 files changed, 396 insertions(+) create mode 100644 headroom/proxy/turn_hooks.py create mode 100644 tests/test_turn_hooks.py diff --git a/headroom/proxy/handlers/anthropic.py b/headroom/proxy/handlers/anthropic.py index 605f9f6be..c3f38abec 100644 --- a/headroom/proxy/handlers/anthropic.py +++ b/headroom/proxy/handlers/anthropic.py @@ -2175,6 +2175,32 @@ class AnthropicHandlerMixin: f"{_ts_saved_tokens}tok" ) + # Turn hooks (opt-in extensions): a registered hook may inspect or + # rewrite the outbound tools/messages before we send upstream — the + # extensible counterpart to the built-in deferral above. A single + # registry check keeps this a no-op when no hook is registered. + from headroom.proxy.turn_hooks import ( + TurnContext, + registered_turn_hooks, + run_request_hooks, + ) + + if registered_turn_hooks(): + _req_ctx = TurnContext( + provider="anthropic", + model=str(model), + messages=optimized_messages, + tools=body.get("tools"), + config=self.config, + ) + run_request_hooks(_req_ctx) + if _req_ctx.messages is not optimized_messages: + optimized_messages = _req_ctx.messages + body["messages"] = optimized_messages + if _req_ctx.tools is not body.get("tools"): + tools = _req_ctx.tools + body["tools"] = tools + # Output shaping (opt-in via HEADROOM_OUTPUT_SHAPER): verbosity # steering appended to the system-prompt tail + effort routing on # mechanical tool_result continuations. Runs after every other @@ -2782,6 +2808,26 @@ class AnthropicHandlerMixin: ) # Update response content with final response resp_json = final_resp_json + # Turn hooks (opt-in extensions) may inspect the turn or + # re-drive the model before we hand back the response. + # Inert when no hook is registered. + from headroom.proxy.turn_hooks import ( + TurnContext, + run_response_hooks, + ) + + final_resp_json = await run_response_hooks( + TurnContext( + provider="anthropic", + model=str(model), + messages=optimized_messages, + tools=tools, + config=self.config, + ), + final_resp_json, + api_call_fn, + ) + resp_json = final_resp_json # Remove encoding headers since content is now uncompressed JSON ccr_response_headers = { k: v diff --git a/headroom/proxy/handlers/openai.py b/headroom/proxy/handlers/openai.py index de8eddd16..47f804470 100644 --- a/headroom/proxy/handlers/openai.py +++ b/headroom/proxy/handlers/openai.py @@ -1709,6 +1709,32 @@ class OpenAIHandlerMixin: modified = True transforms.append("openai:responses:tool_search_deferral") + # Turn hooks (opt-in extensions): a registered hook may inspect or rewrite + # the outbound tools before we send — the extensible counterpart to the + # built-in deferral above. Gated on the registry so it is a no-op (no copy, + # no context construction) when no hook is registered. + from headroom.proxy.turn_hooks import ( + TurnContext, + registered_turn_hooks, + run_request_hooks, + ) + + if registered_turn_hooks(): + if working is payload: + working = copy.deepcopy(payload) + _req_ctx = TurnContext( + provider="openai", + model=str(model), + messages=working.get("input") or working.get("messages") or [], + tools=working.get("tools"), + config=getattr(self, "config", None), + ) + run_request_hooks(_req_ctx) + if _req_ctx.tools is not working.get("tools"): + working["tools"] = _req_ctx.tools + modified = True + transforms.append("openai:responses:turn_hook") + live_units_started = time.perf_counter() ( router_payload, @@ -2815,6 +2841,25 @@ class OpenAIHandlerMixin: api_call_fn, provider="openai", ) + # Turn hooks (opt-in extensions) may inspect the turn + # or re-drive the model before we hand back the + # response. Inert when no hook is registered. + from headroom.proxy.turn_hooks import ( + TurnContext, + run_response_hooks, + ) + + final_resp_json = await run_response_hooks( + TurnContext( + provider="openai", + model=str(model), + messages=optimized_messages, + tools=tools, + config=self.config, + ), + final_resp_json, + api_call_fn, + ) backend_response.body = final_resp_json logger.info( f"[{request_id}] CCR: Retrieval handled " diff --git a/headroom/proxy/turn_hooks.py b/headroom/proxy/turn_hooks.py new file mode 100644 index 000000000..b61d28f1c --- /dev/null +++ b/headroom/proxy/turn_hooks.py @@ -0,0 +1,121 @@ +"""Turn hooks — observe and optionally re-drive a single model turn. + +A *turn hook* wraps the proxy's buffered upstream call: it sees the outbound +request and the model's response, and may call the model again (via +``call_model``) to return a *replacement* response — transparently to the client. +That covers a range of turn-level behaviors an extension might want: resolving an +injected tool call, enforcing a guardrail, retrying on a bad response, serving a +cached answer, or running a small model→proxy→model loop before handing back the +final answer. + +Hooks are registered by opt-in proxy extensions (``proxy/extensions.py``). This +module is **inert unless a hook is registered**: the runner helpers return their +input unchanged, so with no hooks the proxy behaves exactly as if this module did +not exist. Hooks must never raise — a failing hook is logged and skipped so it +cannot take the proxy down. + +Stability: the ``TurnHook`` protocol and the registry/runner functions are part +of the extension surface (see ``proxy/extensions.py``); signature changes follow +the same deprecation policy. +""" + +from __future__ import annotations + +import logging +from collections.abc import Awaitable, Callable +from dataclasses import dataclass +from typing import Any, Protocol, runtime_checkable + +log = logging.getLogger(__name__) + +# Re-drive the model with a message list; returns the provider's response JSON. +# The exact message/response shapes are provider-native (Anthropic / OpenAI / +# Google), matching whatever the surrounding handler already works with. +CallModel = Callable[[list[dict[str, Any]]], Awaitable[dict[str, Any]]] + + +@dataclass +class TurnContext: + """The request side of one model turn, as the proxy forwards it upstream. + + ``tools`` and ``messages`` are the live objects the handler is about to send + (or just sent); a hook's ``on_request`` may mutate them in place. + """ + + provider: str # "anthropic" | "openai" | "google" | ... + model: str + messages: list[dict[str, Any]] + tools: Any = None # provider-native tools value (list, or None) + config: Any = None + + +@runtime_checkable +class TurnHook(Protocol): + """A registered turn observer. Both methods are optional (a hook may define + either); missing methods are simply skipped.""" + + name: str + + def on_request(self, ctx: TurnContext) -> None: + """Inspect / mutate ``ctx`` (e.g. ``ctx.tools``) before it goes upstream.""" + + async def on_response( + self, ctx: TurnContext, response: dict[str, Any], call_model: CallModel + ) -> dict[str, Any] | None: + """Return a replacement response, or ``None`` to leave it unchanged. + + May ``await call_model(messages)`` to re-drive the model (e.g. to resolve + an injected tool call), looping as needed before returning the final.""" + + +_hooks: list[TurnHook] = [] + + +def register_turn_hook(hook: TurnHook) -> None: + """Register a hook. Called by an extension's ``install(app, config)``.""" + _hooks.append(hook) + log.info("registered turn hook: %s", getattr(hook, "name", type(hook).__name__)) + + +def registered_turn_hooks() -> list[TurnHook]: + return list(_hooks) + + +def clear_turn_hooks() -> None: + """Test/reset helper.""" + _hooks.clear() + + +def run_request_hooks(ctx: TurnContext) -> None: + """Run every hook's ``on_request``. Inert when none are registered; never raises.""" + for hook in _hooks: + fn = getattr(hook, "on_request", None) + if fn is None: + continue + try: + fn(ctx) + except Exception: # a hook must never break the proxy + log.exception("turn hook %r on_request failed", getattr(hook, "name", hook)) + + +async def run_response_hooks( + ctx: TurnContext, response: dict[str, Any], call_model: CallModel +) -> dict[str, Any]: + """Run every hook's ``on_response``, chaining any replacements. + + Returns the (possibly replaced) response. Inert when no hooks are registered + (returns ``response`` unchanged); a failing hook is logged and skipped. + """ + current = response + for hook in _hooks: + fn = getattr(hook, "on_response", None) + if fn is None: + continue + try: + replacement = await fn(ctx, current, call_model) + except Exception: + log.exception("turn hook %r on_response failed", getattr(hook, "name", hook)) + continue + if replacement is not None: + current = replacement + return current diff --git a/tests/test_turn_hooks.py b/tests/test_turn_hooks.py new file mode 100644 index 000000000..94be1415e --- /dev/null +++ b/tests/test_turn_hooks.py @@ -0,0 +1,184 @@ +"""Turn-hook registry + runners (headroom/proxy/turn_hooks.py). + +The hook surface is opt-in: with nothing registered the runners must be exact +no-ops (the property the proxy relies on to stay byte-identical for everyone who +has no extension installed). These tests pin that, plus request-mutation, +response-replacement, the re-drive (``call_model``) loop, and the +never-raise guarantee. +""" + +from __future__ import annotations + +import pytest + +from headroom.proxy.turn_hooks import ( + TurnContext, + clear_turn_hooks, + register_turn_hook, + registered_turn_hooks, + run_request_hooks, + run_response_hooks, +) + + +@pytest.fixture(autouse=True) +def _clean_registry(): + clear_turn_hooks() + yield + clear_turn_hooks() + + +def _ctx(**kw): + base = {"provider": "anthropic", "model": "claude-x", "messages": [], "tools": None} + base.update(kw) + return TurnContext(**base) + + +async def _noop_call_model(_messages): # pragma: no cover - never invoked in no-op tests + raise AssertionError("call_model must not be invoked when no hook re-drives") + + +# --- inert-when-empty (the load-bearing guarantee) --------------------------- + + +def test_request_runner_inert_when_empty(): + assert registered_turn_hooks() == [] + ctx = _ctx(tools=[{"name": "a"}]) + before = ctx.tools + run_request_hooks(ctx) # must not raise, must not touch ctx + assert ctx.tools is before + + +@pytest.mark.asyncio +async def test_response_runner_returns_input_unchanged_when_empty(): + resp = {"id": "orig", "content": []} + out = await run_response_hooks(_ctx(), resp, _noop_call_model) + assert out is resp # same object, untouched + + +# --- on_request mutation ----------------------------------------------------- + + +def test_on_request_may_mutate_ctx(): + class Shrink: + name = "shrink" + + def on_request(self, ctx: TurnContext) -> None: + ctx.tools = [t for t in (ctx.tools or []) if t["name"] != "drop_me"] + + register_turn_hook(Shrink()) + ctx = _ctx(tools=[{"name": "keep"}, {"name": "drop_me"}]) + run_request_hooks(ctx) + assert ctx.tools == [{"name": "keep"}] + + +# --- on_response replacement + re-drive loop --------------------------------- + + +@pytest.mark.asyncio +async def test_on_response_can_replace_via_call_model(): + calls: list[list] = [] + + async def call_model(messages): + calls.append(messages) + return {"id": "resolved", "content": [{"type": "text", "text": "done"}]} + + class ResolveOnce: + name = "resolve" + + async def on_response(self, ctx, response, call_model): + if response.get("id") == "needs-work": + return await call_model(ctx.messages + [{"role": "user", "content": "go"}]) + return None + + register_turn_hook(ResolveOnce()) + out = await run_response_hooks( + _ctx(messages=[{"role": "user", "content": "hi"}]), {"id": "needs-work"}, call_model + ) + assert out["id"] == "resolved" + assert len(calls) == 1 + + +@pytest.mark.asyncio +async def test_on_response_none_leaves_response_unchanged(): + class Observer: + name = "observe" + + async def on_response(self, ctx, response, call_model): + return None # observe only + + register_turn_hook(Observer()) + resp = {"id": "orig"} + out = await run_response_hooks(_ctx(), resp, _noop_call_model) + assert out is resp + + +@pytest.mark.asyncio +async def test_replacements_chain_across_hooks(): + class First: + name = "first" + + async def on_response(self, ctx, response, call_model): + return {"id": "after-first", "seen": response["id"]} + + class Second: + name = "second" + + async def on_response(self, ctx, response, call_model): + return {"id": "after-second", "seen": response["id"]} + + register_turn_hook(First()) + register_turn_hook(Second()) + out = await run_response_hooks(_ctx(), {"id": "orig"}, _noop_call_model) + assert out == {"id": "after-second", "seen": "after-first"} # Second saw First's output + + +# --- a failing hook must never break the proxy ------------------------------- + + +def test_failing_on_request_is_swallowed(): + class Boom: + name = "boom" + + def on_request(self, ctx: TurnContext) -> None: + raise RuntimeError("kaboom") + + register_turn_hook(Boom()) + run_request_hooks(_ctx()) # must not raise + + +@pytest.mark.asyncio +async def test_failing_on_response_is_skipped_and_original_survives(): + class Boom: + name = "boom" + + async def on_response(self, ctx, response, call_model): + raise RuntimeError("kaboom") + + class Good: + name = "good" + + async def on_response(self, ctx, response, call_model): + return {"id": "recovered"} + + register_turn_hook(Boom()) + register_turn_hook(Good()) + out = await run_response_hooks(_ctx(), {"id": "orig"}, _noop_call_model) + assert out == {"id": "recovered"} # Boom skipped, Good still ran + + +# --- hooks with only one method defined -------------------------------------- + + +@pytest.mark.asyncio +async def test_hook_without_on_response_is_skipped(): + class OnlyRequest: + name = "only-request" + + def on_request(self, ctx: TurnContext) -> None: + pass + + register_turn_hook(OnlyRequest()) + resp = {"id": "orig"} + out = await run_response_hooks(_ctx(), resp, _noop_call_model) + assert out is resp