mcpforunityserver 9.4.0b20260203025228__py3-none-any.whl

transport/plugin_registry.py

@@ -0,0 +1,182 @@
+"""In-memory registry for connected Unity plugin sessions."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from datetime import datetime, timezone
+
+import asyncio
+
+from core.config import config
+from models.models import ToolDefinitionModel
+
+
+@dataclass(slots=True)
+class PluginSession:
+    """Represents a single Unity plugin connection."""
+
+    session_id: str
+    project_name: str
+    project_hash: str
+    unity_version: str
+    registered_at: datetime
+    connected_at: datetime
+    tools: dict[str, ToolDefinitionModel] = field(default_factory=dict)
+    project_id: str | None = None
+    # Full path to project root (for focus nudging)
+    project_path: str | None = None
+    user_id: str | None = None  # Associated user id (None for local mode)
+
+
+class PluginRegistry:
+    """Stores active plugin sessions in-memory.
+
+    The registry is optimised for quick lookup by either ``session_id`` or
+    ``project_hash`` (which is used as the canonical "instance id" across the
+    HTTP command routing stack).
+
+    In remote-hosted mode, sessions are scoped by (user_id, project_hash) composite key
+    to ensure session isolation between users.
+    """
+
+    def __init__(self) -> None:
+        self._sessions: dict[str, PluginSession] = {}
+        # In local mode: project_hash -> session_id
+        # In remote mode: (user_id, project_hash) -> session_id
+        self._hash_to_session: dict[str, str] = {}
+        self._user_hash_to_session: dict[tuple[str, str], str] = {}
+        self._lock = asyncio.Lock()
+
+    async def register(
+        self,
+        session_id: str,
+        project_name: str,
+        project_hash: str,
+        unity_version: str,
+        project_path: str | None = None,
+        user_id: str | None = None,
+    ) -> PluginSession:
+        """Register (or replace) a plugin session.
+
+        If an existing session already claims the same ``project_hash`` (and ``user_id``
+        in remote-hosted mode) it will be replaced, ensuring that reconnect scenarios
+        always map to the latest WebSocket connection.
+        """
+        if config.http_remote_hosted and not user_id:
+            raise ValueError("user_id is required in remote-hosted mode")
+
+        async with self._lock:
+            now = datetime.now(timezone.utc)
+            session = PluginSession(
+                session_id=session_id,
+                project_name=project_name,
+                project_hash=project_hash,
+                unity_version=unity_version,
+                registered_at=now,
+                connected_at=now,
+                project_path=project_path,
+                user_id=user_id,
+            )
+
+            # Remove old mapping for this hash if it existed under a different session
+            if user_id:
+                # Remote-hosted mode: use composite key (user_id, project_hash)
+                composite_key = (user_id, project_hash)
+                previous_session_id = self._user_hash_to_session.get(
+                    composite_key)
+                if previous_session_id and previous_session_id != session_id:
+                    self._sessions.pop(previous_session_id, None)
+                self._user_hash_to_session[composite_key] = session_id
+            else:
+                # Local mode: use project_hash only
+                previous_session_id = self._hash_to_session.get(project_hash)
+                if previous_session_id and previous_session_id != session_id:
+                    self._sessions.pop(previous_session_id, None)
+                self._hash_to_session[project_hash] = session_id
+
+            self._sessions[session_id] = session
+            return session
+
+    async def touch(self, session_id: str) -> None:
+        """Update the ``connected_at`` timestamp when a heartbeat is received."""
+
+        async with self._lock:
+            session = self._sessions.get(session_id)
+            if session:
+                session.connected_at = datetime.now(timezone.utc)
+
+    async def unregister(self, session_id: str) -> None:
+        """Remove a plugin session from the registry."""
+
+        async with self._lock:
+            session = self._sessions.pop(session_id, None)
+            if session:
+                # Clean up hash mappings
+                if session.project_hash in self._hash_to_session:
+                    mapped = self._hash_to_session.get(session.project_hash)
+                    if mapped == session_id:
+                        del self._hash_to_session[session.project_hash]
+
+                # Clean up user-scoped mappings
+                if session.user_id:
+                    composite_key = (session.user_id, session.project_hash)
+                    if composite_key in self._user_hash_to_session:
+                        mapped = self._user_hash_to_session.get(composite_key)
+                        if mapped == session_id:
+                            del self._user_hash_to_session[composite_key]
+
+    async def register_tools_for_session(self, session_id: str, tools: list[ToolDefinitionModel]) -> None:
+        """Register tools for a specific session."""
+        async with self._lock:
+            session = self._sessions.get(session_id)
+            if session:
+                # Replace existing tools or merge? Usually replace for "set state".
+                # We will replace the dict but keep the field.
+                session.tools.clear()
+                for tool in tools:
+                    session.tools[tool.name] = tool
+
+    async def get_session(self, session_id: str) -> PluginSession | None:
+        """Fetch a session by its ``session_id``."""
+
+        async with self._lock:
+            return self._sessions.get(session_id)
+
+    async def get_session_id_by_hash(self, project_hash: str, user_id: str | None = None) -> str | None:
+        """Resolve a ``project_hash`` (Unity instance id) to a session id."""
+
+        if user_id:
+            async with self._lock:
+                return self._user_hash_to_session.get((user_id, project_hash))
+        else:
+            async with self._lock:
+                return self._hash_to_session.get(project_hash)
+
+    async def list_sessions(self, user_id: str | None = None) -> dict[str, PluginSession]:
+        """Return a shallow copy of sessions.
+
+        Args:
+            user_id: If provided, only return sessions for this user (remote-hosted mode).
+                     If None, return all sessions (local mode only).
+
+        Raises:
+            ValueError: If ``user_id`` is None while running in remote-hosted mode.
+                        This prevents accidentally leaking sessions across users.
+        """
+        if user_id is None and config.http_remote_hosted:
+            raise ValueError(
+                "list_sessions requires user_id in remote-hosted mode"
+            )
+
+        async with self._lock:
+            if user_id is None:
+                return dict(self._sessions)
+            else:
+                return {
+                    sid: session
+                    for sid, session in self._sessions.items()
+                    if session.user_id == user_id
+                }
+
+
+__all__ = ["PluginRegistry", "PluginSession"]

transport/unity_instance_middleware.py

@@ -0,0 +1,262 @@
+"""
+Middleware for managing Unity instance selection per session.
+
+This middleware intercepts all tool calls and injects the active Unity instance
+into the request-scoped state, allowing tools to access it via ctx.get_state("unity_instance").
+"""
+from threading import RLock
+import logging
+
+from fastmcp.server.middleware import Middleware, MiddlewareContext
+
+from core.config import config
+from transport.plugin_hub import PluginHub
+
+logger = logging.getLogger("mcp-for-unity-server")
+
+# Store a global reference to the middleware instance so tools can interact
+# with it to set or clear the active unity instance.
+_unity_instance_middleware = None
+_middleware_lock = RLock()
+
+
+def get_unity_instance_middleware() -> 'UnityInstanceMiddleware':
+    """Get the global Unity instance middleware."""
+    global _unity_instance_middleware
+    if _unity_instance_middleware is None:
+        with _middleware_lock:
+            if _unity_instance_middleware is None:
+                # Auto-initialize if not set (lazy singleton) to handle import order or test cases
+                _unity_instance_middleware = UnityInstanceMiddleware()
+
+    return _unity_instance_middleware
+
+
+def set_unity_instance_middleware(middleware: 'UnityInstanceMiddleware') -> None:
+    """Replace the global middleware instance.
+
+    This is a test seam: production code uses ``get_unity_instance_middleware()``
+    which lazy-initialises the singleton.  Tests call this function to inject a
+    mock or pre-configured middleware before exercising tool/resource code.
+    """
+    global _unity_instance_middleware
+    _unity_instance_middleware = middleware
+
+
+class UnityInstanceMiddleware(Middleware):
+    """
+    Middleware that manages per-session Unity instance selection.
+
+    Stores active instance per session_id and injects it into request state
+    for all tool and resource calls.
+    """
+
+    def __init__(self):
+        super().__init__()
+        self._active_by_key: dict[str, str] = {}
+        self._lock = RLock()
+
+    def get_session_key(self, ctx) -> str:
+        """
+        Derive a stable key for the calling session.
+
+        Prioritizes client_id for stability.
+        In remote-hosted mode, falls back to user_id for session isolation.
+        Otherwise falls back to 'global' (assuming single-user local mode).
+        """
+        client_id = getattr(ctx, "client_id", None)
+        if isinstance(client_id, str) and client_id:
+            return client_id
+
+        # In remote-hosted mode, use user_id so different users get isolated instance selections
+        user_id = ctx.get_state("user_id")
+        if isinstance(user_id, str) and user_id:
+            return f"user:{user_id}"
+
+        # Fallback to global for local dev stability
+        return "global"
+
+    def set_active_instance(self, ctx, instance_id: str) -> None:
+        """Store the active instance for this session."""
+        key = self.get_session_key(ctx)
+        with self._lock:
+            self._active_by_key[key] = instance_id
+
+    def get_active_instance(self, ctx) -> str | None:
+        """Retrieve the active instance for this session."""
+        key = self.get_session_key(ctx)
+        with self._lock:
+            return self._active_by_key.get(key)
+
+    def clear_active_instance(self, ctx) -> None:
+        """Clear the stored instance for this session."""
+        key = self.get_session_key(ctx)
+        with self._lock:
+            self._active_by_key.pop(key, None)
+
+    async def _maybe_autoselect_instance(self, ctx) -> str | None:
+        """
+        Auto-select the sole Unity instance when no active instance is set.
+
+        Note: This method both *discovers* and *persists* the selection via
+        `set_active_instance` as a side-effect, since callers expect the selection
+        to stick for subsequent tool/resource calls in the same session.
+        """
+        try:
+            transport = (config.transport_mode or "stdio").lower()
+            # This implicit behavior works well for solo-users, but is dangerous for multi-user setups
+            if transport == "http" and config.http_remote_hosted:
+                return None
+            if PluginHub.is_configured():
+                try:
+                    sessions_data = await PluginHub.get_sessions()
+                    sessions = sessions_data.sessions or {}
+                    ids: list[str] = []
+                    for session_info in sessions.values():
+                        project = getattr(
+                            session_info, "project", None) or "Unknown"
+                        hash_value = getattr(session_info, "hash", None)
+                        if hash_value:
+                            ids.append(f"{project}@{hash_value}")
+                    if len(ids) == 1:
+                        chosen = ids[0]
+                        self.set_active_instance(ctx, chosen)
+                        logger.info(
+                            "Auto-selected sole Unity instance via PluginHub: %s",
+                            chosen,
+                        )
+                        return chosen
+                except (ConnectionError, ValueError, KeyError, TimeoutError, AttributeError) as exc:
+                    logger.debug(
+                        "PluginHub auto-select probe failed (%s); falling back to stdio",
+                        type(exc).__name__,
+                        exc_info=True,
+                    )
+                except Exception as exc:
+                    if isinstance(exc, (SystemExit, KeyboardInterrupt)):
+                        raise
+                    logger.debug(
+                        "PluginHub auto-select probe failed with unexpected error (%s); falling back to stdio",
+                        type(exc).__name__,
+                        exc_info=True,
+                    )
+
+            if transport != "http":
+                try:
+                    # Import here to avoid circular imports in legacy transport paths.
+                    from transport.legacy.unity_connection import get_unity_connection_pool
+
+                    pool = get_unity_connection_pool()
+                    instances = pool.discover_all_instances(force_refresh=True)
+                    ids = [getattr(inst, "id", None) for inst in instances]
+                    ids = [inst_id for inst_id in ids if inst_id]
+                    if len(ids) == 1:
+                        chosen = ids[0]
+                        self.set_active_instance(ctx, chosen)
+                        logger.info(
+                            "Auto-selected sole Unity instance via stdio discovery: %s",
+                            chosen,
+                        )
+                        return chosen
+                except (ConnectionError, ValueError, KeyError, TimeoutError, AttributeError) as exc:
+                    logger.debug(
+                        "Stdio auto-select probe failed (%s)",
+                        type(exc).__name__,
+                        exc_info=True,
+                    )
+                except Exception as exc:
+                    if isinstance(exc, (SystemExit, KeyboardInterrupt)):
+                        raise
+                    logger.debug(
+                        "Stdio auto-select probe failed with unexpected error (%s)",
+                        type(exc).__name__,
+                        exc_info=True,
+                    )
+        except Exception as exc:
+            if isinstance(exc, (SystemExit, KeyboardInterrupt)):
+                raise
+            logger.debug(
+                "Auto-select path encountered an unexpected error (%s)",
+                type(exc).__name__,
+                exc_info=True,
+            )
+
+        return None
+
+    async def _resolve_user_id(self) -> str | None:
+        """Extract user_id from the current HTTP request's API key."""
+        if not config.http_remote_hosted:
+            return None
+        # Lazy import to avoid circular dependencies (same pattern as _maybe_autoselect_instance).
+        from transport.unity_transport import _resolve_user_id_from_request
+        return await _resolve_user_id_from_request()
+
+    async def _inject_unity_instance(self, context: MiddlewareContext) -> None:
+        """Inject active Unity instance and user_id into context if available."""
+        ctx = context.fastmcp_context
+
+        # Resolve user_id from the HTTP request's API key header
+        user_id = await self._resolve_user_id()
+        if config.http_remote_hosted and user_id is None:
+            raise RuntimeError(
+                "API key authentication required. Provide a valid X-API-Key header."
+            )
+        if user_id:
+            ctx.set_state("user_id", user_id)
+
+        active_instance = self.get_active_instance(ctx)
+        if not active_instance:
+            active_instance = await self._maybe_autoselect_instance(ctx)
+        if active_instance:
+            # If using HTTP transport (PluginHub configured), validate session
+            # But for stdio transport (no PluginHub needed or maybe partially configured),
+            # we should be careful not to clear instance just because PluginHub can't resolve it.
+            # The 'active_instance' (Name@hash) might be valid for stdio even if PluginHub fails.
+
+            session_id: str | None = None
+            # Only validate via PluginHub if we are actually using HTTP transport.
+            # For stdio transport, skip PluginHub entirely - we only need the instance ID.
+            from transport.unity_transport import _is_http_transport
+            if _is_http_transport() and PluginHub.is_configured():
+                try:
+                    # resolving session_id might fail if the plugin disconnected
+                    # We only need session_id for HTTP transport routing.
+                    # For stdio, we just need the instance ID.
+                    # Pass user_id for remote-hosted mode session isolation
+                    session_id = await PluginHub._resolve_session_id(active_instance, user_id=user_id)
+                except (ConnectionError, ValueError, KeyError, TimeoutError) as exc:
+                    # If resolution fails, it means the Unity instance is not reachable via HTTP/WS.
+                    # If we are in stdio mode, this might still be fine if the user is just setting state?
+                    # But usually if PluginHub is configured, we expect it to work.
+                    # Let's LOG the error but NOT clear the instance immediately to avoid flickering,
+                    # or at least debug why it's failing.
+                    logger.debug(
+                        "PluginHub session resolution failed for %s: %s; leaving active_instance unchanged",
+                        active_instance,
+                        exc,
+                        exc_info=True,
+                    )
+                except Exception as exc:
+                    # Re-raise unexpected system exceptions to avoid swallowing critical failures
+                    if isinstance(exc, (SystemExit, KeyboardInterrupt)):
+                        raise
+                    logger.error(
+                        "Unexpected error during PluginHub session resolution for %s: %s",
+                        active_instance,
+                        exc,
+                        exc_info=True
+                    )
+
+            ctx.set_state("unity_instance", active_instance)
+            if session_id is not None:
+                ctx.set_state("unity_session_id", session_id)
+
+    async def on_call_tool(self, context: MiddlewareContext, call_next):
+        """Inject active Unity instance into tool context if available."""
+        await self._inject_unity_instance(context)
+        return await call_next(context)
+
+    async def on_read_resource(self, context: MiddlewareContext, call_next):
+        """Inject active Unity instance into resource context if available."""
+        await self._inject_unity_instance(context)
+        return await call_next(context)

transport/unity_transport.py

@@ -0,0 +1,94 @@
+"""Transport helpers for routing commands to Unity."""
+from __future__ import annotations
+
+import logging
+from typing import Awaitable, Callable, TypeVar
+
+from transport.plugin_hub import PluginHub
+from core.config import config
+from core.constants import API_KEY_HEADER
+from services.api_key_service import ApiKeyService
+from models.models import MCPResponse
+from models.unity_response import normalize_unity_response
+
+logger = logging.getLogger(__name__)
+T = TypeVar("T")
+
+
+def _is_http_transport() -> bool:
+    return config.transport_mode.lower() == "http"
+
+
+async def _resolve_user_id_from_request() -> str | None:
+    """Extract user_id from the current HTTP request's API key header."""
+    if not config.http_remote_hosted:
+        return None
+    if not ApiKeyService.is_initialized():
+        return None
+    try:
+        from fastmcp.server.dependencies import get_http_headers
+        headers = get_http_headers(include_all=True)
+        api_key = headers.get(API_KEY_HEADER.lower())
+        if not api_key:
+            return None
+        service = ApiKeyService.get_instance()
+        result = await service.validate(api_key)
+        return result.user_id if result.valid else None
+    except Exception as e:
+        logger.debug("Failed to resolve user_id from HTTP request: %s", e)
+        return None
+
+
+async def send_with_unity_instance(
+    send_fn: Callable[..., Awaitable[T]],
+    unity_instance: str | None,
+    *args,
+    user_id: str | None = None,
+    **kwargs,
+) -> T:
+    if _is_http_transport():
+        if not args:
+            raise ValueError("HTTP transport requires command arguments")
+        command_type = args[0]
+        params = args[1] if len(args) > 1 else kwargs.get("params")
+        if params is None:
+            params = {}
+        if not isinstance(params, dict):
+            raise TypeError(
+                "Command parameters must be a dict for HTTP transport")
+
+        # Auto-resolve user_id from HTTP request API key (remote-hosted mode)
+        if user_id is None:
+            user_id = await _resolve_user_id_from_request()
+
+        # Auth check
+        if config.http_remote_hosted and not user_id:
+            return normalize_unity_response(
+                MCPResponse(
+                    success=False,
+                    error="auth_required",
+                    message="API key required",
+                ).model_dump()
+            )
+
+        try:
+            raw = await PluginHub.send_command_for_instance(
+                unity_instance,
+                command_type,
+                params,
+                user_id=user_id,
+            )
+            return normalize_unity_response(raw)
+        except Exception as exc:
+            # NOTE: asyncio.TimeoutError has an empty str() by default, which is confusing for clients.
+            err = str(exc) or f"{type(exc).__name__}"
+            # Fail fast with a retry hint instead of hanging for COMMAND_TIMEOUT.
+            # The client can decide whether retrying is appropriate for the command.
+            return normalize_unity_response(
+                MCPResponse(success=False, error=err,
+                            hint="retry").model_dump()
+            )
+
+    if unity_instance:
+        kwargs.setdefault("instance_id", unity_instance)
+    return await send_fn(*args, **kwargs)