akadmin/headroom
1
0
Fork 0
mirror of https://github.com/headroomlabs-ai/headroom.git synced 2026-08-27 14:17:10 -04:00
Code Issues Projects Releases Packages Wiki Activity Actions

feat(proxy): add turn-hook extension point for buffered model turns (#1891)

Browse source 
## 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 <same 4 files>
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
This commit is contained in:
Tejas Chopra 2026-07-08 20:38:45 -04:00 • committed by GitHub
parent 3e85eb1880
commit ec950f7ef1
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
4 changed files with 396 additions and 0 deletions
Download patch file Download diff file Expand all files Collapse all files

46
headroom/proxy/handlers/anthropic.py
View file

@ -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

45
headroom/proxy/handlers/openai.py
View file

@ -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 "

121
headroom/proxy/turn_hooks.py Normal file
View file

@ -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

184
tests/test_turn_hooks.py Normal file
View file

@ -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