fix(bedrock_mantle): send max_completion_tokens for gpt-5.6 and gpt-6 chat completions

litellm/llms/bedrock_mantle/chat/transformation.py

@@ -119,6 +119,29 @@ def get_supported_openai_params(self, model: str) -> list:
         )
         return [*base_params, *extra_params]
 
+    def map_openai_params(
+        self,
+        non_default_params: dict,
+        optional_params: dict,
+        model: str,
+        drop_params: bool,
+        replace_max_completion_tokens_with_max_tokens: bool = True,
+    ) -> dict:
+        if not is_gpt_reasoning_series_name(model):
+            return super().map_openai_params(
+                non_default_params, optional_params, model, drop_params, replace_max_completion_tokens_with_max_tokens
+            )
+        mapped_params: Final = super().map_openai_params(
+            non_default_params, optional_params, model, drop_params, replace_max_completion_tokens_with_max_tokens=False
+        )
+        if "max_tokens" not in mapped_params:
+            return mapped_params
+        return {
+            param: value
+            for param, value in (("max_completion_tokens", mapped_params["max_tokens"]), *mapped_params.items())
+            if param != "max_tokens"
+        }
+
     def _supports_reasoning(self, model: str) -> bool:
         try:
             return litellm.supports_reasoning(model=model, custom_llm_provider=self.custom_llm_provider)

tests/e2e/llm_translation/test_bedrock_mantle_output_limit_e2e.py

@@ -0,0 +1,156 @@
+"""Live e2e for output-token limits on a Bedrock Mantle GPT-5.6 chat deployment.
+
+A deployment registered with model_info.mode "chat" and an explicit Mantle
+OpenAI-compatible api_base is served through Mantle chat completions. Mantle
+rejects max_tokens for the GPT-5.x models there and only accepts
+max_completion_tokens, so a caller that sets an output limit, or a deployment
+that stores one, must still get a completion back instead of the upstream 400
+"Unsupported parameter: 'max_tokens' is not supported with this model."
+"""
+
+from __future__ import annotations
+
+from collections.abc import Iterator
+from typing import Final
+
+import pytest
+from pydantic import BaseModel
+from e2e_config import unique_marker
+from lifecycle import ResourceManager
+from models import (
+    ChatBody,
+    ChatMessage,
+    ChatStreamOptions,
+    ChatTool,
+    ChatToolFunction,
+    LiteLLMParamsBody,
+    ModelInfoBody,
+    ModelNewBody,
+)
+from passthrough_client import PassthroughClient
+
+pytestmark = pytest.mark.e2e
+
+MANTLE_BACKEND: Final = "bedrock_mantle/openai.gpt-5.6-terra"
+MANTLE_CHAT_API_BASE: Final = "https://bedrock-mantle.us-east-1.api.aws/openai/v1"
+OUTPUT_LIMIT: Final = 1000
+WEATHER_TOOL: Final = ChatTool(
+    function=ChatToolFunction(
+        name="get_weather",
+        description="Get the weather for a city",
+        parameters={"type": "object", "properties": {"city": {"type": "string"}}, "required": ["city"]},
+    )
+)
+
+
+def _register_mantle_chat_model(
+    client: PassthroughClient, resources: ResourceManager, max_tokens: int | None = None
+) -> str:
+    model = f"e2e-mantle-chat-output-limit-{unique_marker()}"
+    model_id = client.proxy.register_model(
+        ModelNewBody(
+            model_name=model,
+            litellm_params=LiteLLMParamsBody(
+                model=MANTLE_BACKEND,
+                custom_llm_provider="bedrock_mantle",
+                api_key="os.environ/AWS_BEARER_TOKEN_BEDROCK",
+                api_base=MANTLE_CHAT_API_BASE,
+                reasoning_effort="none",
+                max_tokens=max_tokens,
+            ),
+            model_info=ModelInfoBody(mode="chat", supported_endpoints=["/v1/chat/completions", "/v1/responses"]),
+        )
+    )
+    resources.defer(lambda: client.proxy.delete_model(model_id))
+    return model
+
+
+def _prompt() -> list[ChatMessage]:
+    return [ChatMessage(role="user", content="reply with one word")]
+
+
+def _weather_prompt() -> list[ChatMessage]:
+    return [ChatMessage(role="user", content="What is the weather in Paris? Use the get_weather tool.")]
+
+
+class _StreamToolCallFunction(BaseModel):
+    name: str | None = None
+    arguments: str | None = None
+
+
+class _StreamToolCall(BaseModel):
+    function: _StreamToolCallFunction = _StreamToolCallFunction()
+
+
+class _StreamDelta(BaseModel):
+    tool_calls: list[_StreamToolCall] | None = None
+
+
+class _StreamChoice(BaseModel):
+    delta: _StreamDelta = _StreamDelta()
+
+
+class _StreamChunk(BaseModel):
+    choices: list[_StreamChoice] = []
+
+
+class _WeatherArgs(BaseModel):
+    city: str
+
+
+def _streamed_tool_calls(events: list[str]) -> Iterator[_StreamToolCall]:
+    for event in events:
+        for choice in _StreamChunk.model_validate_json(event).choices:
+            yield from choice.delta.tool_calls or []
+
+
+def _streamed_tool_call(events: list[str]) -> tuple[str, str]:
+    calls = tuple(_streamed_tool_calls(events))
+    name = "".join(call.function.name or "" for call in calls)
+    arguments = "".join(call.function.arguments or "" for call in calls)
+    return name, arguments
+
+
+class TestBedrockMantleChatOutputLimit:
+    def test_client_max_tokens_returns_completion(self, client: PassthroughClient, resources: ResourceManager) -> None:
+        model = _register_mantle_chat_model(client, resources)
+        result = client.proxy.transport.send(
+            "/chat/completions",
+            headers=client.proxy.transport.bearer(resources.key()),
+            json=ChatBody(model=model, messages=_prompt(), max_tokens=OUTPUT_LIMIT),
+        )
+        assert result.ok, f"chat call failed: {result.status_code} {result.body[:400]}"
+
+    def test_client_max_completion_tokens_returns_completion(
+        self, client: PassthroughClient, resources: ResourceManager
+    ) -> None:
+        model = _register_mantle_chat_model(client, resources)
+        result = client.proxy.transport.send(
+            "/chat/completions",
+            headers=client.proxy.transport.bearer(resources.key()),
+            json=ChatBody(model=model, messages=_prompt(), max_completion_tokens=OUTPUT_LIMIT),
+        )
+        assert result.ok, f"chat call failed: {result.status_code} {result.body[:400]}"
+
+    def test_deployment_max_tokens_streams_tool_call_completion(
+        self, client: PassthroughClient, resources: ResourceManager
+    ) -> None:
+        model = _register_mantle_chat_model(client, resources, max_tokens=OUTPUT_LIMIT)
+        result = client.proxy.chat_stream(
+            resources.key(),
+            ChatBody(
+                model=model,
+                messages=_weather_prompt(),
+                stream=True,
+                stream_options=ChatStreamOptions(include_usage=True),
+                tools=[WEATHER_TOOL],
+            ),
+        )
+        assert result.ok and result.stream_error is None and result.stream_done, (
+            f"streamed chat call failed: {result.status_code} {result.stream_error} {result.body[:400]}"
+        )
+        name, arguments = _streamed_tool_call(result.stream_events)
+        assert name == "get_weather", f"streamed tool call named {name!r}: {result.stream_events[:5]}"
+        assert _WeatherArgs.model_validate_json(arguments).city.strip(), (
+            f"streamed tool call arguments missing city: {arguments!r}"
+        )

