From 4e3a8a06371fe9edc8f34de6d368183f809ebb2c Mon Sep 17 00:00:00 2001
From: 0xbyt4 <35742124+0xbyt4@users.noreply.github.com>
Date: Tue, 10 Mar 2026 02:24:53 +0300
Subject: [PATCH] fix: handle empty choices in MCP sampling callback

SamplingHandler.__call__ accessed response.choices[0] without checking
if the list was non-empty. LLM APIs can return empty choices on content
filtering, provider errors, or rate limits, causing an unhandled
IndexError that propagates to the MCP SDK and may crash the connection.

Add a defensive guard that returns a proper ErrorData when choices is
empty, None, or missing. Includes three test cases covering all
variants.
---
 tests/tools/test_mcp_tool.py | 59 ++++++++++++++++++++++++++++++++++++
 tools/mcp_tool.py            |  8 +++++
 2 files changed, 67 insertions(+)

diff --git a/tests/tools/test_mcp_tool.py b/tests/tools/test_mcp_tool.py
index 1acbdfa12..446f80d3e 100644
--- a/tests/tools/test_mcp_tool.py
+++ b/tests/tools/test_mcp_tool.py
@@ -2049,6 +2049,65 @@ class TestSamplingErrors:
             assert "No LLM provider" in result.message
             assert handler.metrics["errors"] == 1
 
+    def test_empty_choices_returns_error(self):
+        """LLM returning choices=[] is handled gracefully, not IndexError."""
+        handler = SamplingHandler("ec", {})
+        fake_client = MagicMock()
+        fake_client.chat.completions.create.return_value = SimpleNamespace(
+            choices=[],
+            model="test-model",
+            usage=SimpleNamespace(total_tokens=0),
+        )
+
+        with patch(
+            "agent.auxiliary_client.get_text_auxiliary_client",
+            return_value=(fake_client, "default-model"),
+        ):
+            result = asyncio.run(handler(None, _make_sampling_params()))
+
+        assert isinstance(result, ErrorData)
+        assert "empty response" in result.message.lower()
+        assert handler.metrics["errors"] == 1
+
+    def test_none_choices_returns_error(self):
+        """LLM returning choices=None is handled gracefully, not TypeError."""
+        handler = SamplingHandler("nc", {})
+        fake_client = MagicMock()
+        fake_client.chat.completions.create.return_value = SimpleNamespace(
+            choices=None,
+            model="test-model",
+            usage=SimpleNamespace(total_tokens=0),
+        )
+
+        with patch(
+            "agent.auxiliary_client.get_text_auxiliary_client",
+            return_value=(fake_client, "default-model"),
+        ):
+            result = asyncio.run(handler(None, _make_sampling_params()))
+
+        assert isinstance(result, ErrorData)
+        assert "empty response" in result.message.lower()
+        assert handler.metrics["errors"] == 1
+
+    def test_missing_choices_attr_returns_error(self):
+        """LLM response without choices attribute is handled gracefully."""
+        handler = SamplingHandler("mc", {})
+        fake_client = MagicMock()
+        fake_client.chat.completions.create.return_value = SimpleNamespace(
+            model="test-model",
+            usage=SimpleNamespace(total_tokens=0),
+        )
+
+        with patch(
+            "agent.auxiliary_client.get_text_auxiliary_client",
+            return_value=(fake_client, "default-model"),
+        ):
+            result = asyncio.run(handler(None, _make_sampling_params()))
+
+        assert isinstance(result, ErrorData)
+        assert "empty response" in result.message.lower()
+        assert handler.metrics["errors"] == 1
+
 
 # ---------------------------------------------------------------------------
 # 10. Model whitelist
diff --git a/tools/mcp_tool.py b/tools/mcp_tool.py
index deb87d483..b0fc35f7f 100644
--- a/tools/mcp_tool.py
+++ b/tools/mcp_tool.py
@@ -538,6 +538,14 @@ class SamplingHandler:
                 f"Sampling LLM call failed: {_sanitize_error(str(exc))}"
             )
 
+        # Guard against empty choices (content filtering, provider errors)
+        if not getattr(response, "choices", None):
+            self.metrics["errors"] += 1
+            return self._error(
+                f"LLM returned empty response (no choices) for server "
+                f"'{self.server_name}'"
+            )
+
         # Track metrics
         choice = response.choices[0]
         self.metrics["requests"] += 1