diff --git a/tensorrt_llm/serve/harmony_adapter.py b/tensorrt_llm/serve/harmony_adapter.py index f326fc74360e..6a37f28a26b5 100644 --- a/tensorrt_llm/serve/harmony_adapter.py +++ b/tensorrt_llm/serve/harmony_adapter.py @@ -417,6 +417,18 @@ def get_debug_info(self) -> dict[str, Any]: self.should_filter_tools } + def active_tool_call_names(self) -> list[str]: + """Names of tool calls still 'active' at end-of-request. + + A tool call that was started but never closed (a recipient or channel + transition did not occur) means the stream ended inside an incomplete + tool invocation, which would otherwise be returned to the caller as an + empty or truncated tool_calls (see #17437 / #16377).""" + return [ + info["name"] for info in self.tool_calls.values() + if info.get("active", True) + ] + def finalize_request(self) -> dict[str, Any] | None: """ Finalize the request and return any remaining closing token delta. @@ -1584,6 +1596,17 @@ def cleanup_stream_state(self, request_id: str) -> None: Call this when a request finishes to free memory. """ if request_id in self._stream_states: + state = self._stream_states[request_id] + incomplete = state.active_tool_call_names() + if incomplete: + logger.warning( + "Request %s finished with incomplete (never-closed) tool " + "call(s): %s. Often caused by speculative-decode " + "acceptance truncating a control-token prelude; the tool " + "call was silently dropped (see #17437 / #16377).", + request_id, + incomplete, + ) del self._stream_states[request_id] logger.debug(f"Cleaned up stream state for request {request_id}") diff --git a/tests/unittest/llmapi/apps/test_harmony_parsing.py b/tests/unittest/llmapi/apps/test_harmony_parsing.py index 8689454a9e14..d2d9500ecdc9 100644 --- a/tests/unittest/llmapi/apps/test_harmony_parsing.py +++ b/tests/unittest/llmapi/apps/test_harmony_parsing.py @@ -1335,5 +1335,84 @@ def test_none_tokenizer_num_postprocess_workers(): _logit_bias_to_embedding_bias({"0": 1.0}, vocab_size=None) +# =========================================================================== +# Fix 3 - incomplete (never-closed) tool call guard (see #17437 / #16377) +# =========================================================================== + + +class TestIncompleteToolCallGuard: + """Verify the end-of-request guardrail for unclosed tool calls. + + A tool call that was started but never closed (for example a + speculative-decode acceptance truncated a control prelude) must be + reported instead of silently dropped. + """ + + def _make_adapter_state(self): + from tensorrt_llm.serve.harmony_adapter import HarmonyAdapter + + adapter = HarmonyAdapter(harmony_input=False, harmony_output=False) + request_id = "test-guard-tool" + state = adapter.create_stream_state( + request_id=request_id, available_tools=None, tool_choice=None + ) + return adapter, request_id, state + + def test_active_tool_call_reported(self): + adapter, request_id, state = self._make_adapter_state() + try: + state.tool_calls["call_1"] = { + "id": "call_1", + "name": "get_weather", + "arguments": '{"city":', + "index": 0, + } + assert state.active_tool_call_names() == ["get_weather"] + finally: + adapter.cleanup_stream_state(request_id) + + def test_closed_tool_call_not_reported(self): + adapter, request_id, state = self._make_adapter_state() + try: + state.tool_calls["call_1"] = { + "id": "call_1", + "name": "get_weather", + "arguments": '{"city":"SF"}', + "index": 0, + "active": False, + } + assert state.active_tool_call_names() == [] + finally: + adapter.cleanup_stream_state(request_id) + + def test_cleanup_warns_on_incomplete_tool_call(self): + adapter, request_id, state = self._make_adapter_state() + state.tool_calls["call_1"] = { + "id": "call_1", + "name": "search", + "arguments": "", + "index": 0, + } + with patch("tensorrt_llm.serve.harmony_adapter.logger.warning") as mock_warn: + adapter.cleanup_stream_state(request_id) + assert mock_warn.call_count == 1 + args = mock_warn.call_args.args + assert request_id in args[0] + assert "search" in args[1] + + def test_cleanup_no_warning_when_complete(self): + adapter, request_id, state = self._make_adapter_state() + state.tool_calls["call_1"] = { + "id": "call_1", + "name": "get_weather", + "arguments": '{"city":"SF"}', + "index": 0, + "active": False, + } + with patch("tensorrt_llm.serve.harmony_adapter.logger.warning") as mock_warn: + adapter.cleanup_stream_state(request_id) + mock_warn.assert_not_called() + + if __name__ == "__main__": pytest.main([__file__, "-v"])