fix(pipeline): pump _stream_with_keepalive source from one task, not one per item

changelog.d/anthropic-stream-keepalive.md

@@ -0,0 +1,6 @@
+---
+category: Fixes
+pr: 804
+---
+
+**Anthropic streaming: re-inject keepalive pings so long generations don't idle out**: The Anthropic SDK's typed stream (`messages.create(stream=True)`) drops the wire `ping` events Anthropic emits during long generations. A model that stays silent before emitting content (e.g. a slow or classifier-gated response) therefore produced no bytes on the proxy→client connection for the entire pre-content phase, so any intermediary with an idle timeout (load balancer, reverse proxy) would cut the healthy stream mid-flight. The gateway now re-emits an Anthropic-style `ping` whenever the upstream is idle longer than `STREAM_KEEPALIVE_SECONDS` (15s), matching the keepalive behavior a direct Anthropic connection relies on.

changelog.d/streaming-otel-context-detach.md

@@ -0,0 +1,29 @@
+---
+category: Fixes
+---
+
+**`_stream_with_keepalive` pumps the upstream generator from one task instead of a fresh task per item, fixing the `opentelemetry.context` "Failed to detach context" ERROR storm**
+  - In production this logger produced roughly 60k ERROR lines per 48h,
+    continuously, since streaming keepalives shipped — about twice per
+    streaming request.
+  - Root cause: `AnthropicClient.stream()` and `_AnthropicPolicyIO._stream()`
+    each hold an OpenTelemetry span open across every chunk of the upstream
+    response (`with tracer.start_as_current_span(...): async for event in
+    ...: yield event`). `_stream_with_keepalive` drove that generator by
+    wrapping every single `__anext__()` call in a brand-new
+    `asyncio.ensure_future(...)` Task. `asyncio.Task` copies
+    `contextvars.Context` at creation, so the span's context-attach token
+    (created in the Task that fetched the first chunk) could not be
+    detached from the different Task that fetched the last one —
+    `contextvars.Context.reset()` raises `ValueError`, which
+    `opentelemetry.context.detach()` catches and logs at ERROR rather than
+    propagating. Reproduced directly against `_stream_with_keepalive` with
+    a real `TracerProvider` (no Sentry involved): the same ERROR line and
+    traceback as production. Sentry was a correlate, not the cause — it
+    was already ruled out by tracing `sentry_sdk`'s default integrations
+    (no OpenTelemetry-backed tracing is enabled by our `init_sentry()`).
+  - **Fix**: `_stream_with_keepalive` now pumps `source` to completion from
+    a single persistent task (`_pump_to_queue`) that feeds a one-slot
+    queue, preserving the existing "never drop a slow item, only inject a
+    keepalive" behavior and the "cancel-on-close" contract, while keeping
+    every span's attach/detach pair inside one task's context.

src/luthien_proxy/pipeline/anthropic_processor.py

@@ -96,6 +96,122 @@ class _StreamErrorEvent(TypedDict):
 tracer = trace.get_tracer(__name__)
 
 
