diff --git a/headroom/proxy/handlers/openai.py b/headroom/proxy/handlers/openai.py index 1ddd047fb..413382166 100644 --- a/headroom/proxy/handlers/openai.py +++ b/headroom/proxy/handlers/openai.py @@ -3471,17 +3471,19 @@ class OpenAIHandlerMixin: tokens_saved = max(0, original_tokens - optimized_tokens) # Turn hooks (opt-in extensions): a registered hook may rewrite the - # outbound tools/messages before we send. Buffered requests only — a - # streamed turn can't be re-driven to resolve whatever the model asks to - # load. Gated on the registry so it is a no-op when none are registered; - # the net tool-schema token delta is recorded so it shows up as a saving. + # outbound tools/messages before we send. on_response re-drive (below, in + # the buffered response path) can't run on a stream, but on_request folds + # can — so on streaming we run only stream-safe (fold-only) hooks, and + # buffered runs all of them. Gated on the registry so it's a no-op when + # none are registered; the net tool-schema token delta is recorded as a + # saving. from headroom.proxy.turn_hooks import ( TurnContext, registered_turn_hooks, run_request_hooks, ) - if registered_turn_hooks() and not stream: + if registered_turn_hooks(): _th_tools_before = body.get("tools") _th_tok_before = ( tokenizer.count_text(json.dumps(_th_tools_before, default=str)) @@ -3502,7 +3504,7 @@ class OpenAIHandlerMixin: _th_msg_before: int | None = tokenizer.count_messages(body["messages"]) except Exception: _th_msg_before = None - run_request_hooks(_th_ctx) + run_request_hooks(_th_ctx, stream_safe_only=stream) # A hook may either replace ctx.messages/ctx.tools or mutate them in # place (the contract allows both). Use object identity only to decide # whether body needs reassignment; measure the saving from the FINAL diff --git a/headroom/proxy/turn_hooks.py b/headroom/proxy/turn_hooks.py index b61d28f1c..1da65fd47 100644 --- a/headroom/proxy/turn_hooks.py +++ b/headroom/proxy/turn_hooks.py @@ -56,6 +56,12 @@ class TurnHook(Protocol): name: str + # Optional: set True if this hook's ``on_request`` needs no later re-drive + # (a fold-only hook). Such hooks run on streaming turns too; hooks that may + # re-drive in ``on_response`` leave it unset and run buffered-only. Absent ⇒ + # False (conservative). See :func:`run_request_hooks`. + stream_safe: bool + def on_request(self, ctx: TurnContext) -> None: """Inspect / mutate ``ctx`` (e.g. ``ctx.tools``) before it goes upstream.""" @@ -86,9 +92,22 @@ def clear_turn_hooks() -> None: _hooks.clear() -def run_request_hooks(ctx: TurnContext) -> None: - """Run every hook's ``on_request``. Inert when none are registered; never raises.""" +def run_request_hooks(ctx: TurnContext, *, stream_safe_only: bool = False) -> None: + """Run each hook's ``on_request``. Inert when none registered; never raises. + + ``on_request`` mutates the outbound request before the upstream send, so it is + safe on a streamed turn *as long as the hook needs no later re-drive*. A hook + that may re-drive the model in ``on_response`` (e.g. defer a tool, then reload + it when the model asks) can't run on a stream — the bytes are already flowing, + there's nothing to re-drive. So on streaming turns the handler passes + ``stream_safe_only=True`` and only hooks that opt in via a truthy ``stream_safe`` + attribute run; fold-only hooks (which never re-drive) set it and thus keep + working on streamed OpenAI-compatible traffic. Default off ⇒ conservative: + a hook is treated as buffered-only unless it declares itself stream-safe. + """ for hook in _hooks: + if stream_safe_only and not getattr(hook, "stream_safe", False): + continue fn = getattr(hook, "on_request", None) if fn is None: continue diff --git a/tests/test_turn_hooks.py b/tests/test_turn_hooks.py index 94be1415e..788ae01ba 100644 --- a/tests/test_turn_hooks.py +++ b/tests/test_turn_hooks.py @@ -59,6 +59,36 @@ async def test_response_runner_returns_input_unchanged_when_empty(): # --- on_request mutation ----------------------------------------------------- +def test_stream_safe_filter_runs_only_optted_in_hooks_on_stream(): + """On a streamed turn only ``stream_safe`` hooks' on_request runs (fold-only, + no re-drive); buffered runs all. A hook that may re-drive stays buffered-only.""" + ran: list[str] = [] + + class Fold: + name = "fold" + stream_safe = True # opts in — safe on streaming + + def on_request(self, ctx: TurnContext) -> None: + ran.append("fold") + + class Redrive: # no stream_safe attr → buffered-only (default) + name = "redrive" + + def on_request(self, ctx: TurnContext) -> None: + ran.append("redrive") + + register_turn_hook(Fold()) + register_turn_hook(Redrive()) + + ran.clear() + run_request_hooks(_ctx(), stream_safe_only=True) # streaming turn + assert ran == ["fold"] + + ran.clear() + run_request_hooks(_ctx()) # buffered turn (default) + assert ran == ["fold", "redrive"] + + def test_on_request_may_mutate_ctx(): class Shrink: name = "shrink"