55into the request-scoped state, allowing tools to access it via ctx.get_state("unity_instance").
66"""
77from threading import RLock
8+ import logging
89
910from fastmcp .server .middleware import Middleware , MiddlewareContext
1011
1112from transport .plugin_hub import PluginHub
1213
14+ logger = logging .getLogger ("mcp-for-unity-server" )
15+
1316# Store a global reference to the middleware instance so tools can interact
1417# with it to set or clear the active unity instance.
1518_unity_instance_middleware = None
1619
1720
1821def get_unity_instance_middleware () -> 'UnityInstanceMiddleware' :
1922 """Get the global Unity instance middleware."""
23+ global _unity_instance_middleware
2024 if _unity_instance_middleware is None :
21- raise RuntimeError (
22- "UnityInstanceMiddleware not initialized. Call set_unity_instance_middleware first." )
25+ # Auto-initialize if not set (lazy singleton) to handle import order or test cases
26+ _unity_instance_middleware = UnityInstanceMiddleware ()
27+
2328 return _unity_instance_middleware
2429
2530
@@ -42,37 +47,36 @@ def __init__(self):
4247 self ._active_by_key : dict [str , str ] = {}
4348 self ._lock = RLock ()
4449
45- def _get_session_key (self , ctx ) -> str :
50+ def get_session_key (self , ctx ) -> str :
4651 """
4752 Derive a stable key for the calling session.
4853
49- Uses ctx.session_id if available, falls back to 'global'.
54+ Prioritizes client_id for stability.
55+ If client_id is missing, falls back to 'global' (assuming single-user local mode),
56+ ignoring session_id which can be unstable in some transports/clients.
5057 """
51- session_id = getattr (ctx , "session_id" , None )
52- if isinstance (session_id , str ) and session_id :
53- return session_id
54-
5558 client_id = getattr (ctx , "client_id" , None )
5659 if isinstance (client_id , str ) and client_id :
5760 return client_id
5861
62+ # Fallback to global for local dev stability
5963 return "global"
6064
6165 def set_active_instance (self , ctx , instance_id : str ) -> None :
6266 """Store the active instance for this session."""
63- key = self ._get_session_key (ctx )
67+ key = self .get_session_key (ctx )
6468 with self ._lock :
6569 self ._active_by_key [key ] = instance_id
6670
6771 def get_active_instance (self , ctx ) -> str | None :
6872 """Retrieve the active instance for this session."""
69- key = self ._get_session_key (ctx )
73+ key = self .get_session_key (ctx )
7074 with self ._lock :
7175 return self ._active_by_key .get (key )
7276
7377 def clear_active_instance (self , ctx ) -> None :
7478 """Clear the stored instance for this session."""
75- key = self ._get_session_key (ctx )
79+ key = self .get_session_key (ctx )
7680 with self ._lock :
7781 self ._active_by_key .pop (key , None )
7882
@@ -82,13 +86,32 @@ async def on_call_tool(self, context: MiddlewareContext, call_next):
8286
8387 active_instance = self .get_active_instance (ctx )
8488 if active_instance :
89+ # If using HTTP transport (PluginHub configured), validate session
90+ # But for stdio transport (no PluginHub needed or maybe partially configured),
91+ # we should be careful not to clear instance just because PluginHub can't resolve it.
92+ # The 'active_instance' (Name@hash) might be valid for stdio even if PluginHub fails.
93+
8594 session_id : str | None = None
95+ # Only validate via PluginHub if we are actually using HTTP transport
96+ # OR if we want to support hybrid mode. For now, let's be permissive.
8697 if PluginHub .is_configured ():
8798 try :
99+ # resolving session_id might fail if the plugin disconnected
100+ # We only need session_id for HTTP transport routing.
101+ # For stdio, we just need the instance ID.
88102 session_id = await PluginHub ._resolve_session_id (active_instance )
89- except Exception :
90- self .clear_active_instance (ctx )
91- return await call_next (context )
103+ except Exception as exc :
104+ # If resolution fails, it means the Unity instance is not reachable via HTTP/WS.
105+ # If we are in stdio mode, this might still be fine if the user is just setting state?
106+ # But usually if PluginHub is configured, we expect it to work.
107+ # Let's LOG the error but NOT clear the instance immediately to avoid flickering,
108+ # or at least debug why it's failing.
109+ logger .debug (
110+ "PluginHub session resolution failed for %s: %s; leaving active_instance unchanged" ,
111+ active_instance ,
112+ exc ,
113+ exc_info = True ,
114+ )
92115
93116 ctx .set_state ("unity_instance" , active_instance )
94117 if session_id is not None :
0 commit comments