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.

tests/test_mcp_tool_registration.py

@@ -0,0 +1,44 @@
+import asyncio
+from pathlib import Path
+
+from src.ai.mcp.base import MCPManager, MCPServer
+
+
+class _DummyMCPServer(MCPServer):
+    def __init__(self):
+        super().__init__(name="dummy", version="1.0.0")
+
+    async def initialize(self):
+        self.register_tool(
+            name="echo",
+            description="Echo text",
+            input_schema={
+                "type": "object",
+                "properties": {"text": {"type": "string"}},
+                "required": ["text"],
+            },
+            handler=self.echo,
+        )
+
+    async def echo(self, text: str) -> str:
+        return text
+
+
+def test_mcp_manager_exports_tool_metadata_for_ai(tmp_path: Path):
+    manager = MCPManager(tmp_path / "mcp.json")
+    asyncio.run(manager.register_server(_DummyMCPServer()))
+
+    tools = asyncio.run(manager.get_all_tools_for_ai())
+    assert len(tools) == 1
+    function_info = tools[0]["function"]
+    assert function_info["name"] == "dummy.echo"
+    assert function_info["description"] == "Echo text"
+    assert function_info["parameters"]["required"] == ["text"]
+
+
+def test_mcp_manager_execute_tool(tmp_path: Path):
+    manager = MCPManager(tmp_path / "mcp.json")
+    asyncio.run(manager.register_server(_DummyMCPServer()))
+
+    result = asyncio.run(manager.execute_tool("dummy.echo", {"text": "hello"}))
+    assert result == "hello"