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:
Tejas Chopra 2026-07-24 21:13:39 -07:00 committed by GitHub
parent c990cfb803
commit a6d4921e82
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
3 changed files with 59 additions and 8 deletions

View file

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

View file

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

View file

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