-
Notifications
You must be signed in to change notification settings - Fork 1.5k
Fix on #837 #1025
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Fix on #837 #1025
Changes from 2 commits
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -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, | ||
| ) | ||
|
Comment on lines
+201
to
+210
|
||
| ) | ||
|
Comment on lines
+201
to
+211
Contributor
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Potential Line 203 uses direct dict access Consider either:
🛡️ Proposed fix to handle missing name gracefully custom_tool_models = []
for td in custom_tool_dicts:
+ tool_name = td.get("name")
+ if not tool_name:
+ logger.warning("Skipping custom tool with missing name: %s", td)
+ continue
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"],
+ name=tool_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,
)
)🤖 Prompt for AI Agents |
||
|
|
||
| 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, | ||
| ) | ||
|
Comment on lines
+213
to
+219
|
||
| 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: | ||
|
|
||
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -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}") | ||
|
Comment on lines
63
to
70
|
||
|
|
||
| # 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. | ||
|
sourcery-ai[bot] marked this conversation as resolved.
|
||
| # 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) | ||
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
ToolParameterModel.required defaults to True in models, but this code uses p.get("required", False). If Unity omits the field (or sends null), parameters will be treated as optional and the generated tool signature will be wrong. Default to True when missing, and consider skipping parameters with missing/empty names instead of using "".