From e4a75f8b7897390e96c3fb059b7f78209fc5ac32 Mon Sep 17 00:00:00 2001 From: JunyiXu-nv <219237550+JunyiXu-nv@users.noreply.github.com> Date: Mon, 2 Nov 2026 16:15:53 +0000 Subject: [PATCH 1/5] [None][fix] Disaggregated Responses: open the stream and keep it alive while the context phase runs On the disaggregated server, a streamed /v1/responses request receives no bytes at all, not even the HTTP status line, until the context phase is done: the route awaits the pipeline before it creates the streaming response, and the pipeline returns its stream only once the context worker has answered, in both schedule orders. Queueing plus prefill of a long agent prompt can take minutes, and a client or gateway that drops a stream without a first event loses the request (Codex behind a gateway gives up after about 120 s); its retry starts the prefill over. An aggregated server sends response.created at once and repeats response.in_progress while the stream is silent. A streamed request now waits for its pipeline for at most _RESPONSES_STREAM_START_WAIT_S (1 s). A pipeline that answers within the wait is answered as before, its failures with their HTTP status. For one still running, the stream starts with response.created and response.in_progress built from the request, and keep_responses_stream_alive repeats response.in_progress whenever the stream is silent. The worker's stream follows without its own opening events, numbered to continue the sequence the client has. The request id is forwarded to the workers, so that they answer under the id the opening named. A pipeline that fails after the stream started ends it with error and response.failed, which carry the message of the HTTP error the request would otherwise have received. A client that disconnects first cancels the pipeline, and a cancelled context phase releases a generation server reserved for the request. Non-streamed requests and the other routes are unchanged. Signed-off-by: JunyiXu-nv <219237550+JunyiXu-nv@users.noreply.github.com> --- tensorrt_llm/serve/openai_disagg_server.py | 126 +++- tensorrt_llm/serve/openai_disagg_service.py | 2 +- tensorrt_llm/serve/responses_utils.py | 39 ++ .../test_disagg_responses_api.py | 630 +++++++++++++++++- .../test_openai_disagg_service.py | 37 + 5 files changed, 825 insertions(+), 9 deletions(-) diff --git a/tensorrt_llm/serve/openai_disagg_server.py b/tensorrt_llm/serve/openai_disagg_server.py index 92a298d4260f..def2f0477e93 100644 --- a/tensorrt_llm/serve/openai_disagg_server.py +++ b/tensorrt_llm/serve/openai_disagg_server.py @@ -20,6 +20,7 @@ import socket import traceback import uuid +from collections.abc import AsyncIterator from contextlib import asynccontextmanager from typing import Callable, Optional @@ -67,10 +68,11 @@ PerfMetricsJsonlWriter, PerfMetricsMiddleware, combine_disagg_metrics) -from tensorrt_llm.serve.responses_utils import (ServerArrivalTimeMiddleware, - get_steady_clock_now_in_seconds, - guard_responses_stream, - stream_error_event) +from tensorrt_llm.serve.responses_utils import ( + STREAM_TERMINATION_UPSTREAM_ERROR, ClosingStreamingResponse, + ResponsesStreamingProcessor, ServerArrivalTimeMiddleware, + get_steady_clock_now_in_seconds, guard_responses_stream, + keep_responses_stream_alive, relay_after_opening, stream_error_event) from tensorrt_llm.serve.router import Router from tensorrt_llm.usage import TerminalOutcome, record_termination_observation from tensorrt_llm.version import __version__ as VERSION @@ -81,6 +83,17 @@ for code in (*range(32), 127) } +# Seconds a streamed /v1/responses request waits for its pipeline before the +# response starts. The pipeline has a stream to return only once its context +# phase is done, and until the response starts the client receives nothing, +# not even the status line, which clients and proxies take for a dead +# connection after a while. A pipeline that answers within the wait is relayed +# or, if it failed, answered with its error status. For one still running, the +# stream starts at once with response.created and response.in_progress, and +# response.in_progress is sent again whenever the stream is silent, as an +# aggregated server does. +_RESPONSES_STREAM_START_WAIT_S = 1.0 + # Most round trips to spend estimating one ctx/gen server's steady-clock # offset. Only the least-delayed sample is kept; see `_sync_server_clock`. # Each probe costs ~0.2 s because the `/steady_clock_offset` handler sleeps @@ -184,6 +197,30 @@ def on_resp_done( )) +class _ResponsesStreamingResponse(ClosingStreamingResponse): + """Cancel the pipeline even when a disconnect precedes body iteration.""" + + def __init__(self, pipeline: asyncio.Future, + content: AsyncIterator[str]) -> None: + super().__init__(content=content, media_type="text/event-stream") + self._pipeline = pipeline + + async def __call__(self, scope, receive, send) -> None: + try: + await super().__call__(scope, receive, send) + finally: + self._pipeline.cancel() + + +def _error_response_message(response: Response) -> str: + """The message of an error response, for a stream to report it.""" + try: + payload = json.loads(response.body) + return str(payload.get("message") or payload.get("detail") or payload) + except (json.JSONDecodeError, AttributeError, TypeError): + return "Internal server error" + + def _upstream_error_message(error: aiohttp.ClientResponseError) -> str: """Recover the worker's own error text without leaking the worker. @@ -348,7 +385,7 @@ def register_routes(self): self.app.add_api_route("/v1/chat/completions", self._wrap_entry_point(self._service.openai_chat_completion, ChatCompletionRequest), methods=["POST"]) self.app.add_api_route("/v1/messages", self.anthropic_messages, methods=["POST"]) self.app.add_api_route("/v1/messages/count_tokens", self.anthropic_count_tokens, methods=["POST"]) - self.app.add_api_route("/v1/responses", self._wrap_entry_point(self._service.openai_responses, ResponsesRequest), methods=["POST"]) + self.app.add_api_route("/v1/responses", self.openai_responses, methods=["POST"]) self.app.add_api_route("/v1/models", self.get_model, methods=["GET"]) self.app.add_api_route("/health", self.health, methods=["GET"]) self.app.add_api_route("/cluster_info", self.cluster_info, methods=["GET"]) @@ -460,6 +497,85 @@ async def wrapper(req: request_type, raw_req: Request) -> Response: return self._handle_exception(e) return wrapper + async def openai_responses(self, request: ResponsesRequest, + raw_request: Request) -> Response: + """Serve Responses through the disaggregated pipeline. + + A streamed request whose pipeline is still running after + _RESPONSES_STREAM_START_WAIT_S is answered from here at once; see + _responses_stream_when_done. + """ + responses = self._wrap_entry_point(self._service.openai_responses) + if not request.stream: + return await responses(request, raw_request) + # The workers name their response after request_id, which is forwarded + # only if set. Assigning it marks it set, so that the workers answer + # under the id that an opening sent from here names. + request.request_id = request.request_id + pipeline = asyncio.ensure_future(responses(request, raw_request)) + try: + await asyncio.wait((pipeline, ), + timeout=_RESPONSES_STREAM_START_WAIT_S) + except asyncio.CancelledError: + # Cancelling the request cancels its pipeline. + pipeline.cancel() + raise + if pipeline.done(): + return pipeline.result() + try: + processor = ResponsesStreamingProcessor( + request=request, + sampling_params=request.to_sampling_params(), + model_name=request.model) + opening = processor.get_initial_responses() + except ValueError: + # The opening is built from the request alone; a request it cannot + # be built from is answered when its pipeline answers. + return await pipeline + # keep_responses_stream_alive repeats response.in_progress whenever the + # stream is silent, as it is while the pipeline runs. + return _ResponsesStreamingResponse( + pipeline=pipeline, + content=keep_responses_stream_alive( + self._responses_stream_when_done(processor, opening, + pipeline))) + + async def _responses_stream_when_done( + self, processor: ResponsesStreamingProcessor, opening: list[str], + pipeline: asyncio.Future) -> AsyncIterator[str]: + """The stream of a Responses request whose pipeline outlasted the start wait. + + ``opening`` is the request's response.created and response.in_progress, + and the pipeline's stream follows it without the worker's own opening + (relay_after_opening). The status line is already out, so a pipeline + that fails ends the stream with error and response.failed, which carry + the message of its error response. A stream abandoned before the + pipeline answers, as when the client disconnects, cancels the pipeline. + """ + try: + for frame in opening: + yield frame + try: + answer = await pipeline + except HTTPException as e: + failure = str(e.detail) + else: + failure = (None if isinstance(answer, StreamingResponse) else + _error_response_message(answer)) + finally: + pipeline.cancel() + if failure is not None: + for frame in processor.get_stream_failed_events( + STREAM_TERMINATION_UPSTREAM_ERROR, failure, len(opening)): + yield frame + return + frames = relay_after_opening(answer.body_iterator, len(opening)) + try: + async for frame in frames: + yield frame + finally: + await frames.aclose() + async def anthropic_messages(self, request: AnthropicMessagesRequest, raw_request: Request) -> Response: """Serve Anthropic Messages through the disaggregated chat pipeline.""" diff --git a/tensorrt_llm/serve/openai_disagg_service.py b/tensorrt_llm/serve/openai_disagg_service.py index c62d62b31f8e..2de1717f2bfa 100644 --- a/tensorrt_llm/serve/openai_disagg_service.py +++ b/tensorrt_llm/serve/openai_disagg_service.py @@ -333,7 +333,7 @@ async def _send_disagg_request_ctx_first( if hooks: hooks.on_disagg_request_id(disagg_request_id) gen_req = self._get_gen_request(request, ctx_response, disagg_request_id) - except Exception: + except (Exception, asyncio.CancelledError): if gen_server: await self._gen_router.finish_request( request, success=False, req_id=gen_reservation_id diff --git a/tensorrt_llm/serve/responses_utils.py b/tensorrt_llm/serve/responses_utils.py index f7d93a1471de..61e1673ab494 100644 --- a/tensorrt_llm/serve/responses_utils.py +++ b/tensorrt_llm/serve/responses_utils.py @@ -3475,6 +3475,45 @@ def until_keepalive() -> Optional[float]: await items.aclose() +_OPENING_FRAME_STARTS = ("event: response.created\n", _IN_PROGRESS_FRAME_START) + + +async def relay_after_opening(stream: AsyncGenerator[Union[str, bytes], None], + events_sent: int) -> AsyncGenerator[str, None]: + """Relay a worker's Responses stream to a client that already has its opening. + + The client holds ``events_sent`` events, ``response.created`` and + ``response.in_progress`` among them, so the worker's own opening events, + and any repeats of them before its first other event, are dropped. Every + other event is numbered to follow on, and is yielded whole, one per item + as ``keep_responses_stream_alive`` takes them, wherever ``stream`` splits + it. ``stream`` is closed when the relay ends. + """ + delimiter = _SSE_EVENT_DELIMITER.encode() + pending = b"" + opening = True + try: + async for chunk in stream: + if isinstance(chunk, str): + chunk = chunk.encode() + *events, pending = (pending + chunk).split(delimiter) + for event in events: + frame = event.decode(errors="replace") + _SSE_EVENT_DELIMITER + if not event or (opening + and frame.startswith(_OPENING_FRAME_STARTS)): + continue + opening = False + yield _stamp_sequence_number(frame, events_sent) + events_sent += 1 + if pending: + # The stream stopped inside an event. + yield pending.decode(errors="replace") + finally: + aclose = getattr(stream, "aclose", None) + if aclose is not None: + await aclose() + + class ServerArrivalTimeMiddleware: """ Custom ASGI middleware to track server arrival time. diff --git a/tests/unittest/disaggregated/test_disagg_responses_api.py b/tests/unittest/disaggregated/test_disagg_responses_api.py index 52a8ad9dcdce..4413c29a9eb8 100644 --- a/tests/unittest/disaggregated/test_disagg_responses_api.py +++ b/tests/unittest/disaggregated/test_disagg_responses_api.py @@ -20,17 +20,28 @@ import asyncio import base64 +import contextlib import json from types import SimpleNamespace -from unittest.mock import patch +from unittest.mock import AsyncMock, Mock, patch import aiohttp +import msgspec import numpy as np import pytest -from fastapi import HTTPException -from openai.types.responses import ResponseCreatedEvent +from fastapi import FastAPI, HTTPException, Request +from openai.types.responses import ( + ResponseCompletedEvent, + ResponseCreatedEvent, + ResponseTextDeltaEvent, +) from pydantic import ValidationError +from tensorrt_llm.executor.utils import context_length_exceeded_message +from tensorrt_llm.llmapi.disagg_utils import DisaggServerConfig +from tensorrt_llm.serve import openai_disagg_server, responses_utils +from tensorrt_llm.serve.disagg_coordinator import DisaggCoordinatorService +from tensorrt_llm.serve.openai_client import _encode_request_body from tensorrt_llm.serve.openai_disagg_server import OpenAIDisaggServer from tensorrt_llm.serve.openai_disagg_service import ( OpenAIDisaggregatedService, @@ -49,10 +60,13 @@ UsageInfo, ) from tensorrt_llm.serve.responses_utils import ( + ResponsesStreamingProcessor, + ServerArrivalTimeMiddleware, create_response_non_store, finish_reason_mapping, responses_done_generator, ) +from tensorrt_llm.serve.router import Router pytestmark = pytest.mark.cpu_only @@ -602,3 +616,613 @@ def test_a_models_failure_is_an_error_response(client, servers, status): response = asyncio.run(server.get_model()) assert response.status_code == status assert json.loads(response.body)["code"] == status + + +# --------------------------------------------------------------------------- +# Streamed requests: a slow context phase must not leave the client silent +# --------------------------------------------------------------------------- + +_MODEL = "test-model" +_SCHEDULE_STYLES = ["context_first", "generation_first"] +_OVER_LENGTH = context_length_exceeded_message(8, 9) + + +def _as_received(request: ResponsesRequest) -> ResponsesRequest: + """``request`` as a worker parses it from the body the service forwards.""" + return ResponsesRequest(**msgspec.msgpack.decode(_encode_request_body(request))) + + +def _response(request: ResponsesRequest, status: str) -> ResponsesResponse: + """A worker's response to ``request``, without output.""" + return ResponsesResponse.from_request( + request=request, + sampling_params=request.to_sampling_params(), + model_name=_MODEL, + created_time=0, + output=[], + status=status, + ) + + +def _sse(event) -> str: + return f"event: {event.type}\ndata: {event.model_dump_json()}\n\n" + + +def _generation_frames(request: ResponsesRequest) -> list[str]: + """A generation worker's stream for ``request``: its opening, "hello", the completion.""" + opening = ResponsesStreamingProcessor( + request, request.to_sampling_params(), model_name=_MODEL + ).get_initial_responses() + delta = ResponseTextDeltaEvent( + type="response.output_text.delta", + sequence_number=2, + item_id="msg_1", + output_index=0, + content_index=0, + delta="hello", + logprobs=[], + ) + completed = ResponseCompletedEvent( + type="response.completed", + sequence_number=3, + response=_response(request, "completed").model_dump(by_alias=True), + ) + return opening + [_sse(delta), _sse(completed)] + + +def _worker_rejection(status: int, message: str) -> aiohttp.ClientResponseError: + """A worker's error response, as the HTTP client raises it.""" + body = json.dumps({"object": "error", "message": message, "code": status}) + return aiohttp.ClientResponseError( + Mock(real_url="http://ctx:8001/v1/responses"), + (), + status=status, + message=f"Error: {body}", + ) + + +class _Workers: + """A context and a generation worker behind the disaggregated service, without a network. + + Each parses its request from the body the service forwards. The context phase runs + ``context_phase`` and then answers with the KV-cache handoff. The generation worker + answers once the context phase has succeeded, as it does once it has the KV cache, and + streams its answer in pieces cut anywhere, as the transport delivers it. + """ + + def __init__(self, context_phase): + self.context_phase = context_phase + self.context_done = asyncio.Event() + self.context_cancelled = asyncio.Event() + self.generation_cancelled = asyncio.Event() + self.generation_requests = [] + # What the generation worker answered: a response, or the frames of a stream. + self.response = None + self.frames = None + + async def send_context_request(self, request, server=None, hooks=None, req_id=None): + received = _as_received(request) + try: + await self.context_phase() + except asyncio.CancelledError: + self.context_cancelled.set() + raise + self.context_done.set() + response = _response(received, "incomplete") + response.finish_reason = "length" + response.prompt_token_ids = [1, 2, 3] + disagg_request_id = received.disaggregated_params.disagg_request_id + response.disaggregated_params = DisaggregatedParams( + request_type="context_only", + disagg_request_id=disagg_request_id, + ctx_request_id=disagg_request_id, + ) + return response + + async def send_generation_request(self, request, server=None, hooks=None, req_id=None): + received = _as_received(request) + self.generation_requests.append(received) + if not received.stream: + await self.context_done.wait() + self.response = _response(received, "completed") + return self.response + + async def stream(): + try: + await self.context_done.wait() + self.frames = _generation_frames(received) + data = "".join(self.frames).encode() + for start in range(0, len(data), 61): + yield data[start : start + 61] + except asyncio.CancelledError: + self.generation_cancelled.set() + raise + + return stream() + + +def _responses_server(schedule_style: str, workers: _Workers) -> OpenAIDisaggServer: + """A disaggregated server in front of ``workers``, with its routes.""" + config = DisaggServerConfig(server_configs=[], schedule_style=schedule_style) + coordinator = DisaggCoordinatorService(config, client_factory=Mock()) + service = OpenAIDisaggregatedService(config, coordinator, client_factory=Mock()) + service.is_ready = AsyncMock(return_value=True) + service._ctx_router = AsyncMock(spec=Router) + service._ctx_router.get_next_server.return_value = ("ctx:8001", {}) + service._gen_router = AsyncMock(spec=Router) + service._gen_router.get_next_server.return_value = ("gen:8002", {}) + service._ctx_client = SimpleNamespace(send_request=workers.send_context_request) + service._gen_client = SimpleNamespace(send_request=workers.send_generation_request) + + server = OpenAIDisaggServer.__new__(OpenAIDisaggServer) + server._config = config + server._coordinator = coordinator + server._disagg_cluster_storage = None + server._service = service + server._perf_metrics_collector = Mock() + server._allow_request_chat_template = False + server._collect_perf_metrics = False + server.app = FastAPI() + server.app.add_middleware(ServerArrivalTimeMiddleware) + with patch("prometheus_client.multiprocess.MultiProcessCollector"): + server.register_routes() + return server + + +def _responses_route(server: OpenAIDisaggServer): + """The handler ``server`` routes /v1/responses to.""" + (route,) = [ + route for route in server.app.routes if getattr(route, "path", "") == "/v1/responses" + ] + return route.endpoint + + +async def _post_responses(app: FastAPI, *, stream=True, on_send=None, leave=None, **fields): + """Send one /v1/responses request to ``app`` as a client and return the messages sent back. + + ``on_send`` sees each message as it is sent. The client disconnects once + ``leave`` is set. + """ + body = {"model": _MODEL, "input": "hi", "stream": stream, **fields} + requested = False + + async def receive(): + nonlocal requested + if not requested: + requested = True + return {"type": "http.request", "body": json.dumps(body).encode(), "more_body": False} + await (leave or asyncio.Event()).wait() + return {"type": "http.disconnect"} + + sent = [] + + async def send(message): + sent.append(message) + if on_send is not None: + on_send(message) + + scope = { + "type": "http", + # As uvicorn declares it: a disconnect cancels a streaming response. + "asgi": {"version": "3.0", "spec_version": "2.3"}, + "http_version": "1.1", + "method": "POST", + "scheme": "http", + "path": "/v1/responses", + "raw_path": b"/v1/responses", + "query_string": b"", + "root_path": "", + "headers": [(b"content-type", b"application/json")], + "client": ("127.0.0.1", 40000), + "server": ("127.0.0.1", 8000), + } + await asyncio.wait_for(app(scope, receive, send), timeout=10) + return sent + + +def _status(sent: list[dict]) -> int: + return next(message["status"] for message in sent if message["type"] == "http.response.start") + + +def _body(sent: list[dict]) -> str: + return b"".join( + message.get("body", b"") for message in sent if message["type"] == "http.response.body" + ).decode() + + +def _sse_events(sent: list[dict]) -> list[tuple[str, dict]]: + """The (event, data) pairs of a streamed response.""" + events = [] + for frame in _body(sent).split("\n\n"): + if frame: + fields = dict(line.split(": ", 1) for line in frame.splitlines()) + events.append((fields["event"], json.loads(fields["data"]))) + return events + + +def _without_number(data: dict) -> dict: + return {key: value for key, value in data.items() if key != "sequence_number"} + + +def _short_waits(monkeypatch): + monkeypatch.setattr(openai_disagg_server, "_RESPONSES_STREAM_START_WAIT_S", 0.05) + monkeypatch.setattr(responses_utils, "RESPONSES_STREAM_KEEPALIVE_INTERVAL_S", 0.05) + + +def _waits_for(event: asyncio.Event): + """A context phase that lasts until ``event`` is set, or five seconds without it.""" + + async def context_phase(): + with contextlib.suppress(asyncio.TimeoutError): + await asyncio.wait_for(event.wait(), 5) + + return context_phase + + +def _on_keepalive(event: asyncio.Event): + """An ``on_send`` that sets ``event`` once response.in_progress has gone out again.""" + in_progress = 0 + + def on_send(message): + nonlocal in_progress + if message.get("body", b"").startswith(b"event: response.in_progress\n"): + in_progress += 1 + if in_progress == 2: + event.set() + + return on_send + + +@pytest.mark.asyncio +@pytest.mark.parametrize("schedule_style", _SCHEDULE_STYLES) +async def test_responses_stream_opens_and_keeps_alive_while_the_context_phase_runs( + monkeypatch, schedule_style +): + _short_waits(monkeypatch) + kept_alive = asyncio.Event() + workers = _Workers(_waits_for(kept_alive)) + keepalive = _on_keepalive(kept_alive) + started_during_context_phase = [] + + def on_send(message): + if message["type"] == "http.response.start": + started_during_context_phase.append(not workers.context_done.is_set()) + keepalive(message) + + sent = await _post_responses(_responses_server(schedule_style, workers).app, on_send=on_send) + + assert _status(sent) == 200 + assert started_during_context_phase == [True] + events = _sse_events(sent) + types = [event for event, _ in events] + first_output = types.index("response.output_text.delta") + # The opening, kept alive while the context phase ran, then the generation + # worker's events without its own opening. + assert types[:2] == ["response.created", "response.in_progress"] + assert first_output > 2 + assert set(types[2:first_output]) == {"response.in_progress"} + relayed = [data for event, data in events[first_output:] if event != "response.in_progress"] + answered = [json.loads(frame.split("data: ", 1)[1]) for frame in workers.frames[2:]] + assert [_without_number(data) for data in relayed] == [ + _without_number(data) for data in answered + ] + # One sequence, for one response: the one the generation worker was asked for. + assert [data["sequence_number"] for _, data in events] == list(range(len(events))) + (generation_request,) = workers.generation_requests + assert {data["response"]["id"] for _, data in events if "response" in data} == { + generation_request.request_id + } + + +@pytest.mark.asyncio +@pytest.mark.parametrize("schedule_style", _SCHEDULE_STYLES) +async def test_responses_stream_after_a_fast_context_phase_is_the_generation_stream( + schedule_style, +): + async def context_phase(): + pass + + workers = _Workers(context_phase) + sent = await _post_responses(_responses_server(schedule_style, workers).app) + + assert _status(sent) == 200 + assert [event for event, _ in _sse_events(sent)] == [ + "response.created", + "response.in_progress", + "response.output_text.delta", + "response.completed", + ] + assert _body(sent) == "".join(workers.frames) + + +_CONTEXT_FAILURES = [ + pytest.param(400, _OVER_LENGTH, id="over_length_400"), + pytest.param(503, "no capacity left", id="503"), +] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("schedule_style", _SCHEDULE_STYLES) +@pytest.mark.parametrize("status, worker_message", _CONTEXT_FAILURES) +async def test_responses_stream_with_a_fast_context_failure_is_an_http_error( + schedule_style, status, worker_message +): + async def context_phase(): + raise _worker_rejection(status, worker_message) + + streamed = await _post_responses(_responses_server(schedule_style, _Workers(context_phase)).app) + not_streamed = await _post_responses( + _responses_server(schedule_style, _Workers(context_phase)).app, stream=False + ) + + assert _status(streamed) == status + assert worker_message in _body(streamed) + # The error response of the same failure without a stream. + assert _status(not_streamed) == status + assert json.loads(_body(streamed)) == json.loads(_body(not_streamed)) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("schedule_style", _SCHEDULE_STYLES) +@pytest.mark.parametrize("status, worker_message", _CONTEXT_FAILURES) +async def test_responses_stream_with_a_late_context_failure_ends_with_response_failed( + monkeypatch, schedule_style, status, worker_message +): + _short_waits(monkeypatch) + kept_alive = asyncio.Event() + wait = _waits_for(kept_alive) + + async def context_phase(): + await wait() + raise _worker_rejection(status, worker_message) + + streamed = await _post_responses( + _responses_server(schedule_style, _Workers(context_phase)).app, + on_send=_on_keepalive(kept_alive), + ) + kept_alive.set() + not_streamed = await _post_responses( + _responses_server(schedule_style, _Workers(context_phase)).app, stream=False + ) + + assert _status(streamed) == 200 + # The stream ends after its events, as any stream does. + assert streamed[-1] == {"type": "http.response.body", "body": b"", "more_body": False} + events = _sse_events(streamed) + types = [event for event, _ in events] + assert types[:2] == ["response.created", "response.in_progress"] + assert set(types[2:-2]) == {"response.in_progress"} + assert types[-2:] == ["error", "response.failed"] + assert [data["sequence_number"] for _, data in events] == list(range(len(events))) + # error and response.failed, as a stream reports any failure, with the + # message of the error response of the same failure without a stream. + error_response = json.loads(_body(not_streamed)) + message = error_response.get("detail") or error_response["message"] + assert worker_message in message + assert (events[-2][1]["code"], events[-2][1]["message"]) == ("upstream_error", message) + failed = events[-1][1]["response"] + assert failed["id"] == events[0][1]["response"]["id"] + assert failed["status"] == "failed" + assert (failed["error"]["code"], failed["error"]["message"]) == ("server_error", message) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("schedule_style", _SCHEDULE_STYLES) +async def test_a_non_streamed_response_still_waits_for_the_whole_pipeline( + monkeypatch, schedule_style +): + _short_waits(monkeypatch) + + async def context_phase(): + # Outlasts the start wait. + await asyncio.sleep(0.2) + + workers = _Workers(context_phase) + sent = await _post_responses(_responses_server(schedule_style, workers).app, stream=False) + + assert _status(sent) == 200 + assert json.loads(_body(sent)) == workers.response.model_dump(mode="json", by_alias=True) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("schedule_style", _SCHEDULE_STYLES) +async def test_a_request_no_opening_can_be_built_for_is_answered_by_its_pipeline( + monkeypatch, schedule_style +): + """The opening is built from the request; one it cannot be built from waits as before.""" + _short_waits(monkeypatch) + rejection = "require 0 <= top_p <= 1, got top_p=2.0" + + async def context_phase(): + # Outlasts the start wait. + await asyncio.sleep(0.2) + raise _worker_rejection(400, rejection) + + sent = await _post_responses( + _responses_server(schedule_style, _Workers(context_phase)).app, top_p=2.0 + ) + + assert _status(sent) == 400 + assert json.loads(_body(sent)) == {"detail": rejection} + + +@pytest.mark.asyncio +@pytest.mark.parametrize("schedule_style", _SCHEDULE_STYLES) +async def test_client_leaving_after_the_responses_stream_starts_cancels_the_context_phase( + monkeypatch, schedule_style +): + _short_waits(monkeypatch) + workers = _Workers(_waits_for(asyncio.Event())) + left = asyncio.Event() + + def on_send(message): + if message["type"] == "http.response.start": + left.set() + + await _post_responses( + _responses_server(schedule_style, workers).app, on_send=on_send, leave=left + ) + + await asyncio.wait_for(workers.context_cancelled.wait(), 2) + + +def _raw_request() -> Request: + raw_request = Request({"type": "http", "headers": []}) + raw_request.state.server_arrival_time = 0 + return raw_request + + +@pytest.mark.asyncio +@pytest.mark.parametrize("schedule_style", _SCHEDULE_STYLES) +async def test_cancelling_the_responses_handler_during_its_wait_cancels_the_context_phase( + monkeypatch, schedule_style +): + monkeypatch.setattr(openai_disagg_server, "_RESPONSES_STREAM_START_WAIT_S", 60) + context_started = asyncio.Event() + release_context = asyncio.Event() + + async def context_phase(): + context_started.set() + await release_context.wait() + + workers = _Workers(context_phase) + handler = _responses_route(_responses_server(schedule_style, workers)) + request = ResponsesRequest(model=_MODEL, input="hi", stream=True) + pending = asyncio.create_task(handler(request, _raw_request())) + try: + await asyncio.wait_for(context_started.wait(), 2) + assert not pending.done() + pending.cancel() + with pytest.raises(asyncio.CancelledError): + await pending + await asyncio.wait_for(workers.context_cancelled.wait(), 2) + assert not workers.context_done.is_set() + finally: + release_context.set() + pending.cancel() + with contextlib.suppress(asyncio.CancelledError): + await pending + + +@pytest.mark.asyncio +@pytest.mark.parametrize("schedule_style", _SCHEDULE_STYLES) +@pytest.mark.parametrize("asgi_version", ["2.3", "2.4"]) +async def test_disconnect_before_the_responses_body_starts_cancels_the_context_phase( + monkeypatch, schedule_style, asgi_version +): + """The pipeline is cancelled even if the body is never iterated.""" + from starlette.requests import ClientDisconnect + + _short_waits(monkeypatch) + release_context = asyncio.Event() + workers = _Workers(_waits_for(release_context)) + handler = _responses_route(_responses_server(schedule_style, workers)) + request = ResponsesRequest(model=_MODEL, input="hi", stream=True) + response = await asyncio.wait_for(handler(request, _raw_request()), 2) + response_started = asyncio.Event() + sent = [] + + async def send(message): + sent.append(message) + response_started.set() + if asgi_version == "2.4": + raise OSError("client disconnected") + await asyncio.Event().wait() + + async def receive(): + await response_started.wait() + return {"type": "http.disconnect"} + + try: + scope = {"type": "http", "asgi": {"version": "3.0", "spec_version": asgi_version}} + if asgi_version == "2.4": + with pytest.raises(ClientDisconnect): + await response(scope, receive, send) + else: + await response(scope, receive, send) + assert [message["type"] for message in sent] == ["http.response.start"] + await asyncio.wait_for(workers.context_cancelled.wait(), 2) + if schedule_style == "generation_first": + await asyncio.wait_for(workers.generation_cancelled.wait(), 2) + assert not workers.context_done.is_set() + finally: + release_context.set() + + +def _frame(event_type: str, sequence_number: int, **fields) -> str: + data = {"type": event_type, "sequence_number": sequence_number, **fields} + return f"event: {event_type}\ndata: {json.dumps(data, separators=(',', ':'))}\n\n" + + +async def _transport(data: str, piece: int): + """``data`` as a relayed stream delivers it, in pieces of ``piece`` bytes.""" + encoded = data.encode() + for start in range(0, len(encoded), piece): + yield encoded[start : start + piece] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("piece", [1, 7, 4096]) +async def test_the_relay_drops_the_worker_opening_and_numbers_the_rest_to_follow(piece): + worker = "".join( + [ + _frame("response.created", 0), + _frame("response.in_progress", 1), + # Kept alive before its first output. + _frame("response.in_progress", 2), + _frame("response.output_text.delta", 3, delta="hi"), + # Kept alive after it. + _frame("response.in_progress", 4), + _frame("response.completed", 5), + ] + ) + + relayed = responses_utils.relay_after_opening(_transport(worker, piece), 2) + + assert [frame async for frame in relayed] == [ + _frame("response.output_text.delta", 2, delta="hi"), + _frame("response.in_progress", 3), + _frame("response.completed", 4), + ] + + +@pytest.mark.asyncio +async def test_the_relay_numbers_a_failure_the_worker_reported_before_its_opening(): + worker = _frame("error", 0, code="upstream_error", message="ClientResponseError") + + relayed = responses_utils.relay_after_opening(_transport(worker, 4096), 3) + + assert [frame async for frame in relayed] == [ + _frame("error", 3, code="upstream_error", message="ClientResponseError") + ] + + +@pytest.mark.asyncio +async def test_the_relay_passes_on_an_event_its_stream_stopped_inside(): + cut = 'event: response.output_text.delta\ndata: {"type":"response.out' + worker = _frame("response.output_text.delta", 2, delta="hi") + cut + + relayed = responses_utils.relay_after_opening(_transport(worker, 7), 2) + + assert [frame async for frame in relayed] == [ + _frame("response.output_text.delta", 2, delta="hi"), + cut, + ] + + +@pytest.mark.asyncio +async def test_the_relay_closes_its_stream_when_it_stops_early(): + closed = asyncio.Event() + + async def worker(): + try: + yield _frame("response.output_text.delta", 2, delta="hi").encode() + yield _frame("response.completed", 3).encode() + finally: + closed.set() + + relayed = responses_utils.relay_after_opening(worker(), 2) + await relayed.__anext__() + await relayed.aclose() + + assert closed.is_set() diff --git a/tests/unittest/disaggregated/test_openai_disagg_service.py b/tests/unittest/disaggregated/test_openai_disagg_service.py index 46b063e71584..ed47990871b7 100644 --- a/tests/unittest/disaggregated/test_openai_disagg_service.py +++ b/tests/unittest/disaggregated/test_openai_disagg_service.py @@ -966,6 +966,43 @@ async def _gen_response(*_args, **_kwargs): ) +@pytest.mark.asyncio +async def test_context_cancellation_releases_conditional_generation_reservation(monkeypatch): + """Cancelling CTX must release the generation server reserved before it.""" + monkeypatch.delenv("TRTLLM_DISAGG_BENCHMARK_GEN_ONLY", raising=False) + service = _make_service("context_first") + service._config.conditional_disagg_config = ConditionalDisaggConfig(max_local_prefill_length=32) + service._coordinator.get_disagg_request_id = AsyncMock(return_value=101) + service._ctx_client = AsyncMock() + service._gen_client = AsyncMock() + router = KvCacheAwareRouter(server_role=ServerRole.GENERATION, servers=[]) + router.get_next_server = AsyncMock( + return_value=("gen:9001", {"match_length": 0, "num_tokens": 96}) + ) + router.finish_request = AsyncMock() + service._gen_router = router + service._ctx_router.get_next_server = AsyncMock(return_value=("ctx:9000", {})) + context_started = asyncio.Event() + + async def context_response(request, **kwargs): + context_started.set() + await asyncio.Event().wait() + + service._ctx_client.send_request = AsyncMock(side_effect=context_response) + request = CompletionRequest(model="model", prompt=[1] * 96, stream=True) + pending = asyncio.create_task(service._send_disagg_request(request)) + try: + await asyncio.wait_for(context_started.wait(), 2) + finally: + pending.cancel() + with pytest.raises(asyncio.CancelledError): + await pending + + router.get_next_server.assert_awaited_once_with(request, req_id=101) + router.finish_request.assert_awaited_once_with(request, success=False, req_id=101) + service._gen_client.send_request.assert_not_awaited() + + @pytest.mark.asyncio async def test_context_retry_preserves_generation_reservation_id(): service = _make_service("context_first") From efc5a811a21e0278ee6ec6faed10830c77c35317 Mon Sep 17 00:00:00 2001 From: JunyiXu-nv <219237550+JunyiXu-nv@users.noreply.github.com> Date: Mon, 2 Nov 2026 16:42:55 +0000 Subject: [PATCH 2/5] [None][fix] Disaggregated Responses: document the early opening Describe in the disaggregated serving guide how a streamed /v1/responses request whose context phase is still running gets its opening from the server, and how a later failure is reported. Import AsyncIterator from typing, with the module's other typing names. Signed-off-by: JunyiXu-nv <219237550+JunyiXu-nv@users.noreply.github.com> --- docs/source/features/disagg-serving.md | 2 +- tensorrt_llm/serve/openai_disagg_server.py | 3 +-- 2 files changed, 2 insertions(+), 3 deletions(-) diff --git a/docs/source/features/disagg-serving.md b/docs/source/features/disagg-serving.md index f0bda3662f1c..27eb09f53a46 100644 --- a/docs/source/features/disagg-serving.md +++ b/docs/source/features/disagg-serving.md @@ -236,7 +236,7 @@ curl http://localhost:8000/v1/completions \ #### Responses API -The disaggregated server also serves `/v1/responses`, with both `context_first` and `generation_first` scheduling, and `/v1/models`, which it answers from a context worker. A streamed response carries the generation worker's events unchanged. Unlike a single `trtllm-serve` instance, the disaggregated server keeps no response history: +The disaggregated server also serves `/v1/responses`, with both `context_first` and `generation_first` scheduling, and `/v1/models`, which it answers from a context worker. A streamed response carries the generation worker's events. If the context phase is still running after a second, the disaggregated server starts the stream itself with `response.created` and `response.in_progress`, repeats `response.in_progress` while the stream is silent, and then relays the generation worker's events without its own opening, numbered to continue the stream; a failure after that point ends the stream with `error` and `response.failed` instead of an HTTP error status. Unlike a single `trtllm-serve` instance, the disaggregated server keeps no response history: - A request with `previous_response_id` is rejected with a 400, because a stored response lives in one worker's memory and a later request may reach another worker. Send the earlier turns in `input` instead. - `store` is ignored: nothing is persisted, and the server logs a warning once when a client asks for it. diff --git a/tensorrt_llm/serve/openai_disagg_server.py b/tensorrt_llm/serve/openai_disagg_server.py index def2f0477e93..5684f6c1a538 100644 --- a/tensorrt_llm/serve/openai_disagg_server.py +++ b/tensorrt_llm/serve/openai_disagg_server.py @@ -20,9 +20,8 @@ import socket import traceback import uuid -from collections.abc import AsyncIterator from contextlib import asynccontextmanager -from typing import Callable, Optional +from typing import AsyncIterator, Callable, Optional import aiohttp import uvicorn From 8106fd8ddfa4b96217ce3427d332460c183acab6 Mon Sep 17 00:00:00 2001 From: JunyiXu-nv <219237550+JunyiXu-nv@users.noreply.github.com> Date: Mon, 2 Nov 2026 20:28:10 +0000 Subject: [PATCH 3/5] [None][fix] Disaggregated Responses: keep an opened stream alive with its opening only Once the server had opened a stream, the relay still passed the generation worker's own keepalives, its repeats of response.in_progress, on to the client. They carry the worker's response object, created on the worker's clock, not the one the client was opened with, while an aggregated server repeats the stream's first response.in_progress exactly. They also came on top of the keepalives this server sends, so a silent worker had the client kept alive twice. The relay now drops every response.created and response.in_progress of the worker's, wherever it comes. That loses nothing, since a worker's keepalive repeats its own opening: the client is kept alive by this server alone, with the opening's response.in_progress and the next sequence number. Signed-off-by: JunyiXu-nv <219237550+JunyiXu-nv@users.noreply.github.com> --- docs/source/features/disagg-serving.md | 2 +- tensorrt_llm/serve/openai_disagg_server.py | 10 +- tensorrt_llm/serve/responses_utils.py | 14 ++- .../test_disagg_responses_api.py | 107 ++++++++++++++++-- 4 files changed, 113 insertions(+), 20 deletions(-) diff --git a/docs/source/features/disagg-serving.md b/docs/source/features/disagg-serving.md index 27eb09f53a46..ad4faf8de4f3 100644 --- a/docs/source/features/disagg-serving.md +++ b/docs/source/features/disagg-serving.md @@ -236,7 +236,7 @@ curl http://localhost:8000/v1/completions \ #### Responses API -The disaggregated server also serves `/v1/responses`, with both `context_first` and `generation_first` scheduling, and `/v1/models`, which it answers from a context worker. A streamed response carries the generation worker's events. If the context phase is still running after a second, the disaggregated server starts the stream itself with `response.created` and `response.in_progress`, repeats `response.in_progress` while the stream is silent, and then relays the generation worker's events without its own opening, numbered to continue the stream; a failure after that point ends the stream with `error` and `response.failed` instead of an HTTP error status. Unlike a single `trtllm-serve` instance, the disaggregated server keeps no response history: +The disaggregated server also serves `/v1/responses`, with both `context_first` and `generation_first` scheduling, and `/v1/models`, which it answers from a context worker. A streamed response carries the generation worker's events. If the context phase is still running after a second, the disaggregated server starts the stream itself with `response.created` and `response.in_progress`, repeats that `response.in_progress` whenever the stream is silent, and relays the generation worker's other events, numbered to continue the stream; a failure after that point ends the stream with `error` and `response.failed` instead of an HTTP error status. Unlike a single `trtllm-serve` instance, the disaggregated server keeps no response history: - A request with `previous_response_id` is rejected with a 400, because a stored response lives in one worker's memory and a later request may reach another worker. Send the earlier turns in `input` instead. - `store` is ignored: nothing is persisted, and the server logs a warning once when a client asks for it. diff --git a/tensorrt_llm/serve/openai_disagg_server.py b/tensorrt_llm/serve/openai_disagg_server.py index 5684f6c1a538..077df6a4079a 100644 --- a/tensorrt_llm/serve/openai_disagg_server.py +++ b/tensorrt_llm/serve/openai_disagg_server.py @@ -531,8 +531,8 @@ async def openai_responses(self, request: ResponsesRequest, # The opening is built from the request alone; a request it cannot # be built from is answered when its pipeline answers. return await pipeline - # keep_responses_stream_alive repeats response.in_progress whenever the - # stream is silent, as it is while the pipeline runs. + # keep_responses_stream_alive repeats the opening's response.in_progress + # whenever the stream is silent, as it is while the pipeline runs. return _ResponsesStreamingResponse( pipeline=pipeline, content=keep_responses_stream_alive( @@ -545,8 +545,10 @@ async def _responses_stream_when_done( """The stream of a Responses request whose pipeline outlasted the start wait. ``opening`` is the request's response.created and response.in_progress, - and the pipeline's stream follows it without the worker's own opening - (relay_after_opening). The status line is already out, so a pipeline + and the pipeline's stream follows it without the worker's own + response.created and response.in_progress events (relay_after_opening), + so that every keepalive the client gets repeats ``opening``'s + response.in_progress. The status line is already out, so a pipeline that fails ends the stream with error and response.failed, which carry the message of its error response. A stream abandoned before the pipeline answers, as when the client disconnects, cancels the pipeline. diff --git a/tensorrt_llm/serve/responses_utils.py b/tensorrt_llm/serve/responses_utils.py index 61e1673ab494..603a90fe4df0 100644 --- a/tensorrt_llm/serve/responses_utils.py +++ b/tensorrt_llm/serve/responses_utils.py @@ -3483,15 +3483,19 @@ async def relay_after_opening(stream: AsyncGenerator[Union[str, bytes], None], """Relay a worker's Responses stream to a client that already has its opening. The client holds ``events_sent`` events, ``response.created`` and - ``response.in_progress`` among them, so the worker's own opening events, - and any repeats of them before its first other event, are dropped. Every + ``response.in_progress`` among them, and is kept alive with that + ``response.in_progress`` by the stream this relay feeds + (``keep_responses_stream_alive``). So every ``response.created`` and + ``response.in_progress`` of the worker's is dropped, wherever it comes: + the worker's opening and its keepalives, which repeat that opening, carry + the worker's response object, not the one the client holds, and a second + source of keepalives would only add to the relaying stream's own. Every other event is numbered to follow on, and is yielded whole, one per item as ``keep_responses_stream_alive`` takes them, wherever ``stream`` splits it. ``stream`` is closed when the relay ends. """ delimiter = _SSE_EVENT_DELIMITER.encode() pending = b"" - opening = True try: async for chunk in stream: if isinstance(chunk, str): @@ -3499,10 +3503,8 @@ async def relay_after_opening(stream: AsyncGenerator[Union[str, bytes], None], *events, pending = (pending + chunk).split(delimiter) for event in events: frame = event.decode(errors="replace") + _SSE_EVENT_DELIMITER - if not event or (opening - and frame.startswith(_OPENING_FRAME_STARTS)): + if not event or frame.startswith(_OPENING_FRAME_STARTS): continue - opening = False yield _stamp_sequence_number(frame, events_sent) events_sent += 1 if pending: diff --git a/tests/unittest/disaggregated/test_disagg_responses_api.py b/tests/unittest/disaggregated/test_disagg_responses_api.py index 4413c29a9eb8..fe8a860a4a0c 100644 --- a/tests/unittest/disaggregated/test_disagg_responses_api.py +++ b/tests/unittest/disaggregated/test_disagg_responses_api.py @@ -649,9 +649,12 @@ def _sse(event) -> str: def _generation_frames(request: ResponsesRequest) -> list[str]: - """A generation worker's stream for ``request``: its opening, "hello", the completion.""" + """A generation worker's stream for ``request``: its opening, "hello", the completion. + + Its response objects are created at 0, so they differ from any built elsewhere. + """ opening = ResponsesStreamingProcessor( - request, request.to_sampling_params(), model_name=_MODEL + request, request.to_sampling_params(), model_name=_MODEL, create_time=0 ).get_initial_responses() delta = ResponseTextDeltaEvent( type="response.output_text.delta", @@ -730,16 +733,45 @@ async def send_generation_request(self, request, server=None, hooks=None, req_id async def stream(): try: await self.context_done.wait() - self.frames = _generation_frames(received) - data = "".join(self.frames).encode() - for start in range(0, len(data), 61): - yield data[start : start + 61] + async for piece in self.generation_stream(received): + yield piece except asyncio.CancelledError: self.generation_cancelled.set() raise return stream() + async def generation_stream(self, received: ResponsesRequest): + """The generation worker's stream, in pieces cut anywhere, as the transport delivers it.""" + self.frames = _generation_frames(received) + data = "".join(self.frames).encode() + for start in range(0, len(data), 61): + yield data[start : start + 61] + + +class _WorkersSilentAfterTheFirstOutput(_Workers): + """Workers whose generation worker is silent for a while after its first output. + + The silence outlasts the keepalive interval, so the worker's own keepalive sends its + response.in_progress again, with its own response object, before the rest of its stream. + """ + + async def generation_stream(self, received: ResponsesRequest): + frames = _generation_frames(received) + + async def produced(): + # The opening and the first output, then the rest after a silence. + for frame in frames[:3]: + yield frame + await asyncio.sleep(0.3) + for frame in frames[3:]: + yield frame + + self.frames = [] + async for frame in responses_utils.keep_responses_stream_alive(produced()): + self.frames.append(frame) + yield frame.encode() + def _responses_server(schedule_style: str, workers: _Workers) -> OpenAIDisaggServer: """A disaggregated server in front of ``workers``, with its routes.""" @@ -914,6 +946,64 @@ def on_send(message): } +def _frame_events(frames: list[str]) -> list[tuple[str, dict]]: + """The (event, data) pairs of SSE frames.""" + return [ + (frame.split("\n", 1)[0].removeprefix("event: "), json.loads(frame.split("data: ", 1)[1])) + for frame in frames + ] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("schedule_style", _SCHEDULE_STYLES) +async def test_every_keepalive_of_an_opened_stream_repeats_its_opening(monkeypatch, schedule_style): + """The worker's own keepalives, which carry its response object, never reach the client.""" + _short_waits(monkeypatch) + kept_alive = asyncio.Event() + workers = _WorkersSilentAfterTheFirstOutput(_waits_for(kept_alive)) + + sent = await _post_responses( + _responses_server(schedule_style, workers).app, on_send=_on_keepalive(kept_alive) + ) + + assert _status(sent) == 200 + events = _sse_events(sent) + types = [event for event, _ in events] + assert types[:2] == ["response.created", "response.in_progress"] + first_output = types.index("response.output_text.delta") + # The worker kept its own stream alive while it was silent after its first output. + answered = _frame_events(workers.frames) + answered_types = [event for event, _ in answered] + assert ( + "response.in_progress" + in answered_types[answered_types.index("response.output_text.delta") :] + ) + # The client was kept alive meanwhile, and every response.in_progress it got is + # the opening's, only numbered on; none is one of the worker's. + assert "response.in_progress" in types[first_output:] + opening = _without_number(events[1][1]) + assert all( + _without_number(data) == opening + for event, data in events + if event == "response.in_progress" + ) + assert opening not in [ + _without_number(data) for event, data in answered if event == "response.in_progress" + ] + # The worker's other events follow, in one sequence, for one response. + relayed = [data for event, data in events[first_output:] if event != "response.in_progress"] + assert [_without_number(data) for data in relayed] == [ + _without_number(data) + for event, data in answered + if event not in ("response.created", "response.in_progress") + ] + assert [data["sequence_number"] for _, data in events] == list(range(len(events))) + (generation_request,) = workers.generation_requests + assert {data["response"]["id"] for _, data in events if "response" in data} == { + generation_request.request_id + } + + @pytest.mark.asyncio @pytest.mark.parametrize("schedule_style", _SCHEDULE_STYLES) async def test_responses_stream_after_a_fast_context_phase_is_the_generation_stream( @@ -1163,7 +1253,7 @@ async def _transport(data: str, piece: int): @pytest.mark.asyncio @pytest.mark.parametrize("piece", [1, 7, 4096]) -async def test_the_relay_drops_the_worker_opening_and_numbers_the_rest_to_follow(piece): +async def test_the_relay_drops_the_worker_opening_and_keepalives_and_numbers_the_rest(piece): worker = "".join( [ _frame("response.created", 0), @@ -1181,8 +1271,7 @@ async def test_the_relay_drops_the_worker_opening_and_numbers_the_rest_to_follow assert [frame async for frame in relayed] == [ _frame("response.output_text.delta", 2, delta="hi"), - _frame("response.in_progress", 3), - _frame("response.completed", 4), + _frame("response.completed", 3), ] From 9563f8ff966e85cef3043c04e7bd645cd14a91a4 Mon Sep 17 00:00:00 2001 From: JunyiXu-nv <219237550+JunyiXu-nv@users.noreply.github.com> Date: Tue, 3 Nov 2026 00:10:08 +0000 Subject: [PATCH 4/5] [None][fix] Disaggregated Responses: fail an opened stream whose worker stream breaks, close it when abandoned A delayed stream guards the generation worker's stream with the processor's failed events, so a read failure before the terminal event ends the stream with one error and one response.failed for the response it opened with, and the stream then ends cleanly. A stream the pipeline answered with inside the start wait keeps the bare error event. The delayed response closes the stream its pipeline answered with however the response ends, also when the body never started. Signed-off-by: JunyiXu-nv <219237550+JunyiXu-nv@users.noreply.github.com> --- tensorrt_llm/serve/openai_disagg_server.py | 62 +++++-- .../test_disagg_responses_api.py | 158 +++++++++++++++++- 2 files changed, 205 insertions(+), 15 deletions(-) diff --git a/tensorrt_llm/serve/openai_disagg_server.py b/tensorrt_llm/serve/openai_disagg_server.py index 077df6a4079a..fe9b71ad5585 100644 --- a/tensorrt_llm/serve/openai_disagg_server.py +++ b/tensorrt_llm/serve/openai_disagg_server.py @@ -24,6 +24,7 @@ from typing import AsyncIterator, Callable, Optional import aiohttp +import anyio import uvicorn from fastapi import FastAPI, HTTPException, Request from fastapi.exceptions import RequestValidationError @@ -197,7 +198,12 @@ def on_resp_done( class _ResponsesStreamingResponse(ClosingStreamingResponse): - """Cancel the pipeline even when a disconnect precedes body iteration.""" + """Release the pipeline however the response ends, body iterated or not. + + A pipeline still running is cancelled, and the stream of one that has + answered is closed: the body may stop, or never start, before it has read + that stream to its end. + """ def __init__(self, pipeline: asyncio.Future, content: AsyncIterator[str]) -> None: @@ -208,7 +214,27 @@ async def __call__(self, scope, receive, send) -> None: try: await super().__call__(scope, receive, send) finally: - self._pipeline.cancel() + pipeline = self._pipeline + pipeline.cancel() + if (pipeline.done() and not pipeline.cancelled() + and pipeline.exception() is None + and isinstance(pipeline.result(), StreamingResponse)): + with anyio.CancelScope(shield=True): + await pipeline.result().body_iterator.aclose() + + +def _guard_relayed_responses_stream(answer: Response) -> Response: + """``answer``, with a stream that fails early ending in an ``error`` event. + + The stream is the generation worker's own, response.created included, so + there is no response of this server's to report as failed. + """ + if not isinstance(answer, StreamingResponse): + return answer + return StreamingResponse( + content=guard_responses_stream(answer.body_iterator, + stream_error_event), + media_type="text/event-stream") def _error_response_message(response: Response) -> str: @@ -478,14 +504,10 @@ async def wrapper(req: request_type, raw_req: Request) -> Response: response_or_generator = await entry_point(req, hooks) self._perf_metrics_collector.total_responses.inc() if req.stream: - stream = response_or_generator - if isinstance(req, ResponsesRequest): - # A relayed Responses stream that fails early ends with - # an `error` event; other protocols end on a sentinel. - stream = guard_responses_stream(stream, - stream_error_event) + # openai_responses ends a Responses stream that fails early + # with error events; other protocols end on a sentinel. return StreamingResponse( - content=stream, + content=response_or_generator, media_type="text/event-stream") # by_alias: e.g. text.format.schema goes out under its wire name. return JSONResponse( @@ -520,7 +542,7 @@ async def openai_responses(self, request: ResponsesRequest, pipeline.cancel() raise if pipeline.done(): - return pipeline.result() + return _guard_relayed_responses_stream(pipeline.result()) try: processor = ResponsesStreamingProcessor( request=request, @@ -530,7 +552,7 @@ async def openai_responses(self, request: ResponsesRequest, except ValueError: # The opening is built from the request alone; a request it cannot # be built from is answered when its pipeline answers. - return await pipeline + return _guard_relayed_responses_stream(await pipeline) # keep_responses_stream_alive repeats the opening's response.in_progress # whenever the stream is silent, as it is while the pipeline runs. return _ResponsesStreamingResponse( @@ -550,8 +572,9 @@ async def _responses_stream_when_done( so that every keepalive the client gets repeats ``opening``'s response.in_progress. The status line is already out, so a pipeline that fails ends the stream with error and response.failed, which carry - the message of its error response. A stream abandoned before the - pipeline answers, as when the client disconnects, cancels the pipeline. + the message of its error response, and so does a pipeline's stream that + fails before its terminal event. A stream abandoned before the pipeline + answers, as when the client disconnects, cancels the pipeline. """ try: for frame in opening: @@ -570,10 +593,21 @@ async def _responses_stream_when_done( STREAM_TERMINATION_UPSTREAM_ERROR, failure, len(opening)): yield frame return - frames = relay_after_opening(answer.body_iterator, len(opening)) + # The guard reads the worker's stream as it comes, so it knows whether + # the terminal event went out; its events are numbered on with the + # worker's by relay_after_opening. + frames = relay_after_opening( + guard_responses_stream(answer.body_iterator, + processor.get_stream_failed_events), + len(opening)) try: async for frame in frames: yield frame + except Exception: # noqa: BLE001 - the guard has reported it + # The guard has logged the failure and sent error and + # response.failed, unless the terminal event had gone out: either + # way the response is over, and so is the stream. + return finally: await frames.aclose() diff --git a/tests/unittest/disaggregated/test_disagg_responses_api.py b/tests/unittest/disaggregated/test_disagg_responses_api.py index fe8a860a4a0c..5db78e9afd06 100644 --- a/tests/unittest/disaggregated/test_disagg_responses_api.py +++ b/tests/unittest/disaggregated/test_disagg_responses_api.py @@ -30,6 +30,7 @@ import numpy as np import pytest from fastapi import FastAPI, HTTPException, Request +from fastapi.responses import StreamingResponse from openai.types.responses import ( ResponseCompletedEvent, ResponseCreatedEvent, @@ -698,6 +699,8 @@ def __init__(self, context_phase): self.context_done = asyncio.Event() self.context_cancelled = asyncio.Event() self.generation_cancelled = asyncio.Event() + # Set once the generation worker's stream has ended, however it ended. + self.generation_closed = asyncio.Event() self.generation_requests = [] # What the generation worker answered: a response, or the frames of a stream. self.response = None @@ -732,14 +735,21 @@ async def send_generation_request(self, request, server=None, hooks=None, req_id async def stream(): try: + yield None await self.context_done.wait() async for piece in self.generation_stream(received): yield piece except asyncio.CancelledError: self.generation_cancelled.set() raise + finally: + self.generation_closed.set() - return stream() + # Started before it is returned, as the HTTP client starts it, so that + # closing it ends it even before it is read. + response = stream() + await response.__anext__() + return response async def generation_stream(self, received: ResponsesRequest): """The generation worker's stream, in pieces cut anywhere, as the transport delivers it.""" @@ -773,6 +783,36 @@ async def produced(): yield frame.encode() +class _WorkersFailingAfterTheFirstOutput(_Workers): + """Workers whose generation worker's connection fails once its first output is relayed.""" + + def __init__(self, context_phase): + super().__init__(context_phase) + self.output_relayed = asyncio.Event() + + async def generation_stream(self, received: ResponsesRequest): + self.frames = _generation_frames(received)[:3] + for frame in self.frames: + yield frame.encode() + await self.output_relayed.wait() + raise aiohttp.ClientPayloadError("Response payload is not completed") + + def on_send(self, message): + """An ``on_send`` that notes the first output reaching the client.""" + if message.get("body", b"").startswith(b"event: response.output_text.delta\n"): + self.output_relayed.set() + + +class _WorkersStillGenerating(_Workers): + """Workers whose generation worker sends its opening and then goes on generating.""" + + async def generation_stream(self, received: ResponsesRequest): + self.frames = _generation_frames(received)[:2] + for frame in self.frames: + yield frame.encode() + await asyncio.Event().wait() + + def _responses_server(schedule_style: str, workers: _Workers) -> OpenAIDisaggServer: """A disaggregated server in front of ``workers``, with its routes.""" config = DisaggServerConfig(server_configs=[], schedule_style=schedule_style) @@ -1096,6 +1136,76 @@ async def context_phase(): assert (failed["error"]["code"], failed["error"]["message"]) == ("server_error", message) +@pytest.mark.asyncio +@pytest.mark.parametrize("schedule_style", _SCHEDULE_STYLES) +async def test_an_opened_stream_whose_generation_stream_fails_ends_with_response_failed( + monkeypatch, schedule_style +): + _short_waits(monkeypatch) + + async def context_phase(): + # Outlasts the start wait. + await asyncio.sleep(0.2) + + workers = _WorkersFailingAfterTheFirstOutput(context_phase) + sent = await _post_responses( + _responses_server(schedule_style, workers).app, on_send=workers.on_send + ) + + assert _status(sent) == 200 + # The stream ends after its events, as any stream does. + assert sent[-1] == {"type": "http.response.body", "body": b"", "more_body": False} + events = _sse_events(sent) + assert events[1][0] == "response.in_progress" + assert [event for event, _ in events if event != "response.in_progress"] == [ + "response.created", + "response.output_text.delta", + "error", + "response.failed", + ] + assert [data["sequence_number"] for _, data in events] == list(range(len(events))) + assert (events[-2][1]["code"], events[-2][1]["message"]) == ( + "upstream_error", + "ClientPayloadError", + ) + # The failed response is the one the stream opened with. + opened, failed = events[0][1]["response"], events[-1][1]["response"] + assert (failed["id"], failed["created_at"]) == (opened["id"], opened["created_at"]) + assert (failed["status"], failed["error"]["code"]) == ("failed", "server_error") + + +@pytest.mark.asyncio +@pytest.mark.parametrize("schedule_style", _SCHEDULE_STYLES) +async def test_a_relayed_generation_stream_that_fails_ends_with_an_error_event(schedule_style): + """After a fast context phase the worker's own opening is relayed, so only error is added.""" + + async def context_phase(): + pass + + workers = _WorkersFailingAfterTheFirstOutput(context_phase) + sent = [] + + def on_send(message): + sent.append(message) + workers.on_send(message) + + with contextlib.suppress(aiohttp.ClientPayloadError): + await _post_responses(_responses_server(schedule_style, workers).app, on_send=on_send) + + events = _sse_events(sent) + assert [event for event, _ in events] == [ + "response.created", + "response.in_progress", + "response.output_text.delta", + "error", + ] + assert [data["sequence_number"] for _, data in events] == [0, 1, 2, 3] + assert (events[-1][1]["code"], events[-1][1]["message"]) == ( + "upstream_error", + "ClientPayloadError", + ) + + @pytest.mark.asyncio @pytest.mark.parametrize("schedule_style", _SCHEDULE_STYLES) async def test_a_non_streamed_response_still_waits_for_the_whole_pipeline( @@ -1239,6 +1349,52 @@ async def receive(): release_context.set() +@pytest.mark.asyncio +@pytest.mark.parametrize("schedule_style", _SCHEDULE_STYLES) +@pytest.mark.parametrize("asgi_version", ["2.3", "2.4"]) +async def test_disconnect_before_the_responses_body_starts_closes_the_stream_it_answered_with( + monkeypatch, schedule_style, asgi_version +): + """The pipeline answers with a stream after the response is returned, which no body reads.""" + from starlette.requests import ClientDisconnect + + _short_waits(monkeypatch) + release_context = asyncio.Event() + workers = _WorkersStillGenerating(_waits_for(release_context)) + handler = _responses_route(_responses_server(schedule_style, workers)) + request = ResponsesRequest(model=_MODEL, input="hi", stream=True) + response = await asyncio.wait_for(handler(request, _raw_request()), 2) + pipeline = response._pipeline + answered = asyncio.Event() + sent = [] + + async def send(message): + sent.append(message) + # The pipeline answers while the response starts, before the body does. + release_context.set() + await asyncio.wait((pipeline,), timeout=2) + answered.set() + if asgi_version == "2.4": + raise OSError("client disconnected") + await asyncio.Event().wait() + + async def receive(): + await answered.wait() + return {"type": "http.disconnect"} + + scope = {"type": "http", "asgi": {"version": "3.0", "spec_version": asgi_version}} + if asgi_version == "2.4": + with pytest.raises(ClientDisconnect): + await response(scope, receive, send) + else: + await response(scope, receive, send) + + assert [message["type"] for message in sent] == ["http.response.start"] + assert isinstance(pipeline.result(), StreamingResponse) + # Closed by the time the response is done, while the generation worker goes on. + assert workers.generation_closed.is_set() + + def _frame(event_type: str, sequence_number: int, **fields) -> str: data = {"type": event_type, "sequence_number": sequence_number, **fields} return f"event: {event_type}\ndata: {json.dumps(data, separators=(',', ':'))}\n\n" From f83682689ae98ff7f53c41c9c98c24f06c2d6f0d Mon Sep 17 00:00:00 2001 From: JunyiXu-nv <219237550+JunyiXu-nv@users.noreply.github.com> Date: Tue, 3 Nov 2026 11:19:04 +0000 Subject: [PATCH 5/5] [None][fix] Disaggregated Responses: keep context_length_exceeded on a late failure; always forward request_id A context-length rejection that arrives after an opened Responses stream has started now ends the stream with error and response.failed carrying context_length_exceeded, the error response's own code, instead of upstream_error and server_error, so agent clients can compact their history and retry. response.failed is built the same way the in-stream over-length rejection builds it, because the SDK's ResponseError does not list that code. The body forwarded to a worker now always includes a Responses request's request_id, replacing a self-assignment that only worked through pydantic's fields-set tracking. Signed-off-by: JunyiXu-nv <219237550+JunyiXu-nv@users.noreply.github.com> --- tensorrt_llm/serve/openai_client.py | 12 ++-- tensorrt_llm/serve/openai_disagg_server.py | 37 ++++++++----- tensorrt_llm/serve/responses_utils.py | 55 +++++++++++++++---- .../test_disagg_responses_api.py | 37 +++++++++++-- 4 files changed, 109 insertions(+), 32 deletions(-) diff --git a/tensorrt_llm/serve/openai_client.py b/tensorrt_llm/serve/openai_client.py index e67ec2e6b371..9aff10b4486e 100644 --- a/tensorrt_llm/serve/openai_client.py +++ b/tensorrt_llm/serve/openai_client.py @@ -109,11 +109,15 @@ def _encode_request_body(request: _SerializableRequest) -> bytes: """Encode a request forwarded to a worker. Fields go out under their wire names, e.g. ``schema``, which the model - holds as ``schema_``. + holds as ``schema_``. Only fields that were set go out, except a Responses + request's ``request_id``, which always does, even when it is the default + generated when the request was parsed: a worker names its response after + it, so every worker answers under the id this server holds. """ - return _msgpack_encoder.encode( - request.model_dump(mode="json", exclude_unset=True, by_alias=True) - ) + body = request.model_dump(mode="json", exclude_unset=True, by_alias=True) + if isinstance(request, ResponsesRequest): + body["request_id"] = request.request_id + return _msgpack_encoder.encode(body) _ResponseT = TypeVar("_ResponseT", bound=BaseModel) diff --git a/tensorrt_llm/serve/openai_disagg_server.py b/tensorrt_llm/serve/openai_disagg_server.py index fe9b71ad5585..576bbaa3f717 100644 --- a/tensorrt_llm/serve/openai_disagg_server.py +++ b/tensorrt_llm/serve/openai_disagg_server.py @@ -69,6 +69,7 @@ PerfMetricsMiddleware, combine_disagg_metrics) from tensorrt_llm.serve.responses_utils import ( + STREAM_TERMINATION_CONTEXT_LENGTH_EXCEEDED, STREAM_TERMINATION_UPSTREAM_ERROR, ClosingStreamingResponse, ResponsesStreamingProcessor, ServerArrivalTimeMiddleware, get_steady_clock_now_in_seconds, guard_responses_stream, @@ -237,13 +238,20 @@ def _guard_relayed_responses_stream(answer: Response) -> Response: media_type="text/event-stream") -def _error_response_message(response: Response) -> str: - """The message of an error response, for a stream to report it.""" +def _error_response_failure(response: Response) -> tuple[str, str]: + """The cause and message a stream reports an error response with. + + An over-length rejection keeps its code, context_length_exceeded, which + clients act on; any other error response is an upstream error. + """ try: payload = json.loads(response.body) - return str(payload.get("message") or payload.get("detail") or payload) + message = str(payload.get("message") or payload.get("detail") or payload) except (json.JSONDecodeError, AttributeError, TypeError): - return "Internal server error" + return STREAM_TERMINATION_UPSTREAM_ERROR, "Internal server error" + if payload.get("code") == CONTEXT_LENGTH_EXCEEDED_CODE: + return STREAM_TERMINATION_CONTEXT_LENGTH_EXCEEDED, message + return STREAM_TERMINATION_UPSTREAM_ERROR, message def _upstream_error_message(error: aiohttp.ClientResponseError) -> str: @@ -529,10 +537,6 @@ async def openai_responses(self, request: ResponsesRequest, responses = self._wrap_entry_point(self._service.openai_responses) if not request.stream: return await responses(request, raw_request) - # The workers name their response after request_id, which is forwarded - # only if set. Assigning it marks it set, so that the workers answer - # under the id that an opening sent from here names. - request.request_id = request.request_id pipeline = asyncio.ensure_future(responses(request, raw_request)) try: await asyncio.wait((pipeline, ), @@ -544,6 +548,8 @@ async def openai_responses(self, request: ResponsesRequest, if pipeline.done(): return _guard_relayed_responses_stream(pipeline.result()) try: + # The opening names request.request_id, the id every worker is sent + # (_encode_request_body) and answers under. processor = ResponsesStreamingProcessor( request=request, sampling_params=request.to_sampling_params(), @@ -572,9 +578,11 @@ async def _responses_stream_when_done( so that every keepalive the client gets repeats ``opening``'s response.in_progress. The status line is already out, so a pipeline that fails ends the stream with error and response.failed, which carry - the message of its error response, and so does a pipeline's stream that - fails before its terminal event. A stream abandoned before the pipeline - answers, as when the client disconnects, cancels the pipeline. + the message of its error response, and for an over-length rejection + its code, context_length_exceeded (_error_response_failure); so does a + pipeline's stream that fails before its terminal event. A stream + abandoned before the pipeline answers, as when the client disconnects, + cancels the pipeline. """ try: for frame in opening: @@ -582,15 +590,16 @@ async def _responses_stream_when_done( try: answer = await pipeline except HTTPException as e: - failure = str(e.detail) + failure = (STREAM_TERMINATION_UPSTREAM_ERROR, str(e.detail)) else: failure = (None if isinstance(answer, StreamingResponse) else - _error_response_message(answer)) + _error_response_failure(answer)) finally: pipeline.cancel() if failure is not None: + cause, detail = failure for frame in processor.get_stream_failed_events( - STREAM_TERMINATION_UPSTREAM_ERROR, failure, len(opening)): + cause, detail, len(opening)): yield frame return # The guard reads the worker's stream as it comes, so it knows whether diff --git a/tensorrt_llm/serve/responses_utils.py b/tensorrt_llm/serve/responses_utils.py index 603a90fe4df0..1c6cc7248d43 100644 --- a/tensorrt_llm/serve/responses_utils.py +++ b/tensorrt_llm/serve/responses_utils.py @@ -17,6 +17,7 @@ Union) import anyio +import pydantic_core from fastapi.responses import StreamingResponse from jinja2.exceptions import TemplateError from openai.types.responses import (ResponseCompletedEvent, @@ -59,6 +60,7 @@ from tensorrt_llm._utils import AdjustedSteadyClock from tensorrt_llm.executor import (EngineDeadError, GenerationResult, RequestError) +from tensorrt_llm.executor.utils import CONTEXT_LENGTH_EXCEEDED_CODE from tensorrt_llm.inputs import PromptInputs, prompt_inputs from tensorrt_llm.inputs.multimodal import MultimodalServerConfig from tensorrt_llm.inputs.utils import resolve_hf_chat_template @@ -3009,6 +3011,35 @@ def _stream_terminal_event( ) +def _response_failed_frame(snapshot: dict, code: str, message: str, + sequence_number: int) -> str: + """A ``response.failed`` frame whose snapshot fails with ``code``. + + ``snapshot`` is a ``ResponsesResponse`` dump, which has no ``error``; the + SDK's Response, which the event carries, does. The SDK's ``ResponseError`` + does not list ``context_length_exceeded``, the code OpenAI fails an + over-length input with and the one clients act on. So the event is + validated with ``server_error`` in ``error.code``; any other code is then + written into its serialized form, with the serializer every frame uses. + """ + event = ResponseFailedEvent( + type="response.failed", + sequence_number=sequence_number, + response={ + **snapshot, "error": { + "code": "server_error", + "message": message + } + }, + ) + if code == "server_error": + return _format_sse_event(event) + payload = event.model_dump(mode="json", by_alias=True) + payload["response"]["error"]["code"] = code + return (f"event: {event.type}\n" + f"data: {pydantic_core.to_json(payload).decode()}\n\n") + + class ResponsesStreamingProcessor: def __init__( @@ -3133,7 +3164,9 @@ def get_stream_failed_events( ``events_sent`` is the number of frames that reached the wire, which the terminal events are numbered after; without it they continue this - processor's own numbering. + processor's own numbering. The failed snapshot's ``error.code`` is + ``context_length_exceeded`` for an over-length rejection and + ``server_error`` for every other cause. """ if events_sent is not None: self.sequence_number = events_sent @@ -3154,15 +3187,12 @@ def get_stream_failed_events( status="failed", usage=None, ).model_dump(by_alias=True) - # Set on the dump: the SDK Response carries `error`, ResponsesResponse - # does not. - snapshot["error"] = {"code": "server_error", "message": detail} - failed_event = self._send_event( - ResponseFailedEvent( - type="response.failed", - sequence_number=-1, - response=snapshot, - )) + code = (CONTEXT_LENGTH_EXCEEDED_CODE + if cause == STREAM_TERMINATION_CONTEXT_LENGTH_EXCEEDED else + "server_error") + failed_event = _response_failed_frame(snapshot, code, detail, + self.sequence_number) + self.sequence_number += 1 return [error_event, failed_event] @@ -3174,6 +3204,11 @@ def get_stream_failed_events( STREAM_TERMINATION_ENGINE_ERROR = "engine_error" STREAM_TERMINATION_UPSTREAM_ERROR = "upstream_error" STREAM_TERMINATION_INTERNAL_ERROR = "internal_error" +# The engine rejected an input longer than the context window. This answers +# the request rather than breaking it: the stream fails with OpenAI's code for +# it, which clients act on (the Codex CLI compacts its history), and then ends +# normally. +STREAM_TERMINATION_CONTEXT_LENGTH_EXCEEDED = CONTEXT_LENGTH_EXCEEDED_CODE _SSE_EVENT_DELIMITER = "\n\n" # A terminal event frame always follows another frame, and JSON payloads diff --git a/tests/unittest/disaggregated/test_disagg_responses_api.py b/tests/unittest/disaggregated/test_disagg_responses_api.py index 5db78e9afd06..a03afd26b8cb 100644 --- a/tests/unittest/disaggregated/test_disagg_responses_api.py +++ b/tests/unittest/disaggregated/test_disagg_responses_api.py @@ -541,6 +541,21 @@ async def send(request, hooks): assert (sent[0].prompt_token_ids, sent[0].prompt_token_ids_b64) == (None, None) +@pytest.mark.parametrize( + "fields", [{}, {"request_id": "resp_set_by_the_client"}], ids=["generated", "client_set"] +) +def test_every_worker_is_sent_the_request_id(fields): + """A worker names its response after request_id, so it goes out even when generated here. + + A stream this server opens itself names that id, and the workers must answer under it. + """ + request = ResponsesRequest(model="m", input="hi", **fields) + # The context request is a copy of the request, as the service makes it. + for forwarded in (request, request.model_copy(update={"stream": False})): + body = msgspec.msgpack.decode(_encode_request_body(forwarded)) + assert body["request_id"] == request.request_id + + def test_a_generation_first_worker_ignores_client_prompt_ids(): final = _responses_response(finish_reason="stop", status="completed") service = _gen_first_service(ctx_response=_responses_response(), gen_response=final) @@ -1094,9 +1109,23 @@ async def context_phase(): @pytest.mark.asyncio @pytest.mark.parametrize("schedule_style", _SCHEDULE_STYLES) -@pytest.mark.parametrize("status, worker_message", _CONTEXT_FAILURES) +@pytest.mark.parametrize( + "status, worker_message, error_code, failed_code", + [ + # An over-length rejection keeps its code in both events: clients act on it, + # as an agent compacts its history and retries. + pytest.param( + 400, + _OVER_LENGTH, + "context_length_exceeded", + "context_length_exceeded", + id="over_length_400", + ), + pytest.param(503, "no capacity left", "upstream_error", "server_error", id="503"), + ], +) async def test_responses_stream_with_a_late_context_failure_ends_with_response_failed( - monkeypatch, schedule_style, status, worker_message + monkeypatch, schedule_style, status, worker_message, error_code, failed_code ): _short_waits(monkeypatch) kept_alive = asyncio.Event() @@ -1129,11 +1158,11 @@ async def context_phase(): error_response = json.loads(_body(not_streamed)) message = error_response.get("detail") or error_response["message"] assert worker_message in message - assert (events[-2][1]["code"], events[-2][1]["message"]) == ("upstream_error", message) + assert (events[-2][1]["code"], events[-2][1]["message"]) == (error_code, message) failed = events[-1][1]["response"] assert failed["id"] == events[0][1]["response"]["id"] assert failed["status"] == "failed" - assert (failed["error"]["code"], failed["error"]["message"]) == ("server_error", message) + assert (failed["error"]["code"], failed["error"]["message"]) == (failed_code, message) @pytest.mark.asyncio