diff --git a/src/graph/nodes.py b/src/graph/nodes.py index 194cdc9..92d30ac 100644 --- a/src/graph/nodes.py +++ b/src/graph/nodes.py @@ -458,16 +458,17 @@ async def _setup_and_execute_agent_step( # Create and execute agent with MCP tools if available if mcp_servers: - async with MultiServerMCPClient(mcp_servers) as client: - loaded_tools = default_tools[:] - for tool in client.get_tools(): - if tool.name in enabled_tools: - tool.description = ( - f"Powered by '{enabled_tools[tool.name]}'.\n{tool.description}" - ) - loaded_tools.append(tool) - agent = create_agent(agent_type, agent_type, loaded_tools, agent_type) - return await _execute_agent_step(state, agent, agent_type) + client = MultiServerMCPClient(mcp_servers) + loaded_tools = default_tools[:] + all_tools = await client.get_tools() + for tool in all_tools: + if tool.name in enabled_tools: + tool.description = ( + f"Powered by '{enabled_tools[tool.name]}'.\n{tool.description}" + ) + loaded_tools.append(tool) + agent = create_agent(agent_type, agent_type, loaded_tools, agent_type) + return await _execute_agent_step(state, agent, agent_type) else: # Use default tools if no MCP servers are configured agent = create_agent(agent_type, agent_type, default_tools, agent_type) diff --git a/tests/integration/test_nodes.py b/tests/integration/test_nodes.py index 73d565e..8c2597f 100644 --- a/tests/integration/test_nodes.py +++ b/tests/integration/test_nodes.py @@ -1110,7 +1110,7 @@ def patch_multiserver_mcp_client(): async def __aexit__(self, exc_type, exc, tb): pass - def get_tools(self): + async def get_tools(self): return [ FakeTool("toolA", "descA"), FakeTool("toolB", "descB"), @@ -1244,7 +1244,7 @@ async def test_setup_and_execute_agent_step_with_mcp_tools_description_update( async def __aexit__(self, exc_type, exc, tb): pass - def get_tools(self): + async def get_tools(self): return [FakeTool("toolA", "descA")] with patch("src.graph.nodes.MultiServerMCPClient", return_value=FakeClient()):