From d0c1761da0993d8a57bb854fb4ddce6a8736e66b Mon Sep 17 00:00:00 2001 From: JunyiXu-nv <219237550+JunyiXu-nv@users.noreply.github.com> Date: Mon, 26 Oct 2026 15:29:10 +0000 Subject: [PATCH 1/5] [None][feat] Anthropic Messages: keepalive pings while a stream is silent The Anthropic streaming format allows any number of ping events, and clients and proxies reset their timers on any event. The server sent none, so a stream stayed silent until the first token (queueing and prefill) and for as long as the model wrote a tool call, which the tool parser releases only once it is complete. A long Write or Edit kept a stream silent for over 30 s, and first tokens took up to 15 s. Until [DONE], a ping now goes out whenever no other event has for STREAM_PING_INTERVAL_S (10 s), before message_start too: the message starts at the first upstream chunk, which carries the prompt's usage. Clients skip pings wherever they appear. One task reads the upstream chat stream to its end and hands its lines over through a queue, so a ping never interrupts a read and the chat stream is never resumed from another task. The reframer closes its line source when it ends, which cancels that task. Signed-off-by: JunyiXu-nv <219237550+JunyiXu-nv@users.noreply.github.com> --- tensorrt_llm/serve/anthropic_adapter.py | 76 ++++++- tensorrt_llm/serve/anthropic_protocol.py | 6 + .../llmapi/apps/test_anthropic_adapter.py | 189 +++++++++++++++++- 3 files changed, 266 insertions(+), 5 deletions(-) diff --git a/tensorrt_llm/serve/anthropic_adapter.py b/tensorrt_llm/serve/anthropic_adapter.py index 489f4d1e0a32..511464b62186 100644 --- a/tensorrt_llm/serve/anthropic_adapter.py +++ b/tensorrt_llm/serve/anthropic_adapter.py @@ -24,12 +24,13 @@ chat templates, or the engine. """ +import asyncio import json import re import traceback import uuid from dataclasses import dataclass, field -from typing import Any, AsyncIterator, Dict, List, Optional, Union +from typing import Any, AsyncIterator, Callable, Dict, List, Optional, Union from fastapi.responses import JSONResponse @@ -50,6 +51,7 @@ AnthropicMessagesResponse, AnthropicMessageStartEvent, AnthropicMessageStopEvent, + AnthropicPingEvent, AnthropicStopReason, AnthropicTextBlock, AnthropicTextDelta, @@ -1073,17 +1075,81 @@ async def _iter_openai_sse_lines( yield buffer.rstrip("\r") +# Seconds a stream may stay silent before it sends a ping. Clients and proxies +# reset their timers on any event; without one the stream is silent until the +# first token (queueing and prefill), and for as long as the model writes a tool +# call, which is sent only once complete. +STREAM_PING_INTERVAL_S = 10.0 + + +async def _with_ticks( + source: AsyncIterator[str], timeout: Callable[[], Optional[float]] +) -> AsyncIterator[Optional[str]]: + """Yield the items of ``source``, and None each time ``timeout()`` seconds pass first. + + ``timeout()`` is asked before every wait; None waits for the next item. One + task reads ``source`` to its end, so a tick never interrupts a read and the + source is never resumed from another task; that task is cancelled when the + consumer stops. + """ + items: asyncio.Queue = asyncio.Queue() + end = object() + + async def read_all() -> None: + try: + async for item in source: + items.put_nowait(item) + finally: + items.put_nowait(end) + + reader = asyncio.create_task(read_all()) + try: + while True: + try: + item = await asyncio.wait_for(items.get(), timeout()) + except asyncio.TimeoutError: + yield None + continue + if item is end: + # Raises what reading the source raised. + await reader + return + yield item + finally: + reader.cancel() + + async def reframe_openai_stream( openai_sse: AsyncIterator[Union[str, bytes]], model: str ) -> AsyncIterator[str]: - """Translate an OpenAI SSE stream into Anthropic SSE frames.""" + """Translate an OpenAI SSE stream into Anthropic SSE frames. + + Until [DONE], a ping goes out whenever no other event has for + ``STREAM_PING_INTERVAL_S`` seconds, before ``message_start`` too: the message + starts at the first upstream chunk, which carries the prompt's usage, so a + slow first token would otherwise leave the stream silent. Clients skip + pings wherever they appear. + """ 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] = [] + loop = asyncio.get_running_loop() + last_event = loop.time() + + def until_ping() -> Optional[float]: + if done: + return None + return max(0.0, STREAM_PING_INTERVAL_S - (loop.time() - last_event)) + + lines = _with_ticks(_iter_openai_sse_lines(openai_sse), until_ping) try: - async for line in _iter_openai_sse_lines(openai_sse): + async for line in lines: + if line is None: + yield anthropic_sse(AnthropicPingEvent()) + last_event = loop.time() + continue 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 @@ -1108,6 +1174,7 @@ async def reframe_openai_stream( raise AnthropicResponseError("Malformed upstream stream chunk") from e for frame in reframer.process_chunk(chunk): yield frame + last_event = loop.time() if done: for frame in closing_frames: yield frame @@ -1137,3 +1204,6 @@ async def reframe_openai_stream( return for frame in reframer.error("Internal server error"): yield frame + finally: + # Stops the task reading the upstream when this stream ends early. + await lines.aclose() diff --git a/tensorrt_llm/serve/anthropic_protocol.py b/tensorrt_llm/serve/anthropic_protocol.py index 0e5f6a0fc638..f257b17030af 100644 --- a/tensorrt_llm/serve/anthropic_protocol.py +++ b/tensorrt_llm/serve/anthropic_protocol.py @@ -491,6 +491,12 @@ class AnthropicMessageStopEvent(AnthropicBaseModel): type: Literal["message_stop"] = "message_stop" +class AnthropicPingEvent(AnthropicBaseModel): + """A keepalive: the stream format allows any number of them after message_start.""" + + type: Literal["ping"] = "ping" + + class AnthropicErrorEvent(AnthropicBaseModel): type: Literal["error"] = "error" error: AnthropicError diff --git a/tests/unittest/llmapi/apps/test_anthropic_adapter.py b/tests/unittest/llmapi/apps/test_anthropic_adapter.py index 82502a55c184..c742f4591ffd 100644 --- a/tests/unittest/llmapi/apps/test_anthropic_adapter.py +++ b/tests/unittest/llmapi/apps/test_anthropic_adapter.py @@ -24,6 +24,7 @@ import pytest from pydantic import ValidationError +import tensorrt_llm.serve.anthropic_adapter as adapter from tensorrt_llm.serve.anthropic_adapter import ( AnthropicRequestError, AnthropicResponseError, @@ -839,8 +840,11 @@ def run_reframer(chunks): def assert_event_invariants(events): - """First-principles invariants of the Anthropic streaming protocol.""" - names = [name for name, _ in events] + """First-principles invariants of the Anthropic streaming protocol. + + Pings may appear anywhere; the invariants hold for the other events. + """ + names = [name for name, _ in events if name != "ping"] assert names[0] == "message_start" assert names.count("message_start") == 1 assert names[-1] == "message_stop" @@ -1704,3 +1708,184 @@ def test_cached_tokens_exceeding_prompt_tokens_is_logged(monkeypatch): assert converted.input_tokens == 0 assert converted.cache_read_input_tokens == 99 assert any("exceeds prompt_tokens" in str(w) for w in warnings), warnings + + +# --------------------------------------------------------------------------- +# Keepalive pings while a stream is silent +# --------------------------------------------------------------------------- + + +def _line(c: ChatCompletionStreamResponse) -> str: + return f"data: {c.model_dump_json()}\n\n" + + +def _reframe_paced(monkeypatch, steps, interval): + """Reframe a source that waits ``pause`` seconds before each of its lines.""" + monkeypatch.setattr(adapter, "STREAM_PING_INTERVAL_S", interval) + + async def source(): + for pause, line in steps: + await asyncio.sleep(pause) + yield line + + async def collect(): + return [frame async for frame in reframe_openai_stream(source(), MODEL)] + + return parse_frames(asyncio.run(collect())) + + +def test_stream_pings_while_silent_after_message_start(monkeypatch): + """A silence longer than the interval is broken by pings, inside the message.""" + events = _reframe_paced( + monkeypatch, + [ + (0, _line(chunk({"content": "Writing the file."}))), + (0.3, _line(chunk({}, finish_reason="stop"))), + (0, "data: [DONE]\n\n"), + ], + interval=0.05, + ) + + assert_event_invariants(events) + names = [name for name, _ in events] + pings = [i for i, name in enumerate(names) if name == "ping"] + assert len(pings) >= 2 + assert names.index("message_start") < pings[0] < pings[-1] < names.index("message_delta") + assert all(events[i][1] == {"type": "ping"} for i in pings) + + +def test_stream_pings_while_waiting_for_the_first_token(monkeypatch): + """A slow first token is covered too: pings come first, then the message starts.""" + events = _reframe_paced( + monkeypatch, + [ + (0.3, _line(chunk({"content": "hello"}))), + (0, _line(chunk({}, finish_reason="stop"))), + (0, "data: [DONE]\n\n"), + ], + interval=0.05, + ) + + assert_event_invariants(events) + assert events[0] == ("ping", {"type": "ping"}) + + +def test_stream_pings_then_fails_before_the_first_token(monkeypatch): + """An upstream that fails before its first chunk still ends the stream with an error.""" + monkeypatch.setattr(adapter, "STREAM_PING_INTERVAL_S", 0.02) + + async def source(): + await asyncio.sleep(0.1) + raise RuntimeError("engine failed") + yield # an async generator + + async def collect(): + return [frame async for frame in reframe_openai_stream(source(), MODEL)] + + names = [name for name, _ in parse_frames(asyncio.run(collect()))] + assert names[0] == "ping" + assert names[-1] == "error" + assert "message_start" not in names + + +def test_stream_pings_when_chunks_bring_no_event(monkeypatch): + """Upstream chunks that produce no event (an empty delta) do not count as activity.""" + steps = [(0, _line(chunk({"content": "Thinking about it."})))] + steps += [(0.02, _line(chunk({}))) for _ in range(15)] + steps += [(0, _line(chunk({}, finish_reason="stop"))), (0, "data: [DONE]\n\n")] + + events = _reframe_paced(monkeypatch, steps, interval=0.05) + + assert_event_invariants(events) + assert "ping" in [name for name, _ in events] + + +def test_stream_sends_no_ping_while_events_flow(monkeypatch): + """Every event restarts the wait: a stream five intervals long but never silent has no ping.""" + steps = [(0.02, _line(chunk({"content": f"word{i} "}))) for i in range(50)] + steps += [(0, _line(chunk({}, finish_reason="stop"))), (0, "data: [DONE]\n\n")] + + events = _reframe_paced(monkeypatch, steps, interval=0.2) + + assert_event_invariants(events) + assert "ping" not in [name for name, _ in events] + + +def test_stopping_the_stream_cancels_its_pending_read(monkeypatch): + """A client that goes away mid-silence stops the upstream read at once.""" + monkeypatch.setattr(adapter, "STREAM_PING_INTERVAL_S", 0.05) + cancelled = [] + + async def source(): + yield _line(chunk({"content": "hello"})) + try: + await asyncio.sleep(30) + except asyncio.CancelledError: + cancelled.append(True) + raise + yield _line(chunk({}, finish_reason="stop")) + + async def stop_at_first_ping(): + stream = reframe_openai_stream(source(), MODEL) + async for frame in stream: + if frame.startswith("event: ping"): + break + await stream.aclose() + # Checked while the loop still runs: asyncio.run cancels every task left at its exit. + for _ in range(100): + if cancelled: + break + await asyncio.sleep(0.01) + return list(cancelled) + + assert asyncio.run(asyncio.wait_for(stop_at_first_ping(), timeout=5)) == [True] + + +def test_stream_reads_its_upstream_from_one_task(monkeypatch): + """Pings never move the upstream between tasks. + + The upstream is the chat stream, and an async generator that holds a + task-bound context (a lock, a timeout or a cancel scope) across a yield + breaks if it is resumed from another task. + """ + monkeypatch.setattr(adapter, "STREAM_PING_INTERVAL_S", 0.02) + readers = set() + + async def source(): + for line in [ + _line(chunk({"content": "one"})), + _line(chunk({"content": "two"})), + _line(chunk({}, finish_reason="stop")), + "data: [DONE]\n\n", + ]: + readers.add(asyncio.current_task()) + await asyncio.sleep(0.1) + yield line + + 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 "ping" in [name for name, _ in events] + assert len(readers) == 1 + + +def test_stream_ends_with_an_error_when_the_upstream_fails_mid_stream(monkeypatch): + """A failure while the stream pings still ends the message with an error event.""" + monkeypatch.setattr(adapter, "STREAM_PING_INTERVAL_S", 0.02) + + async def source(): + yield _line(chunk({"content": "hello"})) + await asyncio.sleep(0.1) + raise RuntimeError("engine failed") + + async def collect(): + return [frame async for frame in reframe_openai_stream(source(), MODEL)] + + events = parse_frames(asyncio.run(collect())) + names = [name for name, _ in events] + assert "ping" in names + assert names[-1] == "error" + # Reported as the upstream's own failure, not as a stream that ended early. + assert events[-1][1]["error"]["message"] == "Internal server error" From 8cd8932ec5576c93850bf38b114db11ebb410866 Mon Sep 17 00:00:00 2001 From: JunyiXu-nv <219237550+JunyiXu-nv@users.noreply.github.com> Date: Mon, 26 Oct 2026 17:13:59 +0000 Subject: [PATCH 2/5] [None][fix] Anthropic Messages: keep a slow client's backpressure on the pinged stream The task that reads the upstream for the keepalive pings now hands over one item at a time, so the upstream is read only as fast as the client takes frames, as before the pings. Tests: a slow client against a fast upstream, which still sees the end or the error, and no ping after [DONE]. Signed-off-by: JunyiXu-nv <219237550+JunyiXu-nv@users.noreply.github.com> --- tensorrt_llm/serve/anthropic_adapter.py | 15 ++++-- .../llmapi/apps/test_anthropic_adapter.py | 54 +++++++++++++++++++ 2 files changed, 65 insertions(+), 4 deletions(-) diff --git a/tensorrt_llm/serve/anthropic_adapter.py b/tensorrt_llm/serve/anthropic_adapter.py index 511464b62186..019db492601d 100644 --- a/tensorrt_llm/serve/anthropic_adapter.py +++ b/tensorrt_llm/serve/anthropic_adapter.py @@ -1092,15 +1092,22 @@ async def _with_ticks( source is never resumed from another task; that task is cancelled when the consumer stops. """ - items: asyncio.Queue = asyncio.Queue() + # One item at a time: the reader waits for the consumer, so a slow client + # holds back the source as reading it directly would. + items: asyncio.Queue = asyncio.Queue(maxsize=1) end = object() async def read_all() -> None: try: async for item in source: - items.put_nowait(item) - finally: - items.put_nowait(end) + await items.put(item) + except asyncio.CancelledError: + # The consumer has stopped; nobody is left to take the end marker. + raise + except Exception: + await items.put(end) + raise + await items.put(end) reader = asyncio.create_task(read_all()) try: diff --git a/tests/unittest/llmapi/apps/test_anthropic_adapter.py b/tests/unittest/llmapi/apps/test_anthropic_adapter.py index c742f4591ffd..4f1bc940c5fd 100644 --- a/tests/unittest/llmapi/apps/test_anthropic_adapter.py +++ b/tests/unittest/llmapi/apps/test_anthropic_adapter.py @@ -1811,6 +1811,60 @@ def test_stream_sends_no_ping_while_events_flow(monkeypatch): assert "ping" not in [name for name, _ in events] +def test_stream_sends_no_ping_after_done(monkeypatch): + """After [DONE] the message is complete: a slow end of the upstream brings no ping.""" + events = _reframe_paced( + monkeypatch, + [ + (0, _line(chunk({"content": "hello"}))), + (0, _line(chunk({}, finish_reason="stop"))), + (0, "data: [DONE]\n\n"), + (0.3, "\n"), + ], + interval=0.05, + ) + + assert_event_invariants(events) + assert "ping" not in [name for name, _ in events] + + +@pytest.mark.parametrize("fails", [False, True]) +def test_a_slow_client_holds_back_the_upstream(monkeypatch, fails): + """The upstream is read only as fast as the client takes frames; its end or error still arrives.""" + monkeypatch.setattr(adapter, "STREAM_PING_INTERVAL_S", 10) + words = [f"w{i} " for i in range(20)] + produced = [] + + async def source(): + for word in words: + produced.append(word) + yield _line(chunk({"content": word})) + if fails: + raise RuntimeError("engine failed") + yield _line(chunk({}, finish_reason="stop")) + yield "data: [DONE]\n\n" + + async def read_slowly(): + frames, lead = [], 0 + async for frame in reframe_openai_stream(source(), MODEL): + frames.append(frame) + taken = sum( + 1 + for name, payload in parse_frames(frames) + if name == "content_block_delta" and payload["delta"].get("text") in words + ) + lead = max(lead, len(produced) - taken) + await asyncio.sleep(0.005) + return parse_frames(frames), lead + + events, lead = asyncio.run(read_slowly()) + + assert lead <= 3, f"the upstream ran {lead} chunks ahead of the client" + texts = [p["delta"]["text"] for n, p in events if n == "content_block_delta"] + assert texts == words + assert events[-1][0] == ("error" if fails else "message_stop") + + def test_stopping_the_stream_cancels_its_pending_read(monkeypatch): """A client that goes away mid-silence stops the upstream read at once.""" monkeypatch.setattr(adapter, "STREAM_PING_INTERVAL_S", 0.05) From e4e149eb56d43d9bdfd920e8e9e133ad1dcc2967 Mon Sep 17 00:00:00 2001 From: JunyiXu-nv <219237550+JunyiXu-nv@users.noreply.github.com> Date: Mon, 26 Oct 2026 20:36:50 +0000 Subject: [PATCH 3/5] [None][fix] Anthropic Messages: close the upstream when a pinged stream's client leaves The reader of the pinged stream closes its source when it is cancelled, the SSE line reader closes the upstream when it is left early, and the stream waits for the reader to stop. A client that leaves while the next frame waits to be handed over thus closes the upstream at once, as it did before the pings. Signed-off-by: JunyiXu-nv <219237550+JunyiXu-nv@users.noreply.github.com> --- tensorrt_llm/serve/anthropic_adapter.py | 47 ++++++++++++------- .../llmapi/apps/test_anthropic_adapter.py | 25 ++++++++++ 2 files changed, 55 insertions(+), 17 deletions(-) diff --git a/tensorrt_llm/serve/anthropic_adapter.py b/tensorrt_llm/serve/anthropic_adapter.py index 019db492601d..dc92fbeda531 100644 --- a/tensorrt_llm/serve/anthropic_adapter.py +++ b/tensorrt_llm/serve/anthropic_adapter.py @@ -1050,24 +1050,30 @@ async def _iter_openai_sse_lines( buffer = "" input_type: Optional[type] = None - async for payload in openai_sse: - if not isinstance(payload, (str, bytes)): - raise TypeError(f"OpenAI SSE payload must be str or bytes, got {type(payload)!r}") - - current_type = type(payload) - if input_type is None: - input_type = current_type - elif current_type is not input_type: - raise TypeError("OpenAI SSE stream cannot mix str and bytes payloads") - - if isinstance(payload, bytes): - buffer += decoder.decode(payload) - else: - buffer += payload + try: + async for payload in openai_sse: + if not isinstance(payload, (str, bytes)): + raise TypeError(f"OpenAI SSE payload must be str or bytes, got {type(payload)!r}") + + current_type = type(payload) + if input_type is None: + input_type = current_type + elif current_type is not input_type: + raise TypeError("OpenAI SSE stream cannot mix str and bytes payloads") + + if isinstance(payload, bytes): + buffer += decoder.decode(payload) + else: + buffer += payload - while "\n" in buffer: - line, buffer = buffer.split("\n", 1) - yield line.rstrip("\r") + while "\n" in buffer: + line, buffer = buffer.split("\n", 1) + yield line.rstrip("\r") + finally: + # Leaving these lines early leaves the upstream too. + aclose = getattr(openai_sse, "aclose", None) + if aclose is not None: + await aclose() if input_type is bytes: buffer += decoder.decode(b"", final=True) @@ -1103,6 +1109,11 @@ async def read_all() -> None: await items.put(item) except asyncio.CancelledError: # The consumer has stopped; nobody is left to take the end marker. + # Stopped while handing over an item, the reader leaves the source + # suspended at its yield, so it closes the source itself. + aclose = getattr(source, "aclose", None) + if aclose is not None: + await aclose() raise except Exception: await items.put(end) @@ -1124,6 +1135,8 @@ async def read_all() -> None: yield item finally: reader.cancel() + # Return only once the reader has stopped and closed the source. + await asyncio.gather(reader, return_exceptions=True) async def reframe_openai_stream( diff --git a/tests/unittest/llmapi/apps/test_anthropic_adapter.py b/tests/unittest/llmapi/apps/test_anthropic_adapter.py index 4f1bc940c5fd..9433155750c7 100644 --- a/tests/unittest/llmapi/apps/test_anthropic_adapter.py +++ b/tests/unittest/llmapi/apps/test_anthropic_adapter.py @@ -1865,6 +1865,31 @@ async def read_slowly(): assert events[-1][0] == ("error" if fails else "message_stop") +def test_a_client_that_leaves_while_a_frame_waits_closes_the_upstream(monkeypatch): + """The upstream is closed by the time the stream is, also while its next frame waits for the client.""" + monkeypatch.setattr(adapter, "STREAM_PING_INTERVAL_S", 10) + closed = [] + + async def source(): + try: + for i in range(20): + yield _line(chunk({"content": f"w{i} "})) + finally: + closed.append(True) + + async def leave_after_one_frame(): + stream = reframe_openai_stream(source(), MODEL) + async for _ in stream: + break + # The reader fills the queue and then waits to hand over the next frame. + for _ in range(5): + await asyncio.sleep(0) + await stream.aclose() + return list(closed) + + assert asyncio.run(leave_after_one_frame()) == [True] + + def test_stopping_the_stream_cancels_its_pending_read(monkeypatch): """A client that goes away mid-silence stops the upstream read at once.""" monkeypatch.setattr(adapter, "STREAM_PING_INTERVAL_S", 0.05) From e1966f6cd38cb43b97ddf777f8e6494d2016ef0b Mon Sep 17 00:00:00 2001 From: JunyiXu-nv <219237550+JunyiXu-nv@users.noreply.github.com> Date: Tue, 27 Oct 2026 04:02:06 +0000 Subject: [PATCH 4/5] [None][fix] Anthropic Messages: tick the pinged stream per payload The keepalive ticks now wrap the upstream's payloads rather than the lines split from them, so a payload costs one hand-off between the reading task and the stream instead of one per line. The reader closes its source in a single finally. The no-ping test uses a 1 s interval against 20 ms gaps, so a loaded worker does not fake a silence. Signed-off-by: JunyiXu-nv <219237550+JunyiXu-nv@users.noreply.github.com> --- tensorrt_llm/serve/anthropic_adapter.py | 31 ++++++++++++------- .../llmapi/apps/test_anthropic_adapter.py | 9 ++++-- 2 files changed, 26 insertions(+), 14 deletions(-) diff --git a/tensorrt_llm/serve/anthropic_adapter.py b/tensorrt_llm/serve/anthropic_adapter.py index bd6ae4de5c77..672e2d0c7285 100644 --- a/tensorrt_llm/serve/anthropic_adapter.py +++ b/tensorrt_llm/serve/anthropic_adapter.py @@ -1075,10 +1075,12 @@ def error(self, message: str) -> List[str]: async def _iter_openai_sse_lines( - openai_sse: AsyncIterator[Union[str, bytes]], -) -> AsyncIterator[str]: + openai_sse: AsyncIterator[Optional[Union[str, bytes]]], +) -> AsyncIterator[Optional[str]]: """Yield complete lines from string frames or arbitrarily chunked bytes. + A None item, a tick from ``_with_ticks``, is passed through as None. + The in-process server produces complete SSE strings, while the disaggregated frontend relays raw ``aiohttp`` byte chunks. Network chunks need not align with either UTF-8 characters or SSE line boundaries, so @@ -1092,6 +1094,9 @@ async def _iter_openai_sse_lines( try: async for payload in openai_sse: + if payload is None: + yield None + continue if not isinstance(payload, (str, bytes)): raise TypeError(f"OpenAI SSE payload must be str or bytes, got {type(payload)!r}") @@ -1129,8 +1134,8 @@ async def _iter_openai_sse_lines( async def _with_ticks( - source: AsyncIterator[str], timeout: Callable[[], Optional[float]] -) -> AsyncIterator[Optional[str]]: + source: AsyncIterator[Any], timeout: Callable[[], Optional[float]] +) -> AsyncIterator[Any]: """Yield the items of ``source``, and None each time ``timeout()`` seconds pass first. ``timeout()`` is asked before every wait; None waits for the next item. One @@ -1149,16 +1154,18 @@ async def read_all() -> None: await items.put(item) except asyncio.CancelledError: # The consumer has stopped; nobody is left to take the end marker. - # Stopped while handing over an item, the reader leaves the source - # suspended at its yield, so it closes the source itself. - aclose = getattr(source, "aclose", None) - if aclose is not None: - await aclose() raise except Exception: await items.put(end) raise - await items.put(end) + else: + await items.put(end) + finally: + # The reader closes the source in every case: stopped while handing + # over an item, it leaves the source suspended at its yield. + aclose = getattr(source, "aclose", None) + if aclose is not None: + await aclose() reader = asyncio.create_task(read_all()) try: @@ -1203,7 +1210,9 @@ def until_ping() -> Optional[float]: return None return max(0.0, STREAM_PING_INTERVAL_S - (loop.time() - last_event)) - lines = _with_ticks(_iter_openai_sse_lines(openai_sse), until_ping) + # The ticks wrap the upstream's payloads, not the lines split from them: + # one hand-off between the reading task and this one per payload. + lines = _iter_openai_sse_lines(_with_ticks(openai_sse, until_ping)) try: async for line in lines: if line is None: diff --git a/tests/unittest/llmapi/apps/test_anthropic_adapter.py b/tests/unittest/llmapi/apps/test_anthropic_adapter.py index 4a85da2fbaae..15f711808a59 100644 --- a/tests/unittest/llmapi/apps/test_anthropic_adapter.py +++ b/tests/unittest/llmapi/apps/test_anthropic_adapter.py @@ -1937,11 +1937,14 @@ def test_stream_pings_when_chunks_bring_no_event(monkeypatch): def test_stream_sends_no_ping_while_events_flow(monkeypatch): - """Every event restarts the wait: a stream five intervals long but never silent has no ping.""" - steps = [(0.02, _line(chunk({"content": f"word{i} "}))) for i in range(50)] + """Every event restarts the wait: a stream longer than the interval but never silent has no ping. + + The 1 s interval is 50 times the 20 ms gap, so a loaded worker does not fake a silence. + """ + steps = [(0.02, _line(chunk({"content": f"word{i} "}))) for i in range(60)] steps += [(0, _line(chunk({}, finish_reason="stop"))), (0, "data: [DONE]\n\n")] - events = _reframe_paced(monkeypatch, steps, interval=0.2) + events = _reframe_paced(monkeypatch, steps, interval=1.0) assert_event_invariants(events) assert "ping" not in [name for name, _ in events] From d8bc3a678458a841ac5a2db48a58790533647005 Mon Sep 17 00:00:00 2001 From: JunyiXu-nv <219237550+JunyiXu-nv@users.noreply.github.com> Date: Tue, 27 Oct 2026 10:28:10 +0000 Subject: [PATCH 5/5] [None][fix] Anthropic Messages: end a pinged stream whose upstream is cancelled on its own The reader task re-raised a CancelledError without handing over the end marker, assuming the consumer had stopped it. A CancelledError raised inside the upstream itself left the consumer waiting, so the stream sent pings forever. The reader now hands over the end marker unless the consumer stopped it, and the stream ends with that cancellation. Signed-off-by: JunyiXu-nv <219237550+JunyiXu-nv@users.noreply.github.com> --- tensorrt_llm/serve/anthropic_adapter.py | 8 ++++++- .../llmapi/apps/test_anthropic_adapter.py | 21 +++++++++++++++++++ 2 files changed, 28 insertions(+), 1 deletion(-) diff --git a/tensorrt_llm/serve/anthropic_adapter.py b/tensorrt_llm/serve/anthropic_adapter.py index 672e2d0c7285..85785555dea0 100644 --- a/tensorrt_llm/serve/anthropic_adapter.py +++ b/tensorrt_llm/serve/anthropic_adapter.py @@ -1147,13 +1147,18 @@ async def _with_ticks( # holds back the source as reading it directly would. items: asyncio.Queue = asyncio.Queue(maxsize=1) end = object() + consumer_stopped = False async def read_all() -> None: try: async for item in source: await items.put(item) except asyncio.CancelledError: - # The consumer has stopped; nobody is left to take the end marker. + if not consumer_stopped: + # Raised inside the source, not by the consumer stopping: the + # consumer still waits, and the end marker ends its stream + # instead of leaving it to pings. + await items.put(end) raise except Exception: await items.put(end) @@ -1181,6 +1186,7 @@ async def read_all() -> None: return yield item finally: + consumer_stopped = True reader.cancel() # Return only once the reader has stopped and closed the source. await asyncio.gather(reader, return_exceptions=True) diff --git a/tests/unittest/llmapi/apps/test_anthropic_adapter.py b/tests/unittest/llmapi/apps/test_anthropic_adapter.py index 15f711808a59..9449c721e1c1 100644 --- a/tests/unittest/llmapi/apps/test_anthropic_adapter.py +++ b/tests/unittest/llmapi/apps/test_anthropic_adapter.py @@ -2107,3 +2107,24 @@ async def collect(): assert names[-1] == "error" # Reported as the upstream's own failure, not as a stream that ended early. assert events[-1][1]["error"]["message"] == "Internal server error" + + +def test_stream_ends_when_the_upstream_is_cancelled_on_its_own(monkeypatch): + """A cancellation raised inside the upstream ends the stream instead of leaving it to pings.""" + monkeypatch.setattr(adapter, "STREAM_PING_INTERVAL_S", 0.02) + + async def source(): + yield _line(chunk({"content": "hello"})) + await asyncio.sleep(0.1) + raise asyncio.CancelledError() + + async def collect(): + frames = [] + with pytest.raises(asyncio.CancelledError): + async for frame in reframe_openai_stream(source(), MODEL): + frames.append(frame) + assert len(frames) < 100, "the stream kept pinging after its upstream ended" + return frames + + events = parse_frames(asyncio.run(asyncio.wait_for(collect(), 5))) + assert "ping" in [name for name, _ in events]