Enhance AIClient and MCPServer to support tool registration with source tracking. Added logging for tool calls and improved error handling. Introduced methods for embedding token limit extraction and budget application in OpenAIModel. Added tests for MCP tool registration and execution.

src/ai/client.py

@@ -43,6 +43,7 @@ class AIClient:
         
         # 初始化工具注册表
         self.tools = ToolRegistry()
+        self._tool_sources: Dict[str, str] = {}
         
         # 初始化人格系统
         self.personality = PersonalitySystem(
@@ -242,8 +243,11 @@ class AIClient:
     ) -> Message:
         """处理工具调用。"""
         messages.append(response)
+        total_calls = len(response.tool_calls or [])
+        if total_calls:
+            logger.info(f"检测到工具调用请求: {total_calls} 个")
         
-        # 鎵ц宸ュ叿璋冪敤
+        # 执行工具调用
         for tool_call in response.tool_calls or []:
             try:
                 tool_name, tool_args, tool_call_id = self._parse_tool_call(tool_call)
@@ -263,6 +267,7 @@ class AIClient:
                 continue
             
             tool_def = self.tools.get(tool_name)
+            tool_source = self._tool_sources.get(tool_name, "custom")
             if not tool_def:
                 error_msg = f"未找到工具: {tool_name}"
                 logger.warning(error_msg)
@@ -275,9 +280,19 @@ class AIClient:
                 continue
             
             try:
+                logger.info(
+                    "工具调用开始: "
+                    f"source={tool_source}, name={tool_name}, "
+                    f"args={self._preview_log_payload(tool_args)}"
+                )
                 result = tool_def.function(**tool_args)
                 if inspect.isawaitable(result):
                     result = await result
+                logger.info(
+                    "工具调用成功: "
+                    f"source={tool_source}, name={tool_name}, "
+                    f"result={self._preview_log_payload(result)}"
+                )
                 messages.append(Message(
                     role="tool",
                     name=tool_name,
@@ -285,6 +300,10 @@ class AIClient:
                     tool_call_id=tool_call_id
                 ))
             except Exception as e:
+                logger.warning(
+                    "工具调用失败: "
+                    f"source={tool_source}, name={tool_name}, error={e}"
+                )
                 messages.append(Message(
                     role="tool",
                     name=tool_name,
@@ -334,6 +353,18 @@ class AIClient:
                 return parsed
         
         raise ValueError(f"不支持的工具参数类型: {type(raw_args)}")
+
+    @staticmethod
+    def _preview_log_payload(payload: Any, max_len: int = 240) -> str:
+        """日志中展示参数/结果时使用的简短预览。"""
+        try:
+            text = json.dumps(payload, ensure_ascii=False, default=str)
+        except Exception:
+            text = str(payload)
+
+        if len(text) > max_len:
+            return text[:max_len] + "..."
+        return text
     
     def set_personality(self, personality_name: str) -> bool:
         """设置人格。"""
@@ -382,7 +413,14 @@ class AIClient:
         """获取任务状态。"""
         return self.task_manager.get_task_status(task_id)
     
-    def register_tool(self, name: str, description: str, parameters: Dict, function: callable):
+    def register_tool(
+        self,
+        name: str,
+        description: str,
+        parameters: Dict,
+        function: callable,
+        source: str = "custom",
+    ):
         """注册工具。"""
         from .base import ToolDefinition
         tool = ToolDefinition(
@@ -392,18 +430,23 @@ class AIClient:
             function=function
         )
         self.tools.register(tool)
-        logger.info(f"已注册工具: {name}")
+        self._tool_sources[name] = source
+        logger.info(f"已注册工具: {name} (source={source})")
 
     def unregister_tool(self, name: str) -> bool:
         """卸载工具。"""
         removed = self.tools.unregister(name)
         if removed:
+            self._tool_sources.pop(name, None)
             logger.info(f"已卸载工具: {name}")
         return removed
 
     def unregister_tools_by_prefix(self, prefix: str) -> int:
         """按前缀批量卸载工具。"""
         removed_count = self.tools.unregister_by_prefix(prefix)
+        for tool_name in list(self._tool_sources.keys()):
+            if tool_name.startswith(prefix):
+                self._tool_sources.pop(tool_name, None)
         if removed_count:
             logger.info(f"Unregistered tools by prefix {prefix}: {removed_count}")
         return removed_count