mirror of
https://github.com/PrefectHQ/fastmcp.git
synced 2026-08-21 13:04:18 +02:00
Fix raise_on_error handling for tool tasks (#3946)
This commit is contained in:
parent
7dd573980a
commit
012f674ee4
3 changed files with 104 additions and 5 deletions
|
|
@ -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,
|
||||
)
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue