mirror of
https://github.com/headroomlabs-ai/headroom.git
synced 2026-08-27 14:17:10 -04:00
## Description Closes #902 Mid-turn user interjections (steering) silently dropped through the Bedrock streaming path. The _stream_response_bedrock code path reconstructs Anthropic SSE events from parsed StreamEvent objects instead of passing raw bytes through, so SSE-level ping keepalives are never forwarded to Claude Code. Claude Code relies on ping events to arm its mid-turn steering / interruptible state; without them, queued interjections are discarded instead of sent. Root cause (confirmed): - Standard direct-Anthropic path does a raw yield-chunk passthrough — pings flow unchanged. - Bedrock path (_stream_response_bedrock.generate()) reconstructs events from litellm/anyllm stream_message() output, which only yields semantic events (message_start, content_block_*, message_delta, message_stop, error). No pings, ever. Fix: emit a synthetic 'event: ping / data: {}' at stream start (before the first message_start) so downstream clients see the same ping-then-content cadence as a real Anthropic stream. Note: periodic pings for very long responses (>~25s) may be needed if steering disarms on a timer. This commit arms it at turn start; follow-up if reporters confirm steering still drops on long turns. The causal link (ping → steering) is the reporter's hypothesis from hands-on debugging. The observable defect (zero pings in stream) is confirmed and fixed. ## Type of Change - [x] Bug fix (non-breaking change that fixes an issue) ## Changes Made - headroom/proxy/handlers/streaming.py: yield ping event before the event loop in _stream_response_bedrock.generate() - tests/test_proxy/test_bedrock_sse_ping.py: 3 new tests asserting ping appears before message_start ## Testing - [x] Unit tests pass (pytest) - [x] Linting passes (ruff check .) - [x] New tests added for new functionality ### Test Output ``` tests/test_proxy/test_bedrock_sse_ping.py::test_bedrock_stream_emits_ping_before_message_start PASSED tests/test_proxy/test_bedrock_sse_ping.py::test_bedrock_stream_ping_has_empty_data PASSED tests/test_proxy/test_bedrock_sse_ping.py::test_bedrock_stream_contains_message_stop PASSED tests/test_backend_streaming_cache_metrics.py (4 tests) PASSED 7 passed in 4.41s ``` ## Real Behavior Proof - Environment: macOS, Python 3.11, headroom unit tests - Exact command / steps: pytest tests/test_proxy/test_bedrock_sse_ping.py -v - Observed result: 3 new tests pass; ping appears before message_start in Bedrock stream - Not tested: end-to-end against live Bedrock + Claude Code (no Bedrock credentials available) ## Review Readiness - [x] I have performed a self-review - [x] This PR is ready for human review ## Checklist - [x] Code follows project style guidelines - [x] Code is commented where non-obvious - [x] No new warnings - [x] Tests added and passing ## Additional Notes The Rust proxy files mentioned in the issue (sse/framing.rs, sse/anthropic.rs) are NOT part of this fix. Those drops are in a telemetry-only tee task that never affects the client byte path — the Rust proxy does a raw bytes passthrough for all responses. The defect is Python-only, confined to the Bedrock path. --------- Co-authored-by: JerrettDavis <mxjerrett@gmail.com>
This commit is contained in:
parent
9fd5ae3d53
commit
4dab254d52
2 changed files with 161 additions and 0 deletions
|
|
@ -1723,6 +1723,14 @@ class StreamingMixin:
|
|||
try:
|
||||
assert self.anthropic_backend is not None
|
||||
|
||||
# Emit a synthetic ping before the first message_start so that
|
||||
# downstream clients (e.g. Claude Code) arm their mid-turn
|
||||
# steering / interruptible state. The Bedrock-to-Anthropic
|
||||
# translation layer never produces SSE-level keepalives; we
|
||||
# synthesise one here to match the real Anthropic wire format
|
||||
# (issue #902).
|
||||
yield b"event: ping\ndata: {}\n\n"
|
||||
|
||||
async for event in self.anthropic_backend.stream_message(body, headers):
|
||||
# Record TTFB on first event
|
||||
if stream_state["ttfb_ms"] is None:
|
||||
|
|
|
|||
153
tests/test_proxy/test_bedrock_sse_ping.py
Normal file
153
tests/test_proxy/test_bedrock_sse_ping.py
Normal file
|
|
@ -0,0 +1,153 @@
|
|||
"""SSE ping passthrough in the Bedrock streaming path (issue #902).
|
||||
|
||||
When Headroom routes through a Bedrock (LiteLLM/AnyLLM) backend the
|
||||
translation layer emits only Anthropic-semantic events — it never produces
|
||||
SSE-level ping keepalives. Claude Code uses ping events to keep a turn in
|
||||
the interruptible / steering-armed state; without them, mid-turn interjections
|
||||
are silently dropped.
|
||||
|
||||
Fix: ``_stream_response_bedrock.generate()`` now emits one synthetic
|
||||
``event: ping\\ndata: {}\\n\\n`` before the first ``message_start`` so the
|
||||
downstream client sees the same ping-then-content cadence as a real Anthropic
|
||||
stream.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from collections.abc import AsyncIterator
|
||||
from typing import Any
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
import pytest
|
||||
|
||||
fastapi = pytest.importorskip("fastapi")
|
||||
|
||||
from fastapi.testclient import TestClient # noqa: E402
|
||||
|
||||
from headroom.backends.base import StreamEvent # noqa: E402
|
||||
from headroom.proxy.server import ProxyConfig, create_app # noqa: E402
|
||||
|
||||
|
||||
def _ev(event_type: str, data: dict[str, Any]) -> StreamEvent:
|
||||
return StreamEvent(event_type=event_type, data=data)
|
||||
|
||||
|
||||
def _minimal_events() -> list[StreamEvent]:
|
||||
return [
|
||||
_ev(
|
||||
"message_start",
|
||||
{
|
||||
"type": "message_start",
|
||||
"message": {
|
||||
"id": "msg_test",
|
||||
"model": "claude-3-5-sonnet-20241022",
|
||||
"role": "assistant",
|
||||
"type": "message",
|
||||
"content": [],
|
||||
"usage": {"input_tokens": 10, "output_tokens": 0},
|
||||
},
|
||||
},
|
||||
),
|
||||
_ev(
|
||||
"content_block_start",
|
||||
{
|
||||
"type": "content_block_start",
|
||||
"index": 0,
|
||||
"content_block": {"type": "text", "text": ""},
|
||||
},
|
||||
),
|
||||
_ev(
|
||||
"content_block_delta",
|
||||
{
|
||||
"type": "content_block_delta",
|
||||
"index": 0,
|
||||
"delta": {"type": "text_delta", "text": "hi"},
|
||||
},
|
||||
),
|
||||
_ev("content_block_stop", {"type": "content_block_stop", "index": 0}),
|
||||
_ev(
|
||||
"message_delta",
|
||||
{
|
||||
"type": "message_delta",
|
||||
"delta": {"stop_reason": "end_turn"},
|
||||
"usage": {"output_tokens": 1},
|
||||
},
|
||||
),
|
||||
_ev("message_stop", {"type": "message_stop"}),
|
||||
]
|
||||
|
||||
|
||||
def _make_bedrock_backend(events: list[StreamEvent]) -> MagicMock:
|
||||
async def fake_stream(body: dict, headers: dict) -> AsyncIterator[StreamEvent]:
|
||||
for evt in events:
|
||||
yield evt
|
||||
|
||||
mock = MagicMock()
|
||||
mock.name = "bedrock"
|
||||
mock.stream_message = fake_stream
|
||||
mock.map_model_id = MagicMock(return_value="claude-3-5-sonnet-20241022")
|
||||
mock.supports_model = MagicMock(return_value=True)
|
||||
return mock
|
||||
|
||||
|
||||
def _run_bedrock_stream(events: list[StreamEvent]) -> str:
|
||||
"""Return the full SSE response body from a Bedrock-backend streaming request."""
|
||||
config = ProxyConfig(
|
||||
optimize=False,
|
||||
cache_enabled=False,
|
||||
rate_limit_enabled=False,
|
||||
backend="anyllm",
|
||||
anyllm_provider="anthropic",
|
||||
)
|
||||
backend = _make_bedrock_backend(events)
|
||||
with patch("headroom.proxy.server.AnyLLMBackend", return_value=backend):
|
||||
app = create_app(config)
|
||||
with TestClient(app) as client:
|
||||
resp = client.post(
|
||||
"/v1/messages",
|
||||
json={
|
||||
"model": "claude-3-5-sonnet-20241022",
|
||||
"messages": [{"role": "user", "content": "hi"}],
|
||||
"max_tokens": 64,
|
||||
"stream": True,
|
||||
},
|
||||
headers={
|
||||
"x-api-key": "sk-ant-test",
|
||||
"anthropic-version": "2023-06-01",
|
||||
},
|
||||
)
|
||||
assert resp.status_code == 200, resp.text[:200]
|
||||
return resp.text
|
||||
|
||||
|
||||
def test_bedrock_stream_emits_ping_before_message_start() -> None:
|
||||
"""Bedrock path must emit a ping event before message_start (issue #902)."""
|
||||
body = _run_bedrock_stream(_minimal_events())
|
||||
|
||||
ping_idx = body.find("event: ping")
|
||||
start_idx = body.find("event: message_start")
|
||||
|
||||
assert ping_idx != -1, "No ping event found in Bedrock stream response"
|
||||
assert start_idx != -1, "No message_start event found in Bedrock stream response"
|
||||
assert ping_idx < start_idx, (
|
||||
f"ping (offset {ping_idx}) must appear before message_start (offset {start_idx})"
|
||||
)
|
||||
|
||||
|
||||
def test_bedrock_stream_ping_has_empty_data() -> None:
|
||||
"""Ping event must carry data: {} to match real Anthropic wire format."""
|
||||
body = _run_bedrock_stream(_minimal_events())
|
||||
|
||||
ping_start = body.find("event: ping")
|
||||
assert ping_start != -1, "No ping event found"
|
||||
|
||||
# The next ~30 bytes after 'event: ping' should contain 'data: {}'
|
||||
ping_block = body[ping_start : ping_start + 40]
|
||||
assert "data: {}" in ping_block, f"Ping block must contain 'data: {{}}', got: {ping_block!r}"
|
||||
|
||||
|
||||
def test_bedrock_stream_contains_message_stop() -> None:
|
||||
"""Smoke test: full event sequence still reaches the client alongside ping."""
|
||||
body = _run_bedrock_stream(_minimal_events())
|
||||
assert "event: message_stop" in body
|
||||
assert "event: message_start" in body
|
||||
Loading…
Add table
Add a link
Reference in a new issue