@@ -42,6 +42,7 @@ def __post_init__(self):
4242 self .port = stdio_port_registry .get_port (self .instance_id )
4343 self ._io_lock = threading .Lock ()
4444 self ._conn_lock = threading .Lock ()
45+ self ._needs_tool_resync = False # Set True after reconnection
4546
4647 def _prepare_socket (self , sock : socket .socket ) -> None :
4748 try :
@@ -65,6 +66,7 @@ def connect(self) -> bool:
6566 self .sock = socket .create_connection (
6667 (self .host , self .port ), connect_timeout )
6768 self ._prepare_socket (self .sock )
69+ self ._needs_tool_resync = True
6870 logger .debug (f"Connected to Unity at { self .host } :{ self .port } " )
6971
7072 # Strict handshake: require FRAMING=1
@@ -933,11 +935,54 @@ async def async_send_command_with_retry(
933935 import asyncio # local import to avoid mandatory asyncio dependency for sync callers
934936 if loop is None :
935937 loop = asyncio .get_running_loop ()
936- return await loop .run_in_executor (
938+ result = await loop .run_in_executor (
937939 None ,
938940 lambda : send_command_with_retry (
939941 command_type , params , instance_id = instance_id , max_retries = max_retries ,
940942 retry_ms = retry_ms , retry_on_reload = retry_on_reload ),
941943 )
944+
945+ # After a successful command, check if the connection was freshly
946+ # established (reconnection after domain reload). If so, re-sync
947+ # tool visibility and custom tool registration from Unity.
948+ # Skip if *this* call is the sync itself to avoid recursion.
949+ if command_type != "get_tool_states" :
950+ try :
951+ pool = get_unity_connection_pool ()
952+ conn = pool .get_connection (instance_id )
953+ if getattr (conn , "_needs_tool_resync" , False ):
954+ conn ._needs_tool_resync = False
955+ logger .info (
956+ "Detected reconnection to Unity; scheduling tool re-sync"
957+ )
958+ asyncio .ensure_future (_resync_tools_after_reconnect (instance_id ))
959+ except Exception :
960+ pass # Best-effort; don't fail the actual command
961+
962+ return result
942963 except Exception as e :
943964 return MCPResponse (success = False , error = str (e ))
965+
966+
967+ async def _resync_tools_after_reconnect (instance_id : str | None ) -> None :
968+ """Background task: re-sync tool visibility and custom tools after reconnection."""
969+ try :
970+ from services .tools import sync_tool_visibility_from_unity
971+ result = await sync_tool_visibility_from_unity (
972+ instance_id = instance_id , notify = True ,
973+ )
974+ if result .get ("synced" ):
975+ logger .info (
976+ "Post-reconnection tool re-sync complete: "
977+ "enabled=[%s], disabled=[%s], custom_tools=%d" ,
978+ ", " .join (result .get ("enabled_groups" , [])),
979+ ", " .join (result .get ("disabled_groups" , [])),
980+ result .get ("custom_tool_count" , 0 ),
981+ )
982+ else :
983+ logger .debug (
984+ "Post-reconnection tool re-sync skipped: %s" ,
985+ result .get ("error" , "unknown" ),
986+ )
987+ except Exception as exc :
988+ logger .debug ("Post-reconnection tool re-sync failed: %s" , exc )
0 commit comments