mcpforunityserver 8.2.3__py3-none-any.whl

transport/models.py

@@ -0,0 +1,62 @@
+from typing import Any
+from pydantic import BaseModel, Field
+from models.models import ToolDefinitionModel
+
+# Outgoing (Server -> Plugin)
+
+
+class WelcomeMessage(BaseModel):
+    type: str = "welcome"
+    serverTimeout: int
+    keepAliveInterval: int
+
+
+class RegisteredMessage(BaseModel):
+    type: str = "registered"
+    session_id: str
+
+
+class ExecuteCommandMessage(BaseModel):
+    type: str = "execute"
+    id: str
+    name: str
+    params: dict[str, Any]
+    timeout: float
+
+# Incoming (Plugin -> Server)
+
+
+class RegisterMessage(BaseModel):
+    type: str = "register"
+    project_name: str = "Unknown Project"
+    project_hash: str
+    unity_version: str = "Unknown"
+
+
+class RegisterToolsMessage(BaseModel):
+    type: str = "register_tools"
+    tools: list[ToolDefinitionModel]
+
+
+class PongMessage(BaseModel):
+    type: str = "pong"
+    session_id: str | None = None
+
+
+class CommandResultMessage(BaseModel):
+    type: str = "command_result"
+    id: str
+    result: dict[str, Any] = Field(default_factory=dict)
+
+# Session Info (API response)
+
+
+class SessionDetails(BaseModel):
+    project: str
+    hash: str
+    unity_version: str
+    connected_at: str
+
+
+class SessionList(BaseModel):
+    sessions: dict[str, SessionDetails]

transport/plugin_hub.py

