Merge pull request #1025 from Scriptwonder/fix-tcp-customtool

Fix on #837

Author: Scriptwonder

MCPForUnity/Editor/Helpers/Response.cs.meta

@@ -2,8 +2,7 @@
     "name": "MCPForUnity.Editor",
     "rootNamespace": "MCPForUnity.Editor",
     "references": [
-        "MCPForUnity.Runtime",
-        "Newtonsoft.Json"
+        "MCPForUnity.Runtime"
     ],
     "includePlatforms": [
         "Editor"

MCPForUnity/Editor/MCPForUnity.Editor.asmdef

@@ -23,11 +23,35 @@ public static object HandleCommand(JObject @params)
                 var toolsArray = new JArray();
                 foreach (var tool in allTools)
                 {
+                    var paramsArray = new JArray();
+                    if (tool.Parameters != null)
+                    {
+                        foreach (var p in tool.Parameters)
+                        {
+                            paramsArray.Add(new JObject
+                            {
+                                ["name"] = p.Name,
+                                ["description"] = p.Description,
+                                ["type"] = p.Type,
+                                ["required"] = p.Required,
+                                ["default_value"] = p.DefaultValue
+                            });
+                        }
+                    }
+
                     toolsArray.Add(new JObject
                     {
                         ["name"] = tool.Name,
                         ["group"] = tool.Group ?? "core",
-                        ["enabled"] = discovery.IsToolEnabled(tool.Name)
+                        ["enabled"] = discovery.IsToolEnabled(tool.Name),
+                        ["description"] = tool.Description,
+                        ["auto_register"] = tool.AutoRegister,
+                        ["is_built_in"] = tool.IsBuiltIn,
+                        ["structured_output"] = tool.StructuredOutput,
+                        ["requires_polling"] = tool.RequiresPolling,
+                        ["poll_action"] = tool.PollAction ?? "status",
+                        ["max_poll_seconds"] = tool.MaxPollSeconds,
+                        ["parameters"] = paramsArray
                     });
                 }
 

MCPForUnity/Editor/Resources/Editor/ToolStates.cs

@@ -48,8 +48,8 @@ public Task<bool> VerifyAsync()
 
         public Task ReregisterToolsAsync()
         {
-            // Stdio transport doesn't support dynamic tool reregistration
-            // Tools are registered at server startup
+            // In stdio mode, Python re-syncs tools automatically on reconnection
+            // after domain reload. No proactive push mechanism exists over TCP.
             return Task.CompletedTask;
         }
 

MCPForUnity/Editor/Services/Transport/Transports/StdioTransportClient.cs

@@ -540,7 +540,7 @@ private async Task SendRegisterToolsAsync(CancellationToken token)
                     ["description"] = tool.Description,
                     ["structured_output"] = tool.StructuredOutput,
                     ["requires_polling"] = tool.RequiresPolling,
-                    ["poll_action"] = tool.PollAction,
+                    ["poll_action"] = tool.PollAction ?? "status",
                     ["max_poll_seconds"] = tool.MaxPollSeconds,
                     ["group"] = string.IsNullOrWhiteSpace(tool.Group) ? "core" : tool.Group
                 };

MCPForUnity/Editor/Services/Transport/Transports/WebSocketTransportClient.cs

