From 8c933cb72f2c41f2c11960f7897342cddb7e2b83 Mon Sep 17 00:00:00 2001 From: Alex Mahrou Date: Sun, 24 Aug 2025 08:29:24 -0400 Subject: [PATCH] Emit object-shaped MCP tool schemas for Codex/Claude compatibility --- pyproject.toml | 5 +- src/main.py | 27 +++++++ src/tool_args.py | 118 +++++++++++++++++++++++++++++ src/tools/account.py | 4 +- src/tools/ai.py | 31 ++++++-- src/tools/backtests.py | 37 +++++++-- src/tools/compile.py | 11 ++- src/tools/files.py | 25 ++++-- src/tools/lean_versions.py | 6 +- src/tools/live.py | 61 ++++++++++++--- src/tools/live_commands.py | 13 +++- src/tools/mcp_server_version.py | 7 +- src/tools/object_store.py | 25 ++++-- src/tools/optimizations.py | 29 +++++-- src/tools/project.py | 32 ++++++-- src/tools/project_collaboration.py | 19 +++-- src/tools/project_nodes.py | 9 ++- tests/tests_tool_schemas.py | 24 ++++++ utils.py | 8 ++ 19 files changed, 422 insertions(+), 69 deletions(-) create mode 100644 src/tool_args.py create mode 100644 tests/tests_tool_schemas.py create mode 100644 utils.py diff --git a/pyproject.toml b/pyproject.toml index b886a72..2c064d2 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -8,8 +8,9 @@ dependencies = [ "python-dotenv>=0.23.0", "httpx>=0.28.1", "mcp[cli]>=1.9.3", - "requests" + "pydantic>=2", + "requests>=2" ] [tool.pytest.ini_options] -pythonpath = "src tests" +pythonpath = ["src"] diff --git a/src/main.py b/src/main.py index a0e2c0f..37f0e2f 100644 --- a/src/main.py +++ b/src/main.py @@ -45,6 +45,33 @@ for f in registration_functions: f(mcp) +# Temporary shim to normalize tool input schemas for Codex/Anthropic +async def _list_tools_with_shim(): + tools = await original_list_tools() + for tool in tools: + schema = tool.inputSchema or {"type": "object", "properties": {}} + if schema.get("properties") and list(schema["properties"].keys()) == ["args"]: + schema = schema["properties"]["args"] + schema.setdefault("type", "object") + schema.setdefault("additionalProperties", False) + + def _fix(node): + if isinstance(node, dict): + if node.get("type") == "integer": + node["type"] = "number" + for v in node.values(): + _fix(v) + elif isinstance(node, list): + for v in node: + _fix(v) + + _fix(schema) + tool.inputSchema = schema + return tools + +original_list_tools = mcp.list_tools +mcp.list_tools = _list_tools_with_shim + if __name__ == "__main__": # Load the organization workspace. OrganizationWorkspace.load() diff --git a/src/tool_args.py b/src/tool_args.py new file mode 100644 index 0000000..613cf39 --- /dev/null +++ b/src/tool_args.py @@ -0,0 +1,118 @@ +from __future__ import annotations + +from typing import Any, Dict, List, Type, Union, get_args, get_origin + +from pydantic import BaseModel, Field, create_model, field_validator + + +class ArgsBaseModel(BaseModel): + """Base class for tool argument models with strict object schema.""" + + model_config = { + "extra": "forbid", + "json_schema_extra": {"additionalProperties": False}, + } + + +_MODEL_CACHE: Dict[type[BaseModel], type[BaseModel]] = {} + + +def _convert_annotation( + annotation: Any, + field_name: str, + validators: Dict[str, classmethod], +) -> Any: + origin = get_origin(annotation) + if origin is None: + if isinstance(annotation, type) and issubclass(annotation, BaseModel): + return get_args_model(annotation) + if annotation is int: + # Replace ints with floats and create validator to coerce to int + @field_validator(field_name, mode="before") + def _coerce(cls, v): + return None if v is None else int(v) + + validators[f"coerce_{field_name}"] = _coerce + return float + return annotation + if origin in (list, List): + inner = _convert_annotation(get_args(annotation)[0], field_name, validators) + return List[inner] + if origin is Union: + converted = [ + _convert_annotation(arg, field_name, validators) for arg in get_args(annotation) + ] + return Union[tuple(converted)] + return annotation + + +def get_args_model(model_cls: Type[BaseModel]) -> Type[BaseModel]: + """Create an Args model for the given request model.""" + if model_cls in _MODEL_CACHE: + return _MODEL_CACHE[model_cls] + + fields: Dict[str, tuple[Any, Any]] = {} + validators: Dict[str, classmethod] = {} + + for name, field in model_cls.model_fields.items(): + ann = _convert_annotation(field.annotation, name, validators) + default = field.default if not field.is_required() else ... + ge = le = None + for meta in getattr(field, "metadata", []): + if hasattr(meta, "ge"): + ge = meta.ge + if hasattr(meta, "le"): + le = meta.le + fields[name] = ( + ann, + Field( + default, + description=getattr(field, "description", None), + ge=ge, + le=le, + ), + ) + + ArgsModel = create_model( + f"{model_cls.__name__}Args", + __base__=ArgsBaseModel, + __validators__=validators, + **fields, + ) + + _MODEL_CACHE[model_cls] = ArgsModel + return ArgsModel + + +from typing import Callable + + +def tool_with_args( + mcp, + request_model: Type[BaseModel] | None = None, + **decorator_kwargs: Any, +) -> Callable: + """Decorator to wrap MCP tools with generated Args models.""" + + def decorator(func: Callable) -> Callable: + if request_model is None: + ArgsModel = ArgsBaseModel + + async def inner(args: ArgsModel): + return await func() + else: + ArgsModel = get_args_model(request_model) + + async def inner(args: ArgsModel): + model = request_model(**args.model_dump()) + return await func(model) + + inner.__name__ = func.__name__ + inner.__doc__ = func.__doc__ + inner.__annotations__ = { + "args": ArgsModel, + "return": func.__annotations__.get("return", Any), + } + return mcp.tool(**decorator_kwargs)(inner) + + return decorator diff --git a/src/tools/account.py b/src/tools/account.py index 0cb8af9..f70782f 100644 --- a/src/tools/account.py +++ b/src/tools/account.py @@ -1,9 +1,11 @@ from api_connection import post from models import AccountResponse +from tool_args import tool_with_args def register_account_tools(mcp): # Read - @mcp.tool( + @tool_with_args( + mcp, annotations={ 'title': 'Read account', 'readOnlyHint': True, diff --git a/src/tools/ai.py b/src/tools/ai.py index 90c066c..d943b88 100644 --- a/src/tools/ai.py +++ b/src/tools/ai.py @@ -14,10 +14,13 @@ SyntaxCheckResponse, SearchResponse ) +from tool_args import tool_with_args def register_ai_tools(mcp): # Get backtest initialization errors - @mcp.tool( + @tool_with_args( + mcp, + BasicFilesRequest, annotations={ 'title': 'Check initialization errors', 'readOnlyHint': True } @@ -29,14 +32,20 @@ async def check_initialization_errors( return await post('/ai/tools/backtest-init', model) # Complete code - @mcp.tool(annotations={'title': 'Complete code', 'readOnlyHint': True}) + @tool_with_args( + mcp, + CodeCompletionRequest, + annotations={'title': 'Complete code', 'readOnlyHint': True} + ) async def complete_code( model: CodeCompletionRequest) -> CodeCompletionResponse: """Show the code completion for a specific text input.""" return await post('/ai/tools/complete', model) # Enchance error message - @mcp.tool( + @tool_with_args( + mcp, + ErrorEnhanceRequest, annotations={'title': 'Enhance error message', 'readOnlyHint': True} ) async def enhance_error_message( @@ -45,7 +54,9 @@ async def enhance_error_message( return await post('/ai/tools/error-enhance', model) # Update code to PEP8 - @mcp.tool( + @tool_with_args( + mcp, + PEP8ConvertRequest, annotations={'title': 'Update code to PEP8', 'readOnlyHint': True} ) async def update_code_to_pep8( @@ -54,13 +65,21 @@ async def update_code_to_pep8( return await post('/ai/tools/pep8-convert', model) # Check syntax - @mcp.tool(annotations={'title': 'Check syntax', 'readOnlyHint': True}) + @tool_with_args( + mcp, + BasicFilesRequest, + annotations={'title': 'Check syntax', 'readOnlyHint': True} + ) async def check_syntax(model: BasicFilesRequest) -> SyntaxCheckResponse: """Check the syntax of a code.""" return await post('/ai/tools/syntax-check', model) # Search - @mcp.tool(annotations={'title': 'Search QuantConnect', 'readOnlyHint': True}) + @tool_with_args( + mcp, + SearchRequest, + annotations={'title': 'Search QuantConnect', 'readOnlyHint': True} + ) async def search_quantconnect(model: SearchRequest) -> SearchResponse: """Search for content in QuantConnect.""" return await post('/ai/tools/search', model) diff --git a/src/tools/backtests.py b/src/tools/backtests.py index 4dbec4a..ddea140 100644 --- a/src/tools/backtests.py +++ b/src/tools/backtests.py @@ -19,10 +19,13 @@ BacktestReportGeneratingResponse, RestResponse ) +from tool_args import tool_with_args def register_backtest_tools(mcp): # Create - @mcp.tool( + @tool_with_args( + mcp, + CreateBacktestRequest, annotations={ 'title': 'Create backtest', 'destructiveHint': False @@ -34,20 +37,30 @@ async def create_backtest( return await post('/backtests/create', model) # Read statistics for a single backtest. - @mcp.tool(annotations={'title': 'Read backtest', 'readOnlyHint': True}) + @tool_with_args( + mcp, + ReadBacktestRequest, + annotations={'title': 'Read backtest', 'readOnlyHint': True} + ) async def read_backtest(model: ReadBacktestRequest) -> BacktestResponse: """Read the results of a backtest.""" return await post('/backtests/read', model) # Read a summary of all the backtests. - @mcp.tool(annotations={'title': 'List backtests', 'readOnlyHint': True}) + @tool_with_args( + mcp, + ListBacktestRequest, + annotations={'title': 'List backtests', 'readOnlyHint': True} + ) async def list_backtests( model: ListBacktestRequest) -> BacktestSummaryResponse: """List all the backtests for the project.""" return await post('/backtests/list', model) # Read the chart of a single backtest. - @mcp.tool( + @tool_with_args( + mcp, + ReadBacktestChartRequest, annotations={'title': 'Read backtest chart', 'readOnlyHint': True} ) async def read_backtest_chart( @@ -56,7 +69,9 @@ async def read_backtest_chart( return await post('/backtests/chart/read', model) # Read the orders of a single backtest. - @mcp.tool( + @tool_with_args( + mcp, + ReadBacktestOrdersRequest, annotations={'title': 'Read backtest orders', 'readOnlyHint': True} ) async def read_backtest_orders( @@ -65,7 +80,9 @@ async def read_backtest_orders( return await post('/backtests/orders/read', model) # Read the insights of a single backtest. - @mcp.tool( + @tool_with_args( + mcp, + ReadBacktestInsightsRequest, annotations={'title': 'Read backtest insights', 'readOnlyHint': True} ) async def read_backtest_insights( @@ -84,7 +101,9 @@ async def read_backtest_insights( # return await post('/backtests/read/report', model) # Update - @mcp.tool( + @tool_with_args( + mcp, + UpdateBacktestRequest, annotations={'title': 'Update backtest', 'idempotentHint': True} ) async def update_backtest(model: UpdateBacktestRequest) -> RestResponse: @@ -92,7 +111,9 @@ async def update_backtest(model: UpdateBacktestRequest) -> RestResponse: return await post('/backtests/update', model) # Delete - @mcp.tool( + @tool_with_args( + mcp, + DeleteBacktestRequest, annotations={'title': 'Delete backtest', 'idempotentHint': True} ) async def delete_backtest(model: DeleteBacktestRequest) -> RestResponse: diff --git a/src/tools/compile.py b/src/tools/compile.py index ae46e91..3145177 100644 --- a/src/tools/compile.py +++ b/src/tools/compile.py @@ -5,10 +5,13 @@ CreateCompileResponse, ReadCompileResponse ) +from tool_args import tool_with_args def register_compile_tools(mcp): # Create - @mcp.tool( + @tool_with_args( + mcp, + CreateCompileRequest, annotations={'title': 'Create compile', 'destructiveHint': False} ) async def create_compile( @@ -17,7 +20,11 @@ async def create_compile( return await post('/compile/create', model) # Read - @mcp.tool(annotations={'title': 'Read compile', 'readOnlyHint': True}) + @tool_with_args( + mcp, + ReadCompileRequest, + annotations={'title': 'Read compile', 'readOnlyHint': True} + ) async def read_compile(model: ReadCompileRequest) -> ReadCompileResponse: """Read a compile packet job result.""" return await post('/compile/read', model) diff --git a/src/tools/files.py b/src/tools/files.py index 0c48fd6..31bea88 100644 --- a/src/tools/files.py +++ b/src/tools/files.py @@ -10,11 +10,14 @@ RestResponse, ProjectFilesResponse ) +from tool_args import tool_with_args def register_file_tools(mcp): # Create - @mcp.tool( + @tool_with_args( + mcp, + CreateProjectFileRequest, annotations={ 'title': 'Create file', 'destructiveHint': False, @@ -27,7 +30,11 @@ async def create_file( return await post('/files/create', add_code_source_id(model)) # Read - @mcp.tool(annotations={'title': 'Read file', 'readOnlyHint': True}) + @tool_with_args( + mcp, + ReadFilesRequest, + annotations={'title': 'Read file', 'readOnlyHint': True} + ) async def read_file(model: ReadFilesRequest) -> ProjectFilesResponse: """Read a file from a project, or all files in the project if no file name is provided. @@ -35,7 +42,9 @@ async def read_file(model: ReadFilesRequest) -> ProjectFilesResponse: return await post('/files/read', add_code_source_id(model)) # Update name - @mcp.tool( + @tool_with_args( + mcp, + UpdateFileNameRequest, annotations={'title': 'Update file name', 'idempotentHint': True} ) async def update_file_name(model: UpdateFileNameRequest) -> RestResponse: @@ -43,7 +52,9 @@ async def update_file_name(model: UpdateFileNameRequest) -> RestResponse: return await post('/files/update', add_code_source_id(model)) # Update contents - @mcp.tool( + @tool_with_args( + mcp, + UpdateFileContentsRequest, annotations={'title': 'Update file contents', 'idempotentHint': True} ) async def update_file_contents( @@ -60,7 +71,11 @@ async def patch_file(model: PatchFileRequest) -> RestResponse: return await post('/files/patch', add_code_source_id(model)) # Delete - @mcp.tool(annotations={'title': 'Delete file', 'idempotentHint': True}) + @tool_with_args( + mcp, + DeleteFileRequest, + annotations={'title': 'Delete file', 'idempotentHint': True} + ) async def delete_file(model: DeleteFileRequest) -> RestResponse: """Delete a file in a project.""" return await post('/files/delete', add_code_source_id(model)) diff --git a/src/tools/lean_versions.py b/src/tools/lean_versions.py index ce34a3c..172dff3 100644 --- a/src/tools/lean_versions.py +++ b/src/tools/lean_versions.py @@ -1,9 +1,11 @@ from api_connection import post -from models import LeanVersionsResponse +from models import LeanVersionsResponse +from tool_args import tool_with_args def register_lean_version_tools(mcp): # Read - @mcp.tool( + @tool_with_args( + mcp, annotations={'title': 'Read LEAN versions', 'readOnlyHint': True} ) async def read_lean_versions() -> LeanVersionsResponse: diff --git a/src/tools/live.py b/src/tools/live.py index b0c161b..41c4c19 100644 --- a/src/tools/live.py +++ b/src/tools/live.py @@ -25,6 +25,7 @@ ReadLiveLogsResponse, RestResponse ) +from tool_args import tool_with_args async def handle_loading_response(response, text): if 'progress' in response: @@ -34,9 +35,11 @@ async def handle_loading_response(response, text): def register_live_trading_tools(mcp): # Authenticate - @mcp.tool( + @tool_with_args( + mcp, + AuthorizeExternalConnectionRequest, annotations={ - 'title': 'Authorize external connection', + 'title': 'Authorize external connection', 'readOnlyHint': False, 'destructiveHint': False, 'idempotentHint': True @@ -69,7 +72,9 @@ async def authorize_connection( return await post('/live/auth0/read', model, 800.0) # Create - @mcp.tool( + @tool_with_args( + mcp, + CreateLiveAlgorithmRequest, annotations={ 'title': 'Create live algorithm', 'destructiveHint': False } @@ -80,21 +85,33 @@ async def create_live_algorithm( return await post('/live/create', model) # Read (singular) - @mcp.tool(annotations={'title': 'Read live algorithm', 'readOnly': True}) + @tool_with_args( + mcp, + ReadLiveAlgorithmRequest, + annotations={'title': 'Read live algorithm', 'readOnly': True} + ) async def read_live_algorithm( model: ReadLiveAlgorithmRequest) -> LiveAlgorithmResults: """Read details of a live algorithm.""" return await post('/live/read', model) # Read (all). - @mcp.tool(annotations={'title': 'List live algorithms', 'readOnly': True}) + @tool_with_args( + mcp, + ListLiveAlgorithmsRequest, + annotations={'title': 'List live algorithms', 'readOnly': True} + ) async def list_live_algorithms( model: ListLiveAlgorithmsRequest) -> LiveAlgorithmListResponse: """List all your past and current live trading deployments.""" return await post('/live/list', model) # Read a chart. - @mcp.tool(annotations={'title': 'Read live chart', 'readOnly': True}) + @tool_with_args( + mcp, + ReadLiveChartRequest, + annotations={'title': 'Read live chart', 'readOnly': True} + ) async def read_live_chart( model: ReadLiveChartRequest) -> ReadChartResponse: """Read a chart from a live algorithm.""" @@ -103,7 +120,11 @@ async def read_live_chart( ) # Read the logs. - @mcp.tool(annotations={'title': 'Read live logs', 'readOnly': True}) + @tool_with_args( + mcp, + ReadLiveLogsRequest, + annotations={'title': 'Read live logs', 'readOnly': True} + ) async def read_live_logs( model: ReadLiveLogsRequest) -> ReadLiveLogsResponse: """Get the logs of a live algorithm. @@ -112,7 +133,11 @@ async def read_live_logs( return await post('/live/logs/read', model) # Read the portfolio state. - @mcp.tool(annotations={'title': 'Read live portfolio', 'readOnly': True}) + @tool_with_args( + mcp, + ReadLivePortfolioRequest, + annotations={'title': 'Read live portfolio', 'readOnly': True} + ) async def read_live_portfolio( model: ReadLivePortfolioRequest) -> LivePortfolioResponse: """Read out the portfolio state of a live algorithm. @@ -121,7 +146,11 @@ async def read_live_portfolio( return await post('/live/portfolio/read', model) # Read the orders. - @mcp.tool(annotations={'title': 'Read live orders', 'readOnly': True}) + @tool_with_args( + mcp, + ReadLiveOrdersRequest, + annotations={'title': 'Read live orders', 'readOnly': True} + ) async def read_live_orders( model: ReadLiveOrdersRequest) -> LiveOrdersResponse: """Read out the orders of a live algorithm. @@ -132,7 +161,11 @@ async def read_live_orders( ) # Read the insights. - @mcp.tool(annotations={'title': 'Read live insights', 'readOnly': True}) + @tool_with_args( + mcp, + ReadLiveInsightsRequest, + annotations={'title': 'Read live insights', 'readOnly': True} + ) async def read_live_insights( model: ReadLiveInsightsRequest) -> LiveInsightsResponse: """Read out the insights of a live algorithm. @@ -141,7 +174,9 @@ async def read_live_insights( return await post('/live/insights/read', model) # Update (stop) - @mcp.tool( + @tool_with_args( + mcp, + StopLiveAlgorithmRequest, annotations={'title': 'Stop live algorithm', 'idempotentHint': True} ) async def stop_live_algorithm( @@ -150,7 +185,9 @@ async def stop_live_algorithm( return await post('/live/update/stop', model) # Update (liquidate) - @mcp.tool( + @tool_with_args( + mcp, + LiquidateLiveAlgorithmRequest, annotations={ 'title': 'Liquidate live algorithm', 'idempotentHint': True } diff --git a/src/tools/live_commands.py b/src/tools/live_commands.py index 8b7d669..6fd48f5 100644 --- a/src/tools/live_commands.py +++ b/src/tools/live_commands.py @@ -4,17 +4,26 @@ BroadcastLiveCommandRequest, RestResponse ) +from tool_args import tool_with_args def register_live_trading_command_tools(mcp): # Create (singular algorithm) - @mcp.tool(annotations={'title': 'Create live command'}) + @tool_with_args( + mcp, + CreateLiveCommandRequest, + annotations={'title': 'Create live command'} + ) async def create_live_command( model: CreateLiveCommandRequest) -> RestResponse: """Send a command to a live trading algorithm.""" return await post('/live/commands/create', model) # Create (multiple algorithms) - Broadcast - @mcp.tool(annotations={'title': 'Broadcast live command'}) + @tool_with_args( + mcp, + BroadcastLiveCommandRequest, + annotations={'title': 'Broadcast live command'} + ) async def broadcast_live_command( model: BroadcastLiveCommandRequest) -> RestResponse: """Broadcast a live command to all live algorithms in an diff --git a/src/tools/mcp_server_version.py b/src/tools/mcp_server_version.py index 24a5156..bad2c31 100644 --- a/src/tools/mcp_server_version.py +++ b/src/tools/mcp_server_version.py @@ -1,10 +1,12 @@ from __init__ import __version__ import requests +from tool_args import tool_with_args def register_mcp_server_version_tools(mcp): # Read current version - @mcp.tool( + @tool_with_args( + mcp, annotations={ 'title': 'Read QC MCP Server version', 'readOnlyHint': True } @@ -14,7 +16,8 @@ async def read_mcp_server_version() -> str: return __version__ # Read latest version - @mcp.tool( + @tool_with_args( + mcp, annotations={ 'title': 'Read latest QC MCP Server version', 'readOnlyHint': True } diff --git a/src/tools/object_store.py b/src/tools/object_store.py index b9ffcf5..a23630a 100644 --- a/src/tools/object_store.py +++ b/src/tools/object_store.py @@ -11,10 +11,13 @@ ListObjectStoreResponse, RestResponse ) +from tool_args import tool_with_args def register_object_store_tools(mcp): # Create - @mcp.tool( + @tool_with_args( + mcp, + ObjectStoreBinaryFile, annotations={ 'title': 'Upload Object Store file', 'idempotentHint': True } @@ -39,7 +42,9 @@ async def upload_object( return response.json() # Read file metadata - @mcp.tool( + @tool_with_args( + mcp, + GetObjectStorePropertiesRequest, annotations={ 'title': 'Read Object Store file properties', 'readOnlyHint': True } @@ -55,7 +60,9 @@ async def read_object_properties( return await post('/object/properties', model) # Read file job Id - @mcp.tool( + @tool_with_args( + mcp, + GetObjectStoreJobIdRequest, annotations={ 'title': 'Read Object Store file job Id', 'destructiveHint': False } @@ -68,7 +75,9 @@ async def read_object_store_file_job_id( return await post('/object/get', model) # Read file download URL - @mcp.tool( + @tool_with_args( + mcp, + GetObjectStoreURLRequest, annotations={ 'title': 'Read Object Store file download URL', 'readOnlyHint': True @@ -80,7 +89,9 @@ async def read_object_store_file_download_url( return await post('/object/get', model) # Read all files - @mcp.tool( + @tool_with_args( + mcp, + ListObjectStoreRequest, annotations={'title': 'List Object Store files', 'readOnlyHint': True} ) async def list_object_store_files( @@ -91,7 +102,9 @@ async def list_object_store_files( return await post('/object/list', model) # Delete - @mcp.tool( + @tool_with_args( + mcp, + DeleteObjectStoreRequest, annotations={ 'title': 'Delete Object Store file', 'idempotentHint': True diff --git a/src/tools/optimizations.py b/src/tools/optimizations.py index 2f9ed0c..d39d500 100644 --- a/src/tools/optimizations.py +++ b/src/tools/optimizations.py @@ -12,10 +12,13 @@ ReadOptimizationResponse, RestResponse ) +from tool_args import tool_with_args def register_optimization_tools(mcp): # Estimate cost - @mcp.tool( + @tool_with_args( + mcp, + EstimateOptimizationRequest, annotations={ 'title': 'Estimate optimization time', 'readOnlyHint': True, @@ -29,7 +32,9 @@ async def estimate_optimization_time( return await post('/optimizations/estimate', model) # Create - @mcp.tool( + @tool_with_args( + mcp, + CreateOptimizationRequest, annotations={ 'title': 'Create optimization', 'destructiveHint': False @@ -41,7 +46,9 @@ async def create_optimization( return await post('/optimizations/create', model) # Read a single optimization job. - @mcp.tool( + @tool_with_args( + mcp, + ReadOptimizationRequest, annotations={'title': 'Read optimization', 'readOnlyHint': True} ) async def read_optimization( @@ -50,7 +57,9 @@ async def read_optimization( return await post('/optimizations/read', model) # Read all optimizations for a project. - @mcp.tool( + @tool_with_args( + mcp, + ListOptimizationRequest, annotations={'title': 'List optimizations', 'readOnlyHint': True} ) async def list_optimizations( @@ -59,7 +68,9 @@ async def list_optimizations( return await post('/optimizations/list', model) # Update the optimization name. - @mcp.tool( + @tool_with_args( + mcp, + UpdateOptimizationRequest, annotations={'title': 'Update optimization', 'idempotentHint': True} ) async def update_optimization( @@ -68,7 +79,9 @@ async def update_optimization( return await post('/optimizations/update', model) # Update the optimization status (stop). - @mcp.tool( + @tool_with_args( + mcp, + AbortOptimizationRequest, annotations={'title': 'Abort optimization', 'idempotentHint': True} ) async def abort_optimization( @@ -77,7 +90,9 @@ async def abort_optimization( return await post('/optimizations/abort', model) # Delete - @mcp.tool( + @tool_with_args( + mcp, + DeleteOptimizationRequest, annotations={'title': 'Delete optimization', 'idempotentHint': True} ) async def delete_optimization(model: DeleteOptimizationRequest) -> RestResponse: diff --git a/src/tools/project.py b/src/tools/project.py index b733a61..dc273ab 100644 --- a/src/tools/project.py +++ b/src/tools/project.py @@ -1,16 +1,19 @@ from api_connection import post from models import ( - CreateProjectRequest, - ReadProjectRequest, + CreateProjectRequest, + ReadProjectRequest, UpdateProjectRequest, DeleteProjectRequest, ProjectListResponse, RestResponse ) +from tool_args import tool_with_args def register_project_tools(mcp): # Create - @mcp.tool( + @tool_with_args( + mcp, + CreateProjectRequest, annotations={ 'title': 'Create project', 'destructiveHint': False, @@ -22,25 +25,40 @@ async def create_project(model: CreateProjectRequest) -> ProjectListResponse: return await post('/projects/create', model) # Read (singular) - @mcp.tool(annotations={'title': 'Read project', 'readOnlyHint': True}) + @tool_with_args( + mcp, + ReadProjectRequest, + annotations={'title': 'Read project', 'readOnlyHint': True} + ) async def read_project(model: ReadProjectRequest) -> ProjectListResponse: """List the details of a project or a set of recent projects.""" return await post('/projects/read', model) # Read (all) - @mcp.tool(annotations={'title': 'List projects', 'readOnlyHint': True}) + @tool_with_args( + mcp, + annotations={'title': 'List projects', 'readOnlyHint': True} + ) async def list_projects() -> ProjectListResponse: """List the details of all projects.""" return await post('/projects/read') # Update - @mcp.tool(annotations={'title': 'Update project', 'idempotentHint': True}) + @tool_with_args( + mcp, + UpdateProjectRequest, + annotations={'title': 'Update project', 'idempotentHint': True} + ) async def update_project(model: UpdateProjectRequest) -> RestResponse: """Update a project's name or description.""" return await post('/projects/update', model) # Delete - @mcp.tool(annotations={'title': 'Delete project', 'idempotentHint': True}) + @tool_with_args( + mcp, + DeleteProjectRequest, + annotations={'title': 'Delete project', 'idempotentHint': True} + ) async def delete_project(model: DeleteProjectRequest) -> RestResponse: """Delete a project.""" return await post('/projects/delete', model) diff --git a/src/tools/project_collaboration.py b/src/tools/project_collaboration.py index e99d1f1..bf3b72a 100644 --- a/src/tools/project_collaboration.py +++ b/src/tools/project_collaboration.py @@ -1,7 +1,7 @@ from api_connection import post from code_source_id import add_code_source_id from models import ( - CreateCollaboratorRequest, + CreateCollaboratorRequest, ReadCollaboratorsRequest, UpdateCollaboratorRequest, DeleteCollaboratorRequest, @@ -12,10 +12,13 @@ DeleteCollaboratorResponse, RestResponse ) +from tool_args import tool_with_args def register_project_collaboration_tools(mcp): # Create - @mcp.tool( + @tool_with_args( + mcp, + CreateCollaboratorRequest, annotations={ 'title': 'Create project collaborator', 'destructiveHint': False, @@ -28,7 +31,9 @@ async def create_project_collaborator( return await post('/projects/collaboration/create', model) # Read - @mcp.tool( + @tool_with_args( + mcp, + ReadCollaboratorsRequest, annotations={ 'title': 'Read project collaborators', 'readOnlyHint': True @@ -40,7 +45,9 @@ async def read_project_collaborators( return await post('/projects/collaboration/read', model) # Update - @mcp.tool( + @tool_with_args( + mcp, + UpdateCollaboratorRequest, annotations={ 'title': 'Update project collaborator', 'idempotentHint': True @@ -52,7 +59,9 @@ async def update_project_collaborator( return await post('/projects/collaboration/update', model) # Delete - @mcp.tool( + @tool_with_args( + mcp, + DeleteCollaboratorRequest, annotations={ 'title': 'Delete project collaborator', 'idempotentHint': True diff --git a/src/tools/project_nodes.py b/src/tools/project_nodes.py index cc34f8b..ec92290 100644 --- a/src/tools/project_nodes.py +++ b/src/tools/project_nodes.py @@ -4,10 +4,13 @@ UpdateProjectNodesRequest, ProjectNodesResponse ) +from tool_args import tool_with_args def register_project_node_tools(mcp): # Read - @mcp.tool( + @tool_with_args( + mcp, + ReadProjectNodesRequest, annotations={'title': 'Read project nodes', 'readOnlyHint': True} ) async def read_project_nodes( @@ -16,7 +19,9 @@ async def read_project_nodes( return await post('/projects/nodes/read', model) # Update - @mcp.tool( + @tool_with_args( + mcp, + UpdateProjectNodesRequest, annotations={ 'title': 'Update project nodes', 'destructiveHint': False, diff --git a/tests/tests_tool_schemas.py b/tests/tests_tool_schemas.py new file mode 100644 index 0000000..bf58ba5 --- /dev/null +++ b/tests/tests_tool_schemas.py @@ -0,0 +1,24 @@ +import asyncio +import pytest + +from main import mcp + + +def _has_integer_type(node): + if isinstance(node, dict): + if node.get("type") == "integer": + return True + return any(_has_integer_type(v) for v in node.values()) + if isinstance(node, list): + return any(_has_integer_type(v) for v in node) + return False + + +@pytest.mark.asyncio +async def test_all_tool_schemas_are_objects(): + tools = await mcp.list_tools() + for tool in tools: + schema = tool.inputSchema or {} + assert schema.get("type") == "object" + assert schema.get("additionalProperties") is False + assert not _has_integer_type(schema) diff --git a/utils.py b/utils.py new file mode 100644 index 0000000..25b2914 --- /dev/null +++ b/utils.py @@ -0,0 +1,8 @@ +import os +import sys + +tests_dir = os.path.join(os.path.dirname(__file__), "tests") +if tests_dir not in sys.path: + sys.path.append(tests_dir) + +from tests.utils import * # noqa: F401,F403