Skip to content

Commit c76d0bb

Browse files
committed
Fix on #837
1 parent 359f2fd commit c76d0bb

6 files changed

Lines changed: 413 additions & 4 deletions

File tree

MCPForUnity/Editor/Resources/Editor/ToolStates.cs

Lines changed: 25 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -23,11 +23,35 @@ public static object HandleCommand(JObject @params)
2323
var toolsArray = new JArray();
2424
foreach (var tool in allTools)
2525
{
26+
var paramsArray = new JArray();
27+
if (tool.Parameters != null)
28+
{
29+
foreach (var p in tool.Parameters)
30+
{
31+
paramsArray.Add(new JObject
32+
{
33+
["name"] = p.Name,
34+
["description"] = p.Description,
35+
["type"] = p.Type,
36+
["required"] = p.Required,
37+
["default_value"] = p.DefaultValue
38+
});
39+
}
40+
}
41+
2642
toolsArray.Add(new JObject
2743
{
2844
["name"] = tool.Name,
2945
["group"] = tool.Group ?? "core",
30-
["enabled"] = discovery.IsToolEnabled(tool.Name)
46+
["enabled"] = discovery.IsToolEnabled(tool.Name),
47+
["description"] = tool.Description,
48+
["auto_register"] = tool.AutoRegister,
49+
["is_built_in"] = tool.IsBuiltIn,
50+
["structured_output"] = tool.StructuredOutput,
51+
["requires_polling"] = tool.RequiresPolling,
52+
["poll_action"] = tool.PollAction,
53+
["max_poll_seconds"] = tool.MaxPollSeconds,
54+
["parameters"] = paramsArray
3155
});
3256
}
3357

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

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -48,8 +48,8 @@ public Task<bool> VerifyAsync()
4848

4949
public Task ReregisterToolsAsync()
5050
{
51-
// Stdio transport doesn't support dynamic tool reregistration
52-
// Tools are registered at server startup
51+
// In stdio mode, Python re-syncs tools automatically on reconnection
52+
// after domain reload. No proactive push mechanism exists over TCP.
5353
return Task.CompletedTask;
5454
}
5555

Server/src/services/custom_tool_service.py

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -118,6 +118,9 @@ async def get_tool_definition(
118118
user_id: str | None = None,
119119
) -> ToolDefinitionModel | None:
120120
tool = self._project_tools.get(project_id, {}).get(tool_name)
121+
if tool:
122+
return tool
123+
tool = self._global_tools.get(tool_name)
121124
if tool:
122125
return tool
123126
return await PluginHub.get_tool_definition(project_id, tool_name, user_id=user_id)

Server/src/services/tools/__init__.py

Lines changed: 67 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -169,6 +169,72 @@ async def sync_tool_visibility_from_unity(
169169

170170
PluginHub._sync_server_tool_visibility(enabled_tools)
171171

172+
# Register custom (non-built-in) tools via CustomToolService.
173+
# The extended get_tool_states response includes is_built_in,
174+
# description, parameters, etc. If those fields are missing
175+
# (older Unity package), we skip custom tool registration.
176+
custom_tool_count = 0
177+
has_extended_metadata = any(
178+
"is_built_in" in t for t in enabled_tools
179+
)
180+
if has_extended_metadata:
181+
custom_tool_dicts = [
182+
t for t in enabled_tools if not t.get("is_built_in", True)
183+
]
184+
if custom_tool_dicts:
185+
try:
186+
from models.models import ToolDefinitionModel, ToolParameterModel
187+
from services.custom_tool_service import CustomToolService
188+
189+
custom_tool_models = []
190+
for td in custom_tool_dicts:
191+
params = [
192+
ToolParameterModel(
193+
name=p.get("name", ""),
194+
description=p.get("description", ""),
195+
type=p.get("type", "string"),
196+
required=p.get("required", False),
197+
default_value=p.get("default_value"),
198+
)
199+
for p in td.get("parameters", [])
200+
]
201+
custom_tool_models.append(
202+
ToolDefinitionModel(
203+
name=td["name"],
204+
description=td.get("description", ""),
205+
structured_output=td.get("structured_output", True),
206+
requires_polling=td.get("requires_polling", False),
207+
poll_action=td.get("poll_action", "status"),
208+
max_poll_seconds=td.get("max_poll_seconds", 0),
209+
parameters=params,
210+
)
211+
)
212+
213+
service = CustomToolService.get_instance()
214+
service.register_global_tools(custom_tool_models)
215+
custom_tool_count = len(custom_tool_models)
216+
logger.info(
217+
"Registered %d custom tool(s) from Unity via stdio sync",
218+
custom_tool_count,
219+
)
220+
except RuntimeError as exc:
221+
logger.debug(
222+
"Skipping custom tool registration: "
223+
"CustomToolService not initialized yet (%s)",
224+
exc,
225+
)
226+
except Exception as exc:
227+
logger.warning(
228+
"Failed to register custom tools from Unity: %s",
229+
exc,
230+
)
231+
else:
232+
logger.debug(
233+
"Unity response does not include extended tool metadata "
234+
"(is_built_in); skipping custom tool registration. "
235+
"Update MCPForUnity to enable custom tool sync in stdio mode."
236+
)
237+
172238
if notify:
173239
await PluginHub._notify_mcp_tool_list_changed()
174240

@@ -191,6 +257,7 @@ async def sync_tool_visibility_from_unity(
191257
"disabled_groups": disabled_groups,
192258
"enabled_tool_count": len(enabled_tools),
193259
"total_tool_count": len(tools),
260+
"custom_tool_count": custom_tool_count,
194261
}
195262

196263
except Exception as exc:

Server/src/transport/legacy/unity_connection.py

Lines changed: 46 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -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

Comments
 (0)