[None][fix] Anthropic Messages: answer an over-length prompt with "prompt is too long", before any stream

tensorrt_llm/executor/utils.py

@@ -155,8 +155,20 @@ def context_length_exceeded_message(max_context_length: int,
 
 
 _CONTEXT_LENGTH_EXCEEDED_RE = re.compile(
-    r"This model's maximum context length is \d+ tokens\. "
-    r"However, your messages resulted in \d+ tokens\.")
+    r"This model's maximum context length is (\d+) tokens\. "
+    r"However, your messages resulted in (\d+) tokens\.")
+
+
+def context_length_exceeded_counts(message: str) -> tuple[int, int] | None:
+    """``(max_context_length, num_prompt_tokens)`` of a ``context_length_exceeded_message``.
+
+    None if ``message`` contains none. Searched for like
+    ``is_context_length_exceeded_message`` does.
+    """
+    match = _CONTEXT_LENGTH_EXCEEDED_RE.search(message)
+    if match is None:
+        return None
+    return int(match.group(1)), int(match.group(2))
 
 
 def is_context_length_exceeded_message(message: str) -> bool:

tensorrt_llm/serve/anthropic_adapter.py

@@ -24,6 +24,7 @@
 chat templates, or the engine.
 """
 
+import asyncio
 import base64
 import hashlib
 import json
@@ -35,6 +36,7 @@
 
 from fastapi.responses import JSONResponse
 
+from tensorrt_llm.executor.utils import RequestError, context_length_exceeded_counts
 from tensorrt_llm.logger import logger
 from tensorrt_llm.serve.anthropic_protocol import (
     AnthropicContentBlockDeltaEvent,
@@ -95,6 +97,99 @@ def anthropic_error_response(
     return JSONResponse(content=envelope.model_dump(exclude_none=True), status_code=status_code)
 
 
+def prompt_too_long_message(message: str) -> Optional[str]:
+    """Anthropic's wording for this server's over-length rejection, or None.
+
+    Anthropic answers a prompt longer than the context window with an
+    ``invalid_request_error`` that reads ``prompt is too long: N tokens > M
+    maximum``, and Claude Code compacts its conversation on that wording. The
+    server's own rejection (``context_length_exceeded_message``) is recognized
+    anywhere in ``message``. M is the longest prompt the server accepts, one
+    token below its context length, so the inequality also holds for a prompt
+    that fills the context exactly.
+    """
+    counts = context_length_exceeded_counts(message)
+    if counts is None:
+        return None
+    max_context_length, num_prompt_tokens = counts
+    return f"prompt is too long: {num_prompt_tokens} tokens > {max_context_length - 1} maximum"
+
+
+def anthropic_chat_error_response(message: str, status_code: int) -> JSONResponse:
+    """The Anthropic error for a failure the chat pipeline reported.
+
+    A 4xx is the client's ``invalid_request_error`` and anything else an
+    ``api_error``. A prompt longer than the context is answered as Anthropic
+    answers it: a 400 ``invalid_request_error`` in Anthropic's wording.
+    """
+    too_long = prompt_too_long_message(message)
+    if too_long is not None:
+        return anthropic_error_response(too_long, "invalid_request_error", 400)
+    error_type = "invalid_request_error" if 400 <= status_code < 500 else "api_error"
+    return anthropic_error_response(message, error_type, status_code)
+
+
+# Seconds a streamed request is held before its response starts, for the
+# engine's verdict on its prompt. The engine rejects a prompt longer than its
+# context when it admits the request, well within this; a first token can take
+# far longer, and the stream then starts without waiting for it.
+STREAM_ADMISSION_WAIT_S = 1.0
+
+
+async def wait_for_stream_start(
+    openai_sse: AsyncIterator[Union[str, bytes]],
+) -> Union[JSONResponse, AsyncIterator[Union[str, bytes]]]:
+    """Wait up to ``STREAM_ADMISSION_WAIT_S`` for the first frame of a chat stream.
+
+    An engine in a worker process rejects a prompt longer than its context only
+    as the failure of the stream's first frame, when a response started at once
+    would already be under way. Anthropic answers such a prompt with HTTP 400
+    before any stream, and Claude Code compacts its conversation on that error,
+    so the rejection is returned here as that 400. Otherwise the stream is
+    returned with its first frame, to be read as before; a rejection that
+    arrives after the wait ends the stream with an error event in the same
+    wording.
+    """
+    first = asyncio.ensure_future(openai_sse.__anext__())
+    try:
+        await asyncio.wait((first,), timeout=STREAM_ADMISSION_WAIT_S)
+    except asyncio.CancelledError:
+        first.cancel()
+        raise
+    if first.done() and not first.cancelled() and isinstance(first.exception(), RequestError):
+        too_long = prompt_too_long_message(str(first.exception()))
+        if too_long is not None:
+            return anthropic_error_response(too_long, "invalid_request_error", 400)
+    return _resume_stream(first, openai_sse)
+
+
+async def _resume_stream(
+    first: asyncio.Future, rest: AsyncIterator[Union[str, bytes]]
+) -> AsyncIterator[Union[str, bytes]]:
+    """Yield the frame ``first`` reads from ``rest``, then the remaining frames of ``rest``.
+
+    Cancelling a read that waits on ``first`` cancels ``first`` too, and
+    closing this stream closes ``rest``.
+    """
+    try:
+        try:
+            frame = await first
+        except StopAsyncIteration:
+            return
+        except RequestError as e:
+            too_long = prompt_too_long_message(str(e))
+            if too_long is None:
+                raise
+            raise AnthropicResponseError(too_long) from e
+        yield frame
+        async for frame in rest:
+            yield frame
+    finally:
+        aclose = getattr(rest, "aclose", None)
+        if aclose is not None:
+            await aclose()
+
+
 # ---------------------------------------------------------------------------
 # Request conversion: Anthropic -> ChatCompletionRequest
 # ---------------------------------------------------------------------------

tensorrt_llm/serve/openai_disagg_server.py

@@ -43,6 +43,7 @@
 from tensorrt_llm.serve._telemetry import create_uvicorn_server
 from tensorrt_llm.serve.anthropic_adapter import (AnthropicRequestError,
                                                   AnthropicResponseError,
+                                                  anthropic_chat_error_response,
                                                   anthropic_error_response,
                                                   convert_anthropic_request,
                                                   convert_chat_response,
@@ -472,10 +473,7 @@ async def anthropic_messages(self, request: AnthropicMessagesRequest,
             openai_response = await self._wrap_entry_point(
                 self._service.openai_chat_completion)(chat_request, raw_request)
         except HTTPException as e:
-            error_type = ("invalid_request_error"
-                          if 400 <= e.status_code < 500 else "api_error")
-            return anthropic_error_response(str(e.detail), error_type,
-                                            e.status_code)
+            return anthropic_chat_error_response(str(e.detail), e.status_code)
 
         if isinstance(openai_response, StreamingResponse):
             return StreamingResponse(
@@ -492,9 +490,10 @@ async def anthropic_messages(self, request: AnthropicMessagesRequest,
                            or payload.get("error") or json.dumps(payload))
             except (json.JSONDecodeError, AttributeError, TypeError):
                 message = "Internal server error"
-            error_type = ("invalid_request_error"
-                          if 400 <= status < 500 else "api_error")
-            return anthropic_error_response(str(message), error_type, status)
+            # The context phase never streams, so a context worker's over-length
+            # rejection arrives here, before any stream, for a streamed request
+            # too.
+            return anthropic_chat_error_response(str(message), status)
 
         try:
             chat_response = ChatCompletionResponse(

tensorrt_llm/serve/openai_server.py

@@ -68,9 +68,10 @@
     get_effective_kv_cache_event_hash_algo
 from tensorrt_llm.sampling_params import GuidedDecodingParams, SamplingParams
 from tensorrt_llm.serve.anthropic_adapter import (
-    AnthropicRequestError, AnthropicResponseError, anthropic_error_response,
+    AnthropicRequestError, AnthropicResponseError,
+    anthropic_chat_error_response, anthropic_error_response,
     convert_anthropic_count_tokens_request, convert_anthropic_request,
-    convert_chat_response, reframe_openai_stream)
+    convert_chat_response, reframe_openai_stream, wait_for_stream_start)
 from tensorrt_llm.serve.anthropic_batches import (AnthropicBatchStore,
                                                   BatchStoreFullError,
                                                   results_to_jsonl)
@@ -2550,9 +2551,13 @@ async def anthropic_messages(self, request: AnthropicMessagesRequest,
                                             "api_error", 500)
 
         if isinstance(response, StreamingResponse):
+            # A prompt the engine rejects as too long is answered with a 400,
+            # not with a stream (see wait_for_stream_start).
+            stream = await wait_for_stream_start(response.body_iterator)
+            if isinstance(stream, JSONResponse):
+                return stream
             return StreamingResponse(
-                content=reframe_openai_stream(response.body_iterator,
-                                              model=self.model),
+                content=reframe_openai_stream(stream, model=self.model),
                 media_type="text/event-stream",
             )
 
@@ -2563,9 +2568,7 @@ async def anthropic_messages(self, request: AnthropicMessagesRequest,
                 message = payload.get("message") or json.dumps(payload)
             except (json.JSONDecodeError, AttributeError):
                 message = "Internal server error"
-            err_type = ("invalid_request_error"
-                        if 400 <= status < 500 else "api_error")
-            return anthropic_error_response(message, err_type, status)
+            return anthropic_chat_error_response(message, status)
 
         try:
             chat_response = ChatCompletionResponse(**json.loads(response.body))