66"""
77from threading import RLock
88import logging
9+ import os
910import time
11+ from urllib .parse import unquote , urlparse
1012
1113from fastmcp .server .middleware import Middleware , MiddlewareContext
1214
@@ -47,6 +49,28 @@ def set_unity_instance_middleware(middleware: 'UnityInstanceMiddleware') -> None
4749 _unity_instance_middleware = middleware
4850
4951
52+ def _file_uri_to_path (uri : str ) -> str | None :
53+ """Convert a ``file://`` URI to a local filesystem path (UNC- and drive-aware)."""
54+ if not uri .startswith ("file://" ):
55+ return None
56+ parsed = urlparse (uri )
57+ host = (parsed .netloc or "" ).strip ()
58+ path = unquote (parsed .path or "" )
59+ if host and host .lower () != "localhost" :
60+ path = f"//{ host } { path } " # UNC: file://server/share/... -> //server/share/...
61+ elif os .name == "nt" and len (path ) >= 3 and path [0 ] == "/" and path [2 ] == ":" :
62+ path = path [1 :] # drive-letter form /C:/... -> C:/...
63+ return path or None
64+
65+
66+ def _strip_assets (project_path : str ) -> str :
67+ """Return the Unity project root (stdio reports ``.../Assets``, HTTP the root)."""
68+ normalized = project_path .replace ("\\ " , "/" ).rstrip ("/" )
69+ if normalized .lower ().endswith ("/assets" ):
70+ return normalized [: - len ("/assets" )]
71+ return normalized
72+
73+
5074class UnityInstanceMiddleware (Middleware ):
5175 """
5276 Middleware that manages per-session Unity instance selection.
@@ -58,6 +82,7 @@ class UnityInstanceMiddleware(Middleware):
5882 def __init__ (self ):
5983 super ().__init__ ()
6084 self ._active_by_key : dict [str , str ] = {}
85+ self ._root_dirs_by_key : dict [str , list [str ]] = {}
6186 self ._lock = RLock ()
6287 self ._metadata_lock = RLock ()
6388 self ._unity_managed_tool_names : set [str ] = set ()
@@ -224,6 +249,101 @@ async def _resolve_instance_value(self, value: str, ctx) -> str:
224249 "Read mcpforunity://instances for current sessions."
225250 )
226251
252+ async def _resolve_launch_dirs (self , ctx ) -> list [str ]:
253+ """Discover the directories the client is working in, client-agnostic.
254+
255+ The package's own ``UNITY_MCP_PROJECT_DIR`` override wins (a single dir);
256+ otherwise the ``file://`` MCP roots the client advertises - the
257+ protocol-native signal for what the session is working on, of which there
258+ may be several. Empty when neither is available.
259+ """
260+ explicit = os .environ .get ("UNITY_MCP_PROJECT_DIR" )
261+ if explicit :
262+ return [explicit ]
263+ return await self ._client_root_dirs (ctx )
264+
265+ async def _client_root_dirs (self , ctx ) -> list [str ]:
266+ """The client's ``file://`` MCP roots, resolved once per session.
267+
268+ Cached because each lookup is a round-trip to the client and a session's
269+ working directory does not change underneath us; this keeps the no-match
270+ path off the wire on every subsequent tool call.
271+ """
272+ key = await self .get_session_key (ctx )
273+ with self ._lock :
274+ cached = self ._root_dirs_by_key .get (key )
275+ if cached is not None :
276+ return cached
277+ dirs = await self ._fetch_client_root_dirs (ctx )
278+ with self ._lock :
279+ self ._root_dirs_by_key [key ] = dirs
280+ return dirs
281+
282+ @staticmethod
283+ async def _fetch_client_root_dirs (ctx ) -> list [str ]:
284+ """Query MCP roots and return their local paths (``file://`` only).
285+
286+ The caller caches the result per session, so a client that does not
287+ support roots costs at most one failed probe, not one per tool call.
288+ """
289+ list_roots = getattr (ctx , "list_roots" , None )
290+ if not callable (list_roots ):
291+ return []
292+ try :
293+ roots = await list_roots ()
294+ except Exception :
295+ # Client does not implement roots, or the request failed; not fatal.
296+ return []
297+ dirs : list [str ] = []
298+ for root in roots or []:
299+ path = _file_uri_to_path (str (getattr (root , "uri" , "" ) or "" ))
300+ if path :
301+ dirs .append (path )
302+ return dirs
303+
304+ @staticmethod
305+ def _select_instance_by_launch_dir (
306+ candidates : list [tuple [str | None , str | None ]],
307+ launch_dirs : list [str ],
308+ ) -> str | None :
309+ """Pick the single connected editor whose project matches a launch dir.
310+
311+ Each candidate's project path is normalized to the project root (Unity
312+ reports the ``Assets`` folder over stdio but the project root over HTTP)
313+ and matched by path lineage against every launch directory - the project
314+ root contains, equals, or is contained by one of them. This routes
315+ per-checkout and git-worktree setups (identical project names, different
316+ paths) without an explicit ``unity_instance``. Returns the id only when
317+ exactly one editor matches; otherwise None, so the caller keeps its
318+ "ask the user to choose" behavior.
319+ """
320+ launch_reals = [
321+ os .path .normcase (os .path .realpath (d )) for d in launch_dirs if d
322+ ]
323+ if not launch_reals :
324+ return None
325+
326+ matches : set [str ] = set ()
327+ for inst_id , project_path in candidates :
328+ if not inst_id or not project_path :
329+ continue
330+ project_root = _strip_assets (project_path )
331+ if not project_root :
332+ continue
333+ project_real = os .path .normcase (os .path .realpath (project_root ))
334+ for launch_real in launch_reals :
335+ try :
336+ shared = os .path .commonpath ([launch_real , project_real ])
337+ except ValueError :
338+ continue # e.g. paths on different Windows drives
339+ if shared in (launch_real , project_real ):
340+ matches .add (inst_id )
341+ break
342+
343+ if len (matches ) == 1 :
344+ return next (iter (matches ))
345+ return None
346+
227347 async def _maybe_autoselect_instance (self , ctx ) -> str | None :
228348 """
229349 Auto-select the sole Unity instance when no active instance is set.
@@ -241,13 +361,17 @@ async def _maybe_autoselect_instance(self, ctx) -> str | None:
241361 try :
242362 sessions_data = await PluginHub .get_sessions ()
243363 sessions = sessions_data .sessions or {}
244- ids : list [str ] = []
364+ candidates : list [tuple [ str | None , str | None ] ] = []
245365 for session_info in sessions .values ():
246366 project = getattr (
247367 session_info , "project" , None ) or "Unknown"
248368 hash_value = getattr (session_info , "hash" , None )
249369 if hash_value :
250- ids .append (f"{ project } @{ hash_value } " )
370+ candidates .append ((
371+ f"{ project } @{ hash_value } " ,
372+ getattr (session_info , "project_path" , None ),
373+ ))
374+ ids = [inst_id for inst_id , _ in candidates ]
251375 if len (ids ) == 1 :
252376 chosen = ids [0 ]
253377 await self .set_active_instance (ctx , chosen )
@@ -257,6 +381,17 @@ async def _maybe_autoselect_instance(self, ctx) -> str | None:
257381 )
258382 return chosen
259383 if len (ids ) > 1 :
384+ launch_dirs = await self ._resolve_launch_dirs (ctx )
385+ chosen = self ._select_instance_by_launch_dir (
386+ candidates , launch_dirs )
387+ if chosen :
388+ await self .set_active_instance (ctx , chosen )
389+ logger .info (
390+ "Auto-selected Unity instance %s via launch directory "
391+ "(of %d running) over PluginHub." ,
392+ chosen , len (ids ),
393+ )
394+ return chosen
260395 logger .info (
261396 "Multiple Unity instances found (%d). Pass unity_instance on any tool call "
262397 "or call set_active_instance to choose one. Available: %s" ,
@@ -284,8 +419,11 @@ async def _maybe_autoselect_instance(self, ctx) -> str | None:
284419
285420 pool = get_unity_connection_pool ()
286421 instances = pool .discover_all_instances (force_refresh = True )
287- ids = [getattr (inst , "id" , None ) for inst in instances ]
288- ids = [inst_id for inst_id in ids if inst_id ]
422+ candidates = [
423+ (getattr (inst , "id" , None ), getattr (inst , "path" , None ))
424+ for inst in instances
425+ ]
426+ ids = [inst_id for inst_id , _ in candidates if inst_id ]
289427 if len (ids ) == 1 :
290428 chosen = ids [0 ]
291429 await self .set_active_instance (ctx , chosen )
@@ -295,6 +433,17 @@ async def _maybe_autoselect_instance(self, ctx) -> str | None:
295433 )
296434 return chosen
297435 if len (ids ) > 1 :
436+ launch_dirs = await self ._resolve_launch_dirs (ctx )
437+ chosen = self ._select_instance_by_launch_dir (
438+ candidates , launch_dirs )
439+ if chosen :
440+ await self .set_active_instance (ctx , chosen )
441+ logger .info (
442+ "Auto-selected Unity instance %s via launch directory "
443+ "(of %d running) via stdio discovery." ,
444+ chosen , len (ids ),
445+ )
446+ return chosen
298447 logger .info (
299448 "Multiple Unity instances found (%d). Pass unity_instance on any tool call "
300449 "or call set_active_instance to choose one. Available: %s" ,
0 commit comments