@@ -0,0 +1,412 @@
+"""WebSocket hub for Unity plugin communication."""
+
+from __future__ import annotations
+
+import asyncio
+import logging
+import time
+import uuid
+from typing import Any
+
+from starlette.endpoints import WebSocketEndpoint
+from starlette.websockets import WebSocket
+
+from core.config import config
+from transport.plugin_registry import PluginRegistry
+from transport.models import (
+    WelcomeMessage,
+    RegisteredMessage,
+    ExecuteCommandMessage,
+    RegisterMessage,
+    RegisterToolsMessage,
+    PongMessage,
+    CommandResultMessage,
+    SessionList,
+    SessionDetails,
+)
+
+logger = logging.getLogger("mcp-for-unity-server")
+
+
+class PluginHub(WebSocketEndpoint):
+    """Manages persistent WebSocket connections to Unity plugins."""
+
+    encoding = "json"
+    KEEP_ALIVE_INTERVAL = 15
+    SERVER_TIMEOUT = 30
+    COMMAND_TIMEOUT = 30
+
+    _registry: PluginRegistry | None = None
+    _connections: dict[str, WebSocket] = {}
+    _pending: dict[str, asyncio.Future] = {}
+    _lock: asyncio.Lock | None = None
+    _loop: asyncio.AbstractEventLoop | None = None
+
+    @classmethod
+    def configure(
+        cls,
+        registry: PluginRegistry,
+        loop: asyncio.AbstractEventLoop | None = None,
+    ) -> None:
+        cls._registry = registry
+        cls._loop = loop or asyncio.get_running_loop()
+        # Ensure coordination primitives are bound to the configured loop
+        cls._lock = asyncio.Lock()
+
+    @classmethod
+    def is_configured(cls) -> bool:
+        return cls._registry is not None and cls._lock is not None
+
+    async def on_connect(self, websocket: WebSocket) -> None:
+        await websocket.accept()
+        msg = WelcomeMessage(
+            serverTimeout=self.SERVER_TIMEOUT,
+            keepAliveInterval=self.KEEP_ALIVE_INTERVAL,
+        )
+        await websocket.send_json(msg.model_dump())
+
+    async def on_receive(self, websocket: WebSocket, data: Any) -> None:
+        if not isinstance(data, dict):
+            logger.warning(f"Received non-object payload from plugin: {data}")
+            return
+
+        message_type = data.get("type")
+        try:
+            if message_type == "register":
+                await self._handle_register(websocket, RegisterMessage(**data))
+            elif message_type == "register_tools":
+                await self._handle_register_tools(websocket, RegisterToolsMessage(**data))
+            elif message_type == "pong":
+                await self._handle_pong(PongMessage(**data))
+            elif message_type == "command_result":
+                await self._handle_command_result(CommandResultMessage(**data))
+            else:
+                logger.debug(f"Ignoring plugin message: {data}")
+        except Exception as e:
+            logger.error(f"Error handling message type {message_type}: {e}")
+
+    async def on_disconnect(self, websocket: WebSocket, close_code: int) -> None:
+        cls = type(self)
+        lock = cls._lock
+        if lock is None:
+            return
+        async with lock:
+            session_id = next(
+                (sid for sid, ws in cls._connections.items() if ws is websocket), None)
+            if session_id:
+                cls._connections.pop(session_id, None)
+                if cls._registry:
+                    await cls._registry.unregister(session_id)
+                logger.info(
+                    f"Plugin session {session_id} disconnected ({close_code})")
+
+    # ------------------------------------------------------------------
+    # Public API
+    # ------------------------------------------------------------------
+    @classmethod
+    async def send_command(cls, session_id: str, command_type: str, params: dict[str, Any]) -> dict[str, Any]:
+        websocket = await cls._get_connection(session_id)
+        command_id = str(uuid.uuid4())
+        future: asyncio.Future = asyncio.get_running_loop().create_future()
+
+        lock = cls._lock
+        if lock is None:
+            raise RuntimeError("PluginHub not configured")
+
+        async with lock:
+            if command_id in cls._pending:
+                raise RuntimeError(
+                    f"Duplicate command id generated: {command_id}")
+            cls._pending[command_id] = future
+
+        try:
+            msg = ExecuteCommandMessage(
+                id=command_id,
+                name=command_type,
+                params=params,
+                timeout=cls.COMMAND_TIMEOUT,
+            )
+            await websocket.send_json(msg.model_dump())
+            result = await asyncio.wait_for(future, timeout=cls.COMMAND_TIMEOUT)
+            return result
+        finally:
+            async with lock:
+                cls._pending.pop(command_id, None)
+
+    @classmethod
+    async def get_sessions(cls) -> SessionList:
+        if cls._registry is None:
+            return SessionList(sessions={})
+        sessions = await cls._registry.list_sessions()
+        return SessionList(
+            sessions={
+                session_id: SessionDetails(
+                    project=session.project_name,
+                    hash=session.project_hash,
+                    unity_version=session.unity_version,
+                    connected_at=session.connected_at.isoformat(),
+                )
+                for session_id, session in sessions.items()
+            }
+        )
+
+    @classmethod
+    async def get_tools_for_project(cls, project_hash: str) -> list[Any]:
+        """Retrieve tools registered for a active project hash."""
+        if cls._registry is None:
+            return []
+
+        session_id = await cls._registry.get_session_id_by_hash(project_hash)
+        if not session_id:
+            return []
+
+        session = await cls._registry.get_session(session_id)
+        if not session:
+            return []
+
+        return list(session.tools.values())
+
+    @classmethod
+    async def get_tool_definition(cls, project_hash: str, tool_name: str) -> Any | None:
+        """Retrieve a specific tool definition for an active project hash."""
+        if cls._registry is None:
+            return None
+
+        session_id = await cls._registry.get_session_id_by_hash(project_hash)
+        if not session_id:
+            return None
+
+        session = await cls._registry.get_session(session_id)
+        if not session:
+            return None
+
+        return session.tools.get(tool_name)
+
+    # ------------------------------------------------------------------
+    # Internal helpers
+    # ------------------------------------------------------------------
+    async def _handle_register(self, websocket: WebSocket, payload: RegisterMessage) -> None:
+        cls = type(self)
+        registry = cls._registry
+        lock = cls._lock
+        if registry is None or lock is None:
+            await websocket.close(code=1011)
+            raise RuntimeError("PluginHub not configured")
+
+        project_name = payload.project_name
+        project_hash = payload.project_hash
+        unity_version = payload.unity_version
+
+        if not project_hash:
+            await websocket.close(code=4400)
+            raise ValueError(
+                "Plugin registration missing project_hash")
+
+        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)
+        async with lock:
+            cls._connections[session.session_id] = websocket
+        logger.info(f"Plugin registered: {project_name} ({project_hash})")
+
+    async def _handle_register_tools(self, websocket: WebSocket, payload: RegisterToolsMessage) -> None:
+        cls = type(self)
+        registry = cls._registry
+        lock = cls._lock
+        if registry is None or lock is None:
+            return
+
+        # Find session_id for this websocket
+        async with lock:
+            session_id = next(
+                (sid for sid, ws in cls._connections.items() if ws is websocket), None)
+
+        if not session_id:
+            logger.warning("Received register_tools from unknown connection")
+            return
+
+        await registry.register_tools_for_session(session_id, payload.tools)
+        logger.info(
+            f"Registered {len(payload.tools)} tools for session {session_id}")
+
+    async def _handle_command_result(self, payload: CommandResultMessage) -> None:
+        cls = type(self)
+        lock = cls._lock
+        if lock is None:
+            return
+        command_id = payload.id
+        result = payload.result
+
+        if not command_id:
+            logger.warning(f"Command result missing id: {payload}")
+            return
+
+        async with lock:
+            future = cls._pending.get(command_id)
+        if future and not future.done():
+            future.set_result(result)
+
+    async def _handle_pong(self, payload: PongMessage) -> None:
+        cls = type(self)
+        registry = cls._registry
+        if registry is None:
+            return
+        session_id = payload.session_id
+        if session_id:
+            await registry.touch(session_id)
+
+    @classmethod
+    async def _get_connection(cls, session_id: str) -> WebSocket:
+        lock = cls._lock
+        if lock is None:
+            raise RuntimeError("PluginHub not configured")
+        async with lock:
+            websocket = cls._connections.get(session_id)
+        if websocket is None:
+            raise RuntimeError(f"Plugin session {session_id} not connected")
+        return websocket
+
+    # ------------------------------------------------------------------
+    # Session resolution helpers
+    # ------------------------------------------------------------------
+    @classmethod
+    async def _resolve_session_id(cls, unity_instance: str | 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.
+        """
+        if cls._registry is None:
+            raise RuntimeError("Plugin registry not configured")
+
+        # Use the same defaults as the stdio transport reload handling so that
+        # HTTP/WebSocket and TCP behave consistently without per-project env.
+        max_retries = max(1, int(getattr(config, "reload_max_retries", 40)))
+        retry_ms = float(getattr(config, "reload_retry_ms", 250))
+        sleep_seconds = max(0.05, retry_ms / 1000.0)
+
+        # Allow callers to provide either just the hash or Name@hash
+        target_hash: str | None = None
+        if unity_instance:
+            if "@" in unity_instance:
+                _, _, suffix = unity_instance.rpartition("@")
+                target_hash = suffix or None
+            else:
+                target_hash = unity_instance
+
+        async def _try_once() -> tuple[str | None, int]:
+            # 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)
+
+            # No target provided: determine if we can auto-select
+            sessions = await cls._registry.list_sessions()
+            count = len(sessions)
+            if count == 0:
+                return None, count
+            if count == 1:
+                return next(iter(sessions.keys())), count
+            # Multiple sessions but no explicit target is ambiguous
+            return None, count
+
+        session_id, session_count = await _try_once()
+        deadline = time.monotonic() + (max_retries * sleep_seconds)
+        wait_started = None
+
+        # If there is no active plugin yet (e.g., Unity starting up or reloading),
+        # 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 unity://instances."
+                )
+            if wait_started is None:
+                wait_started = time.monotonic()
+                logger.debug(
+                    f"No plugin session available (instance={unity_instance or 'default'}); waiting up to {deadline - wait_started:.2f}s",
+                )
+            await asyncio.sleep(sleep_seconds)
+            session_id, session_count = await _try_once()
+
+        if session_id is not None and wait_started is not None:
+            logger.debug(
+                f"Plugin session restored after {time.monotonic() - wait_started:.3f}s (instance={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 unity://instances."
+            )
+
+        if session_id is None:
+            logger.warning(
+                f"No Unity plugin reconnected within {max_retries * sleep_seconds:.2f}s (instance={unity_instance or 'default'})",
+            )
+            # At this point we've given the plugin ample time to reconnect; surface
+            # a clear error so the client can prompt the user to open Unity.
+            raise RuntimeError("No Unity plugins are currently connected")
+
+        return session_id
+
+    @classmethod
+    async def send_command_for_instance(
+        cls,
+        unity_instance: str | None,
+        command_type: str,
+        params: dict[str, Any],
+    ) -> dict[str, Any]:
+        session_id = await cls._resolve_session_id(unity_instance)
+        return await cls.send_command(session_id, command_type, params)
+
+    # ------------------------------------------------------------------
+    # Blocking helpers for synchronous tool code
+    # ------------------------------------------------------------------
+    @classmethod
+    def _run_coroutine_sync(cls, coro: "asyncio.Future[Any]") -> Any:
+        if cls._loop is None:
+            raise RuntimeError("PluginHub event loop not configured")
+        loop = cls._loop
+        if loop.is_running():
+            try:
+                running_loop = asyncio.get_running_loop()
+            except RuntimeError:
+                running_loop = None
+            else:
+                if running_loop is loop:
+                    raise RuntimeError(
+                        "Cannot wait synchronously for PluginHub coroutine from within the event loop"
+                    )
+        future = asyncio.run_coroutine_threadsafe(coro, loop)
+        return future.result()
+
+    @classmethod
+    def send_command_blocking(
+        cls,
+        unity_instance: str | None,
+        command_type: str,
+        params: dict[str, Any],
+    ) -> dict[str, Any]:
+        return cls._run_coroutine_sync(
+            cls.send_command_for_instance(unity_instance, command_type, params)
+        )
+
+    @classmethod
+    def list_sessions_sync(cls) -> SessionList:
+        return cls._run_coroutine_sync(cls.get_sessions())
+
+
+def send_command_to_plugin(
+    *,
+    unity_instance: str | None,
+    command_type: str,
+    params: dict[str, Any],
+) -> dict[str, Any]:
+    return PluginHub.send_command_blocking(unity_instance, command_type, params)

transport/plugin_registry.py

@@ -0,0 +1,123 @@
+"""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 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
+
+
+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).
+    """
+
+    def __init__(self) -> None:
+        self._sessions: dict[str, PluginSession] = {}
+        self._hash_to_session: dict[str, str] = {}
+        self._lock = asyncio.Lock()
+
+    async def register(
+        self,
+        session_id: str,
+        project_name: str,
+        project_hash: str,
+        unity_version: str,
+    ) -> 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.
+        """
+
+        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,
+            )
+
+            # 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)
+
+            self._sessions[session_id] = session
+            self._hash_to_session[project_hash] = session_id
+            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 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]
+
+    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) -> 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)
+
+    async def list_sessions(self) -> dict[str, PluginSession]:
+        """Return a shallow copy of all known sessions."""
+
+        async with self._lock:
+            return dict(self._sessions)
+
+
+__all__ = ["PluginRegistry", "PluginSession"]