tests/e2e/models.py

@@ -1265,6 +1265,10 @@ class LiteLLMParamsBody(BaseModel):
     input_cost_per_token_flex: float | None = None
     output_cost_per_token_flex: float | None = None
     cache_read_input_token_cost_flex: float | None = None
+    max_tokens: int | None = None
+    max_completion_tokens: int | None = None
+    custom_llm_provider: str | None = None
+    reasoning_effort: str | None = None
     extra_headers: dict[str, str] | None = None
     use_in_pass_through: bool | None = None
     complexity_router_config: dict[str, object] | None = None
@@ -1282,7 +1286,7 @@ class LiteLLMParamsBody(BaseModel):
     order: int | None = None
 
 
-ModelMode = Literal["batch", "realtime", "image_generation"]
+ModelMode = Literal["chat", "batch", "realtime", "image_generation"]
 
 
 class ModelInfoBody(BaseModel):
@@ -1292,6 +1296,7 @@ class ModelInfoBody(BaseModel):
     # constraint when a prior run's teardown had not removed the row.
     id: str | None = None
     mode: ModelMode | None = None
+    supported_endpoints: list[str] | None = None
     access_groups: list[str] | None = None
     team_id: str | None = None
     allowed_fails: int | None = None

tests/unit/llms/bedrock_mantle/test_bedrock_mantle_transformation.py

@@ -255,6 +255,41 @@ def test_verbosity_passes_through_for_gpt_5_models(self):
         )
         assert optional_params["verbosity"] == "low"
 
+    @pytest.mark.parametrize(
+        "model", ["openai.gpt-5.6-terra", "us-gov-west-1/openai.gpt-5.6-terra", "openai.gpt-6-sol"]
+    )
+    @pytest.mark.parametrize("limit_param", ["max_tokens", "max_completion_tokens"])
+    def test_gpt_reasoning_series_output_limit_is_sent_as_max_completion_tokens(self, model, limit_param):
+        optional_params = litellm.get_optional_params(
+            model=model,
+            custom_llm_provider="bedrock_mantle",
+            drop_params=False,
+            **{limit_param: 1000},
+        )
+        assert optional_params["max_completion_tokens"] == 1000
+        assert "max_tokens" not in optional_params
+
+    def test_gpt_reasoning_series_explicit_max_completion_tokens_wins_over_max_tokens(self):
+        optional_params = litellm.get_optional_params(
+            model="openai.gpt-5.6-terra",
+            custom_llm_provider="bedrock_mantle",
+            drop_params=False,
+            max_tokens=1000,
+            max_completion_tokens=50,
+        )
+        assert optional_params["max_completion_tokens"] == 50
+        assert "max_tokens" not in optional_params
+
+    def test_non_gpt_reasoning_models_keep_the_openai_compatible_max_tokens(self):
+        optional_params = litellm.get_optional_params(
+            model="openai.gpt-oss-120b",
+            custom_llm_provider="bedrock_mantle",
+            drop_params=False,
+            max_completion_tokens=1000,
+        )
+        assert optional_params["max_tokens"] == 1000
+        assert "max_completion_tokens" not in optional_params
+
 
 class TestBedrockMantleChatAuth:
     """Chat Completions must use the same Bearer-or-SigV4 auth as the Responses