@@ -1,9 +1,7 @@
 {
     "name": "MCPForUnity.Runtime",
     "rootNamespace": "MCPForUnity.Runtime",
-    "references": [
-        "Newtonsoft.Json"
-    ],
+    "references": [],
     "includePlatforms": [],
     "excludePlatforms": [],
     "allowUnsafeCode": false,

MCPForUnity/Runtime/MCPForUnity.Runtime.asmdef

@@ -118,6 +118,9 @@ async def get_tool_definition(
         user_id: str | None = None,
     ) -> ToolDefinitionModel | None:
         tool = self._project_tools.get(project_id, {}).get(tool_name)
+        if tool:
+            return tool
+        tool = self._global_tools.get(tool_name)
         if tool:
             return tool
         return await PluginHub.get_tool_definition(project_id, tool_name, user_id=user_id)

Server/src/services/custom_tool_service.py

@@ -169,6 +169,72 @@ async def sync_tool_visibility_from_unity(
 
         PluginHub._sync_server_tool_visibility(enabled_tools)
 
+        # Register custom (non-built-in) tools via CustomToolService.
+        # The extended get_tool_states response includes is_built_in,
+        # description, parameters, etc.  If those fields are missing
+        # (older Unity package), we skip custom tool registration.
+        custom_tool_count = 0
+        has_extended_metadata = any(
+            "is_built_in" in t for t in enabled_tools
+        )
+        if has_extended_metadata:
+            custom_tool_dicts = [
+                t for t in enabled_tools if not t.get("is_built_in", True)
+            ]
+            if custom_tool_dicts:
+                try:
+                    from models.models import ToolDefinitionModel, ToolParameterModel
+                    from services.custom_tool_service import CustomToolService
+
+                    custom_tool_models = []
+                    for td in custom_tool_dicts:
+                        params = [
+                            ToolParameterModel(
+                                name=p.get("name", ""),
+                                description=p.get("description", ""),
+                                type=p.get("type", "string"),
+                                required=p.get("required", True),
+                                default_value=p.get("default_value"),
+                            )
+                            for p in td.get("parameters", [])
+                        ]
+                        custom_tool_models.append(
+                            ToolDefinitionModel(
+                                name=td["name"],
+                                description=td.get("description", ""),
+                                structured_output=td.get("structured_output", True),
+                                requires_polling=td.get("requires_polling", False),
+                                poll_action=td.get("poll_action") or "status",
+                                max_poll_seconds=td.get("max_poll_seconds", 0),
+                                parameters=params,
+                            )
+                        )
+
+                    service = CustomToolService.get_instance()
+                    service.register_global_tools(custom_tool_models)
+                    custom_tool_count = len(custom_tool_models)
+                    logger.info(
+                        "Registered %d custom tool(s) from Unity via stdio sync",
+                        custom_tool_count,
+                    )
+                except RuntimeError as exc:
+                    logger.debug(
+                        "Skipping custom tool registration: "
+                        "CustomToolService not initialized yet (%s)",
+                        exc,
+                    )
+                except Exception as exc:
+                    logger.warning(
+                        "Failed to register custom tools from Unity: %s",
+                        exc,
+                    )
+        else:
+            logger.debug(
+                "Unity response does not include extended tool metadata "
+                "(is_built_in); skipping custom tool registration. "
+                "Update MCPForUnity to enable custom tool sync in stdio mode."
+            )
+
         if notify:
             await PluginHub._notify_mcp_tool_list_changed()
 
@@ -191,6 +257,7 @@ async def sync_tool_visibility_from_unity(
             "disabled_groups": disabled_groups,
             "enabled_tool_count": len(enabled_tools),
             "total_tool_count": len(tools),
+            "custom_tool_count": custom_tool_count,
         }
 
     except Exception as exc:

Server/src/services/tools/__init__.py

@@ -42,6 +42,7 @@ def __post_init__(self):
             self.port = stdio_port_registry.get_port(self.instance_id)
         self._io_lock = threading.Lock()
         self._conn_lock = threading.Lock()
+        self._needs_tool_resync = False  # Set True after reconnection
 
     def _prepare_socket(self, sock: socket.socket) -> None:
         try:
@@ -65,6 +66,7 @@ def connect(self) -> bool:
                 self.sock = socket.create_connection(
                     (self.host, self.port), connect_timeout)
                 self._prepare_socket(self.sock)
+                self._needs_tool_resync = True
                 logger.debug(f"Connected to Unity at {self.host}:{self.port}")
 
                 # Strict handshake: require FRAMING=1
@@ -933,11 +935,58 @@ async def async_send_command_with_retry(
         import asyncio  # local import to avoid mandatory asyncio dependency for sync callers
         if loop is None:
             loop = asyncio.get_running_loop()
-        return await loop.run_in_executor(
+        result = await loop.run_in_executor(
             None,
             lambda: send_command_with_retry(
                 command_type, params, instance_id=instance_id, max_retries=max_retries,
                 retry_ms=retry_ms, retry_on_reload=retry_on_reload),
         )
+
+        # After a successful command, check if the connection was freshly
+        # established (reconnection after domain reload).  If so, re-sync
+        # tool visibility and custom tool registration from Unity.
+        # Always clear the flag, but only schedule the background resync
+        # when this call is not itself get_tool_states (to avoid recursion).
+        try:
+            pool = get_unity_connection_pool()
+            conn = pool.get_connection(instance_id)
+            if getattr(conn, "_needs_tool_resync", False):
+                conn._needs_tool_resync = False
+                if command_type != "get_tool_states":
+                    logger.info(
+                        "Detected reconnection to Unity; scheduling tool re-sync"
+                    )
+                    asyncio.ensure_future(_resync_tools_after_reconnect(instance_id))
+        except Exception as exc:
+            logger.debug(
+                "Failed to schedule post-reconnection tool re-sync: %s",
+                exc,
+            )
+
+        return result
     except Exception as e:
         return MCPResponse(success=False, error=str(e))
+
+
+async def _resync_tools_after_reconnect(instance_id: str | None) -> None:
+    """Background task: re-sync tool visibility and custom tools after reconnection."""
+    try:
+        from services.tools import sync_tool_visibility_from_unity
+        result = await sync_tool_visibility_from_unity(
+            instance_id=instance_id, notify=True,
+        )
+        if result.get("synced"):
+            logger.info(
+                "Post-reconnection tool re-sync complete: "
+                "enabled=[%s], disabled=[%s], custom_tools=%d",
+                ", ".join(result.get("enabled_groups", [])),
+                ", ".join(result.get("disabled_groups", [])),
+                result.get("custom_tool_count", 0),
+            )
+        else:
+            logger.debug(
+                "Post-reconnection tool re-sync skipped: %s",
+                result.get("error", "unknown"),
+            )
+    except Exception as exc:
+        logger.debug("Post-reconnection tool re-sync failed: %s", exc)