From 12e93d374f38bc1b21ad89643aea6d60daba12d4 Mon Sep 17 00:00:00 2001 From: YBoy Date: Tue, 7 Apr 2026 03:02:06 +0200 Subject: [PATCH] refactor(api): type MCP tool schema and arguments with TypedDict (#34612) --- api/core/mcp/server/streamable_http.py | 19 +++++++++++++++---- 1 file changed, 15 insertions(+), 4 deletions(-) diff --git a/api/core/mcp/server/streamable_http.py b/api/core/mcp/server/streamable_http.py index 278add8cc9..8de002ae55 100644 --- a/api/core/mcp/server/streamable_http.py +++ b/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}