Fix raise_on_error handling for tool tasks (#3946)

This commit is contained in:
Gnani Rahul 2026-04-16 17:37:52 -05:00 committed by GitHub
commit 012f674ee4
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
3 changed files with 104 additions and 5 deletions

View file

@ -288,7 +288,12 @@ class ClientToolsMixin:
if task:
return await self._call_tool_as_task(
name, arguments, task_id, ttl, meta=request_meta or None
name,
arguments,
task_id,
ttl,
raise_on_error=raise_on_error,
meta=request_meta or None,
)
result = await self.call_tool_mcp(
@ -308,6 +313,7 @@ class ClientToolsMixin:
arguments: dict[str, Any] | None = None,
task_id: str | None = None,
ttl: int = 60000,
raise_on_error: bool = True,
meta: dict[str, Any] | None = None,
) -> ToolTask:
"""Call a tool for background execution (SEP-1686).
@ -321,6 +327,7 @@ class ClientToolsMixin:
arguments: Tool arguments
task_id: Optional client-provided task ID (ignored, for backward compatibility)
ttl: Time to keep results available in milliseconds (default 60s)
raise_on_error: Whether task.result() should raise ToolError on errors
meta: Optional request metadata (e.g., version info)
Returns:
@ -356,7 +363,11 @@ class ClientToolsMixin:
self._submitted_task_ids.add(server_task_id)
task_obj = ToolTask(
self, server_task_id, tool_name=name, immediate_result=None
self,
server_task_id,
tool_name=name,
immediate_result=None,
raise_on_error=raise_on_error,
)
self._task_registry[server_task_id] = weakref.ref(task_obj)
return task_obj
@ -369,6 +380,7 @@ class ClientToolsMixin:
synthetic_task_id,
tool_name=name,
immediate_result=parsed_result,
raise_on_error=raise_on_error,
)

View file

@ -15,6 +15,7 @@ import mcp.types
from mcp.types import GetTaskResult, TaskStatusNotification
from fastmcp.client.messages import Message, MessageHandler
from fastmcp.exceptions import ToolError
from fastmcp.utilities.logging import get_logger
logger = get_logger(__name__)
@ -335,6 +336,7 @@ class ToolTask(Task["CallToolResult"]):
task_id: str,
tool_name: str,
immediate_result: CallToolResult | None = None,
raise_on_error: bool = True,
):
"""
Create a ToolTask wrapper.
@ -344,9 +346,11 @@ class ToolTask(Task["CallToolResult"]):
task_id: The task identifier
tool_name: Name of the tool being executed
immediate_result: If server executed synchronously, the immediate result
raise_on_error: Whether task.result() should raise ToolError on errors
"""
super().__init__(client, task_id, immediate_result)
self._tool_name = tool_name
self._raise_on_error = raise_on_error
async def result(self) -> CallToolResult:
"""Wait for and return the tool result.
@ -364,6 +368,14 @@ class ToolTask(Task["CallToolResult"]):
if self._is_immediate:
assert self._immediate_result is not None # Type narrowing
result = self._immediate_result
if result.is_error and self._raise_on_error:
if result.content and isinstance(
result.content[0], mcp.types.TextContent
):
msg = result.content[0].text
else:
msg = f"Tool '{self._tool_name}' returned an error"
raise ToolError(msg)
else:
# Check client connected
self._check_client_connected()
@ -379,12 +391,16 @@ class ToolTask(Task["CallToolResult"]):
# Raw dict from get_task_result - parse as CallToolResult
mcp_result = mcp.types.CallToolResult.model_validate(raw_result)
result = await self._client._parse_call_tool_result(
self._tool_name, mcp_result, raise_on_error=True
self._tool_name,
mcp_result,
raise_on_error=self._raise_on_error,
)
elif isinstance(raw_result, mcp.types.CallToolResult):
# Already a CallToolResult from MCP protocol - parse it
result = await self._client._parse_call_tool_result(
self._tool_name, raw_result, raise_on_error=True
self._tool_name,
raw_result,
raise_on_error=self._raise_on_error,
)
else:
# Legacy ToolResult format - convert to MCP type
@ -397,7 +413,9 @@ class ToolTask(Task["CallToolResult"]):
_meta=raw_result.meta, # type: ignore[call-arg] # _meta is Pydantic alias for meta field # ty:ignore[unknown-argument]
)
result = await self._client._parse_call_tool_result(
self._tool_name, mcp_result, raise_on_error=True
self._tool_name,
mcp_result,
raise_on_error=self._raise_on_error,
)
else:
# Unknown type - just return it

View file

@ -10,6 +10,7 @@ import pytest
from fastmcp import FastMCP
from fastmcp.client import Client
from fastmcp.client.tasks import ToolTask
from fastmcp.exceptions import ToolError
@pytest.fixture
@ -85,3 +86,71 @@ async def test_tool_task_status_and_wait(tool_task_server):
await task.wait(timeout=2.0)
final_status = await task.status()
assert final_status.status == "completed"
async def test_immediate_tool_task_respects_raise_on_error_true():
"""Immediate task fallback should still raise ToolError when requested."""
mcp = FastMCP("immediate-tool-task-error")
@mcp.tool
def failing_tool() -> str:
raise ValueError("immediate task failure")
async with Client(mcp) as client:
task = await client.call_tool("failing_tool", task=True, raise_on_error=True)
assert task.returned_immediately
with pytest.raises(
ToolError, match="does not support task-augmented execution"
):
await task.result()
async def test_immediate_tool_task_respects_raise_on_error_false():
"""Immediate task fallback should return error results when requested."""
mcp = FastMCP("immediate-tool-task-no-raise")
@mcp.tool
def failing_tool() -> str:
raise ValueError("immediate task failure")
async with Client(mcp) as client:
task = await client.call_tool("failing_tool", task=True, raise_on_error=False)
assert task.returned_immediately
result = await task.result()
assert result.is_error is True
assert "does not support task-augmented execution" in str(result)
async def test_background_tool_task_respects_raise_on_error_true():
"""Background tasks should still raise ToolError by default on errors."""
mcp = FastMCP("background-tool-task-error")
@mcp.tool(task=True)
async def failing_tool() -> str:
raise ValueError("background task failure")
async with Client(mcp) as client:
task = await client.call_tool("failing_tool", task=True, raise_on_error=True)
assert not task.returned_immediately
with pytest.raises(ToolError, match="background task failure"):
await task.result()
async def test_background_tool_task_respects_raise_on_error_false():
"""Background tasks should return error results when raise_on_error is disabled."""
mcp = FastMCP("background-tool-task-no-raise")
@mcp.tool(task=True)
async def failing_tool() -> str:
raise ValueError("background task failure")
async with Client(mcp) as client:
task = await client.call_tool("failing_tool", task=True, raise_on_error=False)
assert not task.returned_immediately
result = await task.result()
assert result.is_error is True
assert "background task failure" in str(result)