+# --- Streaming keepalive ------------------------------------------------------
+# Anthropic's wire stream emits `ping` events during long generations to keep the
+# connection alive, but the Anthropic SDK's typed stream (messages.create(
+# stream=True)) drops them. Without re-injecting keepalives, a model that stays
+# silent for a long pre-content phase makes the proxy->client connection idle, and
+# an intermediary (e.g. the ALB idle timeout) cuts the stream mid-flight even
+# though the request is healthy. We emit an Anthropic-style `ping` whenever the
+# upstream produces no event within STREAM_KEEPALIVE_SECONDS.
+STREAM_KEEPALIVE_SECONDS = 15.0
+_KEEPALIVE_SSE = 'event: ping\ndata: {"type": "ping"}\n\n'
+_STREAM_DONE = object()
+
+
+class _Keepalive:
+    """Sentinel yielded by `_stream_with_keepalive` during an upstream gap."""
+
+
+_KEEPALIVE = _Keepalive()
+
+
+class _StreamError:
+    """Sentinel carrying an exception raised while pumping `source` to completion.
+
+    Queued instead of raised directly so the pump task (see `_pump_to_queue`) can
+    report a failure to the consumer without itself needing to be re-awaited from
+    a context that still holds the failing span open.
+    """
+
+    __slots__ = ("exc",)
+
+    def __init__(self, exc: BaseException) -> None:
+        self.exc = exc
+
+
+async def _pump_to_queue(
+    source: AsyncIterator[AnthropicPolicyEmission],
+    queue: "asyncio.Queue[object]",
+) -> None:
+    """Drive `source` to completion from a single persistent task, queuing each item.
+
+    `AnthropicClient.stream()` and `_AnthropicPolicyIO._stream()` each hold an
+    OpenTelemetry span open across every chunk of the upstream response: the
+    span's `with` block attaches its context once, on the first chunk, and
+    detaches it once, when the stream is exhausted. `contextvars.Context` is
+    copied per `asyncio.Task` (see CPython's `asyncio.Task.__init__`), so a
+    context token attached while running in one Task cannot be detached while
+    running in another — `contextvars.Context.reset()` raises `ValueError`, which
+    `opentelemetry.context.detach()` catches and logs at ERROR ("Failed to detach
+    context"). A previous version of `_stream_with_keepalive` created a fresh
+    `asyncio.Task` for every upstream item (`asyncio.ensure_future` after each
+    successful `__anext__`), tripping this twice per streaming request — once for
+    each of the two nested spans above. Pumping `source` from one task, start to
+    finish, keeps every attach/detach pair inside that task's own context.
+
+    A `CancelledError` raised by `source` itself (as opposed to this task being
+    cancelled from the outside, e.g. by `_stream_with_keepalive`'s cleanup on
+    close) is relayed like any other exception: `Task.cancelling()` is 0 in that
+    case, since nothing called `.cancel()` on this task. When this task *was*
+    cancelled from the outside, `cancelling()` is nonzero and the error is
+    re-raised untouched instead of being queued — the consumer has already
+    stopped reading by then, so queuing it could block forever on a full queue.
+    """
+    try:
+        async for item in source:
+            await queue.put(item)
+    except asyncio.CancelledError as exc:
+        pump_task = asyncio.current_task()
+        if pump_task is not None and pump_task.cancelling() > 0:
+            raise
+        await queue.put(_StreamError(exc))
+        return
+    except Exception as exc:
+        await queue.put(_StreamError(exc))
+        return
+    await queue.put(_STREAM_DONE)
+
+
+async def _stream_with_keepalive(
+    source: AsyncIterator[AnthropicPolicyEmission],
+    interval_seconds: float,
+) -> AsyncIterator["AnthropicPolicyEmission | _Keepalive"]:
+    """Yield from `source`, injecting a `_KEEPALIVE` sentinel on slow items.
+
+    A keepalive is injected whenever the next item takes longer than
+    `interval_seconds` to arrive.
+
+    `source` is pumped by a single background task for its entire lifetime (see
+    `_pump_to_queue` for why a fresh task per item is unsafe here). The one-slot
+    queue means the pump is never more than one item ahead of what this generator
+    has yielded, so a slow item is never dropped: the timeout only triggers a
+    keepalive and we keep waiting on the same queue. The pump task is cancelled
+    on close.
+    """
+    queue: "asyncio.Queue[object]" = asyncio.Queue(maxsize=1)
+    pump = asyncio.ensure_future(_pump_to_queue(source, queue))
+    try:
+        while True:
+            try:
+                item = await asyncio.wait_for(queue.get(), interval_seconds)
+            except asyncio.TimeoutError:
+                yield _KEEPALIVE
+                continue
+            if item is _STREAM_DONE:
+                return
+            if isinstance(item, _StreamError):
+                raise item.exc
+            yield cast("AnthropicPolicyEmission", item)
+    finally:
+        if not pump.done():
+            pump.cancel()
+            try:
+                await pump
+            except BaseException:
+                pass
+
+
 class _AnthropicPolicyIO(AnthropicPolicyIOProtocol):
     """Request-scoped I/O helpers for execution-oriented Anthropic policies."""
 
@@ -748,7 +864,15 @@ async def streaming_with_spans() -> AsyncIterator[str]:
                 caught_exception = False
                 try:
                     with tracer.start_as_current_span("policy_execute"):
