fix(router): read chat custom tool calls as tool turns, test the Responses and custom tier shapes

litellm/litellm_core_utils/prompt_templates/factory.py

@@ -6505,6 +6505,19 @@ class _ExchangeChatToolCall(BaseModel):
     function: _ExchangeFunction
 
 
+class _ExchangeCustomInput(BaseModel):
+    model_config = ConfigDict(frozen=True)
+
+    name: str
+    input: str = ""
+
+
+class _ExchangeChatCustomToolCall(BaseModel):
+    model_config = ConfigDict(frozen=True)
+
+    custom: _ExchangeCustomInput
+
+
 class _ExchangeToolUse(BaseModel):
     model_config = ConfigDict(frozen=True)
 
@@ -6522,7 +6535,9 @@ class _ExchangeToolResult(BaseModel):
 
 _EXCHANGE_PARTS: Final = TypeAdapter(tuple[object, ...])
 _EXCHANGE_TEXT_BLOCK: Final = TypeAdapter(_ExchangeTextBlock)
-_EXCHANGE_CHAT_TOOL_CALL: Final = TypeAdapter(_ExchangeChatToolCall)
+_EXCHANGE_CHAT_TOOL_CALL: Final[TypeAdapter[_ExchangeChatToolCall | _ExchangeChatCustomToolCall]] = TypeAdapter(
+    _ExchangeChatToolCall | _ExchangeChatCustomToolCall
+)
 _EXCHANGE_FUNCTION: Final = TypeAdapter(_ExchangeFunction)
 _EXCHANGE_BLOCK: Final[TypeAdapter[_ExchangeToolUse | _ExchangeToolResult]] = TypeAdapter(
     _ExchangeToolUse | _ExchangeToolResult
@@ -6560,6 +6575,14 @@ def _call_text(function: _ExchangeFunction) -> str:
     return f"{function.name}({arguments})"
 
 
+def _chat_call_text(call: _ExchangeChatToolCall | _ExchangeChatCustomToolCall) -> str:
+    match call:
+        case _ExchangeChatToolCall(function=function):
+            return _call_text(function)
+        case _ExchangeChatCustomToolCall(custom=custom):
+            return _call_text(_ExchangeFunction(name=custom.name, arguments=custom.input))
+
+
 def _block_exchange(part: object) -> tuple[ToolExchangeKind, str] | None:
     match _validated_or_none(_EXCHANGE_BLOCK, part):
         case _ExchangeToolUse(name=name, input=arguments):
@@ -6577,9 +6600,9 @@ def tool_exchange_texts(message: Mapping[str, object]) -> tuple[tuple[ToolExchan
     Chat completions puts calls on an assistant turn's `tool_calls` and each result on its own `tool` role,
     while Anthropic Messages puts both inside content blocks, `tool_use` on the assistant turn and
     `tool_result` on the next user turn, so this reads both places. The legacy `function_call` field and
-    `function` role, Anthropic's server tool calls, and its MCP tool blocks read the same way, while server tool
-    results such as `web_search_tool_result` are not read. Responses API input reaches callers already resolved
-    to the chat-completions shape by `resolve_structured_messages`.
+    `function` role, chat `custom` tool calls, Anthropic's server tool calls, and its MCP tool blocks read the
+    same way, while server tool results such as `web_search_tool_result` are not read. Responses API input
+    reaches callers already resolved to the chat-completions shape by `resolve_structured_messages`.
     """
     if message.get("role") in ("tool", "function"):
         result: Final = _exchange_content_text(message.get("content"))
@@ -6588,7 +6611,7 @@ def tool_exchange_texts(message: Mapping[str, object]) -> tuple[tuple[ToolExchan
         exchange for part in _exchange_parts(message.get("content")) if (exchange := _block_exchange(part)) is not None
     )
     chat_calls: Final[tuple[tuple[ToolExchangeKind, str], ...]] = tuple(
-        ("tool call", _call_text(call.function))
+        ("tool call", _chat_call_text(call))
         for part in _exchange_parts(message.get("tool_calls"))
         if (call := _validated_or_none(_EXCHANGE_CHAT_TOOL_CALL, part)) is not None
     )

tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py

@@ -8837,6 +8837,21 @@ def test_wrong_type_template_field_in_requests_leaves_the_deployments_in_rotatio
             (("tool call", 'search({"query": "café"})'),),
             id="chat_object_arguments",
         ),
+        pytest.param(
+            {
+                "role": "assistant",
+                "content": None,
+                "tool_calls": [
+                    {
+                        "id": "call_2",
+                        "type": "custom",
+                        "custom": {"name": "apply_patch", "input": "*** Begin Patch\n*** Update File: ledger.py"},
+                    }
+                ],
+            },
+            (("tool call", "apply_patch(*** Begin Patch\n*** Update File: ledger.py)"),),
+            id="chat_custom_tool_call",
+        ),
         pytest.param(
             {
                 "role": "assistant",

tests/unit/router_strategy/test_complexity_router.py

@@ -11857,6 +11857,81 @@ class TestAgentProgressReachesClassifier:
         assert held_decision.get("classifier_cost") == (ask_decision.get("classifier_cost") if include_tool_turns else None)
         assert held_decision.get("classifier_model") == (ask_decision.get("classifier_model") if include_tool_turns else None)
 
+    @pytest.mark.asyncio
+    @pytest.mark.parametrize(
+        "items",
+        [
+            pytest.param(
+                [
+                    {"type": "function_call", "call_id": "call_1", "name": "read_file", "arguments": '{"path": "TASK.md"}'},
+                    {"type": "function_call_output", "call_id": "call_1", "output": _TASK_FILE},
+                ],
+                id="function_call",
+            ),
+            pytest.param(
+                [
+                    {"type": "custom_tool_call", "call_id": "call_1", "name": "read_file", "input": "TASK.md"},
+                    {"type": "custom_tool_call_output", "call_id": "call_1", "output": _TASK_FILE},
+                ],
+                id="custom_tool_call",
+            ),
+        ],
+    )
+    async def test_responses_continuation_turn_routes_on_what_the_tool_returned(
+        self, items: list[dict[str, object]]
+    ) -> None:
+        """Codex sends the loop as Responses input items, so every_request must read the task file from them too."""
+        client: Final = _TaskAwareJevClient()
+        router: Final = self._jev_router(client, classifier_context_include_tool_turns=True)
+
+        response: Final = await router.async_pre_routing_hook(
+            model="test-router",
+            request_kwargs={"input": [{"role": "user", "content": [{"type": "input_text", "text": _TASK_ASK}]}, *items]},
+        )
+
+        assert response is not None
+        assert response.model == "strong"
+        assert "read_file(" in client.states[-1]
+
+    @pytest.mark.asyncio
+    async def test_user_turn_floor_follows_the_custom_tier_order(self) -> None:
+        """tier_definitions list order is the severity order, so a tool turn moves a user_turn loop up the defined
+        tiers and a cheaper verdict on a later turn leaves it there."""
+        dependency: Final = MagicMock(
+            cache=DualCache(),
+            acompletion=AsyncMock(
+                side_effect=[_llm_response(json.dumps({"tier": tier})) for tier in ("triage", "build", "triage")]
+            ),
+        )
+        router: Final = ComplexityRouter(
+            "test-router",
+            dependency,
+            {
+                "tiers": {"triage": "cheap", "build": "mid", "deep": "strong"},
+                "tier_definitions": [
+                    {"name": "triage", "description": "Reading files and answering questions about them"},
+                    {"name": "build", "description": "Writing and changing code"},
+                    {"name": "deep", "description": "Designing and migrating systems"},
+                ],
+                "fallback_tier": "triage",
+                "classifier_type": "llm",
+                "classifier_llm_config": {"model": "haiku-classifier", "timeout_ms": 400},
+                "classification_mode": "user_turn",
+                "classifier_context_include_tool_turns": True,
+            },
+        )
+        loop: Final = _tool_loop_continuation("chat_completions")
+
+        responses: Final = [
+            await router.async_pre_routing_hook(
+                model="test-router", request_kwargs={"metadata": {"session_id": "custom-loop"}}, messages=turn
+            )
+            for turn in (loop[:1], loop, [*loop, *_RUN_TESTS_TURNS])
+        ]
+
+        assert [response.model for response in responses if response is not None] == ["cheap", "mid", "mid"]
+        assert dependency.acompletion.await_count == 3
+
     @pytest.mark.asyncio
     async def test_user_turn_classifier_failure_keeps_the_held_tier(self) -> None:
         """A failed classifier decided nothing, so its fallback's verdict on the ask alone must not move the loop."""