From a4bd2e62a5bb73f15b3b12e979c69e2b555bee10 Mon Sep 17 00:00:00 2001 From: TenderDeve Date: Wed, 12 Aug 2026 10:06:39 +0530 Subject: [PATCH] fix(proxy): gate mid-turn message coalescing to Claude Code clients (#1643) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## 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 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 --- headroom/proxy/auth_mode.py | 23 ++++++ headroom/proxy/handlers/anthropic.py | 23 ++++-- headroom/proxy/handlers/streaming.py | 14 +++- tests/test_mid_turn_steering.py | 112 +++++++++++++++++++++++++++ 4 files changed, 163 insertions(+), 9 deletions(-) diff --git a/headroom/proxy/auth_mode.py b/headroom/proxy/auth_mode.py index 2c7e97021..97a76259a 100644 --- a/headroom/proxy/auth_mode.py +++ b/headroom/proxy/auth_mode.py @@ -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", ] diff --git a/headroom/proxy/handlers/anthropic.py b/headroom/proxy/handlers/anthropic.py index 059fdb8f0..9e79e3a27 100644 --- a/headroom/proxy/handlers/anthropic.py +++ b/headroom/proxy/handlers/anthropic.py @@ -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) diff --git a/headroom/proxy/handlers/streaming.py b/headroom/proxy/handlers/streaming.py index ef9600381..07d4f70be 100644 --- a/headroom/proxy/handlers/streaming.py +++ b/headroom/proxy/handlers/streaming.py @@ -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} ) diff --git a/tests/test_mid_turn_steering.py b/tests/test_mid_turn_steering.py index 1776190a6..81534ae0d 100644 --- a/tests/test_mid_turn_steering.py +++ b/tests/test_mid_turn_steering.py @@ -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