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

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)
@@ -186,14 +214,16 @@ class UnityInstanceMiddleware(Middleware):
             # 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
-            # OR if we want to support hybrid mode. For now, let's be permissive.
-            if PluginHub.is_configured():
+            # 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.
-                    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?

transport/unity_transport.py

@@ -1,34 +1,49 @@
 """Transport helpers for routing commands to Unity."""
 from __future__ import annotations
 
-import asyncio
-import inspect
-import os
+import logging
 from typing import Awaitable, Callable, TypeVar
 
-from fastmcp import Context
-
 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
-from services.tools import get_unity_instance_from_context
 
+logger = logging.getLogger(__name__)
 T = TypeVar("T")
 
 
 def _is_http_transport() -> bool:
-    return os.environ.get("UNITY_MCP_TRANSPORT", "stdio").lower() == "http"
+    return config.transport_mode.lower() == "http"
 
 
-def _current_transport() -> str:
-    """Expose the active transport mode as a simple string identifier."""
-    return "http" if _is_http_transport() else "stdio"
+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():
@@ -41,11 +56,27 @@ async def send_with_unity_instance(
         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: