From 8664d1928553c6555ccd234cc9ebff9ee2524e1b Mon Sep 17 00:00:00 2001 From: Dmytro Struk <13853051+dmytrostruk@users.noreply.github.com> Date: Thu, 5 Mar 2026 15:16:19 -0800 Subject: [PATCH] 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> --- python/packages/core/agent_framework/_mcp.py | 4 + .../packages/core/agent_framework/_skills.py | 4 +- python/packages/core/tests/core/test_mcp.py | 148 +++++++++++++++++- 3 files changed, 152 insertions(+), 4 deletions(-) diff --git a/python/packages/core/agent_framework/_mcp.py b/python/packages/core/agent_framework/_mcp.py index 0c241cb89a..b07a872204 100644 --- a/python/packages/core/agent_framework/_mcp.py +++ b/python/packages/core/agent_framework/_mcp.py @@ -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 diff --git a/python/packages/core/agent_framework/_skills.py b/python/packages/core/agent_framework/_skills.py index 11de6c3bdb..c7d59d789e 100644 --- a/python/packages/core/agent_framework/_skills.py +++ b/python/packages/core/agent_framework/_skills.py @@ -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) diff --git a/python/packages/core/tests/core/test_mcp.py b/python/packages/core/tests/core/test_mcp.py index 65b4015093..867e7183cf 100644 --- a/python/packages/core/tests/core/test_mcp.py +++ b/python/packages/core/tests/core/test_mcp.py @@ -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