feat: add support to TrustedMCP

Author: tyler-lee

docs/content/1.introduction/1.overview.md

@@ -31,6 +31,7 @@ seo:
 
 - 支持 Identity 管理
 - 支持 API Key 服务鉴权与 OAuth2 用户鉴权能力
+- 支持 AICC 可信部署与 Agent 到 MCP 再到 LLM 服务的端到端加密通信
 
 ### 更完备的可观测性和评估能力
 

docs/content/3.agent/1.agent.md

@@ -265,3 +265,49 @@ print(response)
 ## 多智能体协作
 
 使用 VeADK 可以构建多 Agent 协作， 主 Agent 通过 `sub_agents` 机制协调多个子 Agent 完成复杂任务。
+
+## 可信 Agent
+
+依托机密计算技术（如Jeddak AICC），可以将 Agent 以及大模型服务（如机密豆包）运行在可信环境中，并通过 TrustedMCP 以及端到端加密通信，构建全链路可信的 Agent。其中，TrustedMCP是火山引擎推出的可信MCP，在标准MCP协议的基础上扩展核心组件之间的身份证明及验证能力，提供组件之间端到端的通信安全保障，可以解决MCP应用中服务身份不可信、数据被篡改、流量劫持、数据隐私泄露等安全威胁。
+
+使用 VeADK 可以构建可信 Agent， 并通过 `tools` 机制调用 TrustedMCP 服务。
+
+### 使用示例
+
+```python [agent.py]
+import asyncio
+
+from veadk import Agent
+from veadk.utils.mcp_utils import get_mcp_params
+from veadk.tools.mcp_tool.trusted_mcp_toolset import TrustedMcpToolset
+
+mcp_url = "<TrustedMCP server address>"
+
+# 1. 开启 TrustedMCP 功能以及相关配置
+connection_params = get_mcp_params(mcp_url)
+connection_params.headers = {"x-trusted-mcp": "true"}
+
+# 2. 初始化 TrustedMcpToolset 工具集
+toolset = TrustedMcpToolset(connection_params=connection_params)
+
+# 3. 初始化 Agent
+agent = Agent(tools=[toolset])
+
+# 4. 运行 Agent
+response = asyncio.run(agent.run("北京天气怎么样？"))
+
+print(response) # 北京天气晴朗，气温25°C。
+```
+
+### 选项
+
+TrustedMcpToolset 配置参数：工作流 Agent 采用统一的参数：
+
+::field-group
+  ::field{name="x-trusted-mcp" type="string"}
+  默认为 `true` - 是否开启 TrustedMCP 功能
+  ::
+
+  ::field{name="aicc-config" type="string"}
+  默认为 `./aicc_config.json` - AICC 配置文件路径。配置文件内容请参考[TrustedMCP 配置文件](https://github.com/volcengine/AICC-Trusted-MCP/blob/main/README.md)
+  ::

pyproject.toml

@@ -26,6 +26,7 @@ dependencies = [
     "volcengine>=1.0.193", # For Volcengine sign
     "agent-pilot-sdk>=0.1.2", # Prompt optimization by Volcengine AgentPilot/PromptPilot toolkits
     "fastmcp>=2.11.3", # For running MCP
+    "trustedmcp>=0.0.4", # For running TrustedMCP
     "cookiecutter>=2.6.0", # For cloud deploy
     "omegaconf>=2.3.0", # For agent builder
     "llama-index>=0.14.0",

tests/tools/mcp_tool/test_trusted_mcp_components.py

@@ -0,0 +1,230 @@
+# Copyright (c) 2025 Beijing Volcano Engine Technology Co., Ltd. and/or its affiliates.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import sys
+import unittest
+from unittest import mock
+import asyncio
+
+from veadk.tools.mcp_tool.trusted_mcp_toolset import TrustedMcpToolset
+from veadk.tools.mcp_tool.trusted_mcp_session_manager import (
+    TrustedMcpSessionManager,
+)
+from google.adk.tools.mcp_tool.mcp_session_manager import StreamableHTTPConnectionParams
+
+
+class TestTrustedMcpComponents(unittest.TestCase):
+    def setUp(self):
+        # Create simple mock objects
+        self.mock_stdio_params = mock.MagicMock()
+        # Add serializable headers to stdio_params
+        mock_headers = mock.MagicMock()
+        mock_headers.copy.return_value = {"Content-Type": "application/json"}
+        self.mock_stdio_params.headers = mock_headers
+        # Add getattr method that throws AttributeError by default to simulate non-HTTP connection
+        self.mock_stdio_params.getattr = mock.MagicMock(side_effect=AttributeError)
+
+        self.mock_http_params = mock.MagicMock(spec=StreamableHTTPConnectionParams)
+        self.mock_http_params.url = "http://mock-url.com"
+        self.mock_http_params.timeout = 30
+        self.mock_http_params.sse_read_timeout = 60
+        self.mock_http_params.terminate_on_close = False
+        # Add serializable headers to http_params
+        mock_http_headers = mock.MagicMock()
+        mock_http_headers.copy.return_value = {
+            "Content-Type": "application/json",
+            "x-trusted-mcp": "true",
+        }
+        self.mock_http_params.headers = mock_http_headers
+
+    def test_trusted_mcp_toolset_init(self):
+        """Test the initialization of TrustedMcpToolset"""
+        # Mock TrustedMcpSessionManager and logger
+        with (
+            mock.patch(
+                "veadk.tools.mcp_tool.trusted_mcp_toolset.TrustedMcpSessionManager"
+            ) as mock_session_manager,
+            mock.patch("veadk.tools.mcp_tool.trusted_mcp_toolset.logger"),
+        ):
+            # Create instance directly without mocking parent initialization to automatically set necessary attributes
+            toolset = TrustedMcpToolset(
+                connection_params=self.mock_stdio_params, tool_filter=["read_file"]
+            )
+
+            # Verify TrustedMcpSessionManager was created
+            mock_session_manager.assert_called_once()
+
+            # Verify _trusted_mcp_session_manager attribute is set
+            self.assertIsNotNone(toolset._mcp_session_manager)
+
+    def test_trusted_mcp_session_manager_create_client_trusted(self):
+        """Test TrustedMcpSessionManager._create_client method - TrustedMCP mode"""
+        # Mock trusted_mcp_client_context
+        with mock.patch(
+            "veadk.tools.mcp_tool.trusted_mcp_session_manager.trusted_mcp_client_context"
+        ) as mock_trusted_client:
+            # Create manager instance directly
+            manager = TrustedMcpSessionManager(
+                connection_params=self.mock_http_params, errlog=sys.stderr
+            )
+
+            # Call _create_client with trusted_mcp header
+            merged_headers = {"x-trusted-mcp": "true"}
+            result = manager._create_client(merged_headers)
+
+            # Verify trusted_mcp_client_context was called
+            mock_trusted_client.assert_called_once()
+
+            # Verify result
+            self.assertEqual(result, mock_trusted_client.return_value)
+
+    def test_trusted_mcp_session_manager_create_client_standard(self):
+        """Test TrustedMcpSessionManager._create_client method - Standard mode"""
+        # Mock parent class's _create_client method
+        with mock.patch(
+            "veadk.tools.mcp_tool.trusted_mcp_session_manager.MCPSessionManager._create_client"
+        ) as mock_super_create:
+            expected_client = mock.MagicMock()
+            mock_super_create.return_value = expected_client
+
+            # Create manager instance directly without mocking parent initialization to automatically set necessary attributes
+            manager = TrustedMcpSessionManager(
+                connection_params=self.mock_stdio_params, errlog=sys.stderr
+            )
+
+            # Call _create_client with normal headers
+            merged_headers = {"content-type": "application/json"}
+            result = manager._create_client(merged_headers)
+
+            # Verify parent's _create_client was called
+            mock_super_create.assert_called_once_with(merged_headers)
+
+            # Verify result
+            self.assertEqual(result, expected_client)
+
+    def test_trusted_mcp_session_manager_trusted_create_session(self):
+        """Test TrustedMcpSessionManager.create_session method - TrustedMCP mode"""
+
+        # Use coroutine test helper to run async test
+        async def run_test():
+            # Create a mock async context manager
+            class MockAsyncContext:
+                def __init__(self):
+                    self.session = mock.MagicMock()
+
+                async def __aenter__(self):
+                    return self.session
+
+                async def __aexit__(self, exc_type, exc_val, exc_tb):
+                    return False
+
+            # Create a mock AsyncExitStack that supports async methods
+            class MockAsyncExitStack:
+                def __init__(self):
+                    self.entered_contexts = []
+
+                async def enter_async_context(self, context):
+                    self.entered_contexts.append(context)
+                    return await context.__aenter__()
+
+                async def aclose(self):
+                    pass
+
+            # Mock required functions and classes
+            mock_trusted_context = MockAsyncContext()
+            with (
+                mock.patch(
+                    "veadk.tools.mcp_tool.trusted_mcp_session_manager.trusted_mcp_client",
+                    return_value=mock_trusted_context,
+                ),
+                mock.patch(
+                    "veadk.tools.mcp_tool.trusted_mcp_session_manager.MCPSessionManager._merge_headers",
+                    return_value={"x-trusted-mcp": "true"},
+                ),
+                mock.patch(
+                    "veadk.tools.mcp_tool.trusted_mcp_session_manager.MCPSessionManager._generate_session_key",
+                    return_value="session-key",
+                ),
+                mock.patch(
+                    "veadk.tools.mcp_tool.trusted_mcp_session_manager.MCPSessionManager._is_session_disconnected",
+                    return_value=False,
+                ),
+                mock.patch(
+                    "veadk.tools.mcp_tool.trusted_mcp_session_manager.AsyncExitStack",
+                    return_value=MockAsyncExitStack(),
+                ),
+            ):
+                # Create manager instance and set necessary attributes
+                manager = TrustedMcpSessionManager(
+                    connection_params=self.mock_http_params, errlog=sys.stderr
+                )
+                manager._sessions = {}
+                manager._session_lock = asyncio.Lock()
+
+                # Call create_session
+                headers = {"x-trusted-mcp": "true"}
+                result = await manager.create_session(headers)
+
+                # Verify result is the mock session
+                self.assertEqual(result, mock_trusted_context.session)
+
+        # Run async test
+        asyncio.run(run_test())
+
+    def test_trusted_mcp_session_manager_session_reuse(self):
+        """Test TrustedMcpSessionManager.create_session method - Session reuse"""
+
+        # Use coroutine test helper to run async test
+        async def run_test():
+            # Mock necessary methods
+            with (
+                mock.patch(
+                    "veadk.tools.mcp_tool.trusted_mcp_session_manager.MCPSessionManager._merge_headers",
+                    return_value={"header": "value"},
+                ),
+                mock.patch(
+                    "veadk.tools.mcp_tool.trusted_mcp_session_manager.MCPSessionManager._generate_session_key",
+                    return_value="session-key",
+                ),
+                mock.patch(
+                    "veadk.tools.mcp_tool.trusted_mcp_session_manager.MCPSessionManager._is_session_disconnected",
+                    return_value=False,
+                ),
+            ):
+                # Create manager instance
+                manager = TrustedMcpSessionManager(
+                    connection_params=self.mock_http_params, errlog=sys.stderr
+                )
+                manager._sessions = {}
+                manager._session_lock = asyncio.Lock()
+
+                # Set up an existing session
+                existing_session = mock.MagicMock()
+                existing_exit_stack = mock.MagicMock()
+                manager._sessions = {
+                    "session-key": (existing_session, existing_exit_stack)
+                }
+
+                # Call create_session
+                result = await manager.create_session({"header": "value"})
+
+                # Verify existing session was returned
+                self.assertEqual(result, existing_session)
+
+        # Run async test
+        asyncio.run(run_test())
+
+
+if __name__ == "__main__":
+    unittest.main()

veadk/tools/builtin_tools/web_search.py

@@ -27,7 +27,7 @@
 logger = get_logger(__name__)
 
 
-def web_search(query: str, tool_context: ToolContext) -> list[str]:
+def web_search(query: str, tool_context: ToolContext | None = None) -> list[str]:
     """Search a query in websites.
 
     Args:
@@ -36,8 +36,11 @@ def web_search(query: str, tool_context: ToolContext) -> list[str]:
     Returns:
         A list of result documents.
     """
-    ak = tool_context.state.get("VOLCENGINE_ACCESS_KEY")
-    sk = tool_context.state.get("VOLCENGINE_SECRET_KEY")
+    ak = None
+    sk = None
+    if tool_context:
+        ak = tool_context.state.get("VOLCENGINE_ACCESS_KEY")
+        sk = tool_context.state.get("VOLCENGINE_SECRET_KEY")
     session_token = ""
 
     if not (ak and sk):

veadk/tools/mcp_tool/__init__.py

@@ -0,0 +1,13 @@
+# Copyright (c) 2025 Beijing Volcano Engine Technology Co., Ltd. and/or its affiliates.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.