[None][fix] Anthropic Messages: read the chat stream to its end so its metrics are recorded

tensorrt_llm/serve/anthropic_adapter.py

@@ -1062,24 +1062,40 @@ async def reframe_openai_stream(
 ) -> AsyncIterator[str]:
     """Translate an OpenAI SSE stream into Anthropic SSE frames."""
     reframer = AnthropicStreamReframer(model=model)
+    done = False
+    # The frames that close the message, held from [DONE] until the upstream
+    # has been read to its end.
+    closing_frames: List[str] = []
     try:
         async for line in _iter_openai_sse_lines(openai_sse):
+            if done:
+                # The upstream is read to its end before the message is closed:
+                # the chat stream records the request's metrics after its last
+                # frame, and a consumer that stops reading at message_stop would
+                # leave it suspended at [DONE], to be closed there (GeneratorExit)
+                # without doing that work. A relayed byte stream likewise ends
+                # cleanly at the worker's EOF.
+                continue
             line = line.strip()
             if not line.startswith("data:"):
                 continue
             data = line[len("data:") :].strip()
             if not data:
                 continue
             if data == "[DONE]":
-                for frame in reframer.finish():
-                    yield frame
-                return
+                closing_frames = reframer.finish()
+                done = True
+                continue
             try:
                 chunk = ChatCompletionStreamResponse(**json.loads(data))
             except (json.JSONDecodeError, ValueError) as e:
                 raise AnthropicResponseError("Malformed upstream stream chunk") from e
             for frame in reframer.process_chunk(chunk):
                 yield frame
+        if done:
+            for frame in closing_frames:
+                yield frame
+            return
         # Upstream ended without [DONE]. Both streaming producers emit it only
         # after the response is complete, so its absence means the stream was
         # cut short. Calling finish() here would close the message with
