[None][fix] Responses: fail over-length streamed requests in-stream with context_length_exceeded

tensorrt_llm/serve/openai_disagg_server.py

@@ -69,10 +69,10 @@
                                              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_CONTEXT_LENGTH_EXCEEDED, ResponsesStreamingProcessor,
+    ServerArrivalTimeMiddleware, get_steady_clock_now_in_seconds,
+    guard_responses_stream, rejected_responses_stream, 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
@@ -213,6 +213,43 @@ def _upstream_error_message(error: aiohttp.ClientResponseError) -> str:
     return f"context worker rejected the request with status {error.status}"
 
 
+def _context_length_exceeded_message(
+        response: Optional[Response]) -> Optional[str]:
+    """The message of a worker's ``context_length_exceeded`` envelope, else None.
+
+    ``_handle_exception`` answers that rejection, and only that one, with the
+    worker's 4xx error envelope as a JSONResponse; every other failure raises,
+    or returns None for a ``CppExecutorError``.
+    """
+    if not isinstance(response, JSONResponse):
+        return None
+    if not 400 <= response.status_code < 500:
+        return None
+    body = json.loads(response.body)
+    if isinstance(body, dict) and body.get(
+            "code") == CONTEXT_LENGTH_EXCEEDED_CODE:
+        return str(body.get("message", ""))
+    return None
+
+
+def _context_length_exceeded_stream(request: ResponsesRequest,
+                                    message: str) -> StreamingResponse:
+    """Fail a streamed Responses request in-stream, as an aggregated server does.
+
+    The context phase never streams, so its worker rejects an over-length
+    prompt with a 400. The snapshot names the model the client asked for; the
+    frontend does not know the served name.
+    """
+    processor = ResponsesStreamingProcessor(
+        request=request,
+        sampling_params=request.to_sampling_params(),
+        model_name=request.model,
+        use_harmony=False)
+    return StreamingResponse(content=rejected_responses_stream(
+        processor, STREAM_TERMINATION_CONTEXT_LENGTH_EXCEEDED, message),
+                             media_type="text/event-stream")
+
+
 class OpenAIDisaggServer:
     def __init__(self,
                  config: DisaggServerConfig,
@@ -461,7 +498,16 @@ async def wrapper(req: request_type, raw_req: Request) -> Response:
             except Exception as e:
                 # Usually raises; returns a Response for worker errors that
                 # carry a machine-readable code (context_length_exceeded).
-                return self._handle_exception(e)
+                error_response = self._handle_exception(e)
+                if isinstance(req, ResponsesRequest) and req.stream:
+                    # The handler, which a subclass may override, has ruled
+                    # on the failure. Only its context-length envelope is
+                    # re-shaped: a streamed Responses request is answered
+                    # in-stream, as on an aggregated server.
+                    message = _context_length_exceeded_message(error_response)
+                    if message is not None:
+                        return _context_length_exceeded_stream(req, message)
+                return error_response
         return wrapper
 
     async def anthropic_messages(self, request: AnthropicMessagesRequest,

tensorrt_llm/serve/openai_server.py

@@ -39,7 +39,7 @@
 from tensorrt_llm._torch.async_llm import AsyncLLM
 from tensorrt_llm._utils import EnergyMonitor
 # yapf: disable
-from tensorrt_llm.executor import CppExecutorError
+from tensorrt_llm.executor import CppExecutorError, RequestError
 from tensorrt_llm.executor.postproc_worker import PostprocArgs, PostprocParams
 from tensorrt_llm.executor.request import DEFAULT_REQUEST_PRIORITY
 from tensorrt_llm.executor.rpc import RPCError
@@ -129,14 +129,16 @@
                                        prepare_chat_request, render_chat,
                                        render_chat_for_length,
                                        render_conversation)
-from tensorrt_llm.serve.responses_utils import (ClosingStreamingResponse,
-                                                ConversationHistoryStore,
-                                                ResponsesStreamingProcessor,
-                                                ServerArrivalTimeMiddleware)
+from tensorrt_llm.serve.responses_utils import (
+    STREAM_TERMINATION_CONTEXT_LENGTH_EXCEEDED, ClosingStreamingResponse,
+    ConversationHistoryStore, ResponsesStreamingProcessor,
+    ServerArrivalTimeMiddleware, classify_stream_termination)
 from tensorrt_llm.serve.responses_utils import \
     create_response as responses_api_create_response
-from tensorrt_llm.serve.responses_utils import (guard_responses_stream,
-                                                keep_responses_stream_alive)
+from tensorrt_llm.serve.responses_utils import (describe_stream_termination,
+                                                guard_responses_stream,
+                                                keep_responses_stream_alive,
+                                                rejected_responses_stream)
 from tensorrt_llm.serve.responses_utils import \
     request_preprocess as responses_api_request_preprocess
 from tensorrt_llm.serve.responses_web_search import web_search_rejection_reason
@@ -3512,16 +3514,31 @@ async def create_streaming_generator(promise: RequestOutput,
                 request.disaggregated_params)
             conversation_params = to_llm_conversation_params(
                 request.conversation_params)
-            promise = self.generator.generate_async(
-                inputs=input_tokens,
-                sampling_params=sampling_params,
-                streaming=request.stream,
-                disaggregated_params=disaggregated_params,
-                conversation_params=conversation_params,
-                kv_hints=request.kv_hints,
-                _postproc_params=postproc_params
-                if self.postproc_worker_enabled else None,
-            )
+            try:
+                promise = self.generator.generate_async(
+                    inputs=input_tokens,
+                    sampling_params=sampling_params,
+                    streaming=request.stream,
+                    disaggregated_params=disaggregated_params,
+                    conversation_params=conversation_params,
+                    kv_hints=request.kv_hints,
+                    _postproc_params=postproc_params
+                    if self.postproc_worker_enabled else None,
+                )
+            except RequestError as error:
+                # An executor that runs in this process rejects an over-length
+                # prompt here; one in a worker process rejects it while the
+                # stream runs (guard_responses_stream). A streamed request
+                # fails in-stream either way.
+                cause = classify_stream_termination(error)
+                if (streaming_processor is None
+                        or cause != STREAM_TERMINATION_CONTEXT_LENGTH_EXCEEDED):
+                    raise
+                detail = describe_stream_termination(error, cause)
+                logger.error(f"Responses request rejected ({cause}): {detail}")
+                return StreamingResponse(content=rejected_responses_stream(
+                    streaming_processor, cause, detail),
+                                         media_type="text/event-stream")
             if not self.postproc_worker_enabled:
                 # Set by the executor only for postprocessing it runs itself.
                 postproc_args.num_prompt_tokens = len(promise.prompt_token_ids)

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,8 @@
 from tensorrt_llm._utils import AdjustedSteadyClock
 from tensorrt_llm.executor import (EngineDeadError, GenerationResult,
                                    RequestError)
+from tensorrt_llm.executor.utils import (CONTEXT_LENGTH_EXCEEDED_CODE,
+                                         is_context_length_exceeded_message)
 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
@@ -3015,6 +3018,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__(
@@ -3139,7 +3171,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
@@ -3160,15 +3194,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]
 
 
@@ -3180,6 +3211,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
@@ -3200,6 +3236,11 @@ def _sse_delimiter_and_markers(sample: Any) -> Tuple[Any, tuple]:
 
 def classify_stream_termination(exc: Exception) -> str:
     """Attribute a stream-ending exception to a cause."""
+    # An over-length rejection crosses the executor IPC boundary as message
+    # text, so it is recognized by that text.
+    if isinstance(exc, RequestError) and is_context_length_exceeded_message(
+            str(exc)):
+        return STREAM_TERMINATION_CONTEXT_LENGTH_EXCEEDED
     if isinstance(exc, (RequestError, EngineDeadError)):
         return STREAM_TERMINATION_ENGINE_ERROR
     # By module, to keep the HTTP client library out of this layer.
@@ -3212,8 +3253,11 @@ def describe_stream_termination(exc: Exception, cause: str) -> str:
     """The client-facing detail of a stream-ending exception.
 
     Only engine errors carry their message; transport and internal errors can
-    name internal hosts.
+    name internal hosts. An over-length rejection is described by its message
+    alone, the text a non-streamed request is rejected with.
     """
+    if cause == STREAM_TERMINATION_CONTEXT_LENGTH_EXCEEDED:
+        return str(exc)
     if cause == STREAM_TERMINATION_ENGINE_ERROR and str(exc):
         return f"{type(exc).__name__}: {exc}"
     return type(exc).__name__
@@ -3240,8 +3284,11 @@ async def guard_responses_stream(
 ) -> AsyncGenerator[Any, None]:
     """End a stream that fails before its terminal event with events saying so.
 
-    The exception is re-raised and frames pass through untouched. A client hangup (CancelledError or
-    GeneratorExit) is not caught: there is nobody left to tell.
+    Frames pass through untouched, and the exception is re-raised after the
+    terminal events, except for an over-length rejection: those events answer
+    the request in full, so the stream ends normally. A client hangup
+    (CancelledError or GeneratorExit) is not caught: there is nobody left to
+    tell.
     """
     events_sent = 0
     tail = None
@@ -3290,9 +3337,29 @@ async def guard_responses_stream(
                 yield delimiter
         for frame in frames:
             yield frame
+        if frames and cause == STREAM_TERMINATION_CONTEXT_LENGTH_EXCEEDED:
+            # Re-raising would abort the connection, which a client reads as
+            # a transport error: one that settles on response.failed only
+            # when the stream ends (the Codex CLI does) never sees its code.
+            return
         raise
 
 
+async def rejected_responses_stream(processor: ResponsesStreamingProcessor,
+                                    cause: str,
+                                    detail: str) -> AsyncGenerator[str, None]:
+    """The stream of a request rejected before generation started.
+
+    It opens as any other stream does and fails at once, so a client reads the
+    rejection exactly as it reads one that ``guard_responses_stream`` reports
+    while a stream runs.
+    """
+    opening = processor.get_initial_responses()
+    for frame in opening + processor.get_stream_failed_events(
+            cause, detail, events_sent=len(opening)):
+        yield frame
+
+
 # --------------------------------------------------------------------------
 # Keepalive while a stream is silent
 # --------------------------------------------------------------------------