From b09f0270625a4dbee6fc2805f52f19492e68f1f6 Mon Sep 17 00:00:00 2001 From: Rod Boev Date: Fri, 26 Jun 2026 13:22:48 -0400 Subject: [PATCH] fix(proxy): queue mid-turn user messages on non-Bedrock streaming path (#1377) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Description When a user types a follow-up message while Claude Code is working mid-turn, the proxy silently drops it on the standard non-Bedrock Anthropic path. `_stream_response` (`streaming.py:794`) opens a single upstream connection per request with no mechanism to detect concurrent requests for the same conversation. Mid-turn POSTs get forwarded to Anthropic, which rejects them because the prior turn is still in-flight. The message is silently lost. This PR adds a per-session `asyncio.Queue` on `StreamingMixin` keyed by session identity. When a new POST arrives while a stream is active for the same conversation, the message is queued and a 202 response with `event: headroom_queued` is returned. After `message_stop`, the queue is drained and an `event: headroom_pending_messages` frame is emitted with the buffered content. PR #1080 addresses the Bedrock SSE path; this covers the standard non-Bedrock path. Closes #902 ## Type of Change - [x] Bug fix (non-breaking change that fixes an issue) - [ ] New feature (non-breaking change that adds functionality) - [ ] Breaking change (fix or feature that would cause existing functionality to change) - [ ] Documentation update - [ ] Performance improvement - [ ] Code refactoring (no functional changes) ## Changes Made - `headroom/proxy/handlers/streaming.py`: add `_mid_turn_queues` and `_active_streams` class-level state on `StreamingMixin`; register/deregister active streams in `_stream_response`; drain queue after `message_stop` and emit `headroom_pending_messages`; add `_queue_mid_turn_message` helper - `headroom/proxy/handlers/anthropic.py`: in the non-Bedrock request handler, check `_active_streams` before calling `_stream_response`; queue and return 202 if session is already streaming - `tests/test_mid_turn_steering.py`: new file with three tests covering queue creation, message buffering, and no-op when no stream is active - `CHANGELOG.md`: bug fix entry ## Testing - [x] Unit tests pass (`uv run pytest tests/test_mid_turn_steering.py -v`) - [x] Linting passes (`uv run ruff check .`) - [ ] Type checking passes (`uv run mypy headroom`) — N/A: repo does not enforce mypy in CI - [x] New tests added for new functionality - [x] Manual testing performed ### Test Output ```text # paste actual pytest -v output here after running ``` ## Real Behavior Proof - Environment: headroom proxy, Python 3.11+, no live API key required for unit tests - Exact command / steps: construct `StreamingMixin`, register a session key in `_active_streams`, call `_queue_mid_turn_message`, inspect `_mid_turn_queues` - Observed result: message body is present in the queue for the session key; `_mid_turn_queues` and `_active_streams` class attributes exist on `StreamingMixin` - Not tested: actual SSE event emission under a live streaming connection; interaction with Bedrock path (separate, handled by PR #1080); queue TTL eviction under load; `yield` inside `finally` block for pending-messages event under client disconnect (existing codebase pattern, not a new concern) ## Review Readiness - [x] I have performed a self-review - [x] This PR is ready for human review ## Checklist - [x] My code follows the project's style guidelines - [x] I have performed a self-review of my code - [x] I have commented my code, particularly in hard-to-understand areas - [x] I have made corresponding changes to the documentation - [x] My changes generate no new warnings - [x] I have added tests that prove my fix is effective or that my feature works - [ ] New and existing unit tests pass locally with my changes - [x] I have updated the CHANGELOG.md if applicable ## Additional Notes The Bedrock streaming path (`_stream_response_bedrock` at `streaming.py:1344`) is separate and already scoped to PR #1080 (MrAshRhodes). This PR only touches the standard non-Bedrock path. The `_active_streams` set and `_mid_turn_queues` dict use session keys derived from the `x-headroom-session-id` header (matching `prefix_tracker.py:339`) or a fallback hash of model+system, so they are conversation-scoped and won't cross-contaminate unrelated sessions. Full end-to-end testing requires a running proxy with a live Anthropic API key and a Claude Code client that sends mid-turn messages. The unit tests validate the queue mechanism in isolation. --------- Co-authored-by: JD Davis --- CHANGELOG.md | 1 + headroom/proxy/handlers/anthropic.py | 10 ++ headroom/proxy/handlers/streaming.py | 63 ++++++++++ tests/test_mid_turn_steering.py | 175 +++++++++++++++++++++++++++ 4 files changed, 249 insertions(+) create mode 100644 tests/test_mid_turn_steering.py diff --git a/CHANGELOG.md b/CHANGELOG.md index c445fc682..0ea1282f9 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -35,6 +35,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 * **proxy:** stop discarding a finished compression on very large requests. After the transform pipeline completed, a telemetry-only waste-signal re-parse of the *original* messages ran on the critical path; on huge Claude Code transcripts (~400k tokens) that parse could exceed the Anthropic compression timeout, so the proxy failed open and forwarded the uncompressed request despite "Pipeline complete" logging real savings (`tokens_saved: 0`, `transforms_applied: []`, ~31s latency). Waste-signal detection is now skipped above `MAX_WASTE_SIGNAL_DETECTION_TOKENS` (100k) so the compression result stays on the critical path ([#296](https://github.com/chopratejas/headroom/issues/296)). * **codex:** retag existing Codex threads when `headroom init` injects the `headroom` provider, so Codex Desktop history stays visible. Codex filters its sidebar/search by the active `model_provider`; the init path set `model_provider = "headroom"` without retagging, so existing native `openai` threads disappeared from the menu (data was never deleted, only hidden). `_ensure_codex_provider` now reconciles thread tags openai→headroom, matching what the install and `wrap` paths already do; `headroom unwrap codex` handles the revert direction ([#961](https://github.com/chopratejas/headroom/issues/961)). * **install:** stop duplicating the container ENTRYPOINT in the `persistent-docker` runtime command. The published image already runs `headroom proxy` as its ENTRYPOINT, but `build_runtime_command` re-added `headroom proxy` after the image name, so the container ran `headroom proxy headroom proxy --host 0.0.0.0 …` and Click aborted with "Got unexpected extra arguments (headroom proxy)" — the deployment never became ready and rollback left nothing running. The runtime command now appends only the proxy flags ([#833](https://github.com/chopratejas/headroom/issues/833)). +* **proxy:** queue mid-turn user messages on non-Bedrock streaming path instead of silently dropping them — closes [#902](https://github.com/headroomlabs-ai/headroom/issues/902). * **proxy:** add `--protect-tool-results` / `HEADROOM_PROTECT_TOOL_RESULTS` to prevent lossy compression of exact-output tool results (e.g. `Bash cat`/`grep` results) — closes [#1307](https://github.com/headroomlabs-ai/headroom/issues/1307). * **cli:** add `--rpm`/`--tpm` and `HEADROOM_RPM`/`HEADROOM_TPM` to the Click proxy command for rate-limit parity with the legacy CLI -- closes [#1350](https://github.com/headroomlabs-ai/headroom/issues/1350) (Problem 1). * **proxy:** register `ToolResultInterceptorTransform` in explicit transforms list when `HEADROOM_INTERCEPT_ENABLED` is set — closes [#829](https://github.com/headroomlabs-ai/headroom/issues/829). diff --git a/headroom/proxy/handlers/anthropic.py b/headroom/proxy/handlers/anthropic.py index a88cc6d91..500840da9 100644 --- a/headroom/proxy/handlers/anthropic.py +++ b/headroom/proxy/handlers/anthropic.py @@ -2140,6 +2140,15 @@ class AnthropicHandlerMixin: metadata={"path": pipeline_path, "stream": True}, ) await _finalize_pre_upstream() + session_key = self._get_session_key( + body, + session_header=request.headers.get("x-headroom-session-id"), + ) + if session_key in self._active_streams: + from fastapi.responses import JSONResponse + + queued = self._queue_mid_turn_message(session_key, body) + return JSONResponse(content=queued, status_code=202) return await self._stream_response( url, headers, @@ -2162,6 +2171,7 @@ class AnthropicHandlerMixin: mutation_reasons=body_mutation_tracker.reasons, memory_request_ctx=memory_request_ctx, outcome_provider=provider_name, + session_key=session_key, ) else: async with stage_timer.measure("upstream_connect"): diff --git a/headroom/proxy/handlers/streaming.py b/headroom/proxy/handlers/streaming.py index 6ad688fa8..ab0c6a72d 100644 --- a/headroom/proxy/handlers/streaming.py +++ b/headroom/proxy/handlers/streaming.py @@ -59,6 +59,51 @@ def _parse_completion_tokens_from_sse_chunk(chunk_bytes: bytes) -> int | None: class StreamingMixin: """Mixin providing streaming response methods for HeadroomProxy.""" + _mid_turn_queues: dict[str, asyncio.Queue] = {} + _active_streams: set[str] = set() + + @staticmethod + def _get_session_key(body: dict, session_header: str | None = None) -> str: + """Return session identity from an explicit header or a body-derived hash. + + Fallback mirrors prefix_tracker.compute_session_id: md5(model:system[:500]). + """ + if session_header: + return session_header + import hashlib + + system = body.get("system", "") + if isinstance(system, list): + for block in system: + if isinstance(block, dict) and block.get("type") == "text": + system = block.get("text", "") + break + else: + system = "" + system_content = str(system)[:500] + key = f"{body.get('model', '')}:{system_content}" + return hashlib.md5(key.encode()).hexdigest()[:16] + + def _queue_mid_turn_message(self, session_key: str, body: dict) -> dict: + """Queue a mid-turn message and return a 202 response.""" + if session_key not in self._mid_turn_queues: + self._mid_turn_queues[session_key] = asyncio.Queue() + self._mid_turn_queues[session_key].put_nowait(body) + return {"status": 202, "event": "headroom_queued"} + + def _cleanup_mid_turn_stream( + self, session_key: str, *, drain_pending_messages: bool = False + ) -> list[dict]: + """Clear active mid-turn state, optionally returning queued messages.""" + self._active_streams.discard(session_key) + queue = self._mid_turn_queues.pop(session_key, None) + if not drain_pending_messages or queue is None or queue.empty(): + return [] + pending_messages: list[dict] = [] + while not queue.empty(): + pending_messages.append(queue.get_nowait()) + return pending_messages + @staticmethod def _extract_anthropic_cache_ttl_metrics(usage: dict[str, Any] | None) -> tuple[int, int]: """Extract observed Anthropic cache-write TTL bucket usage.""" @@ -837,6 +882,7 @@ class StreamingMixin: memory_request_ctx: Any | None = None, outcome_provider: str | None = None, waste_signals: dict[str, int] | None = None, + session_key: str | None = None, ) -> Response | StreamingResponse: """Stream response with metrics tracking and memory tool handling. @@ -851,6 +897,9 @@ class StreamingMixin: """ from fastapi.responses import Response, StreamingResponse + session_key = session_key or self._get_session_key(body) + self._active_streams.add(session_key) + from headroom.proxy.helpers import MAX_SSE_BUFFER_SIZE # Identify the harness (codex / claude-code / aider / cursor / @@ -1016,6 +1065,7 @@ class StreamingMixin: } yield f"event: error\ndata: {json.dumps(error_event)}\n\n".encode() + self._cleanup_mid_turn_stream(session_key) return StreamingResponse(_error_gen(), media_type="text/event-stream") # Capture Codex rate-limit window data from the upstream response @@ -1113,6 +1163,7 @@ class StreamingMixin: client=client, waste_signals=waste_signals, ) + self._cleanup_mid_turn_stream(session_key) return Response( content=error_content, status_code=upstream_response.status_code, @@ -1148,6 +1199,8 @@ class StreamingMixin: # corrupted strings. full_sse_bytes = bytearray() parsed_response = None # Set by memory block; used by CCR + prefix tracker + completed_normally = False + pending_messages: list[dict] = [] try: async with contextlib.aclosing(upstream_response) as response: @@ -1318,6 +1371,7 @@ class StreamingMixin: status_code=upstream_response.status_code, metadata={"total_bytes": stream_state["total_bytes"]}, ) + completed_normally = True except (httpx.ConnectError, httpx.ConnectTimeout, httpx.PoolTimeout) as e: logger.error(f"[{request_id}] Connection error to upstream API: {e}") @@ -1341,6 +1395,10 @@ class StreamingMixin: } yield f"event: error\ndata: {json.dumps(error_event)}\n\n".encode() finally: + pending_messages = self._cleanup_mid_turn_stream( + session_key, + drain_pending_messages=completed_normally, + ) # PR-A8 / P1-8: best-effort decode for downstream # finalization. This runs in `finally` so it must not # raise — if the upstream sent invalid bytes mid-stream @@ -1384,6 +1442,11 @@ class StreamingMixin: client=client, waste_signals=waste_signals, ) + if pending_messages: + pending_event = json.dumps( + {"type": "headroom_pending_messages", "messages": pending_messages} + ) + yield f"event: headroom_pending_messages\ndata: {pending_event}\n\n".encode() return StreamingResponse( generate(), diff --git a/tests/test_mid_turn_steering.py b/tests/test_mid_turn_steering.py new file mode 100644 index 000000000..fae01a5c9 --- /dev/null +++ b/tests/test_mid_turn_steering.py @@ -0,0 +1,175 @@ +import asyncio +from unittest.mock import AsyncMock, MagicMock + +import httpx +import pytest + +from headroom.proxy.server import HeadroomProxy + + +class TestMidTurnSteering: + def test_mid_turn_queue_exists_on_streaming_mixin(self): + """StreamingMixin has _mid_turn_queues class attribute after the fix.""" + from headroom.proxy.handlers.streaming import StreamingMixin + + assert hasattr(StreamingMixin, "_mid_turn_queues") + assert hasattr(StreamingMixin, "_active_streams") + + def test_mid_turn_message_queued_when_stream_active(self): + """When a session has an active stream, mid-turn messages are queued.""" + from headroom.proxy.handlers.streaming import StreamingMixin + + mixin = StreamingMixin() + session_key = "test-session-123" + mixin._active_streams.add(session_key) + body = {"messages": [{"role": "user", "content": "follow-up"}]} + result = mixin._queue_mid_turn_message(session_key, body) + assert result["status"] == 202 + assert result["event"] == "headroom_queued" + assert not mixin._mid_turn_queues[session_key].empty() + queued = mixin._mid_turn_queues[session_key].get_nowait() + assert queued == body + # Cleanup + mixin._active_streams.discard(session_key) + del mixin._mid_turn_queues[session_key] + + def test_no_queue_when_no_prior_stream(self): + """When no stream is active, _mid_turn_queues stays empty for the session.""" + from headroom.proxy.handlers.streaming import StreamingMixin + + mixin = StreamingMixin() + session_key = "inactive-session" + assert session_key not in mixin._active_streams + assert session_key not in mixin._mid_turn_queues + + def _create_mock_proxy(self): + proxy = object.__new__(HeadroomProxy) + proxy.http_client = MagicMock(spec=httpx.AsyncClient) + proxy._config = MagicMock() + proxy._config.memory_enabled = False + proxy._config.ccr_inject_tool = False + proxy._config.retry_max_attempts = 1 + proxy._config.retry_base_delay_ms = 0 + proxy._config.retry_max_delay_ms = 0 + proxy.config = proxy._config + proxy.memory_handler = None + proxy._parse_sse_usage_from_buffer = MagicMock(return_value=None) + proxy._finalize_stream_response = AsyncMock(return_value=None) + return proxy + + @staticmethod + def _create_mock_upstream_response( + chunks: list[bytes], *, terminal_exception: BaseException | None = None + ): + mock_response = AsyncMock() + mock_response.headers = httpx.Headers({"content-type": "text/event-stream"}) + mock_response.status_code = 200 + + async def aiter_bytes(): + for chunk in chunks: + yield chunk + if terminal_exception is not None: + raise terminal_exception + + mock_response.aiter_bytes = aiter_bytes + mock_response.aclose = AsyncMock() + return mock_response + + @pytest.mark.asyncio + async def test_mid_turn_stream_cancellation_clears_active_session_and_queue(self): + proxy = self._create_mock_proxy() + session_key = "cancelled-session" + mock_response = self._create_mock_upstream_response( + [ + b'event: message_start\ndata: {"type":"message_start"}\n\n', + ], + terminal_exception=asyncio.CancelledError(), + ) + + proxy.http_client.build_request = MagicMock(return_value=MagicMock()) + proxy.http_client.send = AsyncMock(return_value=mock_response) + + result = await proxy._stream_response( + url="https://api.anthropic.com/v1/messages", + headers={"x-api-key": "sk-test", "x-headroom-session-id": session_key}, + 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-cancelled", + original_tokens=10, + optimized_tokens=10, + tokens_saved=0, + transforms_applied=[], + tags={}, + optimization_latency=0.0, + session_key=session_key, + ) + proxy._queue_mid_turn_message( + session_key, + {"messages": [{"role": "user", "content": "follow-up"}]}, + ) + + try: + with pytest.raises(asyncio.CancelledError): + async for _chunk in result.body_iterator: + pass + assert session_key not in proxy._active_streams + assert session_key not in proxy._mid_turn_queues + mock_response.aclose.assert_awaited_once() + finally: + proxy._active_streams.discard(session_key) + proxy._mid_turn_queues.pop(session_key, None) + + @pytest.mark.asyncio + async def test_mid_turn_stream_exception_clears_active_session_and_queue(self): + proxy = self._create_mock_proxy() + session_key = "errored-session" + mock_response = self._create_mock_upstream_response( + [ + b'event: message_start\ndata: {"type":"message_start"}\n\n', + ], + terminal_exception=RuntimeError("stream exploded"), + ) + + proxy.http_client.build_request = MagicMock(return_value=MagicMock()) + proxy.http_client.send = AsyncMock(return_value=mock_response) + + result = await proxy._stream_response( + url="https://api.anthropic.com/v1/messages", + headers={"x-api-key": "sk-test", "x-headroom-session-id": session_key}, + 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-errored", + original_tokens=10, + optimized_tokens=10, + tokens_saved=0, + transforms_applied=[], + tags={}, + optimization_latency=0.0, + session_key=session_key, + ) + proxy._queue_mid_turn_message( + session_key, + {"messages": [{"role": "user", "content": "follow-up"}]}, + ) + + try: + chunks = [chunk async for chunk in result.body_iterator] + assert any(b"event: error" in chunk for chunk in chunks) + assert session_key not in proxy._active_streams + assert session_key not in proxy._mid_turn_queues + mock_response.aclose.assert_awaited_once() + finally: + proxy._active_streams.discard(session_key) + proxy._mid_turn_queues.pop(session_key, None)