@@ -1097,5 +1113,11 @@ async def reframe_openai_stream(
             yield frame
     except Exception as e:  # noqa: BLE001 - stream must end with an event
         logger.error(f"Anthropic stream reframing failed: {e}\n{traceback.format_exc()}")
+        if done:
+            # The response was complete at [DONE], so a failure while reading
+            # the rest of the upstream is logged and the message still closes.
+            for frame in closing_frames:
+                yield frame
+            return
         for frame in reframer.error("Internal server error"):
             yield frame

tensorrt_llm/serve/openai_server.py

@@ -2061,7 +2061,10 @@ async def chat_stream_generator(
                 yield "data: [DONE]\n\n"
                 await self._extract_metrics(res, raw_request)
                 nvtx_mark("generation ends")
-            except:
+            except Exception:
+                # Only failures are logged: a consumer that stops reading closes
+                # the stream with GeneratorExit, and a client that disconnects
+                # cancels it, and neither is an error of this stream.
                 logger.error(traceback.format_exc())
                 raise
 

tests/unittest/llmapi/apps/test_anthropic_adapter.py

@@ -1132,6 +1132,122 @@ async def collect():
     assert "".join(text_deltas) == "hello"
 
 
+def test_stream_reframer_reads_the_upstream_to_its_end_after_done():
+    """The upstream runs past [DONE], so the work it does there runs.
+
+    The chat stream records the request's metrics after its last frame. A
+    reframer that stopped reading at [DONE] left it suspended there, to be
+    closed with GeneratorExit: every streamed /v1/messages request lost its
+    metrics, and the server logged a traceback for it.
+    """
+    after_done = []
+
+    async def source():
+        yield f"data: {chunk({'content': 'hello'}).model_dump_json()}\n\n"
+        yield f"data: {chunk({}, finish_reason='stop').model_dump_json()}\n\n"
+        yield "data: [DONE]\n\n"
+        after_done.append("ran")
+
+    async def collect():
+        return [frame async for frame in reframe_openai_stream(source(), MODEL)]
+
+    events = parse_frames(asyncio.run(collect()))
+    assert_event_invariants(events)
+    assert after_done == ["ran"]
+
+
+def test_stream_reframer_reads_the_upstream_to_its_end_before_message_stop():
+    """A consumer that stops reading at message_stop still lets the upstream end.
+
+    A client may disconnect as soon as it has message_stop. Had message_stop
+    gone out before the upstream was drained, the chat stream would be left
+    suspended at [DONE] and its metrics never recorded.
+    """
+    after_done = []
+
+    async def source():
+        yield f"data: {chunk({'content': 'hello'}).model_dump_json()}\n\n"
+        yield f"data: {chunk({}, finish_reason='stop').model_dump_json()}\n\n"
+        yield "data: [DONE]\n\n"
+        after_done.append("ran")
+
+    async def read_until_message_stop():
+        stream = reframe_openai_stream(source(), MODEL)
+        frames = []
+        async for frame in stream:
+            frames.append(frame)
+            if frame.startswith("event: message_stop"):
+                break
+        await stream.aclose()
+        return frames
+
+    events = parse_frames(asyncio.run(read_until_message_stop()))
+    assert_event_invariants(events)
+    assert after_done == ["ran"]
+
+
+def test_stream_reframer_drops_anything_after_done():
+    """Lines after [DONE] are read to reach the end, never reframed."""
+
+    async def source():
+        yield f"data: {chunk({'content': 'hello'}).model_dump_json()}\n\n"
+        yield f"data: {chunk({}, finish_reason='stop').model_dump_json()}\n\n"
+        yield "data: [DONE]\n\n"
+        yield f"data: {chunk({'content': 'late'}).model_dump_json()}\n\n"
+
+    async def collect():
+        return [frame async for frame in reframe_openai_stream(source(), MODEL)]
+
+    events = parse_frames(asyncio.run(collect()))
+    assert_event_invariants(events)
+    text_deltas = [
+        payload["delta"]["text"]
+        for name, payload in events
+        if name == "content_block_delta" and payload["delta"]["type"] == "text_delta"
+    ]
+    assert "".join(text_deltas) == "hello"
+
+
+def test_stream_reframer_adds_no_event_when_the_upstream_fails_after_done():
+    """message_stop stays the last event when reading past [DONE] fails."""
+
+    async def source():
+        yield f"data: {chunk({'content': 'hello'}).model_dump_json()}\n\n"
+        yield f"data: {chunk({}, finish_reason='stop').model_dump_json()}\n\n"
+        yield "data: [DONE]\n\n"
+        raise RuntimeError("failed after the last frame")
+
+    async def collect():
+        return [frame async for frame in reframe_openai_stream(source(), MODEL)]
+
+    events = parse_frames(asyncio.run(collect()))
+    assert_event_invariants(events)
+    assert "error" not in [name for name, _ in events]
+
+
+def test_stream_reframer_reads_a_byte_stream_to_its_end_when_done_is_split():
+    """A relayed byte stream is read to its EOF however [DONE] is chunked."""
+    body = (
+        f"data: {chunk({'content': 'hello'}).model_dump_json()}\n\n"
+        f"data: {chunk({}, finish_reason='stop').model_dump_json()}\n\n"
+        "data: [DONE]\n\n"
+    ).encode()
+    split = body.index(b"[DONE]") + len(b"[DO")
+    after_last_chunk = []
+
+    async def source():
+        yield body[:split]
+        yield body[split:]
+        after_last_chunk.append("ran")
+
+    async def collect():
+        return [frame async for frame in reframe_openai_stream(source(), MODEL)]
+
+    events = parse_frames(asyncio.run(collect()))
+    assert_event_invariants(events)
+    assert after_last_chunk == ["ran"]
+
+
 def test_stream_reframer_splits_multibyte_characters():
     """A chunk boundary landing inside a UTF-8 sequence must not corrupt output.
 

tests/unittest/llmapi/apps/test_anthropic_routes.py

@@ -27,8 +27,10 @@
 from fastapi.testclient import TestClient
 
 from tensorrt_llm.inputs.chat_template_guard import UnusedChatTemplateKwargsError
+from tensorrt_llm.serve import openai_server as openai_server_module
 from tensorrt_llm.serve.openai_disagg_server import OpenAIDisaggServer
 from tensorrt_llm.serve.openai_protocol import (
+    ChatCompletionRequest,
     ChatCompletionResponse,
     ChatCompletionResponseChoice,
     ChatCompletionStreamResponse,
@@ -38,6 +40,7 @@
     UsageInfo,
 )
 from tensorrt_llm.serve.openai_server import OpenAIServer
+from tensorrt_llm.serve.render import chat as render_chat_module
 
 # These tests are CPU-only (no GPU, engine or sockets) and run in the
 # CPU-Generic CI stage, which selects with `-m cpu_only`.
@@ -271,6 +274,168 @@ def test_messages_route_reframes_streaming_response(server_kind):
     assert backend.await_args.args[0].stream
 
 
+class _StreamedGeneration:
+    """A streamed generation whose steps are already post-processed frames.
+
+    After its steps it ends, raises ``then``, or (``hang=True``) waits forever
+    for the next step, as a generation does while the model is still running.
+    """
+
+    prompt_token_ids = [1, 2, 3]
+
+    def __init__(self, steps, *, then=None, hang=False):
+        self._steps = list(steps)
+        self._then = then
+        self._hang = hang
+
+    def __aiter__(self):
+        return self
+
+    async def __anext__(self):
+        if self._steps:
+            frames = self._steps.pop(0)
+            return SimpleNamespace(
+                finished=not self._steps,
+                outputs=[SimpleNamespace(_postprocess_result=frames)],
+            )
+        if self._then is not None:
+            raise self._then
+        if self._hang:
+            await asyncio.Event().wait()
+        raise StopAsyncIteration
+
+
+class _ModelConfig:
+    model_type = "llama"
+    vocab_size = 1024
+
+
+async def _chat_stream(monkeypatch, generation):
+    """The body of a streamed /v1/chat/completions response, and the server.
+
+    The real ``openai_chat`` handler runs on a server built with
+    ``object.__new__``; prompt rendering and the engine are stubbed, and the
+    engine's output is already post-processed into SSE frames.
+    """
+
+    async def render(**kwargs):
+        return "rendered prompt"
+
+    monkeypatch.setattr(render_chat_module, "async_apply_chat_template", render)
+    server = object.__new__(OpenAIServer)
+    server.model = MODEL
+    server.allow_request_chat_template = False
+    server.model_config = _ModelConfig()
+    server.processor = None
+    server.tokenizer = SimpleNamespace(tokenizer=SimpleNamespace(vocab_size=1024))
+    server.chat_template = None
+    server.tool_parser = None
+    server.tool_call_id_type = "random"
+    server.multimodal_server_config = None
+    server._is_visual_gen = False
+    server._adjusted_steady_clock = SimpleNamespace(now=time.monotonic)
+    server.generator = SimpleNamespace(
+        args=SimpleNamespace(
+            gather_generation_logits=False,
+            reasoning_parser=None,
+            backend="pytorch",
+            guided_decoding_backend=None,
+            num_postprocess_workers=1,
+        ),
+        generate_async=Mock(return_value=generation),
+    )
+    server.await_disconnected = AsyncMock()
+    server._extract_metrics = AsyncMock()
+    request = ChatCompletionRequest(
+        model=MODEL, messages=[{"role": "user", "content": "hi"}], stream=True
+    )
+    raw_request = SimpleNamespace(state=SimpleNamespace(), headers={})
+    response = await server.openai_chat(request, raw_request)
+    assert isinstance(response, StreamingResponse), response.body
+    return response.body_iterator, server
+
+
+@pytest.fixture
+def server_logger(monkeypatch):
+    log = Mock()
+    monkeypatch.setattr(openai_server_module, "logger", log)
+    return log
+
+
+def test_chat_stream_read_to_its_end_records_the_metrics(monkeypatch, server_logger):
+    """Past [DONE] the chat stream records the request's metrics."""
+
+    async def run():
+        generation = _StreamedGeneration([["data: one\n\n"], ["data: two\n\n"]])
+        stream, server = await _chat_stream(monkeypatch, generation)
+        frames = [frame async for frame in stream]
+        return frames, server
+
+    frames, server = asyncio.run(run())
+
+    assert frames == ["data: one\n\n", "data: two\n\n", "data: [DONE]\n\n"]
+    server._extract_metrics.assert_awaited_once()
+    server_logger.error.assert_not_called()
+
+
+def test_chat_stream_closed_by_its_consumer_logs_no_error(monkeypatch, server_logger):
+    """A consumer that stops reading (GeneratorExit) is not a stream error."""
+
+    async def run():
+        generation = _StreamedGeneration([["data: one\n\n"], ["data: two\n\n"]])
+        stream, server = await _chat_stream(monkeypatch, generation)
+        assert await stream.__anext__() == "data: one\n\n"
+        await stream.aclose()
+        return server
+
+    server = asyncio.run(run())
+
+    server._extract_metrics.assert_not_awaited()
+    server_logger.error.assert_not_called()
+
+
+def test_chat_stream_cancelled_by_a_disconnect_logs_no_error(monkeypatch, server_logger):
+    """A client that disconnects cancels the stream; that is not a stream error."""
+
+    async def run():
+        generation = _StreamedGeneration([["data: one\n\n"]], hang=True)
+        stream, _ = await _chat_stream(monkeypatch, generation)
+        first_frame = asyncio.Event()
+
+        async def consume():
+            async for _ in stream:
+                first_frame.set()
+
+        consumer = asyncio.create_task(consume())
+        # Bounded, so a stream that fails before its first frame fails the
+        # test instead of hanging it.
+        await asyncio.wait_for(first_frame.wait(), timeout=5.0)
+        # The stream is now waiting on the generation for its next step.
+        consumer.cancel()
+        with pytest.raises(asyncio.CancelledError):
+            await consumer
+
+    asyncio.run(run())
+
+    server_logger.error.assert_not_called()
+
+
+def test_chat_stream_logs_a_failure_and_raises_it(monkeypatch, server_logger):
+    """A failure of the generation itself is still logged and raised."""
+
+    async def run():
+        generation = _StreamedGeneration([["data: one\n\n"]], then=RuntimeError("engine failed"))
+        stream, _ = await _chat_stream(monkeypatch, generation)
+        with pytest.raises(RuntimeError, match="engine failed"):
+            async for _ in stream:
+                pass
+
+    asyncio.run(run())
+
+    server_logger.error.assert_called_once()
+    assert "engine failed" in server_logger.error.call_args.args[0]
+
+
 def test_standard_and_disagg_register_messages_route(monkeypatch, tmp_path):
     # The disagg server's register_routes mounts a prometheus multiprocess
     # collector, which errors unless PROMETHEUS_MULTIPROC_DIR points at a real