"""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"]