From 3e4e5516131b251fffc38f0ba85329968d7e22d3 Mon Sep 17 00:00:00 2001
From: He Tao <taohe@bytedance.com>
Date: Wed, 23 Apr 2025 18:02:58 +0800
Subject: [PATCH] fix: add server name in desc

---
 src/graph/nodes.py | 15 ++++++++++-----
 1 file changed, 10 insertions(+), 5 deletions(-)

diff --git a/src/graph/nodes.py b/src/graph/nodes.py
index 12e2db6..00cd9a7 100644
--- a/src/graph/nodes.py
+++ b/src/graph/nodes.py
@@ -294,7 +294,7 @@ async def _setup_and_execute_agent_step(
     """
     configurable = Configuration.from_runnable_config(config)
     mcp_servers = {}
-    enabled_tools = set()
+    enabled_tools = {}
 
     # Extract MCP server configuration for this agent type
     if configurable.mcp_settings:
@@ -308,14 +308,19 @@ async def _setup_and_execute_agent_step(
                     for k, v in server_config.items()
                     if k in ("transport", "command", "args", "url", "env")
                 }
-                enabled_tools.update(server_config["enabled_tools"])
+                for tool_name in server_config["enabled_tools"]:
+                    enabled_tools[tool_name] = server_name
 
     # Create and execute agent with MCP tools if available
     if mcp_servers:
         async with MultiServerMCPClient(mcp_servers) as client:
-            loaded_tools = [
-                tool for tool in client.get_tools() if tool.name in enabled_tools
-            ] + default_tools
+            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)
     else: