fix bugs

Author: twishabansal

packages/toolbox-core/src/toolbox_core/mcp_transport.py

@@ -13,13 +13,18 @@
 # limitations under the License.
 
 import uuid
-import asyncio
-from typing import Any, Mapping, Optional
+from typing import Any, Mapping, Optional, Union
 
 from aiohttp import ClientSession
 
 from .itransport import ITransport
-from .protocol import ManifestSchema, Protocol
+from .protocol import (
+    AdditionalPropertiesSchema,
+    ManifestSchema,
+    ParameterSchema,
+    Protocol,
+    ToolSchema,
+)
 
 
 class McpHttpTransport(ITransport):
@@ -50,6 +55,49 @@ def __init__(
     def base_url(self) -> str:
         return self.__base_url
 
+    def _convert_tool_schema(self, tool_data: dict) -> ToolSchema:
+        parameters = []
+        input_schema = tool_data.get("inputSchema", {})
+        properties = input_schema.get("properties", {})
+        required = input_schema.get("required", [])
+
+        for name, schema in properties.items():
+            additional_props_value = schema.get("additionalProperties")
+            final_additional_properties: Union[bool, AdditionalPropertiesSchema] = True
+
+            if isinstance(additional_props_value, dict):
+                final_additional_properties = AdditionalPropertiesSchema(
+                    type=additional_props_value["type"]
+                )
+            parameters.append(
+                ParameterSchema(
+                    name=name,
+                    type=schema["type"],
+                    description=schema.get("description", ""),
+                    required=name in required,
+                    additionalProperties=final_additional_properties,
+                )
+            )
+
+        return ToolSchema(description=tool_data["description"], parameters=parameters)
+
+    async def _list_tools(
+        self,
+        toolset_name: Optional[str] = None,
+        headers: Optional[Mapping[str, str]] = None,
+    ) -> Any:
+        """Private helper to fetch the raw tool list from the server."""
+        # TODO: Do not use lazy initialisation
+        if not self.__mcp_initialized:
+            await self._initialize_session()
+        if toolset_name:
+            url = f"{self.__base_url}/mcp/{toolset_name}"
+        else:
+            url = f"{self.__base_url}/mcp/"
+        return await self._send_request(
+            url=url, method="tools/list", params={}, headers=headers
+        )
+
     async def tool_get(
         self, tool_name: str, headers: Optional[Mapping[str, str]] = None
     ) -> ManifestSchema:
@@ -63,6 +111,12 @@ async def tool_get(
         else:
             raise ValueError(f"Tool '{tool_name}' not found.")
 
+        tool_details = ManifestSchema(
+            serverVersion=self.__server_version,
+            tools={tool_name: tool_def},
+        )
+        return tool_details
+
     async def tools_list(
         self,
         toolset_name: Optional[str] = None,
@@ -80,15 +134,21 @@ async def tools_list(
         return ManifestSchema(**result)
 
     async def tool_invoke(
-        self, tool_name: str, arguments: dict, headers: Mapping[str, str]
-    ) -> dict:
+        self, tool_name: str, arguments: dict, headers: Optional[Mapping[str, str]]
+    ) -> str:
         """Invokes a specific tool on the server using the MCP protocol."""
         url = f"{self.__base_url}/mcp/"
         params = {"name": tool_name, "arguments": arguments}
         result = await self._send_request(
             url=url, method="tools/call", params=params, headers=headers
         )
-        return result
+        all_content = result.get("content", result)
+        content_str = "".join(
+            content.get("text", "")
+            for content in all_content
+            if isinstance(content, dict)
+        )
+        return content_str or "null"
 
     async def close(self):
         if self.__manage_session and not self.__session.closed:
@@ -109,7 +169,13 @@ async def _send_request(
             "params": params,
             "id": request_id,
         }
-        async with self.__session.post(url, json=payload, headers=headers) as response:
+
+        if not method.startswith("notifications/"):
+            payload["id"] = str(uuid.uuid4())
+
+        async with self.__session.post(
+            url, json=payload, headers=req_headers
+        ) as response:
             if not response.ok:
                 error_text = await response.text()
                 raise RuntimeError(