mcpforunityserver 9.4.0b20260203025228__py3-none-any.whl

services/custom_tool_service.py

@@ -0,0 +1,499 @@
+import asyncio
+import inspect
+import logging
+import time
+from hashlib import sha256
+from typing import Optional
+
+from fastmcp import Context, FastMCP
+from pydantic import BaseModel, Field, ValidationError
+from starlette.requests import Request
+from starlette.responses import JSONResponse
+
+from models.models import MCPResponse, ToolDefinitionModel, ToolParameterModel
+from core.logging_decorator import log_execution
+from core.telemetry_decorator import telemetry_tool
+from transport.unity_transport import send_with_unity_instance
+from transport.legacy.unity_connection import (
+    async_send_command_with_retry,
+    get_unity_connection_pool,
+)
+from transport.plugin_hub import PluginHub
+from services.tools import get_unity_instance_from_context
+from services.registry import get_registered_tools
+
+logger = logging.getLogger("mcp-for-unity-server")
+
+_DEFAULT_POLL_INTERVAL = 1.0
+_MAX_POLL_SECONDS = 600
+
+
+class RegisterToolsPayload(BaseModel):
+    project_id: str
+    project_hash: str | None = None
+    tools: list[ToolDefinitionModel]
+
+
+class ToolRegistrationResponse(BaseModel):
+    success: bool
+    registered: list[str]
+    replaced: list[str]
+    message: str
+
+
+class CustomToolService:
+    _instance: "CustomToolService | None" = None
+
+    def __init__(self, mcp: FastMCP, project_scoped_tools: bool = True):
+        CustomToolService._instance = self
+        self._mcp = mcp
+        self._project_scoped_tools = project_scoped_tools
+        self._project_tools: dict[str, dict[str, ToolDefinitionModel]] = {}
+        self._hash_to_project: dict[str, str] = {}
+        self._global_tools: dict[str, ToolDefinitionModel] = {}
+        self._register_http_routes()
+
+    @classmethod
+    def get_instance(cls) -> "CustomToolService":
+        if cls._instance is None:
+            raise RuntimeError("CustomToolService has not been initialized")
+        return cls._instance
+
+    # --- HTTP Routes -----------------------------------------------------
+    def _register_http_routes(self) -> None:
+        @self._mcp.custom_route("/register-tools", methods=["POST"])
+        async def register_tools(request: Request) -> JSONResponse:
+            try:
+                payload = RegisterToolsPayload.model_validate(await request.json())
+            except ValidationError as exc:
+                return JSONResponse({"success": False, "error": exc.errors()}, status_code=400)
+
+            registered, replaced = self._register_project_tools(
+                payload.project_id, payload.tools, project_hash=payload.project_hash)
+
+            message = f"Registered {len(registered)} tool(s)"
+            if replaced:
+                message += f" (replaced: {', '.join(replaced)})"
+
+            response = ToolRegistrationResponse(
+                success=True,
+                registered=registered,
+                replaced=replaced,
+                message=message,
+            )
+            return JSONResponse(response.model_dump())
+
+    # --- Public API for MCP tools ---------------------------------------
+    async def list_registered_tools(self, project_id: str) -> list[ToolDefinitionModel]:
+        legacy = list(self._project_tools.get(project_id, {}).values())
+        hub_tools = await PluginHub.get_tools_for_project(project_id)
+        return legacy + hub_tools
+
+    async def get_tool_definition(self, project_id: str, tool_name: str) -> ToolDefinitionModel | None:
+        tool = self._project_tools.get(project_id, {}).get(tool_name)
+        if tool:
+            return tool
+        return await PluginHub.get_tool_definition(project_id, tool_name)
+
+    async def execute_tool(
+        self,
+        project_id: str,
+        tool_name: str,
+        unity_instance: str | None,
+        params: dict[str, object] | None = None,
+    ) -> MCPResponse:
+        params = params or {}
+        logger.info(
+            f"Executing tool '{tool_name}' for project '{project_id}' (instance={unity_instance}) with params: {params}"
+        )
+
+        definition = await self.get_tool_definition(project_id, tool_name)
+        if definition is None:
+            return MCPResponse(
+                success=False,
+                message=f"Tool '{tool_name}' not found for project {project_id}",
+            )
+
+        response = await send_with_unity_instance(
+            async_send_command_with_retry,
+            unity_instance,
+            tool_name,
+            params,
+        )
+
+        if not definition.requires_polling:
+            result = self._normalize_response(response)
+            logger.info(f"Tool '{tool_name}' immediate response: {result}")
+            return result
+
+        result = await self._poll_until_complete(
+            tool_name,
+            unity_instance,
+            params,
+            response,
+            definition.poll_action or "status",
+        )
+        logger.info(f"Tool '{tool_name}' polled response: {result}")
+        return result
+
+    # --- Internal helpers ------------------------------------------------
+    def _is_registered(self, project_id: str, tool_name: str) -> bool:
+        return tool_name in self._project_tools.get(project_id, {})
+
+    def _register_tool(self, project_id: str, definition: ToolDefinitionModel) -> None:
+        self._project_tools.setdefault(project_id, {})[
+            definition.name] = definition
+
+    def get_project_id_for_hash(self, project_hash: str | None) -> str | None:
+        if not project_hash:
+            return None
+        return self._hash_to_project.get(project_hash.lower())
+
+    async def _poll_until_complete(
+        self,
+        tool_name: str,
+        unity_instance,
+        initial_params: dict[str, object],
+        initial_response,
+        poll_action: str,
+    ) -> MCPResponse:
+        poll_params = dict(initial_params)
+        poll_params["action"] = poll_action or "status"
+
+        deadline = time.time() + _MAX_POLL_SECONDS
+        response = initial_response
+
+        while True:
+            status, poll_interval = self._interpret_status(response)
+
+            if status in ("complete", "error", "final"):
+                return self._normalize_response(response)
+
+            if time.time() > deadline:
+                return MCPResponse(
+                    success=False,
+                    message=f"Timeout waiting for {tool_name} to complete",
+                    data=self._safe_response(response),
+                )
+
+            await asyncio.sleep(poll_interval)
+
+            try:
+                response = await send_with_unity_instance(
+                    async_send_command_with_retry, unity_instance, tool_name, poll_params
+                )
+            except Exception as exc:  # pragma: no cover - network/domain reload variability
+                logger.debug(f"Polling {tool_name} failed, will retry: {exc}")
+                # Back off modestly but stay responsive.
+                response = {
+                    "_mcp_status": "pending",
+                    "_mcp_poll_interval": min(max(poll_interval * 2, _DEFAULT_POLL_INTERVAL), 5.0),
+                    "message": f"Retrying after transient error: {exc}",
+                }
+
+    def _interpret_status(self, response) -> tuple[str, float]:
+        if response is None:
+            return "pending", _DEFAULT_POLL_INTERVAL
+
+        if not isinstance(response, dict):
+            return "final", _DEFAULT_POLL_INTERVAL
+
+        status = response.get("_mcp_status")
+        if status is None:
+            if len(response.keys()) == 0:
+                return "pending", _DEFAULT_POLL_INTERVAL
+            return "final", _DEFAULT_POLL_INTERVAL
+
+        if status == "pending":
+            interval_raw = response.get(
+                "_mcp_poll_interval", _DEFAULT_POLL_INTERVAL)
+            try:
+                interval = float(interval_raw)
+            except (TypeError, ValueError):
+                interval = _DEFAULT_POLL_INTERVAL
+
+            interval = max(0.1, min(interval, 5.0))
+            return "pending", interval
+
+        if status == "complete":
+            return "complete", _DEFAULT_POLL_INTERVAL
+
+        if status == "error":
+            return "error", _DEFAULT_POLL_INTERVAL
+
+        return "final", _DEFAULT_POLL_INTERVAL
+
+    def _normalize_response(self, response) -> MCPResponse:
+        if isinstance(response, MCPResponse):
+            return response
+        if isinstance(response, dict):
+            return MCPResponse(
+                success=response.get("success", True),
+                message=response.get("message"),
+                error=response.get("error"),
+                data=response.get(
+                    "data", response) if "data" not in response else response["data"],
+            )
+
+        success = True
+        message = None
+        error = None
+        data = None
+
+        if isinstance(response, dict):
+            success = response.get("success", True)
+            if "_mcp_status" in response and response["_mcp_status"] == "error":
+                success = False
+            message = str(response.get("message")) if response.get(
+                "message") else None
+            error = str(response.get("error")) if response.get(
+                "error") else None
+            data = response.get("data")
+            if "success" not in response and "_mcp_status" not in response:
+                data = response
+        else:
+            success = False
+            message = str(response)
+
+        return MCPResponse(success=success, message=message, error=error, data=data)
+
+    def _safe_response(self, response):
+        if isinstance(response, dict):
+            return response
+        if response is None:
+            return None
+        return {"message": str(response)}
+
+    def _register_project_tools(
+        self,
+        project_id: str,
+        tools: list[ToolDefinitionModel],
+        project_hash: str | None = None,
+    ) -> tuple[list[str], list[str]]:
+        registered: list[str] = []
+        replaced: list[str] = []
+        for tool in tools:
+            if self._is_registered(project_id, tool.name):
+                replaced.append(tool.name)
+            self._register_tool(project_id, tool)
+            registered.append(tool.name)
+            if not self._project_scoped_tools:
+                self._register_global_tool(tool)
+
+        if project_hash:
+            self._hash_to_project[project_hash.lower()] = project_id
+
+        return registered, replaced
+
+    def register_global_tools(self, tools: list[ToolDefinitionModel]) -> None:
+        if self._project_scoped_tools:
+            return
+        builtin_names = self._get_builtin_tool_names()
+        for tool in tools:
+            if tool.name in builtin_names:
+                logger.info(
+                    "Skipping global custom tool registration for built-in tool '%s'",
+                    tool.name,
+                )
+                continue
+            self._register_global_tool(tool)
+
+    def _get_builtin_tool_names(self) -> set[str]:
+        return {tool["name"] for tool in get_registered_tools()}
+
+    def _register_global_tool(self, definition: ToolDefinitionModel) -> None:
+        existing = self._global_tools.get(definition.name)
+        if existing:
+            if existing.model_dump() != definition.model_dump():
+                logger.warning(
+                    "Custom tool '%s' already registered with a different schema; keeping existing definition.",
+                    definition.name,
+                )
+            return
+
+        handler = self._build_global_tool_handler(definition)
+        wrapped = log_execution(definition.name, "Tool")(handler)
+        wrapped = telemetry_tool(definition.name)(wrapped)
+
+        try:
+            wrapped = self._mcp.tool(
+                name=definition.name,
+                description=definition.description,
+            )(wrapped)
+        except Exception as exc:  # pragma: no cover - defensive against tool conflicts
+            logger.warning(
+                "Failed to register custom tool '%s' globally: %s",
+                definition.name,
+                exc,
+            )
+            return
+
+        self._global_tools[definition.name] = definition
+
+    def _build_global_tool_handler(self, definition: ToolDefinitionModel):
+        async def _handler(ctx: Context, **kwargs) -> MCPResponse:
+            unity_instance = get_unity_instance_from_context(ctx)
+            if not unity_instance:
+                return MCPResponse(
+                    success=False,
+                    message="No active Unity instance. Call set_active_instance with Name@hash from mcpforunity://instances.",
+                )
+
+            project_id = resolve_project_id_for_unity_instance(unity_instance)
+            if project_id is None:
+                return MCPResponse(
+                    success=False,
+                    message=f"Could not resolve project id for {unity_instance}. Ensure Unity is running and reachable.",
+                )
+
+            params = {k: v for k, v in kwargs.items() if v is not None}
+            service = CustomToolService.get_instance()
+            return await service.execute_tool(project_id, definition.name, unity_instance, params)
+
+        _handler.__name__ = f"custom_tool_{definition.name}"
+        _handler.__doc__ = definition.description or ""
+        _handler.__signature__ = self._build_signature(definition)
+        _handler.__annotations__ = self._build_annotations(definition)
+        return _handler
+
+    def _build_signature(self, definition: ToolDefinitionModel) -> inspect.Signature:
+        params: list[inspect.Parameter] = [
+            inspect.Parameter(
+                "ctx",
+                inspect.Parameter.POSITIONAL_OR_KEYWORD,
+                annotation=Context,
+            )
+        ]
+        for param in definition.parameters:
+            if not param.name.isidentifier():
+                logger.warning(
+                    "Custom tool '%s' has non-identifier parameter '%s'; exposing via kwargs only.",
+                    definition.name,
+                    param.name,
+                )
+                continue
+            default = inspect._empty if param.required else self._coerce_default(
+                param.default_value, param.type)
+            params.append(
+                inspect.Parameter(
+                    param.name,
+                    inspect.Parameter.POSITIONAL_OR_KEYWORD,
+                    default=default,
+                    annotation=self._map_param_type(param),
+                )
+            )
+        return inspect.Signature(parameters=params)
+
+    def _build_annotations(self, definition: ToolDefinitionModel) -> dict[str, object]:
+        annotations: dict[str, object] = {"ctx": Context}
+        for param in definition.parameters:
+            if not param.name.isidentifier():
+                continue
+            annotations[param.name] = self._map_param_type(param)
+        return annotations
+
+    def _map_param_type(self, param: ToolParameterModel):
+        ptype = (param.type or "string").lower()
+        if ptype in ("integer", "int"):
+            return int
+        if ptype in ("number", "float", "double"):
+            return float
+        if ptype in ("bool", "boolean"):
+            return bool
+        if ptype in ("array", "list"):
+            return list
+        if ptype in ("object", "dict"):
+            return dict
+        return str
+
+    def _coerce_default(self, value: str | None, param_type: str | None):
+        if value is None:
+            return None
+        try:
+            ptype = (param_type or "string").lower()
+            if ptype in ("integer", "int"):
+                return int(value)
+            if ptype in ("number", "float", "double"):
+                return float(value)
+            if ptype in ("bool", "boolean"):
+                return str(value).lower() in ("1", "true", "yes", "on")
+            return value
+        except Exception:
+            return value
+
+
+def compute_project_id(project_name: str, project_path: str) -> str:
+    """
+    DEPRECATED: Computes a SHA256-based project ID.
+    This function is no longer used as of the multi-session fix.
+    Unity instances now use their native project_hash (SHA1-based) for consistency
+    across stdio and WebSocket transports.
+    """
+    combined = f"{project_name}:{project_path}"
+    return sha256(combined.encode("utf-8")).hexdigest().upper()[:16]
+
+
+def resolve_project_id_for_unity_instance(unity_instance: str | None) -> str | None:
+    if unity_instance is None:
+        return None
+
+    # stdio transport: resolve via discovered instances with name+path
+    try:
+        pool = get_unity_connection_pool()
+        instances = pool.discover_all_instances()
+        target = None
+        if "@" in unity_instance:
+            name_part, _, hash_hint = unity_instance.partition("@")
+            target = next(
+                (
+                    inst for inst in instances
+                    if inst.name == name_part and inst.hash.startswith(hash_hint)
+                ),
+                None,
+            )
+        else:
+            target = next(
+                (
+                    inst for inst in instances
+                    if inst.id == unity_instance or inst.hash.startswith(unity_instance)
+                ),
+                None,
+            )
+
+        if target:
+            # Return the project_hash from Unity (not a computed SHA256 hash).
+            # This matches the hash Unity uses when registering tools via WebSocket.
+            if target.hash:
+                return target.hash
+            logger.warning(
+                f"Unity instance {target.id} has empty hash; cannot resolve project ID")
+            return None
+    except Exception:
+        logger.debug(
+            f"Failed to resolve project id via connection pool for {unity_instance}")
+
+    # HTTP/WebSocket transport: resolve via PluginHub using project_hash
+    try:
+        hash_part: Optional[str] = None
+        if "@" in unity_instance:
+            _, _, suffix = unity_instance.partition("@")
+            hash_part = suffix or None
+        else:
+            hash_part = unity_instance
+
+        if hash_part:
+            lowered = hash_part.lower()
+            mapped: Optional[str] = None
+            try:
+                service = CustomToolService.get_instance()
+                mapped = service.get_project_id_for_hash(lowered)
+            except RuntimeError:
+                mapped = None
+            if mapped:
+                return mapped
+            return lowered
+    except Exception:
+        logger.debug(
+            f"Failed to resolve project id via plugin hub for {unity_instance}")
+
+    return None

services/registry/__init__.py

@@ -0,0 +1,22 @@
+"""
+Registry package for MCP tool auto-discovery.
+"""
+from .tool_registry import (
+    mcp_for_unity_tool,
+    get_registered_tools,
+    clear_tool_registry,
+)
+from .resource_registry import (
+    mcp_for_unity_resource,
+    get_registered_resources,
+    clear_resource_registry,
+)
+
+__all__ = [
+    'mcp_for_unity_tool',
+    'get_registered_tools',
+    'clear_tool_registry',
+    'mcp_for_unity_resource',
+    'get_registered_resources',
+    'clear_resource_registry'
+]

services/registry/resource_registry.py

@@ -0,0 +1,53 @@
+"""
+Resource registry for auto-discovery of MCP resources.
+"""
+from typing import Callable, Any
+
+# Global registry to collect decorated resources
+_resource_registry: list[dict[str, Any]] = []
+
+
+def mcp_for_unity_resource(
+    uri: str,
+    name: str | None = None,
+    description: str | None = None,
+    **kwargs
+) -> Callable:
+    """
+    Decorator for registering MCP resources in the server's resources directory.
+
+    Resources are registered in the global resource registry.
+
+    Args:
+        name: Resource name (defaults to function name)
+        description: Resource description
+        **kwargs: Additional arguments passed to @mcp.resource()
+
+    Example:
+        @mcp_for_unity_resource("mcpforunity://resource", description="Gets something interesting")
+        async def my_custom_resource(ctx: Context, ...):
+            pass
+    """
+    def decorator(func: Callable) -> Callable:
+        resource_name = name if name is not None else func.__name__
+        _resource_registry.append({
+            'func': func,
+            'uri': uri,
+            'name': resource_name,
+            'description': description,
+            'kwargs': kwargs
+        })
+
+        return func
+
+    return decorator
+
+
+def get_registered_resources() -> list[dict[str, Any]]:
+    """Get all registered resources"""
+    return _resource_registry.copy()
+
+
+def clear_resource_registry():
+    """Clear the resource registry (useful for testing)"""
+    _resource_registry.clear()

services/registry/tool_registry.py

@@ -0,0 +1,51 @@
+"""
+Tool registry for auto-discovery of MCP tools.
+"""
+from typing import Callable, Any
+
+# Global registry to collect decorated tools
+_tool_registry: list[dict[str, Any]] = []
+
+
+def mcp_for_unity_tool(
+    name: str | None = None,
+    description: str | None = None,
+    **kwargs
+) -> Callable:
+    """
+    Decorator for registering MCP tools in the server's tools directory.
+
+    Tools are registered in the global tool registry.
+
+    Args:
+        name: Tool name (defaults to function name)
+        description: Tool description
+        **kwargs: Additional arguments passed to @mcp.tool()
+
+    Example:
+        @mcp_for_unity_tool(description="Does something cool")
+        async def my_custom_tool(ctx: Context, ...):
+            pass
+    """
+    def decorator(func: Callable) -> Callable:
+        tool_name = name if name is not None else func.__name__
+        _tool_registry.append({
+            'func': func,
+            'name': tool_name,
+            'description': description,
+            'kwargs': kwargs
+        })
+
+        return func
+
+    return decorator
+
+
+def get_registered_tools() -> list[dict[str, Any]]:
+    """Get all registered tools"""
+    return _tool_registry.copy()
+
+
+def clear_tool_registry():
+    """Clear the tool registry (useful for testing)"""
+    _tool_registry.clear()