mirror of
https://github.com/headroomlabs-ai/headroom.git
synced 2026-08-27 14:17:10 -04:00
feat(proxy/hooks): run fold-only (stream-safe) turn hooks on streaming OpenAI chat (#2549)
## Description Fixes the last harness gap in the turn-hook seam (the "B4" finding from the savings audit). The OpenAI chat handler gated hooks on `not stream`, so **streamed** `/v1/chat/completions` requests ran **no** turn hooks — the lossless-guard plugin's on_request fold and tool-schema shrink were skipped, unlike the Anthropic path (hooks run unconditionally). Affects opencode / Cursor / older OpenAI SDKs / some Copilot flows; **not** Claude Code (Anthropic path). The gate existed for a real reason: hooks that **re-drive** the model in `on_response` (defer a tool, reload it when asked) can't run mid-stream. But an **on_request fold** mutates the outbound request before the send — safe on a stream. ## Change - Add an opt-in `stream_safe` hook attribute (fold-only hooks set it). `run_request_hooks(ctx, stream_safe_only=…)` filters to stream-safe hooks when set. - OpenAI chat handler runs `on_request` on streaming with `stream_safe_only=stream`; buffered runs all hooks; the `on_response` re-drive (buffered response path) is untouched. - **Default off = conservative:** a hook is buffered-only unless it declares `stream_safe`, so **no behavior change** until a hook opts in. ## Type of Change - [x] Bug fix / feature (opt-in, backward-compatible) ## Testing ```text pytest tests/test_turn_hooks.py tests/test_openai_chat_turn_hooks.py -q → 25 passed ruff + mypy → clean ``` New test pins the filter: streaming runs only stream-safe hooks' on_request; buffered runs all. ## Notes The companion plugin PR (headroom-lossless-guard) sets `stream_safe = True` on its fold-only hook to actually claim the streaming savings. Anthropic path already ran hooks on streaming, so it's unaffected. ## Checklist - [x] Self-reviewed; tests pass; no CHANGELOG edit
This commit is contained in:
parent
c990cfb803
commit
a6d4921e82
3 changed files with 59 additions and 8 deletions
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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"
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue