feat: publish tool schemas to dataplane

mcpgateway/services/dataplane_publisher.py

@@ -10,6 +10,7 @@
 
 # Standard
 import asyncio
+from collections.abc import Sequence
 from collections import defaultdict
 import logging
 import os
@@ -67,11 +68,23 @@ class BackendConfig(TypedDict):
     remove_headers: list[str]
     capabilities: dict[str, Any]
     allowed_tool_names: list[str]
+    tool_schemas: dict[str, dict[str, Any]]
     allowed_resource_names: list[str]
     allowed_resource_uris: list[str]
     allowed_prompt_names: list[str]
 
 
+class GatewayBaseConfig(TypedDict):
+    """Gateway connection fields shared by every virtual-host backend."""
+
+    name: str
+    url: str
+    passthrough_headers: list[str]
+    add_headers: dict[str, str]
+    remove_headers: list[str]
+    capabilities: dict[str, Any]
+
+
 class VirtualHostConfig(TypedDict):
     """Virtual host configuration mapping backend IDs to their configs."""
 
@@ -84,7 +97,27 @@ class UserConfig(TypedDict):
     virtual_hosts: dict[str, VirtualHostConfig]
 
 
-BackendItems = dict[str, list[str]]
+class BackendItems(TypedDict):
+    """Database identifiers associated with one server backend."""
+
+    tools: list[str]
+    resources: list[str]
+    prompts: list[str]
+
+
+class PublishedBackendItems(BackendItems):
+    """User-filtered backend items enriched with tool input schemas."""
+
+    tool_schemas: dict[str, dict[str, Any]]
+
+
+class ToolMetadata(TypedDict):
+    """Tool fields needed to publish dataplane routing configuration."""
+
+    name: str
+    input_schema: dict[str, Any]
+
+
 BackendItemsByServer = dict[str, dict[str, BackendItems]]
 
 
@@ -129,7 +162,7 @@ async def shutdown(self) -> None:
         self.task = None
         logger.info("Dataplane publisher stopped.")
 
-    async def fetch_payload(self) -> dict[str, dict[str, dict[str, Any]]] | None:
+    async def fetch_payload(self) -> dict[str, UserConfig] | None:
         """Fetch the payload to publish to Redis. Returns None on error."""
         user_data = await self.get_data_from_db()
         if user_data is None:
@@ -231,7 +264,7 @@ def create_payload(
             # The dataplane proxies streamable-HTTP upstreams only. Exclude
             # every other transport before building its transport-agnostic
             # backend config.
-            gateway_base = {
+            gateway_base: dict[str, GatewayBaseConfig] = {
                 gateway["id"]: {
                     "name": gateway["name"],
                     "url": gateway["url"],
@@ -261,8 +294,14 @@ def create_payload(
                         continue
 
                     backends[gateway_id] = {
-                        **gateway_config,
+                        "name": gateway_config["name"],
+                        "url": gateway_config["url"],
+                        "passthrough_headers": gateway_config["passthrough_headers"],
+                        "add_headers": gateway_config["add_headers"],
+                        "remove_headers": gateway_config["remove_headers"],
+                        "capabilities": gateway_config["capabilities"],
                         "allowed_tool_names": backend_items["tools"],
+                        "tool_schemas": backend_items["tool_schemas"],
                         "allowed_resource_names": allowed_resource_names,
                         "allowed_resource_uris": allowed_resource_uris,
                         "allowed_prompt_names": allowed_prompt_names,
@@ -326,7 +365,7 @@ async def get_data_from_db(self) -> dict[str, Any] | None:
                         DbResource.uri_template.is_(None),
                     )
                 ).all()
-                tool_rows = db.execute(select(DbTool.id, DbTool.original_name, DbTool.owner_email, DbTool.team_id, DbTool.visibility).where(DbTool.enabled.is_(True))).all()
+                tool_rows = db.execute(select(DbTool.id, DbTool.original_name, DbTool.input_schema, DbTool.owner_email, DbTool.team_id, DbTool.visibility).where(DbTool.enabled.is_(True))).all()
                 backend_items_by_server = self._get_backend_items_by_server(db)
 
                 return {
@@ -353,21 +392,32 @@ def _build_user_data(
         user_email: str,
         team_ids: set[str],
         is_admin: bool,
-        server_rows: list[Any],
-        gateway_rows: list[Any],
-        prompt_rows: list[Any],
-        resource_rows: list[Any],
-        tool_rows: list[Any],
+        server_rows: Sequence[Any],
+        gateway_rows: Sequence[Any],
+        prompt_rows: Sequence[Any],
+        resource_rows: Sequence[Any],
+        tool_rows: Sequence[Any],
         backend_items_by_server: BackendItemsByServer,
     ) -> dict[str, Any]:
         """Build already-filtered dataplane data for one user."""
-        tool_name_by_id = {tool.id: tool.original_name for tool in tool_rows if self._filter_for_user(tool, user_email, team_ids, is_admin=is_admin)}
+        visible_tools = [tool for tool in tool_rows if self._filter_for_user(tool, user_email, team_ids, is_admin=is_admin)]
+        for tool in visible_tools:
+            if not isinstance(tool.input_schema, dict):
+                raise ValueError(f"Tool {tool.id} has a non-object input schema")
+
+        tool_by_id: dict[str, ToolMetadata] = {
+            tool.id: {
+                "name": tool.original_name,
+                "input_schema": tool.input_schema,
+            }
+            for tool in visible_tools
+        }
 
         return {
             "servers": [
                 {
                     "id": server.id,
-                    "backend_items": self._filter_backend_items_for_user(backend_items_by_server.get(server.id, {}), tool_name_by_id),
+                    "backend_items": self._filter_backend_items_for_user(backend_items_by_server.get(server.id, {}), tool_by_id),
                 }
                 for server in server_rows
                 if self._filter_for_user(server, user_email, team_ids, is_admin=is_admin)
@@ -403,11 +453,12 @@ def _filter_for_user(row: Any, user_email: str, team_ids: set[str], is_admin: bo
         return row.team_id in team_ids and visibility == "team"
 
     @staticmethod
-    def _filter_backend_items_for_user(backend_items_by_gateway: dict[str, BackendItems], tool_name_by_id: dict[str, str]) -> dict[str, BackendItems]:
-        """Filter backend tool IDs for one user and convert visible tools to names."""
+    def _filter_backend_items_for_user(backend_items_by_gateway: dict[str, BackendItems], tool_by_id: dict[str, ToolMetadata]) -> dict[str, PublishedBackendItems]:
+        """Filter backend tool IDs for one user and publish visible names with schemas."""
         return {
             gateway_id: {
-                "tools": [tool_name_by_id[tool_id] for tool_id in backend_items["tools"] if tool_id in tool_name_by_id],
+                "tools": [tool_by_id[tool_id]["name"] for tool_id in backend_items["tools"] if tool_id in tool_by_id],
+                "tool_schemas": {tool_by_id[tool_id]["name"]: tool_by_id[tool_id]["input_schema"] for tool_id in backend_items["tools"] if tool_id in tool_by_id},
                 "resources": list(backend_items["resources"]),
                 "prompts": list(backend_items["prompts"]),
             }

tests/unit/mcpgateway/services/test_dataplane_publisher.py

@@ -11,7 +11,6 @@
 
 import pytest
 
-
 USER1_ID = "11111111-1111-1111-1111-111111111111"
 USER2_ID = "22222222-2222-2222-2222-222222222222"
 USER3_ID = "33333333-3333-3333-3333-333333333333"
@@ -214,6 +213,10 @@ async def test_full_payload_generation_with_mock_db():
     tool1.id = "t1"
     tool1.name = "gw1-public_tool"
     tool1.original_name = "public_tool"
+    tool1.input_schema = {
+        "type": "object",
+        "properties": {"region": {"type": "string", "x-mcp-header": "Region"}},
+    }
     tool1.owner_email = "user1@example.com"
     tool1.team_id = "team1"
     tool1.visibility = "public"
@@ -223,6 +226,7 @@ async def test_full_payload_generation_with_mock_db():
     tool2.id = "t2"
     tool2.name = "gw1-private_tool"
     tool2.original_name = "private_tool"
+    tool2.input_schema = {}
     tool2.owner_email = "user1@example.com"
     tool2.team_id = "team1"
     tool2.visibility = "private"
@@ -232,6 +236,7 @@ async def test_full_payload_generation_with_mock_db():
     tool3.id = "t3"
     tool3.name = "gw1-team2_tool"
     tool3.original_name = "team2_tool"
+    tool3.input_schema = {"type": "object", "properties": {"count": {"type": "integer"}}}
     tool3.owner_email = "user2@example.com"
     tool3.team_id = "team2"
     tool3.visibility = "team"
@@ -289,6 +294,10 @@ async def test_full_payload_generation_with_mock_db():
             "remove_headers": ["Cookie"],
             "capabilities": {"resources": {"subscribe": True}},
             "allowed_tool_names": ["public_tool", "private_tool"],
+            "tool_schemas": {
+                "public_tool": tool1.input_schema,
+                "private_tool": {},
+            },
             "allowed_resource_names": ["Resource 1"],
             "allowed_resource_uris": ["resource://one"],
             "allowed_prompt_names": ["Prompt 1"],
@@ -302,6 +311,11 @@ async def test_full_payload_generation_with_mock_db():
         assert "add_headers" in selected_keys, "Gateway SELECT must include add_headers"
         assert "remove_headers" in selected_keys, "Gateway SELECT must include remove_headers"
 
+        tool_execute_call = mock_db.execute.call_args_list[6]
+        tool_stmt = tool_execute_call[0][0]
+        selected_tool_keys = {col.key for col in tool_stmt.selected_columns}
+        assert "input_schema" in selected_tool_keys, "Tool SELECT must include input_schema"
+
         # Verify user2 sees public server but not private server from user1
         user2_config = payload[USER2_ID]
         assert "s1" in user2_config["virtual_hosts"]  # public
@@ -310,13 +324,30 @@ async def test_full_payload_generation_with_mock_db():
         assert "s2" not in user2_config["virtual_hosts"]
         user2_backend = user2_config["virtual_hosts"]["s1"]["backends"]["g1"]
         assert user2_backend["allowed_tool_names"] == ["public_tool", "team2_tool"]
+        assert user2_backend["tool_schemas"] == {
+            "public_tool": tool1.input_schema,
+            "team2_tool": tool3.input_schema,
+        }
 
         # Verify active users with no team membership still get public-only config.
         user3_config = payload[USER3_ID]
         assert "s1" in user3_config["virtual_hosts"]
         assert "s2" not in user3_config["virtual_hosts"]
         user3_backend = user3_config["virtual_hosts"]["s1"]["backends"]["g1"]
         assert user3_backend["allowed_tool_names"] == ["public_tool"]
+        assert user3_backend["tool_schemas"] == {"public_tool": tool1.input_schema}
+
+
+def test_build_user_data_rejects_non_object_tool_schema():
+    """Dataplane snapshots fail closed when an enabled tool schema is malformed."""
+    from unittest.mock import Mock
+
+    from mcpgateway.services.dataplane_publisher import DataplanePublisherService
+
+    tool = Mock(id="bad-tool", original_name="bad", input_schema=None, visibility="public")
+
+    with pytest.raises(ValueError, match="Tool bad-tool has a non-object input schema"):
+        DataplanePublisherService()._build_user_data("user@example.com", set(), False, [], [], [], [], [tool], {})
 
 
 # ============================================================================
@@ -394,7 +425,7 @@ def test_create_payload_filters_empty_backends():
                 {
                     "id": "server1",
                     "backend_items": {
-                        "gateway1": {"tools": [], "resources": [], "prompts": []},
+                        "gateway1": {"tools": [], "tool_schemas": {}, "resources": [], "prompts": []},
                     },
                 }
             ],
@@ -423,7 +454,12 @@ def test_create_payload_excludes_non_streamable_gateways(transport: str):
                 {
                     "id": "server1",
                     "backend_items": {
-                        "gateway_non_streamable": {"tools": ["tool1"], "resources": [], "prompts": []},
+                        "gateway_non_streamable": {
+                            "tools": ["tool1"],
+                            "tool_schemas": {},
+                            "resources": [],
+                            "prompts": [],
+                        },
                     },
                 }
             ],
@@ -450,7 +486,7 @@ def test_create_payload_normalizes_null_passthrough_headers():
                 {
                     "id": "server1",
                     "backend_items": {
-                        "gateway1": {"tools": ["tool1"], "resources": [], "prompts": []},
+                        "gateway1": {"tools": ["tool1"], "tool_schemas": {}, "resources": [], "prompts": []},
                     },
                 }
             ],
@@ -480,7 +516,12 @@ def test_create_payload_handles_missing_references():
                 {
                     "id": "server1",
                     "backend_items": {
-                        "missing_gateway": {"tools": ["tool1"], "resources": ["missing_res"], "prompts": ["missing_prompt"]},
+                        "missing_gateway": {
+                            "tools": ["tool1"],
+                            "tool_schemas": {},
+                            "resources": ["missing_res"],
+                            "prompts": ["missing_prompt"],
+                        },
                     },
                 }
             ],