refactor(api): type MCP tool schema and arguments with TypedDict (#34612)

api/core/mcp/server/streamable_http.py

@@ -1,7 +1,7 @@
 import json
 import logging
 from collections.abc import Mapping
-from typing import Any, cast
+from typing import Any, NotRequired, TypedDict, cast
 
 from graphon.variables.input_entities import VariableEntity, VariableEntityType
 
@@ -15,6 +15,17 @@ from services.app_generate_service import AppGenerateService
 logger = logging.getLogger(__name__)
 
 
+class ToolParameterSchemaDict(TypedDict):
+    type: str
+    properties: dict[str, Any]
+    required: list[str]
+
+
+class ToolArgumentsDict(TypedDict):
+    query: NotRequired[str]
+    inputs: dict[str, Any]
+
+
 def handle_mcp_request(
     app: App,
     request: mcp_types.ClientRequest,
@@ -119,7 +130,7 @@ def handle_list_tools(
             mcp_types.Tool(
                 name=app_name,
                 description=description,
-                inputSchema=parameter_schema,
+                inputSchema=cast(dict[str, Any], parameter_schema),
             )
         ],
     )
@@ -154,7 +165,7 @@ def build_parameter_schema(
     app_mode: str,
     user_input_form: list[VariableEntity],
     parameters_dict: dict[str, str],
-) -> dict[str, Any]:
+) -> ToolParameterSchemaDict:
     """Build parameter schema for the tool"""
     parameters, required = convert_input_form_to_parameters(user_input_form, parameters_dict)
 
@@ -174,7 +185,7 @@ def build_parameter_schema(
     }
 
 
-def prepare_tool_arguments(app: App, arguments: dict[str, Any]) -> dict[str, Any]:
+def prepare_tool_arguments(app: App, arguments: dict[str, Any]) -> ToolArgumentsDict:
     """Prepare arguments based on app mode"""
     if app.mode == AppMode.WORKFLOW:
         return {"inputs": arguments}