-                        async for emitted in emissions:
+                        async for emitted in _stream_with_keepalive(emissions, STREAM_KEEPALIVE_SECONDS):
+                            if isinstance(emitted, _Keepalive):
+                                # Upstream gap: forward an Anthropic-style ping so
+                                # the connection doesn't idle out mid-generation.
+                                # Only after message_start (emitted_any) so the
+                                # wire event ordering stays intact.
+                                if emitted_any:
+                                    yield _KEEPALIVE_SSE
+                                continue
                             if _is_anthropic_response_emission(emitted):
                                 raise TypeError(
                                     "Streaming Anthropic execution policies must emit streaming events, "

tests/luthien_proxy/unit_tests/pipeline/test_anthropic_processor.py

@@ -2,7 +2,9 @@
 
 import asyncio
 import json
+import logging
 from collections.abc import AsyncIterator
+from typing import Iterator
 from unittest.mock import AsyncMock, MagicMock, patch
 
 import pytest
@@ -23,6 +25,11 @@
 from fastapi.responses import StreamingResponse as FastAPIStreamingResponse
 from httpx import Request as HttpxRequest
 from httpx import Response as HttpxResponse
+from opentelemetry import trace
+from opentelemetry.sdk.trace import TracerProvider
+from opentelemetry.sdk.trace.export import SimpleSpanProcessor
+from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter
+from opentelemetry.trace import StatusCode
 from tests.constants import DEFAULT_TEST_MODEL
 from tests.luthien_proxy.fixtures.policy_context import make_policy_context
 
@@ -45,6 +52,30 @@
 from luthien_proxy.policy_core.policy_context import PolicyContext
 
 
+@pytest.fixture
+def span_exporter(monkeypatch: pytest.MonkeyPatch) -> Iterator[InMemorySpanExporter]:
+    """Route `anthropic_processor`'s module-level tracer to an in-memory exporter."""
+    previous_provider = trace._TRACER_PROVIDER
+    previous_once_done = trace._TRACER_PROVIDER_SET_ONCE._done
+    trace._TRACER_PROVIDER_SET_ONCE._done = False
+
+    provider = TracerProvider()
+    exporter = InMemorySpanExporter()
+    provider.add_span_processor(SimpleSpanProcessor(exporter))
+    trace.set_tracer_provider(provider)
+    monkeypatch.setattr(
+        "luthien_proxy.pipeline.anthropic_processor.tracer",
+        provider.get_tracer("luthien_proxy.pipeline.anthropic_processor"),
+    )
+
+    try:
+        yield exporter
+    finally:
+        provider.shutdown()
+        trace._TRACER_PROVIDER = previous_provider
+        trace._TRACER_PROVIDER_SET_ONCE._done = previous_once_done
+
+
 class TestFormatSSEEvent:
     """Tests for _format_sse_event helper function."""
 
@@ -2736,3 +2767,136 @@ async def emissions():
 
         webhook.fire_and_forget.assert_called_once()
         recorder.flush.assert_called()  # cleanup completed despite webhook failure
+
+
+class TestStreamWithKeepalive:
+    """`_stream_with_keepalive` injects SSE keepalives during upstream gaps without
+    dropping or reordering real events. Regression: the Anthropic SDK drops upstream
+    `ping` events, so a long silent generation idled out at the ALB timeout."""
+
+    async def test_injects_keepalive_during_gap(self):
+        from luthien_proxy.pipeline.anthropic_processor import (
+            _Keepalive,
+            _stream_with_keepalive,
+        )
+
+        async def source():
+            yield "A"
+            await asyncio.sleep(0.12)  # > interval -> keepalives expected in this gap
+            yield "B"
+
+        out = [item async for item in _stream_with_keepalive(source(), 0.02)]
+        reals = [x for x in out if not isinstance(x, _Keepalive)]
+        keepalives = [x for x in out if isinstance(x, _Keepalive)]
+        assert reals == ["A", "B"]  # order preserved, nothing dropped
+        assert len(keepalives) >= 1  # the gap produced at least one keepalive
+        assert out[0] == "A" and out[-1] == "B"  # keepalives sit between real events
+
+    async def test_no_keepalive_when_fast(self):
+        from luthien_proxy.pipeline.anthropic_processor import (
+            _Keepalive,
+            _stream_with_keepalive,
+        )
+
+        async def source():
+            yield "A"
+            yield "B"
+            yield "C"
+
+        out = [item async for item in _stream_with_keepalive(source(), 0.5)]
+        assert out == ["A", "B", "C"]
+        assert not any(isinstance(x, _Keepalive) for x in out)
+
+    async def test_empty_source_terminates(self):
+        from luthien_proxy.pipeline.anthropic_processor import _stream_with_keepalive
+
+        async def source():
+            return
+            yield  # pragma: no cover - makes this an async generator
+
+        out = [item async for item in _stream_with_keepalive(source(), 0.02)]
+        assert out == []
+
+    async def test_close_cancels_pending(self):
+        from luthien_proxy.pipeline.anthropic_processor import (
+            _Keepalive,
+            _stream_with_keepalive,
+        )
+
+        cancelled = asyncio.Event()
+
+        async def source():
+            yield "A"
+            try:
+                await asyncio.sleep(10)  # long-pending item, in flight at close
+            except asyncio.CancelledError:
+                cancelled.set()
+                raise
+            yield "B"  # pragma: no cover - never reached
+
+        gen = _stream_with_keepalive(source(), 0.02)
+        assert await gen.__anext__() == "A"
+        nxt = await gen.__anext__()  # pending sleep in flight -> keepalive
+        assert isinstance(nxt, _Keepalive)
+        await gen.aclose()  # must cancel the pending __anext__, not hang
+        await asyncio.sleep(0.01)
+        assert cancelled.is_set()
+
+    async def test_source_exception_propagates(self):
+        """An exception raised by `source` surfaces to the consumer, not just a hang.
+
+        `_pump_to_queue` catches it and hands it back through the queue as a
+        `_StreamError` sentinel; `_stream_with_keepalive` must unwrap and re-raise it.
+        """
+        from luthien_proxy.pipeline.anthropic_processor import _stream_with_keepalive
+
+        async def source():
+            yield "A"
+            raise RuntimeError("upstream exploded")
+
+        gen = _stream_with_keepalive(source(), 0.02)
+        assert await gen.__anext__() == "A"
+        with pytest.raises(RuntimeError, match="upstream exploded"):
+            await gen.__anext__()
+
+    async def test_span_spanning_multiple_yields_detaches_cleanly(
+        self,
+        span_exporter: InMemorySpanExporter,
+        caplog: pytest.LogCaptureFixture,
+    ):
+        """Regression for the production `Failed to detach context` ERROR storm.
+
+        `AnthropicClient.stream()` and `_AnthropicPolicyIO._stream()` each hold a
+        span open across every chunk of the upstream response — exactly this
+        shape, reproduced directly against `_stream_with_keepalive` instead of the
+        full pipeline. A keepalive gap forces at least one wait-for-timeout cycle
+        mid-span, which used to move the underlying generator's `__anext__` onto a
+        fresh `asyncio.Task` (and therefore a fresh `contextvars.Context`) on every
+        item, so the span's `__exit__` detached a token created in a different
+        Context than the one it ran in. That raised `ValueError` inside
+        `opentelemetry.context.detach()`, which logs it rather than propagating it
+        — so the only observable symptom is the ERROR log line asserted against
+        below.
+        """
+        from luthien_proxy.pipeline import anthropic_processor as mod
+        from luthien_proxy.pipeline.anthropic_processor import _Keepalive, _stream_with_keepalive
+
+        async def source():
+            with mod.tracer.start_as_current_span("fake_upstream"):
+                yield "A"
+                await asyncio.sleep(0.12)  # > interval -> forces a keepalive mid-span
+                yield "B"
+
+        with caplog.at_level(logging.ERROR, logger="opentelemetry.context"):
+            out = [item async for item in _stream_with_keepalive(source(), 0.02)]
+
+        reals = [x for x in out if not isinstance(x, _Keepalive)]
+        assert reals == ["A", "B"]
+
+        failed_detaches = [r for r in caplog.records if "Failed to detach context" in r.message]
+        assert failed_detaches == []
+
+        finished = span_exporter.get_finished_spans()
+        assert len(finished) == 1
+        assert finished[0].name == "fake_upstream"
+        assert finished[0].status.status_code == StatusCode.UNSET