From 95318acd8ac982491efb04e155a004ea8e69db35 Mon Sep 17 00:00:00 2001 From: kavmur Date: Thu, 3 Jul 2025 11:08:40 -0500 Subject: [PATCH] feat(experimental): Add mcp to deadline client tools Signed-off-by: kavmur --- .github/workflows/code_quality.yml | 4 +- README.md | 14 +- docs/mcp_guide.md | 131 +++++ pyproject.toml | 5 + requirements-integ-testing.txt | 3 + requirements-testing.txt | 2 + src/deadline/_mcp/__init__.py | 1 + src/deadline/_mcp/registry.py | 70 +++ src/deadline/_mcp/server.py | 13 + src/deadline/_mcp/tools/__init__.py | 1 + src/deadline/_mcp/tools/job.py | 198 ++++++++ src/deadline/_mcp/utils.py | 187 +++++++ src/deadline/client/api/__init__.py | 1 + src/deadline/client/api/_telemetry.py | 9 +- src/deadline/client/cli/_deadline_cli.py | 3 + src/deadline/client/cli/_mcp_server.py | 51 ++ test/integ/deadline_mcp/__init__.py | 1 + .../test_mcp_server_integration.py | 456 ++++++++++++++++++ .../unit/deadline_client/cli/test_cli_farm.py | 4 +- test/unit/deadline_mcp/test_mcp.py | 159 ++++++ test/unit/deadline_mcp/test_mcp_telemetry.py | 358 ++++++++++++++ 21 files changed, 1658 insertions(+), 13 deletions(-) create mode 100644 docs/mcp_guide.md create mode 100644 src/deadline/_mcp/__init__.py create mode 100644 src/deadline/_mcp/registry.py create mode 100644 src/deadline/_mcp/server.py create mode 100644 src/deadline/_mcp/tools/__init__.py create mode 100644 src/deadline/_mcp/tools/job.py create mode 100644 src/deadline/_mcp/utils.py create mode 100644 src/deadline/client/cli/_mcp_server.py create mode 100644 test/integ/deadline_mcp/__init__.py create mode 100644 test/integ/deadline_mcp/test_mcp_server_integration.py create mode 100644 test/unit/deadline_mcp/test_mcp.py create mode 100644 test/unit/deadline_mcp/test_mcp_telemetry.py diff --git a/.github/workflows/code_quality.yml b/.github/workflows/code_quality.yml index fe08afe16..6e1114ba3 100644 --- a/.github/workflows/code_quality.yml +++ b/.github/workflows/code_quality.yml @@ -36,5 +36,5 @@ jobs: os: ${{ matrix.os }} python-version: ${{ matrix.python-version }} ref: ${{inputs.tag}} - # The incremental output download feature doesn't run on Python 3.8, so test coverage is lower - cov-fail-under: "75" + # The incremental output download and mcp feature doesn't run on Python 3.8, so test coverage is lower + cov-fail-under: "74" diff --git a/README.md b/README.md index 0a7249ca1..a55ed9962 100644 --- a/README.md +++ b/README.md @@ -15,6 +15,7 @@ Notable features include: * A library of functions that implement AWS Deadline Cloud's Job Attachments functionality. * A library of functions for creating a job submission UI within any content creation tool that supports Python 3.8+ based plugins and the Qt GUI framework. +* A Model Context Protocol (MCP) server for AI assistant integration, enabling natural language interaction with AWS Deadline Cloud resources. [cas]: https://en.wikipedia.org/wiki/Content-addressable_storage [deadline-cloud]: https://docs.aws.amazon.com/deadline-cloud/latest/userguide/what-is-deadline-cloud.html @@ -62,6 +63,11 @@ or if you want the optional gui dependencies: $ pip install "deadline[gui]" ``` +if you want the optional mcp dependencies: +```sh +$ pip install "deadline[mcp]" +``` + ## Usage After installation it can then be used as a command line tool: @@ -116,8 +122,6 @@ deadline config gui ``` -Shared storage is possible with customer-managed fleets (CMF) but not service-managed fleets (SMF). See [shared storage][shared-storage] for more information. - ## Job Bundles A job bundle is one of the tools that you can use to define jobs for AWS Deadline Cloud. They group an [Open Job Description (OpenJD)][openjd] template with @@ -308,6 +312,12 @@ Available modes: - `USER`: Credentials with full queue-role permissions. - `READ`: Credentials with read-only permissions for queue logs +## Model Context Protocol (MCP) Server + +The AWS Deadline Cloud client includes an MCP server that enables AI assistants to interact with AWS Deadline Cloud resources through natural language. The MCP server uses the [MCP Python SDK](https://github.com/modelcontextprotocol/python-sdk) for simplified tool registration while maintaining full protocol compliance. + +See [MCP Guide](https://github.com/aws-deadline/deadline-cloud/blob/release/docs/mcp_guide.md) for more information. + ## Code of Conduct diff --git a/docs/mcp_guide.md b/docs/mcp_guide.md new file mode 100644 index 000000000..c41705d5a --- /dev/null +++ b/docs/mcp_guide.md @@ -0,0 +1,131 @@ +# AWS Deadline Cloud MCP Server Guide + +**MCP ([Model Context Protocol](https://modelcontextprotocol.io/docs/getting-started/intro))** is an open standard that enables AI assistants to securely connect to external data sources and tools. It acts a bridge that allows Large Language Models (LLMs) like Claude, GPT, or other AI assistants to interact with applications and services through natural language. + +With the AWS Deadline Cloud MCP Server, you can use natural language for various Deadline Cloud workflows such as submitting a job, reading job/farm/queue information etc. + +## User Guide + +1. **Install the server** + ```bash + pip install 'deadline[mcp]' + ``` +2. **Verify the MCP server command is available:** + ```bash + deadline mcp-server --help + ``` + You should see: + ``` + Usage: deadline mcp-server [OPTIONS] + + Start the AWS Deadline Cloud MCP (Model Context Protocol) server. + + The MCP server provides LLM tools with access to AWS Deadline Cloud + operations through the Model Context Protocol. This allows AI assistants to + interact with Deadline Cloud services on your behalf. + + The server will run until interrupted with Ctrl+C or Ctrl+D. + ``` +3. **Configure AWS credentials:** + - Standard AWS credentials (AWS Profiles, environment variables) + - Or Deadline Cloud monitor credentials: + ```bash + deadline auth login + ``` + - Verify authentication status: + ```bash + deadline auth status + ``` +4. **Configure your MCP client** to connect to the server + Add to your MCP configuration (e.g., `./settings/mcp.json`): + ```json + { + "mcpServers": { + "deadline-cloud": { + "command": "deadline", + "args": ["mcp-server"], + "disabled": false, + "autoApprove": [] + } + } + } + ``` +5. **Start having conversations** with your AI assistant about your rendering workflows + +## Example prompts + +``` +- "List all my AWS Deadline Cloud farms" +- "Show me the queues in my farm" +- "List the jobs in my queue" +- "Submit the render job in /path/to/my-job-bundle" +- "Submit a job with priority 80 to my render queue" +- "Show me the status of job job-3a907bac684841f69fc344867ee166de" +- "Download output from job job-3a907bac684841f69fc344867ee166de" +- "Download output from step step-render in job job-3a907bac684841f69fc344867ee166de" +``` + +## Available Tools + +The MCP server provides access to all allowlisted/configured Deadline Cloud API functions through automatic registration: + +- `deadline_list_farms()`: List available farms +- `deadline_list_queues()`: List queues in a farm +- `deadline_list_jobs()`: List jobs in a queue +- `deadline_list_fleets()`: List fleets in a farm +- `deadline_list_storage_profiles_for_queue()`: List storage profiles for a queue + +- `deadline_check_authentication_status()`: Check current authentication status + +- `deadline_submit_job()`: Submit an Open Job Description job bundle to AWS Deadline Cloud +- `deadline_download_job_output()`: Download job output files from AWS Deadline Cloud + + +## Developer Guide + +The MCP server exposes public Deadline Cloud operations as tools that AI assistants can call directly. Tools are defined in `config.py` where each entry maps a tool name to its corresponding function and parameters. + +## Project Structure + +``` +src/deadline/mcp/ +├── server.py # Main server with FastMCP setup and auto-registration +├── config.py # Tool configuration definitions +├── utils.py # Auto-registration utilities +└── tools/ # Tool modules + └── job.py # Job management tools (submit, download) +``` + +### Adding New MCP Tools + +#### Step 1: Use public api operations + +Prefer consistency with existing CLI patterns by using functions from the public API module `deadline.client.api.*` when available. Use direct boto3 calls when no wrapper exists or when the wrapper doesn't provide the needed functionality. Tools like `submit_job` and `download_job_output` are exceptions because they require job attachment functionality that are not available in the public API. + +```python +# ✅ Good: Use existing public API +from deadline.client.api import list_farms, list_queues + +# ❌ Bad: Don't use internal modules +from deadline.client._internal.some_module import internal_function +``` + +#### Step 2: Configure the MCP Tool + +Add your new tool to the `API_TOOLS_CONFIG` in `src/deadline/mcp/config.py`: + +```python +API_TOOLS_CONFIG: dict[str, ToolConfig] = { + # ... existing tools ... + "your_new_tool": { + "func": api.your_new_function, # Must be from deadline.client.api + "params": ["param1", "param2", "optional_param"], # List all parameters, or None if no params + }, +} +``` + +#### Step 3: Test the Integration + +1. **Unit Tests**: Add tests to `test/unit/deadline_mcp/test_mcp.py` +2. **Integration Tests**: Add tests to `test/integ/deadline_mcp/test_mcp_server_integration.py` +3. **Manual Testing**: Test with real MCP clients diff --git a/pyproject.toml b/pyproject.toml index af5a7de52..6c46e9b26 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -56,6 +56,9 @@ gui = [ # If the version changes, update the version in deadline/client/ui/__init__.py "PySide6-essentials >= 6.6,< 6.10", ] +mcp = [ + "mcp >= 1.13.0", +] [project.scripts] deadline = "deadline.client.cli:main" @@ -132,6 +135,7 @@ known-first-party = ["deadline"] # XML, HTML, and terminal reports. [tool.pytest.ini_options] xfail_strict = true + addopts = [ "--durations=5", "--cov=src/deadline", @@ -147,6 +151,7 @@ markers = [ "integ: tests that run against AWS resources", "docker: marks tests to be run only in a Docker environment", "cross_account: tests that run against other aws accounts", + "asyncio: mark test as async", ] # looponfailroots is deprecated, this removes the deprecation from the test output filterwarnings = ["ignore::DeprecationWarning"] diff --git a/requirements-integ-testing.txt b/requirements-integ-testing.txt index 7502b75a0..fb6259ccd 100644 --- a/requirements-integ-testing.txt +++ b/requirements-integ-testing.txt @@ -1 +1,4 @@ deadline-cloud-test-fixtures == 0.18.* +# MCP (Model Context Protocol) library for MCP server integration tests +mcp >= 1.13.0 +pytest-asyncio == 0.* \ No newline at end of file diff --git a/requirements-testing.txt b/requirements-testing.txt index b216519d5..9d8070dfd 100644 --- a/requirements-testing.txt +++ b/requirements-testing.txt @@ -16,3 +16,5 @@ ruff == 0.13.* moto == 5.* jsondiff == 2.* pyinstrument == 5.* +# MCP (Model Context Protocol) library for MCP unit tests +mcp >= 1.13.0; python_version >= '3.10' diff --git a/src/deadline/_mcp/__init__.py b/src/deadline/_mcp/__init__.py new file mode 100644 index 000000000..8d929cc86 --- /dev/null +++ b/src/deadline/_mcp/__init__.py @@ -0,0 +1 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. diff --git a/src/deadline/_mcp/registry.py b/src/deadline/_mcp/registry.py new file mode 100644 index 000000000..f269d927a --- /dev/null +++ b/src/deadline/_mcp/registry.py @@ -0,0 +1,70 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. + +"""Tool registry and mapping definitions for MCP tools in Deadline Cloud.""" + +from typing import Any, Callable, List, Optional, TypedDict, Dict + +from ..client import api +from .tools import job + + +class ToolDefinition(TypedDict): + """Definition of a single MCP tool including its function and parameters.""" + + func: Callable[..., Any] + param_names: Optional[List[str]] + + +def get_tool_definition(tool_name: str) -> ToolDefinition: + """Get the definition for a specific tool.""" + if tool_name not in TOOL_REGISTRY: + raise ValueError(f"Tool '{tool_name}' not found in registry") + return TOOL_REGISTRY[tool_name] + + +def get_all_tool_names() -> List[str]: + """Get all registered tool names.""" + return list(TOOL_REGISTRY.keys()) + + +TOOL_REGISTRY: Dict[str, ToolDefinition] = { + "list_farms": { + "func": api.list_farms, + "param_names": ["nextToken", "principalId", "maxResults"], + }, + "list_queues": { + "func": api.list_queues, + "param_names": ["farmId", "principalId", "status", "nextToken", "maxResults"], + }, + "list_jobs": { + "func": api.list_jobs, + "param_names": ["farmId", "queueId", "principalId", "nextToken", "maxResults"], + }, + "list_fleets": { + "func": api.list_fleets, + "param_names": [ + "farmId", + "principalId", + "displayName", + "status", + "nextToken", + "maxResults", + ], + }, + "list_storage_profiles_for_queue": { + "func": api.list_storage_profiles_for_queue, + "param_names": ["farmId", "queueId", "nextToken", "maxResults"], + }, + "check_authentication_status": { + "func": api.check_authentication_status, + "param_names": None, + }, + "submit_job": { + "func": job.submit_job, + "param_names": None, + }, + "download_job_output": { + "func": job.download_job_output, + "param_names": None, + }, +} diff --git a/src/deadline/_mcp/server.py b/src/deadline/_mcp/server.py new file mode 100644 index 000000000..d4d60b688 --- /dev/null +++ b/src/deadline/_mcp/server.py @@ -0,0 +1,13 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. + +from mcp.server.fastmcp import FastMCP + +from .utils import register_api_tools + +app = FastMCP("deadline-cloud") + +register_api_tools(app, prefix="deadline_") + + +def main(): + app.run() diff --git a/src/deadline/_mcp/tools/__init__.py b/src/deadline/_mcp/tools/__init__.py new file mode 100644 index 000000000..8d929cc86 --- /dev/null +++ b/src/deadline/_mcp/tools/__init__.py @@ -0,0 +1 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. diff --git a/src/deadline/_mcp/tools/job.py b/src/deadline/_mcp/tools/job.py new file mode 100644 index 000000000..37b3c45ef --- /dev/null +++ b/src/deadline/_mcp/tools/job.py @@ -0,0 +1,198 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. + +""" +Deadline Cloud Job tools. +""" + +import io +import json +import os +import time +from contextlib import redirect_stdout +from typing import Any, Dict, Optional + +from ...client.api import create_job_from_job_bundle +from ...client.cli._groups.job_group import _download_job_output +from ...client.config import config_file + +# TODO: Make submit_job tool async once progress reporting feature is supported in clients + + +def submit_job( + job_bundle_dir: str, + job_parameters: Optional[str] = None, + name: Optional[str] = None, + farm_id: Optional[str] = None, + queue_id: Optional[str] = None, + storage_profile_id: Optional[str] = None, + priority: Optional[int] = 50, + max_failed_tasks_count: Optional[int] = None, + max_retries_per_task: Optional[int] = None, + max_worker_count: Optional[int] = None, + job_attachments_file_system: Optional[str] = None, + require_paths_exist: bool = False, + submitter_name: Optional[str] = None, + known_asset_paths: Optional[str] = None, +) -> Dict[str, Any]: + """ + Submit an Open Job Description job bundle to AWS Deadline Cloud. + + Args: + job_bundle_dir: Path to the job bundle directory containing template.json/yaml + job_parameters: JSON string of job parameters in format [{"name": "param_name", "value": "param_value"}] + name: Job name to override the one in the bundle + farm_id: Farm ID to submit to (uses default if not provided) + queue_id: Queue ID to submit to (uses default if not provided) + storage_profile_id: Storage profile ID to use + priority: Job priority (1-100, default 50) + max_failed_tasks_count: Maximum failed tasks before job fails + max_retries_per_task: Maximum retries per task + max_worker_count: Maximum worker count for the job + job_attachments_file_system: File system type (COPIED or VIRTUAL) + require_paths_exist: Return error if input files are missing + submitter_name: Name of the submitting application + known_asset_paths: JSON array of paths that shouldn't generate warnings + + Returns: + Dictionary containing job_id and submission status + """ + start_time = time.time() + + if not os.path.exists(job_bundle_dir): + raise ValueError(f"Job bundle directory does not exist: {job_bundle_dir}") + + if not os.path.isdir(job_bundle_dir): + raise ValueError(f"Path is not a directory: {job_bundle_dir}") + + # Parse job parameters + parsed_job_parameters = [] + if job_parameters: + parsed_job_parameters = json.loads(job_parameters) + if not isinstance(parsed_job_parameters, list): + raise ValueError( + "job_parameters must be a JSON array of objects with 'name' and 'value' keys" + ) + + # Parse known asset paths + parsed_known_asset_paths = [] + if known_asset_paths: + parsed_known_asset_paths = json.loads(known_asset_paths) + if not isinstance(parsed_known_asset_paths, list): + raise ValueError("known_asset_paths must be a JSON array of strings") + + # Read config and use parameter values if provided, otherwise fall back to config defaults + config = config_file.read_config() + + farm_id = farm_id or config.get("defaults", "farm_id", fallback=None) + queue_id = queue_id or config.get("defaults", "queue_id", fallback=None) + storage_profile_id = storage_profile_id or config.get( + "defaults", "storage_profile_id", fallback=None + ) + + if not farm_id: + raise ValueError("farm_id is required") + if not queue_id: + raise ValueError("queue_id is required") + + config.set("defaults", "farm_id", farm_id) + config.set("defaults", "queue_id", queue_id) + if storage_profile_id: + config.set("defaults", "storage_profile_id", storage_profile_id) + + # Submit the job + job_id = create_job_from_job_bundle( + job_bundle_dir=job_bundle_dir, + job_parameters=parsed_job_parameters, + name=name, + config=config, + priority=priority, + max_failed_tasks_count=max_failed_tasks_count, + max_retries_per_task=max_retries_per_task, + max_worker_count=max_worker_count, + job_attachments_file_system=job_attachments_file_system, + require_paths_exist=require_paths_exist, + submitter_name=submitter_name or "MCP", + known_asset_paths=parsed_known_asset_paths, + ) + + total_time = time.time() - start_time + + return { + "status": "success", + "job_id": job_id, + "message": f"Successfully submitted job bundle from {job_bundle_dir}", + "total_time_seconds": round(total_time, 1), + } + + +def download_job_output( + farm_id: Optional[str] = None, + queue_id: Optional[str] = None, + job_id: Optional[str] = None, + step_id: Optional[str] = None, + task_id: Optional[str] = None, + conflict_resolution: Optional[str] = None, +) -> Dict[str, Any]: + """ + Download job output files from AWS Deadline Cloud. + + Args: + farm_id: Farm ID (uses default if not provided) + queue_id: Queue ID (uses default if not provided) + job_id: Job ID to download output from + step_id: Optional step ID to download output from specific step + task_id: Optional task ID to download output from specific task (requires step_id) + conflict_resolution: How to handle file conflicts - SKIP, OVERWRITE, or CREATE_COPY (default) + + Returns: + Dictionary containing download status and summary + """ + start_time = time.time() + + if task_id and not step_id: + raise ValueError("step_id is required when task_id is provided") + if not job_id: + raise ValueError("job_id is required") + if conflict_resolution and conflict_resolution.upper() not in [ + "SKIP", + "OVERWRITE", + "CREATE_COPY", + ]: + raise ValueError( + f"Invalid conflict_resolution: {conflict_resolution}. Must be SKIP, OVERWRITE, or CREATE_COPY" + ) + + config = config_file.read_config() + config.set("defaults", "farm_id", farm_id or config.get("defaults", "farm_id", fallback="")) + config.set("defaults", "queue_id", queue_id or config.get("defaults", "queue_id", fallback="")) + config.set("defaults", "job_id", job_id) + + if not config.has_section("settings"): + config.add_section("settings") + config.set("settings", "auto_accept", "true") + if conflict_resolution: + config.set("settings", "conflict_resolution", conflict_resolution.upper()) + + captured_output = io.StringIO() + with redirect_stdout(captured_output): + _download_job_output( + config, + config.get("defaults", "farm_id"), + config.get("defaults", "queue_id"), + job_id, + step_id, + task_id, + is_json_format=False, + ) + + output_text = captured_output.getvalue() + total_time = time.time() - start_time + + return { + "status": "success", + "job_id": job_id, + "step_id": step_id, + "task_id": task_id, + "total_time_seconds": round(total_time, 1), + "output": output_text.strip(), + } diff --git a/src/deadline/_mcp/utils.py b/src/deadline/_mcp/utils.py new file mode 100644 index 000000000..80427fca5 --- /dev/null +++ b/src/deadline/_mcp/utils.py @@ -0,0 +1,187 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. + +"""Utility functions for MCP tool registration in Deadline Cloud.""" + +import inspect +import json +import logging +import time +from typing import Any, Callable, Dict, List, Optional + +from mcp.server.fastmcp import FastMCP + +from .registry import get_all_tool_names, get_tool_definition, ToolDefinition +from deadline.client.api._telemetry import get_deadline_cloud_library_telemetry_client + +logger = logging.getLogger(__name__) + + +def _default_serializer(obj: Any) -> Any: + """Default serializer for API responses.""" + if hasattr(obj, "__dict__"): + return obj.__dict__ + elif hasattr(obj, "to_dict"): + return obj.to_dict() + return str(obj) + + +def _default_error_handler(e: Exception) -> Dict: + """Default error handler for API calls.""" + error_info = {"error": str(e), "type": type(e).__name__} + if hasattr(e, "response") and hasattr(e.response, "status_code"): + error_info["status_code"] = e.response.status_code + logger.error(f"API tool error: {error_info}", exc_info=True) + return error_info + + +def _create_wrapper( + config: ToolDefinition, serializer: Callable, error_handler: Callable +) -> Callable: + """Create a wrapper function based on the tool configuration.""" + func = config["func"] + param_names = config["param_names"] + + # Inspect signature once and reuse + func_sig = inspect.signature(func) + + if param_names is None: + param_names = [p for p in func_sig.parameters.keys() if p != "config"] + + # Check if function accepts 'config' parameter + has_config_param = "config" in func_sig.parameters + + signature = inspect.Signature( + [ + inspect.Parameter( + name, + inspect.Parameter.KEYWORD_ONLY, + default=None, + annotation=func_sig.parameters[name].annotation + if name in func_sig.parameters + else str, + ) + for name in param_names + ], + return_annotation=dict, + ) + + def wrapper(**kwargs) -> dict: + start_t = time.perf_counter_ns() + success = True + error_type = None + + try: + # Filter out empty/null values that MCP clients might send + filtered_kwargs = {k: v for k, v in kwargs.items() if v not in (None, "", "null")} + + # Type conversion now handled by preserving original function annotations + + # Add config parameter if function accepts it + if has_config_param: + filtered_kwargs["config"] = None + + result = json.loads(json.dumps(func(**filtered_kwargs), default=serializer)) + except Exception as e: + success = False + error_type = type(e).__name__ + result = error_handler(e) + + # Record telemetry data + try: + telemetry_client = get_deadline_cloud_library_telemetry_client() + latency = time.perf_counter_ns() - start_t + + telemetry_client.record_event( + event_type="com.amazon.rum.deadline.mcp.latency", + event_details={ + "latency": latency, + "tool_name": wrapper.__name__, + "usage_mode": "MCP", + }, + ) + + telemetry_client.record_event( + event_type="com.amazon.rum.deadline.mcp.usage", + event_details={ + "tool_name": wrapper.__name__, + "is_success": success, + "error_type": error_type, + "usage_mode": "MCP", + }, + ) + except Exception as telemetry_error: + logger.debug( + f"Failed to record telemetry for MCP tool {func.__name__}: {telemetry_error}" + ) + + return result + + wrapper.__signature__ = signature # type: ignore[attr-defined] + return wrapper + + +def register_api_tools( + app: FastMCP, + tools: Optional[List[Callable]] = None, + prefix: str = "", + error_handler: Optional[Callable] = None, + serializer: Optional[Callable] = None, +) -> None: + """Register API tools with the MCP server. + + Args: + app: FastMCP application instance + tools: Optional list of specific functions to register. If None, registers all configured tools. + prefix: Prefix to add to tool names + error_handler: Optional custom error handler + serializer: Optional custom serializer + """ + error_handler = error_handler or _default_error_handler + serializer = serializer or _default_serializer + + if tools is not None: + # Register only the specified tools + for func in tools: + if not callable(func): + raise ValueError(f"Tool {func} is not callable") + + if hasattr(func, "_mcp_tool_registered"): + continue # Already registered, skip + + # Find the tool name for this function in the registry + tool_name = None + for name in get_all_tool_names(): + config = get_tool_definition(name) + if config["func"] == func: + tool_name = name + break + + if tool_name is None: + raise ValueError(f"Function {func.__name__} not found in tool registry") + + config = get_tool_definition(tool_name) + wrapper = _create_wrapper(config, serializer, error_handler) + wrapper.__name__ = tool_name + + description = func.__doc__ + wrapper.__doc__ = description + + app.tool(name=f"{prefix}{tool_name}", description=description)(wrapper) + func._mcp_tool_registered = True # type: ignore[attr-defined] + else: + # Register all configured tools + for tool_name in get_all_tool_names(): + config = get_tool_definition(tool_name) + func = config["func"] + + if not callable(func) or hasattr(func, "_mcp_tool_registered"): + continue + + wrapper = _create_wrapper(config, serializer, error_handler) + wrapper.__name__ = tool_name + + description = func.__doc__ + wrapper.__doc__ = description + + app.tool(name=f"{prefix}{tool_name}", description=description)(wrapper) + func._mcp_tool_registered = True # type: ignore[attr-defined] diff --git a/src/deadline/client/api/__init__.py b/src/deadline/client/api/__init__.py index c6cef4785..380182546 100644 --- a/src/deadline/client/api/__init__.py +++ b/src/deadline/client/api/__init__.py @@ -72,6 +72,7 @@ list_fleets, list_storage_profiles_for_queue, ) + from ._queue_parameters import get_queue_parameter_definitions from ._queue_credentials import ( assume_queue_role_for_user, diff --git a/src/deadline/client/api/_telemetry.py b/src/deadline/client/api/_telemetry.py index 0c62d1a1c..14c59198c 100644 --- a/src/deadline/client/api/_telemetry.py +++ b/src/deadline/client/api/_telemetry.py @@ -1,7 +1,7 @@ # Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. import atexit -from functools import lru_cache +from functools import lru_cache, wraps import json import logging import os @@ -429,12 +429,8 @@ def record_function_latency_telemetry_event(**decorator_kwargs: Any) -> Callable """ def inner(function: F) -> F: + @wraps(function) def wrapper(*args: Any, **kwargs: Any) -> Any: - """ - Wrapper to time a function for latency telemetry - :param * Python variable argument. See https://docs.python.org/3/glossary.html#term-parameter - :param ** Python variable argument. See https://docs.python.org/3/glossary.html#term-parameter - """ start_t = time.perf_counter_ns() ret_val = function(*args, **kwargs) end_t = time.perf_counter_ns() @@ -449,7 +445,6 @@ def wrapper(*args: Any, **kwargs: Any) -> Any: return ret_val - wrapper.__doc__ = function.__doc__ return cast(F, wrapper) return inner diff --git a/src/deadline/client/cli/_deadline_cli.py b/src/deadline/client/cli/_deadline_cli.py index 5d8847489..8b82e5d63 100644 --- a/src/deadline/client/cli/_deadline_cli.py +++ b/src/deadline/client/cli/_deadline_cli.py @@ -20,3 +20,6 @@ attachment_group, manifest_group, ) + +# MCP server command +from ._mcp_server import cli_mcp_server # noqa: F401 diff --git a/src/deadline/client/cli/_mcp_server.py b/src/deadline/client/cli/_mcp_server.py new file mode 100644 index 000000000..3b72bef38 --- /dev/null +++ b/src/deadline/client/cli/_mcp_server.py @@ -0,0 +1,51 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. + +""" +The `deadline mcp-server` command. +""" + +import sys +import click + +from ._common import _handle_error +from ._main import main +from deadline.client.api._telemetry import get_deadline_cloud_library_telemetry_client + + +@main.command(name="mcp-server") +@_handle_error +def cli_mcp_server(): + """ + Start the AWS Deadline Cloud MCP (Model Context Protocol) server. + + The MCP server provides LLM tools with access to AWS Deadline Cloud operations + through the Model Context Protocol. This allows AI assistants to interact with + Deadline Cloud services on your behalf. + + The server will run until interrupted with Ctrl+C or Ctrl+D. + + Note: This command requires MCP dependencies. Install them with: + pip install 'deadline[mcp]' + """ + try: + from ..._mcp.server import main as mcp_main + except ImportError: + click.echo( + "Error: MCP dependencies not installed.\n" + "Please install them with: pip install 'deadline[mcp]'", + err=True, + ) + sys.exit(1) + + # Record server startup telemetry + try: + telemetry_client = get_deadline_cloud_library_telemetry_client() + telemetry_client.record_event( + event_type="com.amazon.rum.deadline.mcp.server_startup", + event_details={"usage_mode": "MCP", "startup_method": "cli"}, + ) + except Exception: + # Don't let telemetry errors affect server startup + pass + + mcp_main() diff --git a/test/integ/deadline_mcp/__init__.py b/test/integ/deadline_mcp/__init__.py new file mode 100644 index 000000000..8d929cc86 --- /dev/null +++ b/test/integ/deadline_mcp/__init__.py @@ -0,0 +1 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. diff --git a/test/integ/deadline_mcp/test_mcp_server_integration.py b/test/integ/deadline_mcp/test_mcp_server_integration.py new file mode 100644 index 000000000..ca9119a47 --- /dev/null +++ b/test/integ/deadline_mcp/test_mcp_server_integration.py @@ -0,0 +1,456 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. + +"""Integration test for the Deadline Cloud MCP Server.""" + +import json +import pytest +import boto3 +import asyncio +import sys +from unittest.mock import patch, MagicMock + +# Skip all tests in this module if MCP dependencies are not available +pytest.importorskip("mcp", reason="MCP dependencies not available") + +from deadline._mcp import server +from deadline._mcp.registry import TOOL_REGISTRY, get_tool_definition + +from deadline.client.api import list_farms +from deadline.client.cli._mcp_server import cli_mcp_server +from click.testing import CliRunner + + +@pytest.fixture(scope="session") +def get_boto_session(): + """ + Fixture to establish authenticated session for MCP integration tests. + """ + session = boto3.Session() + sts_client = session.client("sts") + _ = sts_client.get_caller_identity() + return session + + +@pytest.mark.asyncio +@pytest.mark.integ +async def test_mcp_server_integration(get_boto_session): + """Test that MCP server starts, registers tools, and handles tool calls.""" + with patch("deadline.client.api._session.get_boto3_session") as mock_session: + mock_session.return_value = get_boto_session + + app = server.app + + # 1. Check if all tools in TOOL_REGISTRY are available + tools_result = await app.list_tools() + tools = tools_result if isinstance(tools_result, list) else tools_result.tools + available_tool_names = [tool.name for tool in tools] + expected_tools = [f"deadline_{tool_name}" for tool_name in TOOL_REGISTRY.keys()] + missing_tools = [tool for tool in expected_tools if tool not in available_tool_names] + assert not missing_tools, f"Missing expected tools from TOOL_REGISTRY: {missing_tools}" + + # 2. Call deadline_check_authentication_status tool and verify response + auth_result = await app.call_tool("deadline_check_authentication_status", {}) + assert auth_result is not None, "Authentication tool call should return a result" + + if auth_result and hasattr(auth_result[0], "text"): # type: ignore[index] + try: + auth_data = json.loads(auth_result[0].text) # type: ignore[index] + assert "error" not in auth_data, ( + f"Authentication check failed with error: {auth_data.get('error')}" + ) + except json.JSONDecodeError as e: + pytest.fail( + f"Check authentication tool returned non-JSON response: {auth_result[0].text[:200]}... (JSONDecodeError: {e})" # type: ignore[index] + ) + + # 3. Call deadline_list_farms tool and verify response + farms_result = await app.call_tool("deadline_list_farms", {}) + assert farms_result is not None, "List farms tool call should return a result" + if farms_result and hasattr(farms_result[0], "text"): # type: ignore[index] + try: + farms_data = json.loads(farms_result[0].text) # type: ignore[index] + assert "error" not in farms_data, ( + f"List farms failed with error: {farms_data.get('error')}" + ) + except json.JSONDecodeError as e: + pytest.fail( + f"List farms returned non-JSON response: {farms_result[0].text[:200]}... (JSONDecodeError: {e})" # type: ignore[index] + ) + + # 4. Check metadata on deadline_list_queues tool + list_queues_tool = None + for tool in tools: + if tool.name == "deadline_list_queues": + list_queues_tool = tool + break + + assert list_queues_tool is not None, "deadline_list_queues tool not found" + + expected_params = get_tool_definition("list_queues")["param_names"] + assert expected_params is not None, ( + "list_queues should have parameters defined in TOOL_REGISTRY" + ) + + assert hasattr(list_queues_tool, "inputSchema"), ( + "deadline_list_queues tool should have an inputSchema" + ) + + input_schema = list_queues_tool.inputSchema + assert input_schema is not None, "inputSchema should not be None" + + available_params = [] + if isinstance(input_schema, dict) and "properties" in input_schema: + available_params = list(input_schema["properties"].keys()) + elif hasattr(input_schema, "properties") and input_schema.properties: + available_params = list(input_schema.properties.keys()) + elif hasattr(input_schema, "model_fields"): + available_params = list(input_schema.model_fields.keys()) + + assert "farmId" in available_params, ( + f"farmId parameter should be available in list_queues tool. " + f"Available params: {available_params}, Expected params: {expected_params}" + ) + + missing_params = [param for param in expected_params if param not in available_params] + assert not missing_params, ( + f"Missing expected parameters in list_queues tool: {missing_params}. " + f"Available: {available_params}, Expected: {expected_params}" + ) + + return True + + +@pytest.mark.asyncio +@pytest.mark.integ +async def test_tool_description_extraction(get_boto_session): + """Test that tool descriptions are properly extracted from original function docstrings.""" + app = server.app + + tools_result = await app.list_tools() + tools = tools_result if isinstance(tools_result, list) else tools_result.tools + + list_farms_tool = None + for tool in tools: + if tool.name == "deadline_list_farms": + list_farms_tool = tool + break + + assert list_farms_tool is not None, "deadline_list_farms tool not found" + + original_docstring = list_farms.__doc__ + assert original_docstring is not None, "Original docstring should be available" + + assert list_farms_tool.description == original_docstring, ( + f"Tool description should match original docstring.\n" + f"Expected: {original_docstring}\n" + f"Got: {list_farms_tool.description}" + ) + + assert "Calls the deadline:ListFarms API call" in list_farms_tool.description, ( + "Tool description should come from the function's doc string" + ) + + return True + + +class MCPClient: + """Simple MCP client for testing the server process via stdio.""" + + def __init__(self, process): + self.process = process + self.request_id = 0 + + async def send_request(self, method, params=None): + """Send a JSON-RPC request to the MCP server.""" + self.request_id += 1 + request = { + "jsonrpc": "2.0", + "id": self.request_id, + "method": method, + "params": params or {}, + } + + request_json = json.dumps(request) + "\n" + self.process.stdin.write(request_json.encode()) + await self.process.stdin.drain() + + response_line = await self.process.stdout.readline() + if not response_line: + raise RuntimeError("No response from MCP server") + + return json.loads(response_line.decode().strip()) + + async def initialize(self): + """Initialize the MCP session.""" + response = await self.send_request( + "initialize", + { + "protocolVersion": "2024-11-05", + "capabilities": {"tools": {}}, + "clientInfo": {"name": "test-client", "version": "1.0.0"}, + }, + ) + + if "error" in response: + raise RuntimeError(f"Initialize failed: {response['error']}") + + # Send initialized notification (no response expected) + await self.send_notification("notifications/initialized") + return response + + async def send_notification(self, method, params=None): + """Send a JSON-RPC notification (no response expected).""" + request = {"jsonrpc": "2.0", "method": method, "params": params or {}} + + request_json = json.dumps(request) + "\n" + self.process.stdin.write(request_json.encode()) + await self.process.stdin.drain() + + async def list_tools(self): + """List available tools.""" + response = await self.send_request("tools/list", {}) + if "error" in response: + raise RuntimeError(f"List tools failed: {response['error']}") + return response["result"]["tools"] + + async def call_tool(self, name, arguments=None): + """Call a tool.""" + response = await self.send_request( + "tools/call", {"name": name, "arguments": arguments or {}} + ) + if "error" in response: + raise RuntimeError(f"Tool call failed: {response['error']}") + return response["result"] + + +@pytest.mark.asyncio +@pytest.mark.integ +async def test_mcp_server_process_integration(get_boto_session): + """Test MCP server as a separate process via stdio communication.""" + with patch("deadline.client.api._session.get_boto3_session") as mock_session: + mock_session.return_value = get_boto_session + process = await asyncio.create_subprocess_exec( + sys.executable, + "-c", + "from deadline._mcp.server import main; main()", + stdin=asyncio.subprocess.PIPE, + stdout=asyncio.subprocess.PIPE, + ) + + try: + client = MCPClient(process) + init_response = await client.initialize() + + assert "result" in init_response, "Initialize should return a result" + assert "capabilities" in init_response["result"], "Server should return capabilities" + + tools = await client.list_tools() + assert len(tools) > 0, "Server should have tools available" + + tool_names = [tool["name"] for tool in tools] + expected_tools = [f"deadline_{tool_name}" for tool_name in TOOL_REGISTRY.keys()] + missing_tools = [tool for tool in expected_tools if tool not in tool_names] + assert not missing_tools, f"Missing expected tools: {missing_tools}" + + farms_result = await client.call_tool("deadline_list_farms") + assert "content" in farms_result, "List farms should return content" + assert len(farms_result["content"]) > 0, "List farms should return non-empty content" + + farms_content = farms_result["content"][0]["text"] + try: + farms_data = json.loads(farms_content) + if "error" in farms_data: + print(f"List farms returned error (may be expected): {farms_data['error']}") + except json.JSONDecodeError: + pytest.fail(f"List farms tool returned invalid JSON: {farms_content[:200]}...") + + return True + + finally: + if process and process.returncode is None: + process.terminate() + try: + await asyncio.wait_for(process.wait(), timeout=5.0) + except asyncio.TimeoutError: + process.kill() + await process.wait() + + +@pytest.mark.asyncio +@pytest.mark.integ +async def test_mcp_tool_telemetry(get_boto_session): + """Test that MCP tools record telemetry events.""" + with patch("deadline.client.api._session.get_boto3_session") as mock_session: + mock_session.return_value = get_boto_session + + mock_telemetry_client = MagicMock() + with patch( + "deadline._mcp.utils.get_deadline_cloud_library_telemetry_client" + ) as mock_get_client: + mock_get_client.return_value = mock_telemetry_client + + app = server.app + + auth_result = await app.call_tool("deadline_check_authentication_status", {}) + assert auth_result is not None, "Authentication tool call should return a result" + assert mock_telemetry_client.record_event.call_count >= 2, ( + "Should record at least 2 telemetry events (latency and usage)" + ) + + recorded_calls = mock_telemetry_client.record_event.call_args_list + event_types = [call[1]["event_type"] for call in recorded_calls] + + assert "com.amazon.rum.deadline.mcp.latency" in event_types, ( + "Should record latency telemetry event" + ) + assert "com.amazon.rum.deadline.mcp.usage" in event_types, ( + "Should record usage telemetry event" + ) + + latency_events = [ + call + for call in recorded_calls + if call[1]["event_type"] == "com.amazon.rum.deadline.mcp.latency" + ] + usage_events = [ + call + for call in recorded_calls + if call[1]["event_type"] == "com.amazon.rum.deadline.mcp.usage" + ] + + assert len(latency_events) >= 1, "Should have at least one latency event" + assert len(usage_events) >= 1, "Should have at least one usage event" + + # Assert latency event + latency_event_details = latency_events[0][1]["event_details"] + expected_latency_keys = {"latency", "tool_name", "usage_mode"} + assert set(latency_event_details.keys()) == expected_latency_keys, ( + f"Latency event should contain exactly {expected_latency_keys}, " + f"but got {set(latency_event_details.keys())}" + ) + assert isinstance(latency_event_details["latency"], int), "Latency should be an integer" + assert latency_event_details["tool_name"] == "check_authentication_status" + assert latency_event_details["usage_mode"] == "MCP" + + # Assert usage event + usage_event_details = usage_events[0][1]["event_details"] + expected_usage_keys = {"tool_name", "is_success", "error_type", "usage_mode"} + assert set(usage_event_details.keys()) == expected_usage_keys, ( + f"Usage event should contain exactly {expected_usage_keys}, " + f"but got {set(usage_event_details.keys())}" + ) + assert usage_event_details["tool_name"] == "check_authentication_status" + assert isinstance(usage_event_details["is_success"], bool), ( + "is_success should be a boolean" + ) + assert usage_event_details["usage_mode"] == "MCP" + assert usage_event_details["error_type"] is None, ( + "error_type should be None for successful calls" + ) + + return True + + +@pytest.mark.integ +def test_mcp_server_startup_telemetry(): + """Test that MCP server startup records telemetry.""" + + mock_telemetry_client = MagicMock() + + with patch( + "deadline.client.cli._mcp_server.get_deadline_cloud_library_telemetry_client" + ) as mock_get_client, patch("deadline._mcp.server.main") as mock_mcp_main: + mock_get_client.return_value = mock_telemetry_client + + runner = CliRunner() + result = runner.invoke(cli_mcp_server, []) + + assert result.exit_code == 0, f"Command failed with output: {result.output}" + + mock_telemetry_client.record_event.assert_called_once_with( + event_type="com.amazon.rum.deadline.mcp.server_startup", + event_details={"usage_mode": "MCP", "startup_method": "cli"}, + ) + + mock_mcp_main.assert_called_once() + + +@pytest.mark.asyncio +@pytest.mark.integ +async def test_mcp_tool_error_handling(): + """Test that MCP tools properly handle and return errors using real AWS service validation.""" + + def parse_tool_result(result): + """Helper to parse and validate tool result JSON.""" + assert result is not None, "Tool call should return a result" + if result and hasattr(result[0], "text"): + return json.loads(result[0].text) + return {} + + mock_session = MagicMock() + + deadline_client = boto3.client("deadline", region_name="us-west-2") + + mock_session.client.side_effect = lambda service, **kwargs: {"deadline": deadline_client}.get( + service, MagicMock() + ) + + with patch("deadline.client.api._session.get_boto3_session") as mock_get_session: + mock_get_session.return_value = mock_session + app = server.app + + # Test 1: Invalid farm ID format + queues_data = parse_tool_result( + await app.call_tool("deadline_list_queues", {"farmId": "invalid-farm-id-format"}) + ) + assert "error" in queues_data, f"Expected AWS error, got: {queues_data}" + assert "type" in queues_data, "Error response should contain 'type' field" + # Check for specific AWS exception types that indicate authentication/authorization issues + expected_error_types = [ + "AccessDeniedException", + "ValidationException", + "UnauthorizedException", + "ExpiredTokenException", + ] + assert any(error_type in queues_data["type"] for error_type in expected_error_types), ( + f"Expected one of {expected_error_types}, got: {queues_data['type']}" + ) + assert "error occurred" in queues_data["error"].lower() + print(f"✅ Test 1 - Got expected error type: {queues_data['type']}") + + # Test 2: Invalid queue ID format + jobs_data = parse_tool_result( + await app.call_tool( + "deadline_list_jobs", + { + "farmId": "farm-1234567890abcdef1234567890abcdef", # Valid format but likely non-existent + "queueId": "invalid-queue-format", # Invalid format will trigger validation error + }, + ) + ) + assert "error" in jobs_data, f"Expected AWS error, got: {jobs_data}" + assert "type" in jobs_data, "Error response should contain 'type' field" + assert any(error_type in jobs_data["type"] for error_type in expected_error_types), ( + f"Expected one of {expected_error_types}, got: {jobs_data['type']}" + ) + assert "error occurred" in jobs_data["error"].lower() + print(f"✅ Test 2 - Got expected error type: {jobs_data['type']}") + + # Test 3: Non-existent IDs + fleets_data = parse_tool_result( + await app.call_tool( + "deadline_list_fleets", + { + "farmId": "farm-0000000000000000000000000000000" # Properly formatted but non-existent + }, + ) + ) + assert "error" in fleets_data, f"Expected AWS error, got: {fleets_data}" + assert "type" in fleets_data, "Error response should contain 'type' field" + resource_error_types = expected_error_types + [ + "ResourceNotFoundException", + "ForbiddenException", + ] + assert any(error_type in fleets_data["type"] for error_type in resource_error_types), ( + f"Expected one of {resource_error_types}, got: {fleets_data['type']}" + ) diff --git a/test/unit/deadline_client/cli/test_cli_farm.py b/test/unit/deadline_client/cli/test_cli_farm.py index a57ce6865..e8f67c401 100644 --- a/test/unit/deadline_client/cli/test_cli_farm.py +++ b/test/unit/deadline_client/cli/test_cli_farm.py @@ -89,7 +89,7 @@ def test_cli_farm_list_client_error(fresh_deadline_config): assert result.exit_code != 0 -def test_cli_farm_get(fresh_deadline_config): +def test_cli_farm_get(fresh_deadline_config, mock_telemetry): """ Confirm that the CLI interface prints out the expected farm, given mock data. """ @@ -151,7 +151,7 @@ def test_cli_farm_get_no_default_set(fresh_deadline_config): assert result.exit_code != 0 -def test_cli_farm_get_explicit_farm_id(fresh_deadline_config): +def test_cli_farm_get_explicit_farm_id(fresh_deadline_config, mock_telemetry): """ Confirm that the CLI interface prints out the expected farm, given mock data. """ diff --git a/test/unit/deadline_mcp/test_mcp.py b/test/unit/deadline_mcp/test_mcp.py new file mode 100644 index 000000000..396eda776 --- /dev/null +++ b/test/unit/deadline_mcp/test_mcp.py @@ -0,0 +1,159 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. + +"""Unit tests for MCP server""" + +from unittest.mock import MagicMock + +import pytest + +try: + from deadline.client import api + from deadline._mcp.registry import TOOL_REGISTRY, ToolDefinition + from deadline._mcp.utils import ( + _create_wrapper, + _default_error_handler, + _default_serializer, + register_api_tools, + ) + + MCP_AVAILABLE = True +except ImportError: + MCP_AVAILABLE = False + +if not MCP_AVAILABLE: + pytest.skip("MCP dependencies not available", allow_module_level=True) + + +class TestToolRegistry: + """Test the TOOL_REGISTRY structure.""" + + def test_registry_structure(self): + """Test that TOOL_REGISTRY has valid structure.""" + assert isinstance(TOOL_REGISTRY, dict) + assert len(TOOL_REGISTRY) > 0 + + for definition in TOOL_REGISTRY.values(): + assert "func" in definition + assert "param_names" in definition + assert callable(definition["func"]) + + def test_expected_functions(self): + """Test that expected functions are in the registry.""" + expected = ["list_farms", "submit_job", "download_job_output"] + for func_name in expected: + assert func_name in TOOL_REGISTRY + + +class TestUtilityFunctions: + """Test utility functions.""" + + def test_default_serializer(self): + """Test serializer with different object types.""" + assert _default_serializer("string") == "string" + assert _default_serializer(123) == "123" + + class TestObj: + def __init__(self): + self.attr = "value" + + assert _default_serializer(TestObj()) == {"attr": "value"} + + def test_default_error_handler(self): + """Test error handler formats errors correctly.""" + error = Exception("test error") + result = _default_error_handler(error) + + assert result["error"] == "test error" + assert result["type"] == "Exception" + + def test_create_wrapper(self): + """Test wrapper creation and execution.""" + + def mock_func(param1: str, param2: int = 10): + return {"param1": param1, "param2": param2} + + config = ToolDefinition( + func=mock_func, + param_names=["param1", "param2"], + ) + + wrapper = _create_wrapper(config, _default_serializer, _default_error_handler) + + assert callable(wrapper) + assert hasattr(wrapper, "__signature__") + + result = wrapper(param1="test", param2=20) + assert isinstance(result, dict) + + +class TestRegisterAPITools: + """Test API tools registration.""" + + def test_basic_registration(self): + """Test registering valid tools.""" + mock_app = MagicMock() + + # Clean up any existing markers + if hasattr(api.list_farms, "_mcp_tool_registered"): + delattr(api.list_farms, "_mcp_tool_registered") + + register_api_tools(mock_app, [api.list_farms]) + + assert mock_app.tool.call_count == 1 + assert hasattr(api.list_farms, "_mcp_tool_registered") + + # Clean up + delattr(api.list_farms, "_mcp_tool_registered") + + def test_invalid_tools_raise_exception(self): + """Test that invalid tools raise exceptions.""" + mock_app = MagicMock() + + with pytest.raises(ValueError, match="Tool not_a_function is not callable"): + register_api_tools(mock_app, ["not_a_function"]) # type: ignore[list-item] + + mock_app.tool.assert_not_called() + + def test_unregistered_function_raises_exception(self): + """Test that callable functions not in registry raise exceptions.""" + mock_app = MagicMock() + + def unregistered_function(): + """A function not in the registry.""" + pass + + with pytest.raises( + ValueError, match="Function unregistered_function not found in tool registry" + ): + register_api_tools(mock_app, [unregistered_function]) + + mock_app.tool.assert_not_called() + + +class TestParameterTypeConversion: + """Test parameter type conversion functionality.""" + + def test_signature_preserves_original_annotations(self): + """Test that MCP wrapper preserves original function type annotations.""" + from deadline._mcp.utils import _create_wrapper + from deadline._mcp.tools.job import submit_job + from deadline._mcp.registry import ToolDefinition + + # Create a mock config + config: ToolDefinition = {"func": submit_job, "param_names": None} + + # Create wrapper with our improved approach + wrapper = _create_wrapper(config, lambda x: x, lambda e: {"error": str(e)}) + + # Check that the wrapper signature preserves original annotations + wrapper_sig = getattr(wrapper, "__signature__") + + # Test a few key parameters that should have preserved annotations + assert "priority" in wrapper_sig.parameters + assert "max_failed_tasks_count" in wrapper_sig.parameters + + # The annotations should match the original function (Optional[int]) + # Note: We can't directly compare annotations due to how Optional types work, + # but we can verify they're not all strings anymore + priority_annotation = wrapper_sig.parameters["priority"].annotation + assert priority_annotation is not str # Should not be forced to str anymore diff --git a/test/unit/deadline_mcp/test_mcp_telemetry.py b/test/unit/deadline_mcp/test_mcp_telemetry.py new file mode 100644 index 000000000..7fe89b881 --- /dev/null +++ b/test/unit/deadline_mcp/test_mcp_telemetry.py @@ -0,0 +1,358 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. + +"""Unit tests for MCP telemetry functionality.""" + +import pytest +from unittest.mock import patch, MagicMock + +# Skip all tests in this module if MCP dependencies are not available +pytest.importorskip("mcp", reason="MCP dependencies not available") + +from deadline._mcp.utils import _create_wrapper +from deadline._mcp.registry import ToolDefinition + + +@pytest.fixture(scope="function", name="mock_telemetry_client") +def fixture_mock_telemetry_client(): + """Fixture to provide a mock telemetry client for MCP tests.""" + mock_client = MagicMock() + mock_client.is_initialized = True + mock_client.event_queue = MagicMock() + return mock_client + + +def test_mcp_tool_telemetry_success(mock_telemetry_client): + """Test that MCP tools record telemetry on successful execution.""" + mock_func = MagicMock() + mock_func.__name__ = "test_function" + mock_func.return_value = {"result": "success"} + + config: ToolDefinition = {"func": mock_func, "param_names": ["param1", "param2"]} + + mock_serializer = MagicMock(side_effect=lambda x: x) + mock_error_handler = MagicMock() + + with patch( + "deadline._mcp.utils.get_deadline_cloud_library_telemetry_client" + ) as mock_get_client: + mock_get_client.return_value = mock_telemetry_client + + wrapper = _create_wrapper(config, mock_serializer, mock_error_handler) + result = wrapper(param1="value1", param2="value2") + + mock_func.assert_called_once_with(param1="value1", param2="value2") + assert result == {"result": "success"} + assert mock_telemetry_client.record_event.call_count == 2 + + calls = mock_telemetry_client.record_event.call_args_list + + latency_call = next( + call for call in calls if call[1]["event_type"] == "com.amazon.rum.deadline.mcp.latency" + ) + latency_details = latency_call[1]["event_details"] + assert "latency" in latency_details + assert latency_details["tool_name"] == "wrapper" # In unit tests, wrapper name is used + assert latency_details["usage_mode"] == "MCP" + + usage_call = next( + call for call in calls if call[1]["event_type"] == "com.amazon.rum.deadline.mcp.usage" + ) + usage_details = usage_call[1]["event_details"] + assert usage_details["tool_name"] == "wrapper" # In unit tests, wrapper name is used + assert usage_details["is_success"] is True + assert usage_details["error_type"] is None + assert usage_details["usage_mode"] == "MCP" + + +def test_mcp_tool_telemetry_failure(mock_telemetry_client): + """Test that MCP tools record telemetry on failed execution.""" + mock_func = MagicMock() + mock_func.__name__ = "test_function" + mock_func.side_effect = ValueError("Test error") + + config: ToolDefinition = {"func": mock_func, "param_names": ["param1"]} + + mock_serializer = MagicMock() + mock_error_handler = MagicMock(return_value={"error": "Test error", "type": "ValueError"}) + + with patch( + "deadline._mcp.utils.get_deadline_cloud_library_telemetry_client" + ) as mock_get_client: + mock_get_client.return_value = mock_telemetry_client + + wrapper = _create_wrapper(config, mock_serializer, mock_error_handler) + result = wrapper(param1="value1") + + mock_error_handler.assert_called_once() + assert result == {"error": "Test error", "type": "ValueError"} + assert mock_telemetry_client.record_event.call_count == 2 + + calls = mock_telemetry_client.record_event.call_args_list + + latency_call = next( + call for call in calls if call[1]["event_type"] == "com.amazon.rum.deadline.mcp.latency" + ) + latency_details = latency_call[1]["event_details"] + assert "latency" in latency_details + assert latency_details["tool_name"] == "wrapper" # In unit tests, wrapper name is used + assert latency_details["usage_mode"] == "MCP" + + usage_call = next( + call for call in calls if call[1]["event_type"] == "com.amazon.rum.deadline.mcp.usage" + ) + usage_details = usage_call[1]["event_details"] + assert usage_details["tool_name"] == "wrapper" # In unit tests, wrapper name is used + assert usage_details["is_success"] is False + assert usage_details["error_type"] == "ValueError" + assert usage_details["usage_mode"] == "MCP" + + +def test_mcp_tool_telemetry_error_handling(mock_telemetry_client): + """Test that telemetry errors don't affect tool execution.""" + mock_func = MagicMock() + mock_func.__name__ = "test_function" + mock_func.return_value = {"result": "success"} + + mock_telemetry_client.record_event.side_effect = Exception("Telemetry error") + + config: ToolDefinition = {"func": mock_func, "param_names": ["param1"]} + + mock_serializer = MagicMock(side_effect=lambda x: x) + mock_error_handler = MagicMock() + + with patch( + "deadline._mcp.utils.get_deadline_cloud_library_telemetry_client" + ) as mock_get_client: + mock_get_client.return_value = mock_telemetry_client + + wrapper = _create_wrapper(config, mock_serializer, mock_error_handler) + # Should not raise exception despite telemetry error + result = wrapper(param1="value1") + + mock_func.assert_called_once_with(param1="value1") + assert result == {"result": "success"} + assert mock_telemetry_client.record_event.call_count > 0 + + +def test_mcp_tool_parameter_filtering(): + """Test that empty/null parameters are filtered correctly.""" + mock_func = MagicMock() + mock_func.__name__ = "test_function" + mock_func.return_value = {"result": "success"} + + mock_telemetry_client = MagicMock() + + config: ToolDefinition = {"func": mock_func, "param_names": ["param1", "param2", "param3"]} + + mock_serializer = MagicMock(side_effect=lambda x: x) + mock_error_handler = MagicMock() + + with patch( + "deadline._mcp.utils.get_deadline_cloud_library_telemetry_client" + ) as mock_get_client: + mock_get_client.return_value = mock_telemetry_client + + wrapper = _create_wrapper(config, mock_serializer, mock_error_handler) + result = wrapper(param1="value1", param2=None, param3="", param4="null") + + mock_func.assert_called_once_with(param1="value1") + assert result == {"result": "success"} + + +def test_mcp_tool_telemetry_client_initialization_error(): + """Test behavior when telemetry client fails to initialize.""" + mock_func = MagicMock() + mock_func.__name__ = "test_function" + mock_func.return_value = {"result": "success"} + + config: ToolDefinition = {"func": mock_func, "param_names": ["param1"]} + + mock_serializer = MagicMock(side_effect=lambda x: x) + mock_error_handler = MagicMock() + + with patch( + "deadline._mcp.utils.get_deadline_cloud_library_telemetry_client" + ) as mock_get_client: + mock_get_client.side_effect = Exception("Telemetry client initialization failed") + + wrapper = _create_wrapper(config, mock_serializer, mock_error_handler) + # Should not raise exception despite telemetry client error + result = wrapper(param1="value1") + + mock_func.assert_called_once_with(param1="value1") + assert result == {"result": "success"} + mock_get_client.assert_called() + + +def test_mcp_tool_aws_error_handling(): + """Test that MCP tools properly handle AWS service errors.""" + mock_func = MagicMock() + mock_func.__name__ = "list_farms" + + from botocore.exceptions import ClientError + + error_response = { + "Error": {"Code": "ValidationException", "Message": "Invalid farm ID format"}, + "ResponseMetadata": {"HTTPStatusCode": 400}, + } + mock_func.side_effect = ClientError(error_response, "ListFarms") + + mock_telemetry_client = MagicMock() + + config: ToolDefinition = {"func": mock_func, "param_names": ["farmId"]} + + from deadline._mcp.utils import _default_error_handler, _default_serializer + + with patch( + "deadline._mcp.utils.get_deadline_cloud_library_telemetry_client" + ) as mock_get_client: + mock_get_client.return_value = mock_telemetry_client + + wrapper = _create_wrapper(config, _default_serializer, _default_error_handler) + result = wrapper(farmId="invalid-farm-id") + + mock_func.assert_called_once_with(farmId="invalid-farm-id") + + assert "error" in result, "Error response should contain 'error' field" + assert "type" in result, "Error response should contain 'type' field" + + assert result["type"] == "ClientError" + assert "ValidationException" in result["error"] or "Invalid farm ID" in result["error"] + + if "status_code" in result: + assert result["status_code"] == 400 + + assert mock_telemetry_client.record_event.call_count == 2 + + calls = mock_telemetry_client.record_event.call_args_list + usage_call = next( + call for call in calls if call[1]["event_type"] == "com.amazon.rum.deadline.mcp.usage" + ) + usage_details = usage_call[1]["event_details"] + assert usage_details["is_success"] is False + assert usage_details["error_type"] == "ClientError" + + +def test_mcp_tool_network_error_handling(): + """Test that MCP tools properly handle network-related errors.""" + mock_func = MagicMock() + mock_func.__name__ = "list_queues" + + from botocore.exceptions import ConnectTimeoutError + + mock_func.side_effect = ConnectTimeoutError( + endpoint_url="https://deadline.us-west-2.amazonaws.com" + ) + + mock_telemetry_client = MagicMock() + + config: ToolDefinition = {"func": mock_func, "param_names": ["farmId"]} + + from deadline._mcp.utils import _default_error_handler, _default_serializer + + with patch( + "deadline._mcp.utils.get_deadline_cloud_library_telemetry_client" + ) as mock_get_client: + mock_get_client.return_value = mock_telemetry_client + + wrapper = _create_wrapper(config, _default_serializer, _default_error_handler) + result = wrapper(farmId="farm-123") + + mock_func.assert_called_once_with(farmId="farm-123") + + assert "error" in result, "Error response should contain 'error' field" + assert "type" in result, "Error response should contain 'type' field" + + assert result["type"] == "ConnectTimeoutError" + assert "timeout" in result["error"].lower() or "connect" in result["error"].lower() + + assert mock_telemetry_client.record_event.call_count == 2 + + calls = mock_telemetry_client.record_event.call_args_list + usage_call = next( + call for call in calls if call[1]["event_type"] == "com.amazon.rum.deadline.mcp.usage" + ) + usage_details = usage_call[1]["event_details"] + assert usage_details["is_success"] is False + assert usage_details["error_type"] == "ConnectTimeoutError" + + +def test_mcp_tool_telemetry_event_structure(mock_telemetry_client): + """Test that MCP telemetry events have the correct structure.""" + mock_func = MagicMock() + mock_func.__name__ = "list_farms" + mock_func.return_value = {"farms": []} + + config: ToolDefinition = {"func": mock_func, "param_names": ["farmId"]} + + mock_serializer = MagicMock(side_effect=lambda x: x) + mock_error_handler = MagicMock() + + with patch( + "deadline._mcp.utils.get_deadline_cloud_library_telemetry_client" + ) as mock_get_client: + mock_get_client.return_value = mock_telemetry_client + + wrapper = _create_wrapper(config, mock_serializer, mock_error_handler) + result = wrapper(farmId="farm-0000001234567adba49becbca1fce5f6") + + mock_func.assert_called_once_with(farmId="farm-0000001234567adba49becbca1fce5f6") + assert result == {"farms": []} + assert mock_telemetry_client.record_event.call_count == 2 + + calls = mock_telemetry_client.record_event.call_args_list + + for call in calls: + assert "event_type" in call[1] + assert "event_details" in call[1] + + event_type = call[1]["event_type"] + event_details = call[1]["event_details"] + + assert event_type.startswith("com.amazon.rum.deadline.mcp.") + assert event_type in [ + "com.amazon.rum.deadline.mcp.latency", + "com.amazon.rum.deadline.mcp.usage", + ] + + assert "tool_name" in event_details + assert "usage_mode" in event_details + assert event_details["usage_mode"] == "MCP" + + if event_type == "com.amazon.rum.deadline.mcp.latency": + assert "latency" in event_details + assert isinstance(event_details["latency"], (int, float)) + elif event_type == "com.amazon.rum.deadline.mcp.usage": + assert "is_success" in event_details + assert isinstance(event_details["is_success"], bool) + assert "error_type" in event_details + + +def test_mcp_tool_telemetry_queue_integration(mock_telemetry_client): + """Test that telemetry events are properly queued.""" + mock_func = MagicMock() + mock_func.__name__ = "submit_job" + mock_func.return_value = {"jobId": "job-0000001234567adba49becbca1fce5f6"} + + config: ToolDefinition = {"func": mock_func, "param_names": ["job_bundle_dir"]} + + mock_serializer = MagicMock(side_effect=lambda x: x) + mock_error_handler = MagicMock() + + with patch( + "deadline._mcp.utils.get_deadline_cloud_library_telemetry_client" + ) as mock_get_client: + mock_get_client.return_value = mock_telemetry_client + + wrapper = _create_wrapper(config, mock_serializer, mock_error_handler) + result = wrapper(job_bundle_dir="/path/to/bundle") + + mock_func.assert_called_once_with(job_bundle_dir="/path/to/bundle") + assert result == {"jobId": "job-0000001234567adba49becbca1fce5f6"} + assert mock_telemetry_client.record_event.call_count == 2 + + calls = mock_telemetry_client.record_event.call_args_list + for call in calls: + event_details = call[1]["event_details"] + # In unit tests, the wrapper function name is used + assert event_details["tool_name"] == "wrapper"