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