mirror of
https://github.com/headroomlabs-ai/headroom.git
synced 2026-08-27 14:17:10 -04:00
fix(proxy): gate mid-turn message coalescing to Claude Code clients (#1643)
## Description
`headroom wrap opencode` (and any other `@ai-sdk/anthropic` client)
can't use subagents. The subagent is spawned, receives the prompt, and
never responds; OpenCode throws `invalid_union / "No matching
discriminator" / discriminator: "type"`.
Root cause is headroom's mid-turn message coalescing. It keys concurrent
streaming requests by `md5(model:system[:500])` (`_get_session_key`,
`handlers/streaming.py`). An OpenCode subagent runs concurrently with
the main agent on the same model and same first-500-char system prefix,
so it produces the **same** session key and collides with the
still-active main stream. Two things then break it:
1. `handlers/anthropic.py` sees the key in `_active_streams` and answers
the subagent's request with a bare `202 headroom_queued` instead of
forwarding it — so the subagent never gets a response.
2. When the main stream ends, `handlers/streaming.py` emits a
non-standard `event: headroom_pending_messages` SSE event.
`@ai-sdk/anthropic`'s SSE parser keys its Zod union on `type`, and
`headroom_pending_messages` isn't a valid Anthropic event type — hence
the error.
The 202 reply and the `headroom_pending_messages` event are a Claude
Code-only protocol (nothing else consumes them). This gates coalescing
to Claude Code clients; every other harness streams normally.
Closes #1608
## Type of Change
- [x] Bug fix (non-breaking change that fixes an issue)
## Changes Made
- `handlers/streaming.py`: only register a stream in `_active_streams`
when `classify_client(headers) == "claude-code"`, and only emit the
`headroom_pending_messages` SSE event for Claude Code.
- `handlers/anthropic.py`: only take the queue-and-`202` branch when the
client is Claude Code (in addition to the existing `session_key in
_active_streams` check).
- Regression tests in `tests/test_mid_turn_steering.py` for all four
cases (active-stream registration and pending-event emission, each for a
Claude Code vs. a non-Claude-Code client).
## 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
### Test Output
```text
$ pytest tests/test_mid_turn_steering.py -q
9 passed in 0.46s
$ ruff check headroom/proxy/handlers/streaming.py headroom/proxy/handlers/anthropic.py tests/test_mid_turn_steering.py
All checks passed!
$ ruff format --check <same files>
3 files already formatted
$ mypy headroom/proxy/handlers/streaming.py headroom/proxy/handlers/anthropic.py --ignore-missing-imports
Success: no issues found in 2 source files
```
## Real Behavior Proof
- Environment: macOS (arm64), Python 3.14 venv, editable install of this
branch.
- Exact command / steps: ran `pytest tests/test_mid_turn_steering.py` —
the new tests drive `_stream_response` with a queued mid-turn message
under an `opencode/1.0` User-Agent vs. a `claude-code/1.2.3` User-Agent
and assert the streamed bytes. Also ran the streaming + anthropic
handler suites (`pytest tests/test_mid_turn_steering.py
tests/test_proxy_streaming_* tests/test_anthropic_*
tests/test_streaming_usage_parser.py`).
- Observed result: with the `opencode/1.0` client the session is never
added to `_active_streams` and the response contains no
`headroom_pending_messages` event; with `claude-code/1.2.3` both still
happen (protocol preserved). Handler suites: 155 passed, 3 skipped.
Before this change the non-Claude client received the
`headroom_pending_messages` event (the exact byte string the OpenCode
parser rejects).
- Not tested: end-to-end against a live OpenCode + real subagent run —
reproduced deterministically at the proxy layer instead (the emitted SSE
bytes are the direct source of the reported `invalid_union` error).
## Review Readiness
- [x] I have performed a self-review
- [x] This PR is ready for human review
## Additional Notes
Gating on `classify_client == "claude-code"` (User-Agent `claude-code/`
/ `claude-cli/`) is the same client identification used elsewhere in the
proxy. Unidentified clients (no recognized User-Agent) are treated as
non-Claude-Code and stream normally, which is the safe default for this
feature.
## Maintainer Update (2026-07-21)
- Removed the manual `CHANGELOG.md` entry so release-please remains the
source of changelog updates; pushed `d8e36540`.
- Validation: `tests/test_mid_turn_steering.py` passed (12 tests), the
related streaming/Anthropic suite passed (72 tests), Ruff check passed
for touched files, Ruff format check passed, and `git diff --check
upstream/main...HEAD` passed.
---------
Co-authored-by: JerrettDavis <mxjerrett@gmail.com>
This commit is contained in:
parent
89493714d2
commit
a4bd2e62a5
4 changed files with 163 additions and 9 deletions
|
|
@ -162,6 +162,28 @@ def should_stamp_codex_client(path: str, headers: Mapping[str, Any] | Any) -> bo
|
|||
return should_stamp_codex_client_signals(path, _auth_signals(headers))
|
||||
|
||||
|
||||
# Client harnesses that can consume Headroom's mid-turn message-coalescing
|
||||
# protocol: the 202 ``headroom_queued`` reply and the synthetic
|
||||
# ``headroom_pending_messages`` SSE event. This is a custom protocol only Claude
|
||||
# Code parses today; other harnesses would receive events they can't decode, so
|
||||
# a concurrent same-session request from them must stream normally instead of
|
||||
# being queued and replayed. Keeping the capability in one named place — rather
|
||||
# than scattering ``client == "claude-code"`` string checks across the request
|
||||
# handlers — makes it a single, documented decision to revisit as more clients
|
||||
# learn the protocol (#1608).
|
||||
_COALESCING_CAPABLE_CLIENTS = frozenset({"claude-code"})
|
||||
|
||||
|
||||
def supports_mid_turn_coalescing(client: str | None) -> bool:
|
||||
"""Whether ``client`` can consume the mid-turn coalescing protocol.
|
||||
|
||||
``client`` is a value returned by :func:`classify_client`. See
|
||||
:data:`_COALESCING_CAPABLE_CLIENTS` for why the set is currently limited to
|
||||
Claude Code.
|
||||
"""
|
||||
return client in _COALESCING_CAPABLE_CLIENTS
|
||||
|
||||
|
||||
__all__ = [
|
||||
"AuthMode",
|
||||
"CLIENT_UA_MAP",
|
||||
|
|
@ -170,4 +192,5 @@ __all__ = [
|
|||
"classify_auth_mode",
|
||||
"classify_client",
|
||||
"should_stamp_codex_client",
|
||||
"supports_mid_turn_coalescing",
|
||||
]
|
||||
|
|
|
|||
|
|
@ -27,7 +27,11 @@ from headroom.agent_savings import proxy_pipeline_kwargs
|
|||
from headroom.ccr.context_tracker import looks_like_claude_code_compact_summary
|
||||
from headroom.copilot_auth import build_copilot_upstream_url
|
||||
from headroom.pipeline import PipelineStage, summarize_routing_markers
|
||||
from headroom.proxy.auth_mode import classify_auth_mode, classify_client
|
||||
from headroom.proxy.auth_mode import (
|
||||
classify_auth_mode,
|
||||
classify_client,
|
||||
supports_mid_turn_coalescing,
|
||||
)
|
||||
from headroom.proxy.compression_decision import CompressionDecision
|
||||
from headroom.proxy.forwarded_headers import resolve_client_ip
|
||||
from headroom.proxy.handlers._debug_dump import _debug_dump_mode, _redact_debug_value
|
||||
|
|
@ -3024,11 +3028,18 @@ class AnthropicHandlerMixin:
|
|||
body,
|
||||
session_header=explicit_session_header,
|
||||
)
|
||||
# Only opt-in (header-bearing) callers participate in
|
||||
# mid-turn steering; see StreamingMixin._should_queue_mid_turn
|
||||
# for why the coarse md5 fallback must not queue concurrent
|
||||
# independent streams (it wrongly 202s a streaming caller).
|
||||
if self._should_queue_mid_turn(session_key, explicit_session_header):
|
||||
# Coalesce mid-turn messages only for Claude Code, the sole
|
||||
# client that understands the 202 `headroom_queued` reply and
|
||||
# the `headroom_pending_messages` SSE event. Other harnesses
|
||||
# (e.g. OpenCode subagents sharing a body-derived session key)
|
||||
# would otherwise have their request swallowed and never
|
||||
# answered. (#1608) `_should_queue_mid_turn` further restricts
|
||||
# this to opt-in (header-bearing) callers with an active
|
||||
# stream, so the coarse md5 fallback can't 202 a streaming
|
||||
# caller.
|
||||
if supports_mid_turn_coalescing(
|
||||
classify_client(request.headers)
|
||||
) and self._should_queue_mid_turn(session_key, explicit_session_header):
|
||||
from fastapi.responses import JSONResponse
|
||||
|
||||
queued = self._queue_mid_turn_message(session_key, body)
|
||||
|
|
|
|||
|
|
@ -12,7 +12,7 @@ import logging
|
|||
import time
|
||||
from typing import TYPE_CHECKING, Any
|
||||
|
||||
from headroom.proxy.auth_mode import classify_client
|
||||
from headroom.proxy.auth_mode import classify_client, supports_mid_turn_coalescing
|
||||
from headroom.proxy.helpers import (
|
||||
RETRYABLE_OVERLOAD_STATUSES,
|
||||
jitter_delay_ms,
|
||||
|
|
@ -1033,7 +1033,6 @@ class StreamingMixin:
|
|||
4. Streams the final response to the client
|
||||
"""
|
||||
session_key = session_key or self._get_session_key(body)
|
||||
self._active_streams.add(session_key)
|
||||
|
||||
# Guard everything up to the generator's own try/finally (which owns
|
||||
# cleanup once streaming starts): any exception here — including
|
||||
|
|
@ -1106,6 +1105,15 @@ class StreamingMixin:
|
|||
# ...) from the *client's* User-Agent before copilot-auth
|
||||
# potentially rewrites headers for upstream.
|
||||
client = classify_client(headers)
|
||||
# Mid-turn message coalescing (queueing a concurrent same-session
|
||||
# request and later replaying it via a `headroom_pending_messages`
|
||||
# SSE event) is a Claude Code-only protocol. Only register the stream
|
||||
# as active for coalescing when the client can consume that protocol,
|
||||
# so concurrent requests from other harnesses (e.g. OpenCode subagents
|
||||
# that share a body-derived session key) are streamed normally instead
|
||||
# of being swallowed. (#1608)
|
||||
if supports_mid_turn_coalescing(client):
|
||||
self._active_streams.add(session_key)
|
||||
headers = await apply_copilot_api_auth(headers, url=url)
|
||||
start_time = time.time()
|
||||
|
||||
|
|
@ -1651,7 +1659,7 @@ class StreamingMixin:
|
|||
client=client,
|
||||
waste_signals=waste_signals,
|
||||
)
|
||||
if pending_messages:
|
||||
if supports_mid_turn_coalescing(client) and pending_messages:
|
||||
pending_event = json.dumps(
|
||||
{"type": "headroom_pending_messages", "messages": pending_messages}
|
||||
)
|
||||
|
|
|
|||
|
|
@ -201,3 +201,115 @@ class TestMidTurnSteering:
|
|||
finally:
|
||||
proxy._active_streams.discard(session_key)
|
||||
proxy._mid_turn_queues.pop(session_key, None)
|
||||
|
||||
# --- #1608: mid-turn coalescing must be gated to Claude Code clients ---
|
||||
|
||||
def _normal_stream(self):
|
||||
return self._create_mock_upstream_response(
|
||||
[
|
||||
b'event: message_start\ndata: {"type":"message_start"}\n\n',
|
||||
b'event: message_stop\ndata: {"type":"message_stop"}\n\n',
|
||||
]
|
||||
)
|
||||
|
||||
async def _run_stream(self, proxy, session_key, user_agent):
|
||||
mock_response = self._normal_stream()
|
||||
proxy.http_client.build_request = MagicMock(return_value=MagicMock())
|
||||
proxy.http_client.send = AsyncMock(return_value=mock_response)
|
||||
return await proxy._stream_response(
|
||||
url="https://api.anthropic.com/v1/messages",
|
||||
headers={"x-api-key": "sk-test", "user-agent": user_agent},
|
||||
body={
|
||||
"model": "claude-sonnet-4-20250514",
|
||||
"max_tokens": 100,
|
||||
"stream": True,
|
||||
"messages": [{"role": "user", "content": "hi"}],
|
||||
},
|
||||
provider="anthropic",
|
||||
model="claude-sonnet-4-20250514",
|
||||
request_id="test-1608",
|
||||
original_tokens=10,
|
||||
optimized_tokens=10,
|
||||
tokens_saved=0,
|
||||
transforms_applied=[],
|
||||
tags={},
|
||||
optimization_latency=0.0,
|
||||
session_key=session_key,
|
||||
)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_non_claude_client_not_registered_active(self):
|
||||
# An OpenCode subagent shares the main agent's body-derived session
|
||||
# key; if it registered as active, the concurrent request would be
|
||||
# swallowed. Non-Claude-Code clients must never be registered.
|
||||
proxy = self._create_mock_proxy()
|
||||
session_key = "opencode-session"
|
||||
result = await self._run_stream(proxy, session_key, "opencode/1.0")
|
||||
try:
|
||||
assert session_key not in proxy._active_streams
|
||||
finally:
|
||||
async for _chunk in result.body_iterator:
|
||||
pass
|
||||
proxy._active_streams.discard(session_key)
|
||||
proxy._mid_turn_queues.pop(session_key, None)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_claude_code_client_registered_active(self):
|
||||
proxy = self._create_mock_proxy()
|
||||
session_key = "claude-session"
|
||||
result = await self._run_stream(proxy, session_key, "claude-code/1.2.3")
|
||||
try:
|
||||
assert session_key in proxy._active_streams
|
||||
finally:
|
||||
async for _chunk in result.body_iterator:
|
||||
pass
|
||||
proxy._active_streams.discard(session_key)
|
||||
proxy._mid_turn_queues.pop(session_key, None)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_pending_event_not_emitted_for_non_claude(self):
|
||||
# Even with a queued message, a non-Claude-Code stream must not emit the
|
||||
# custom `headroom_pending_messages` SSE event — @ai-sdk/anthropic can't
|
||||
# parse it and throws "invalid_union / No matching discriminator".
|
||||
proxy = self._create_mock_proxy()
|
||||
session_key = "opencode-pending"
|
||||
proxy._queue_mid_turn_message(
|
||||
session_key, {"messages": [{"role": "user", "content": "queued"}]}
|
||||
)
|
||||
result = await self._run_stream(proxy, session_key, "opencode/1.0")
|
||||
try:
|
||||
body = b"".join([chunk async for chunk in result.body_iterator])
|
||||
assert b"headroom_pending_messages" not in body
|
||||
finally:
|
||||
proxy._active_streams.discard(session_key)
|
||||
proxy._mid_turn_queues.pop(session_key, None)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_pending_event_emitted_for_claude_code(self):
|
||||
proxy = self._create_mock_proxy()
|
||||
session_key = "claude-pending"
|
||||
proxy._queue_mid_turn_message(
|
||||
session_key, {"messages": [{"role": "user", "content": "queued"}]}
|
||||
)
|
||||
result = await self._run_stream(proxy, session_key, "claude-code/1.2.3")
|
||||
try:
|
||||
body = b"".join([chunk async for chunk in result.body_iterator])
|
||||
assert b"headroom_pending_messages" in body
|
||||
finally:
|
||||
proxy._active_streams.discard(session_key)
|
||||
proxy._mid_turn_queues.pop(session_key, None)
|
||||
|
||||
|
||||
class TestCoalescingCapability:
|
||||
"""The capability predicate that gates the mid-turn coalescing protocol."""
|
||||
|
||||
def test_claude_code_supports_coalescing(self):
|
||||
from headroom.proxy.auth_mode import supports_mid_turn_coalescing
|
||||
|
||||
assert supports_mid_turn_coalescing("claude-code") is True
|
||||
|
||||
def test_other_clients_do_not_support_coalescing(self):
|
||||
from headroom.proxy.auth_mode import supports_mid_turn_coalescing
|
||||
|
||||
for client in ("opencode", "codex", "cursor", "aider", "", None):
|
||||
assert supports_mid_turn_coalescing(client) is False
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue