mcpforunityserver 9.3.0b20260129121506__py3-none-any.whl → 9.3.0b20260131003150__py3-none-any.whl

services/resources/tests.py

@@ -4,6 +4,7 @@ from pydantic import BaseModel, Field
 from fastmcp import Context
 
 from models import MCPResponse
+from models.unity_response import parse_resource_response
 from services.registry import mcp_for_unity_resource
 from services.tools import get_unity_instance_from_context
 from transport.unity_transport import send_with_unity_instance
@@ -53,7 +54,7 @@ async def get_tests(ctx: Context) -> GetTestsResponse | MCPResponse:
         "get_tests",
         {},
     )
-    return GetTestsResponse(**response) if isinstance(response, dict) else response
+    return parse_resource_response(response, GetTestsResponse)
 
 
 @mcp_for_unity_resource(
@@ -84,4 +85,4 @@ async def get_tests_for_mode(
         "get_tests_for_mode",
         {"mode": mode},
     )
-    return GetTestsResponse(**response) if isinstance(response, dict) else response
+    return parse_resource_response(response, GetTestsResponse)

services/resources/unity_instances.py

@@ -7,7 +7,7 @@ from fastmcp import Context
 from services.registry import mcp_for_unity_resource
 from transport.legacy.unity_connection import get_unity_connection_pool
 from transport.plugin_hub import PluginHub
-from transport.unity_transport import _current_transport
+from core.config import config
 
 
 @mcp_for_unity_resource(
@@ -36,10 +36,13 @@ async def unity_instances(ctx: Context) -> dict[str, Any]:
     await ctx.info("Listing Unity instances")
 
     try:
-        transport = _current_transport()
+        transport = (config.transport_mode or "stdio").lower()
         if transport == "http":
             # HTTP/WebSocket transport: query PluginHub
-            sessions_data = await PluginHub.get_sessions()
+            # In remote-hosted mode, filter sessions by user_id
+            user_id = ctx.get_state(
+                "user_id") if config.http_remote_hosted else None
+            sessions_data = await PluginHub.get_sessions(user_id=user_id)
             sessions = sessions_data.sessions
 
             instances = []

services/resources/windows.py

@@ -2,6 +2,7 @@ from pydantic import BaseModel
 from fastmcp import Context
 
 from models import MCPResponse
+from models.unity_response import parse_resource_response
 from services.registry import mcp_for_unity_resource
 from services.tools import get_unity_instance_from_context
 from transport.unity_transport import send_with_unity_instance
@@ -44,4 +45,4 @@ async def get_windows(ctx: Context) -> WindowsResponse | MCPResponse:
         "get_windows",
         {}
     )
-    return WindowsResponse(**response) if isinstance(response, dict) else response
+    return parse_resource_response(response, WindowsResponse)

services/tools/set_active_instance.py

@@ -8,7 +8,7 @@ from services.registry import mcp_for_unity_tool
 from transport.legacy.unity_connection import get_unity_connection_pool
 from transport.unity_instance_middleware import get_unity_instance_middleware
 from transport.plugin_hub import PluginHub
-from transport.unity_transport import _current_transport
+from core.config import config
 
 
 @mcp_for_unity_tool(
@@ -21,11 +21,14 @@ async def set_active_instance(
         ctx: Context,
         instance: Annotated[str, "Target instance (Name@hash or hash prefix)"]
 ) -> dict[str, Any]:
-    transport = _current_transport()
+    transport = (config.transport_mode or "stdio").lower()
 
     # Discover running instances based on transport
     if transport == "http":
-        sessions_data = await PluginHub.get_sessions()
+        # In remote-hosted mode, filter sessions by user_id
+        user_id = ctx.get_state(
+            "user_id") if config.http_remote_hosted else None
+        sessions_data = await PluginHub.get_sessions(user_id=user_id)
         sessions = sessions_data.sessions
         instances = []
         for session_id, session in sessions.items():

transport/plugin_hub.py

@@ -13,8 +13,10 @@ from starlette.endpoints import WebSocketEndpoint
 from starlette.websockets import WebSocket
 
 from core.config import config
+from core.constants import API_KEY_HEADER
 from models.models import MCPResponse
 from transport.plugin_registry import PluginRegistry
+from services.api_key_service import ApiKeyService
 from transport.models import (
     WelcomeMessage,
     RegisteredMessage,
@@ -38,6 +40,22 @@ class NoUnitySessionError(RuntimeError):
     """Raised when no Unity plugins are available."""
 
 
+class InstanceSelectionRequiredError(RuntimeError):
+    """Raised when the caller must explicitly select a Unity instance."""
+
+    _SELECTION_REQUIRED = (
+        "Unity instance selection is required. "
+        "Call set_active_instance with Name@hash from mcpforunity://instances."
+    )
+    _MULTIPLE_INSTANCES = (
+        "Multiple Unity instances are connected. "
+        "Call set_active_instance with Name@hash from mcpforunity://instances."
+    )
+
+    def __init__(self, message: str | None = None):
+        super().__init__(message or self._SELECTION_REQUIRED)
+
+
 class PluginHub(WebSocketEndpoint):
     """Manages persistent WebSocket connections to Unity plugins."""
 
@@ -77,6 +95,50 @@ class PluginHub(WebSocketEndpoint):
         return cls._registry is not None and cls._lock is not None
 
     async def on_connect(self, websocket: WebSocket) -> None:
+        # Validate API key in remote-hosted mode (fail closed)
+        if config.http_remote_hosted:
+            if not ApiKeyService.is_initialized():
+                logger.debug(
+                    "WebSocket connection rejected: auth service not initialized")
+                await websocket.close(code=1013, reason="Try again later")
+                return
+
+            api_key = websocket.headers.get(API_KEY_HEADER)
+
+            if not api_key:
+                logger.debug("WebSocket connection rejected: API key required")
+                await websocket.close(code=4401, reason="API key required")
+                return
+
+            service = ApiKeyService.get_instance()
+            result = await service.validate(api_key)
+
+            if not result.valid:
+                # Transient auth failures are retryable (1013)
+                if result.error and any(
+                    indicator in result.error.lower()
+                    for indicator in ("unavailable", "timeout", "service error")
+                ):
+                    logger.debug(
+                        "WebSocket connection rejected: auth service unavailable")
+                    await websocket.close(code=1013, reason="Try again later")
+                    return
+
+                logger.debug("WebSocket connection rejected: invalid API key")
+                await websocket.close(code=4403, reason="Invalid API key")
+                return
+
+            # Both valid and user_id must be present to accept
+            if not result.user_id:
+                logger.debug(
+                    "WebSocket connection rejected: validated key missing user_id")
+                await websocket.close(code=4403, reason="Invalid API key")
+                return
+
+            # Store user_id in websocket state for later use during registration
+            websocket.state.user_id = result.user_id
+            websocket.state.api_key_metadata = result.metadata
+
         await websocket.accept()
         msg = WelcomeMessage(
             serverTimeout=self.SERVER_TIMEOUT,
@@ -217,10 +279,15 @@ class PluginHub(WebSocketEndpoint):
                 cls._pending.pop(command_id, None)
 
     @classmethod
-    async def get_sessions(cls) -> SessionList:
+    async def get_sessions(cls, user_id: str | None = None) -> SessionList:
+        """Get all active plugin sessions.
+
+        Args:
+            user_id: If provided (remote-hosted mode), only return sessions for this user.
+        """
         if cls._registry is None:
             return SessionList(sessions={})
-        sessions = await cls._registry.list_sessions()
+        sessions = await cls._registry.list_sessions(user_id=user_id)
         return SessionList(
             sessions={
                 session_id: SessionDetails(
@@ -286,15 +353,23 @@ class PluginHub(WebSocketEndpoint):
             raise ValueError(
                 "Plugin registration missing project_hash")
 
+        # Get user_id from websocket state (set during API key validation)
+        user_id = getattr(websocket.state, "user_id", None)
+
         session_id = str(uuid.uuid4())
         # Inform the plugin of its assigned session ID
         response = RegisteredMessage(session_id=session_id)
         await websocket.send_json(response.model_dump())
 
-        session = await registry.register(session_id, project_name, project_hash, unity_version, project_path)
+        session = await registry.register(session_id, project_name, project_hash, unity_version, project_path, user_id=user_id)
         async with lock:
             cls._connections[session.session_id] = websocket
-        logger.info(f"Plugin registered: {project_name} ({project_hash})")
+
+        if user_id:
+            logger.info(
+                f"Plugin registered: {project_name} ({project_hash}) for user {user_id}")
+        else:
+            logger.info(f"Plugin registered: {project_name} ({project_hash})")
 
     async def _handle_register_tools(self, websocket: WebSocket, payload: RegisterToolsMessage) -> None:
         cls = type(self)
@@ -375,13 +450,17 @@ class PluginHub(WebSocketEndpoint):
     # Session resolution helpers
     # ------------------------------------------------------------------
     @classmethod
-    async def _resolve_session_id(cls, unity_instance: str | None) -> str:
+    async def _resolve_session_id(cls, unity_instance: str | None, user_id: str | None = None) -> str:
         """Resolve a project hash (Unity instance id) to an active plugin session.
 
         During Unity domain reloads the plugin's WebSocket session is torn down
         and reconnected shortly afterwards. Instead of failing immediately when
         no sessions are available, we wait for a bounded period for a plugin
         to reconnect so in-flight MCP calls can succeed transparently.
+
+        Args:
+            unity_instance: Target instance (Name@hash or hash)
+            user_id: User ID from API key validation (for remote-hosted mode session isolation)
         """
         if cls._registry is None:
             raise RuntimeError("Plugin registry not configured")
@@ -411,24 +490,35 @@ class PluginHub(WebSocketEndpoint):
             else:
                 target_hash = unity_instance
 
-        async def _try_once() -> tuple[str | None, int]:
+        async def _try_once() -> tuple[str | None, int, bool]:
+            explicit_required = config.http_remote_hosted
             # Prefer a specific Unity instance if one was requested
             if target_hash:
-                session_id = await cls._registry.get_session_id_by_hash(target_hash)
-                sessions = await cls._registry.list_sessions()
-                return session_id, len(sessions)
+                # In remote-hosted mode with user_id, use user-scoped lookup
+                if config.http_remote_hosted and user_id:
+                    session_id = await cls._registry.get_session_id_by_hash(target_hash, user_id)
+                    sessions = await cls._registry.list_sessions(user_id=user_id)
+                else:
+                    session_id = await cls._registry.get_session_id_by_hash(target_hash)
+                    sessions = await cls._registry.list_sessions(user_id=user_id)
+                return session_id, len(sessions), explicit_required
 
             # No target provided: determine if we can auto-select
-            sessions = await cls._registry.list_sessions()
+            # In remote-hosted mode, filter sessions by user_id
+            sessions = await cls._registry.list_sessions(user_id=user_id)
             count = len(sessions)
             if count == 0:
-                return None, count
+                return None, count, explicit_required
+            if explicit_required:
+                return None, count, explicit_required
             if count == 1:
-                return next(iter(sessions.keys())), count
+                return next(iter(sessions.keys())), count, explicit_required
             # Multiple sessions but no explicit target is ambiguous
-            return None, count
+            return None, count, explicit_required
 
-        session_id, session_count = await _try_once()
+        session_id, session_count, explicit_required = await _try_once()
+        if session_id is None and explicit_required and not target_hash and session_count > 0:
+            raise InstanceSelectionRequiredError()
         deadline = time.monotonic() + max_wait_s
         wait_started = None
 
@@ -436,10 +526,10 @@ class PluginHub(WebSocketEndpoint):
         # wait politely for a session to appear before surfacing an error.
         while session_id is None and time.monotonic() < deadline:
             if not target_hash and session_count > 1:
-                raise RuntimeError(
-                    "Multiple Unity instances are connected. "
-                    "Call set_active_instance with Name@hash from mcpforunity://instances."
-                )
+                raise InstanceSelectionRequiredError(
+                    InstanceSelectionRequiredError._MULTIPLE_INSTANCES)
+            if session_id is None and explicit_required and not target_hash and session_count > 0:
+                raise InstanceSelectionRequiredError()
             if wait_started is None:
                 wait_started = time.monotonic()
                 logger.debug(
@@ -448,7 +538,7 @@ class PluginHub(WebSocketEndpoint):
                     max_wait_s,
                 )
             await asyncio.sleep(sleep_seconds)
-            session_id, session_count = await _try_once()
+            session_id, session_count, explicit_required = await _try_once()
 
         if session_id is not None and wait_started is not None:
             logger.debug(
@@ -457,10 +547,11 @@ class PluginHub(WebSocketEndpoint):
                 unity_instance or "default",
             )
         if session_id is None and not target_hash and session_count > 1:
-            raise RuntimeError(
-                "Multiple Unity instances are connected. "
-                "Call set_active_instance with Name@hash from mcpforunity://instances."
-            )
+            raise InstanceSelectionRequiredError(
+                InstanceSelectionRequiredError._MULTIPLE_INSTANCES)
+
+        if session_id is None and explicit_required and not target_hash and session_count > 0:
+            raise InstanceSelectionRequiredError()
 
         if session_id is None:
             logger.warning(
@@ -481,9 +572,18 @@ class PluginHub(WebSocketEndpoint):
         unity_instance: str | None,
         command_type: str,
         params: dict[str, Any],
+        user_id: str | None = None,
     ) -> dict[str, Any]:
+        """Send a command to a Unity instance.
+
+        Args:
+            unity_instance: Target instance (Name@hash or hash)
+            command_type: Command type to execute
+            params: Command parameters
+            user_id: User ID for session isolation in remote-hosted mode
+        """
         try:
-            session_id = await cls._resolve_session_id(unity_instance)
+            session_id = await cls._resolve_session_id(unity_instance, user_id=user_id)
         except NoUnitySessionError:
             logger.debug(
                 "Unity session unavailable; returning retry: command=%s instance=%s",

transport/plugin_registry.py

@@ -7,6 +7,7 @@ from datetime import datetime, timezone
 
 import asyncio
 
+from core.config import config
 from models.models import ToolDefinitionModel
 
 
@@ -22,7 +23,9 @@ class PluginSession:
     connected_at: datetime
     tools: dict[str, ToolDefinitionModel] = field(default_factory=dict)
     project_id: str | None = None
-    project_path: str | None = None  # Full path to project root (for focus nudging)
+    # 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:
@@ -31,11 +34,17 @@ class PluginRegistry:
     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(
@@ -45,13 +54,16 @@ class PluginRegistry:
         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`` it will be
-        replaced, ensuring that reconnect scenarios always map to the latest
-        WebSocket connection.
+        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)
@@ -63,15 +75,26 @@ class PluginRegistry:
                 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
-            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)
+            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
-            self._hash_to_session[project_hash] = session_id
             return session
 
     async def touch(self, session_id: str) -> None:
@@ -87,11 +110,20 @@ class PluginRegistry:
 
         async with self._lock:
             session = self._sessions.pop(session_id, None)
-            if session and session.project_hash in self._hash_to_session:
-                # Only delete the mapping if it still points at the removed session.
-                mapped = self._hash_to_session.get(session.project_hash)
-                if mapped == session_id:
-                    del self._hash_to_session[session.project_hash]
+            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."""
@@ -110,17 +142,41 @@ class PluginRegistry:
         async with self._lock:
             return self._sessions.get(session_id)
 
-    async def get_session_id_by_hash(self, project_hash: str) -> str | None:
+    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."""
 
-        async with self._lock:
-            return self._hash_to_session.get(project_hash)
+        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.
 
-    async def list_sessions(self) -> dict[str, PluginSession]:
-        """Return a shallow copy of all known 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:
-            return dict(self._sessions)
+            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

@@ -9,6 +9,7 @@ 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")
@@ -32,7 +33,12 @@ def get_unity_instance_middleware() -> 'UnityInstanceMiddleware':
 
 
 def set_unity_instance_middleware(middleware: 'UnityInstanceMiddleware') -> None:
-    """Set the global Unity instance middleware (called during server initialization)."""
+    """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
 
@@ -55,13 +61,18 @@ class UnityInstanceMiddleware(Middleware):
         Derive a stable key for the calling session.
 
         Prioritizes client_id for stability.
-        If client_id is missing, falls back to 'global' (assuming single-user local mode),
-        ignoring session_id which can be unstable in some transports/clients.
+        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"
 
@@ -92,10 +103,10 @@ class UnityInstanceMiddleware(Middleware):
         to stick for subsequent tool/resource calls in the same session.
         """
         try:
-            # Import here to avoid circular dependencies / optional transport modules.
-            from transport.unity_transport import _current_transport
-
-            transport = _current_transport()
+            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()
@@ -172,10 +183,27 @@ class UnityInstanceMiddleware(Middleware):
 
         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 into context if available."""
+        """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)
@@ -193,7 +221,8 @@ class UnityInstanceMiddleware(Middleware):
                     # 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.
-                    session_id = await PluginHub._resolve_session_id(active_instance)
+                    # 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?