Python: Propagated MCP isError flag through function middleware pipeline (#4511)

* Propagated MCP isError flag through function middleware pipeline

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

* Small update

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

* Fix CI

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

---------

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
This commit is contained in:
Dmytro Struk
2026-03-05 15:16:19 -08:00
committed by GitHub
Unverified
parent ce7b5b17c1
commit 8664d19285
3 changed files with 152 additions and 4 deletions
@@ -901,7 +901,11 @@ class MCPTool:
for attempt in range(2):
try:
result = await self.session.call_tool(tool_name, arguments=filtered_kwargs, meta=otel_meta) # type: ignore
if result.isError:
raise ToolExecutionException(parser(result))
return parser(result)
except ToolExecutionException:
raise
except ClosedResourceError as cl_ex:
if attempt == 0:
# First attempt failed, try reconnecting
@@ -563,10 +563,10 @@ class SkillsProvider(BaseContextProvider):
try:
if inspect.iscoroutinefunction(resource.function):
result = (
await resource.function(**kwargs) if resource._accepts_kwargs else await resource.function()
await resource.function(**kwargs) if resource._accepts_kwargs else await resource.function() # pyright: ignore[reportPrivateUsage]
)
else:
result = resource.function(**kwargs) if resource._accepts_kwargs else resource.function()
result = resource.function(**kwargs) if resource._accepts_kwargs else resource.function() # pyright: ignore[reportPrivateUsage]
return str(result)
except Exception as exc:
logger.exception("Failed to read resource '%s' from skill '%s'", resource_name, skill_name)
+146 -2
View File
@@ -14,6 +14,8 @@ from pydantic import AnyUrl, BaseModel
from agent_framework import (
Content,
FunctionInvocationContext,
FunctionMiddleware,
MCPStdioTool,
MCPStreamableHTTPTool,
MCPWebsocketTool,
@@ -30,6 +32,7 @@ from agent_framework._mcp import (
_prepare_message_for_mcp,
logger,
)
from agent_framework._middleware import FunctionMiddlewarePipeline
from agent_framework.exceptions import ToolException, ToolExecutionException
# Integration test skip condition
@@ -898,6 +901,147 @@ async def test_local_mcp_server_function_execution_error():
await func.invoke(param="test_value")
async def test_mcp_tool_call_tool_raises_on_is_error():
"""Test that call_tool raises ToolExecutionException when MCP returns isError=True."""
class TestServer(MCPTool):
async def connect(self):
self.session = Mock(spec=ClientSession)
self.session.list_tools = AsyncMock(
return_value=types.ListToolsResult(
tools=[
types.Tool(
name="test_tool",
description="Test tool",
inputSchema={
"type": "object",
"properties": {"param": {"type": "string"}},
"required": ["param"],
},
)
]
)
)
self.session.call_tool = AsyncMock(
return_value=types.CallToolResult(
content=[types.TextContent(type="text", text="Something went wrong")],
isError=True,
)
)
def get_mcp_client(self) -> _AsyncGeneratorContextManager[Any, None]:
return None
server = TestServer(name="test_server")
async with server:
await server.load_tools()
func = server.functions[0]
with pytest.raises(ToolExecutionException, match="Something went wrong"):
await func.invoke(param="test_value")
async def test_mcp_tool_call_tool_succeeds_when_is_error_false():
"""Test that call_tool returns normally when MCP returns isError=False."""
class TestServer(MCPTool):
async def connect(self):
self.session = Mock(spec=ClientSession)
self.session.list_tools = AsyncMock(
return_value=types.ListToolsResult(
tools=[
types.Tool(
name="test_tool",
description="Test tool",
inputSchema={
"type": "object",
"properties": {"param": {"type": "string"}},
"required": ["param"],
},
)
]
)
)
self.session.call_tool = AsyncMock(
return_value=types.CallToolResult(
content=[types.TextContent(type="text", text="Success")],
isError=False,
)
)
def get_mcp_client(self) -> _AsyncGeneratorContextManager[Any, None]:
return None
server = TestServer(name="test_server")
async with server:
await server.load_tools()
func = server.functions[0]
result = await func.invoke(param="test_value")
assert result == "Success"
async def test_mcp_tool_is_error_propagates_through_function_middleware():
"""Test that MCP isError=True propagates as ToolExecutionException through function middleware."""
error_seen_in_middleware = False
class ErrorCheckMiddleware(FunctionMiddleware):
async def process(self, context: FunctionInvocationContext, call_next):
nonlocal error_seen_in_middleware
try:
await call_next()
except ToolExecutionException:
error_seen_in_middleware = True
raise
class TestServer(MCPTool):
async def connect(self):
self.session = Mock(spec=ClientSession)
self.session.list_tools = AsyncMock(
return_value=types.ListToolsResult(
tools=[
types.Tool(
name="test_tool",
description="Test tool",
inputSchema={
"type": "object",
"properties": {"param": {"type": "string"}},
"required": ["param"],
},
)
]
)
)
self.session.call_tool = AsyncMock(
return_value=types.CallToolResult(
content=[types.TextContent(type="text", text="MCP error occurred")],
isError=True,
)
)
def get_mcp_client(self) -> _AsyncGeneratorContextManager[Any, None]:
return None
server = TestServer(name="test_server")
async with server:
await server.load_tools()
func = server.functions[0]
middleware_pipeline = FunctionMiddlewarePipeline(ErrorCheckMiddleware())
middleware_context = FunctionInvocationContext(
function=func,
arguments={"param": "test_value"},
)
with pytest.raises(ToolExecutionException, match="MCP error occurred"):
await middleware_pipeline.execute(
middleware_context,
lambda ctx: func.invoke(arguments=ctx.arguments),
)
assert error_seen_in_middleware, "Middleware should have seen the ToolExecutionException"
async def test_local_mcp_server_prompt_execution():
"""Test prompt execution through MCP server."""
@@ -2098,7 +2242,7 @@ async def test_mcp_tool_connection_properly_invalidated_after_closed_resource_er
tool._tools_loaded = True
# First call should work - connection is valid
mock_session.call_tool.return_value = MagicMock(content=[])
mock_session.call_tool.return_value = types.CallToolResult(content=[])
result = await tool.call_tool("test_tool", arg1="value1")
assert result is not None
@@ -2111,7 +2255,7 @@ async def test_mcp_tool_connection_properly_invalidated_after_closed_resource_er
call_count += 1
if call_count == 1:
raise ClosedResourceError
return MagicMock(content=[])
return types.CallToolResult(content=[])
mock_session.call_tool = call_tool_with_error