001 state aware tool filtering#864
Conversation
- Added `filter_middleware.py` to filter tools based on editor state. - Enhanced `editor_state.py` to include selection state attributes. - Updated tool management scripts to utilize prerequisite checks for visibility. - Created integration tests for tool filtering logic in `test_tool_filtering_integration.py`. - Developed unit tests for `tool_filter_decorator.py` to validate prerequisite logic. - Introduced `EditorStateSelectionTests.cs` for Unity to verify selection state handling.
…hot handling and async decorator tests
…hatevertogo/unity-mcp into 001-state-aware-tool-filtering
Fixes CoplayDev#709 - Remove logic in ToolDiscoveryService.EnsurePreferenceInitialized that forced built-in tools to be re-enabled on every Unity launch - Add ReregisterToolsAsync() to IMcpTransportClient for dynamic tool list updates when toggles change in the UI - Implement tool reregistration in WebSocketTransportClient - Add ReregisterToolsAsync() to McpToolsSection to trigger updates when tool toggles are flipped - Add unit tests for ToolDiscoveryService preference persistence Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
Reviewer's GuideImplements state-aware tool filtering by extending the Unity editor state snapshot with selection and advice metadata, adding server-side prerequisite/decorator infrastructure and middleware to hide or block tools based on editor state, wiring tool re‑registration on Unity when user toggles tools, and adding tests for both the Unity side and Python server side. Sequence diagram for tool re-registration on Unity tool togglesequenceDiagram
actor User
participant McpToolsSection
participant ThreadPool
participant MCPServiceLocator
participant TransportManager
participant HttpClient as WebSocketTransportClient
User->>McpToolsSection: ToggleTool(tool, enabled)
McpToolsSection->>McpToolsSection: HandleToggleChange(tool, enabled, updateSummary)
McpToolsSection->>McpToolsSection: UpdateSummary()
McpToolsSection->>ThreadPool: QueueUserWorkItem(ReregisterToolsAsync)
activate ThreadPool
ThreadPool->>MCPServiceLocator: TransportManager
MCPServiceLocator-->>ThreadPool: TransportManager
ThreadPool->>TransportManager: GetClient(TransportMode.Http)
TransportManager-->>ThreadPool: HttpClient
alt HttpClient is not null and IsConnected
ThreadPool->>HttpClient: ReregisterToolsAsync()
activate HttpClient
HttpClient->>HttpClient: SendRegisterToolsAsync(_lifecycleCts.Token)
HttpClient-->>ThreadPool: Completed
deactivate HttpClient
else No client or not connected
ThreadPool->>ThreadPool: Skip reregistration
end
ThreadPool-->>McpToolsSection: Return
deactivate ThreadPool
Flow diagram for tool list filtering based on editor stateflowchart TD
A["FastMCP Server
all_tools list"] --> B["get_tools_matching_state(ctx, all_tools)"]
B --> C["get_editor_state(ctx)"]
C --> D{Editor state acquired?}
D -- No --> E["Log warning
return all_tools"]
D -- Yes --> F[For each tool in all_tools]
F --> G{Tool has
ToolPrerequisite?}
G -- No --> H["Add tool to
filtered_tools"]
G -- Yes --> I["prereq = tool_prerequisites[name]
prereq.is_met(state_data)"]
I --> J{is_met?}
J -- Yes --> H
J -- No --> K["Skip tool
log hidden with blocking_reason"]
H --> L[Next tool]
K --> L
L -->|Loop until done| M["Return filtered_tools
to FastMCP"]
File-Level Changes
Tips and commandsInteracting with Sourcery
Customizing Your ExperienceAccess your dashboard to:
Getting Help
|
|
@codex review it |
|
Caution Review failedPull request was closed or merged during review 📝 WalkthroughWalkthroughThis PR introduces a cross-cutting state-aware prerequisite system for MCP tools. It expands EditorStateCache with selection state tracking and advisory metadata, implements a server-side prerequisite decorator and filter middleware to gate tool visibility and execution based on editor conditions (compilation, selection, play mode, tests), and wires reregistration flows. Prerequisites are checked at execution time and during tool discovery, with comprehensive test coverage. Changes
Estimated code review effort🎯 4 (Complex) | ⏱️ ~60 minutes The PR introduces a substantial feature spanning client and server with multiple new abstractions (ToolPrerequisite, prerequisite_check decorator, filter_middleware), integrates state-aware filtering at both execution and discovery time, requires understanding of decorator mechanics, async/sync wrapper patterns, thread-safe registry management, and includes dense test coverage. The changes affect tool registration, transport, and editor state tracking across heterogeneous systems, demanding careful review of control flow, error handling paths, and prerequisite evaluation logic. Possibly related PRs
Suggested labels
Suggested reviewers
Poem
🚥 Pre-merge checks | ❌ 3❌ Failed checks (2 warnings, 1 inconclusive)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 Generate unit tests (beta)
Tip Try Coding Plans. Let us write the prompt for your AI agent so you can ship faster (with fewer bugs). Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
|
To use Codex here, create a Codex account and connect to github. |
|
@codex review it |
|
To use Codex here, create a Codex account and connect to github. |
![sourcery-ai[bot]](https://avatars.githubusercontent.com/in/48477?s=60&v=4){kind=small}
There was a problem hiding this comment.
Hey - I've found 3 issues, and left some high level feedback:
- In
McpToolsSection.ReregisterToolsAsync, you spin up a ThreadPool work item and then synchronously block onclient.ReregisterToolsAsync().Wait(), which risks deadlocks or long-lived background threads if the implementation ever needs the calling context; consider either making the call chain fully async (and awaited from the UI event) or running the async operation without blocking (e.g.,ReregisterToolsAsync().ConfigureAwait(false)inside the worker). - The synchronous wrapper in
prerequisite_checkmanually creates/sets an asyncio event loop and may callrun_until_completeon a loop that could already be running, which can conflict with the hosting server; it would be safer to avoid loop management in the decorator (or to restrict prerequisites checking to the async path) rather than mutating the global event loop from library code.
Prompt for AI Agents
Please address the comments from this code review:
## Overall Comments
- In `McpToolsSection.ReregisterToolsAsync`, you spin up a ThreadPool work item and then synchronously block on `client.ReregisterToolsAsync().Wait()`, which risks deadlocks or long-lived background threads if the implementation ever needs the calling context; consider either making the call chain fully async (and awaited from the UI event) or running the async operation without blocking (e.g., `ReregisterToolsAsync().ConfigureAwait(false)` inside the worker).
- The synchronous wrapper in `prerequisite_check` manually creates/sets an asyncio event loop and may call `run_until_complete` on a loop that could already be running, which can conflict with the hosting server; it would be safer to avoid loop management in the decorator (or to restrict prerequisites checking to the async path) rather than mutating the global event loop from library code.
## Individual Comments
### Comment 1
<location path="MCPForUnity/Editor/Windows/Components/Tools/McpToolsSection.cs" line_range="241-250" />
<code_context>
+ ReregisterToolsAsync();
+ }
+
+ private void ReregisterToolsAsync()
+ {
+ // Fire and forget - don't block UI
+ ThreadPool.QueueUserWorkItem(_ =>
+ {
+ try
+ {
+ var transportManager = MCPServiceLocator.TransportManager;
+ var client = transportManager.GetClient(TransportMode.Http);
+ if (client != null && client.IsConnected)
+ {
+ client.ReregisterToolsAsync().Wait();
+ }
+ }
</code_context>
<issue_to_address>
**suggestion (bug_risk):** Avoid blocking on async ReregisterToolsAsync() using .Wait() inside ThreadPool work item.
Since this runs on a ThreadPool thread, synchronously waiting on `ReregisterToolsAsync()` can still cause deadlocks or thread pool starvation if the async method ever captures a context or awaits other blocking work. Instead, keep this path fully async by using an async delegate, e.g. `ThreadPool.QueueUserWorkItem(async _ => { ... await client.ReregisterToolsAsync().ConfigureAwait(false); });` so you never block a worker thread.
Suggested implementation:
```csharp
private void ReregisterToolsAsync()
{
// Fire and forget - don't block UI
ThreadPool.QueueUserWorkItem(async _ =>
{
```
```csharp
var transportManager = MCPServiceLocator.TransportManager;
var client = transportManager.GetClient(TransportMode.Http);
if (client != null && client.IsConnected)
{
await client.ReregisterToolsAsync().ConfigureAwait(false);
}
```
</issue_to_address>
### Comment 2
<location path="Server/src/core/tool_filter_decorator.py" line_range="138-147" />
<code_context>
+ @functools.wraps(func)
</code_context>
<issue_to_address>
**suggestion (bug_risk):** The sync wrapper’s event loop management is complex and may leak loops or conflict with existing ones.
In `_sync_wrapper`, the logic that obtains/creates an event loop and calls `run_until_complete` is fragile: when a loop is already running it just warns and skips the check, and when a new loop is created it’s never closed. If sync usage is required, either use a dedicated `asyncio.run` path (to get proper loop setup/teardown) or explicitly treat sync as a best-effort mode (e.g., skip the prereq check) rather than partially managing event loops in a way that may vary across runtimes.
Suggested implementation:
```python
# Check prerequisites before executing.
# NOTE: sync invocation is treated as a *best-effort* mode.
# For direct tool calls from async code, prefer the async wrapper.
# For tool list filtering, the middleware handles prerequisite checks.
```
The core of the requested change is in `_sync_wrapper`’s event-loop handling, which is below the snippet you shared. Because I can’t see the full function, you’ll need to replace its body directly. Here is a concrete implementation that addresses the comment by:
1. Completely removing ad‑hoc loop creation / `run_until_complete`.
2. Using `asyncio.run` when there is **no** running loop (proper setup/teardown).
3. Treating sync as **best‑effort** when a loop is already running: we **skip** the async prerequisite check rather than trying to reuse or manipulate the existing loop.
You can replace the entire `_sync_wrapper` definition with:
```python
@functools.wraps(func)
def _sync_wrapper(*args, **kwargs) -> Any:
"""
Synchronous wrapper for tools.
Sync usage is treated as best-effort: we avoid managing asyncio event
loops directly to prevent interference with existing runtimes.
- When there is no running event loop, we use `asyncio.run` to perform
prerequisite checks and then invoke the tool.
- When there is already a running loop (e.g. inside an async framework),
we skip the async prerequisite checks here; such cases should use the
async wrapper or middleware-level checks instead.
"""
import asyncio
import warnings
ctx = None
if args and hasattr(args[0], "get_state"):
ctx = args[0]
elif "ctx" in kwargs:
ctx = kwargs["ctx"]
# If there is already a running loop, we do *not* try to manage it here.
try:
running_loop = asyncio.get_running_loop()
except RuntimeError:
running_loop = None
if running_loop is not None:
# Best-effort mode: do not attempt to run async code from sync via
# run_until_complete on an already-running loop (which would error).
# Rely on the async path / middleware to enforce prerequisites.
warnings.warn(
"Tool called via synchronous wrapper while an event loop is already "
"running; skipping async prerequisite checks. Prefer using the async "
"wrapper in this context.",
RuntimeWarning,
stacklevel=2,
)
return func(*args, **kwargs)
# No running loop: we can safely perform async prerequisites via asyncio.run
async def _run_with_prereq() -> Any:
# This assumes you have an async prerequisite checker already used by
# the async wrapper. Replace the following call with the actual one,
# or no-op if checks are purely synchronous.
if ctx is not None:
# e.g. await _check_prerequisites(ctx, tool_name=func.__name__)
await _check_prerequisites(ctx, tool_name=func.__name__) # type: ignore[name-defined]
return func(*args, **kwargs)
return asyncio.run(_run_with_prereq())
```
You will need to:
1. Ensure `asyncio` and `warnings` are imported at the top of the file if they are not already:
```python
import asyncio
import warnings
```
2. Replace `await _check_prerequisites(...)` with the actual async prerequisite check your async wrapper uses. If your current `_sync_wrapper` duplicated that logic inline, move that logic into a shared async helper (e.g., `_check_prerequisites(ctx, tool_name)`).
3. Remove any previous logic in `_sync_wrapper` that:
- Calls `asyncio.get_event_loop()` / `asyncio.new_event_loop()`
- Calls `loop.run_until_complete(...)`
- Creates loops that are not later closed
After these changes, the sync wrapper will no longer create or partially manage event loops itself, avoiding loop leaks and conflicts, while making it explicit that full prerequisite enforcement is guaranteed only on the async path.
</issue_to_address>
### Comment 3
<location path="Server/tests/unit/test_tool_filter_decorator.py" line_range="283" />
<code_context>
+ del tool_prerequisites[tool_name]
+
+
+class TestConcurrentAccess:
+ """Test thread-safe concurrent access to tool_prerequisites dictionary."""
+
</code_context>
<issue_to_address>
**suggestion (testing):** The concurrent access test doesn’t really exercise the lock-based behavior and may give a false sense of thread-safety.
In `TestConcurrentAccess`, `register_tools` patches `core.tool_filter_decorator._prerequisites_lock`, so the real lock is never used and `read_tools` just iterates the dict without any synchronization. This only verifies that concurrent dict access doesn’t crash, not that your locking protocol in `tool_filter_decorator` is correct.
To actually validate the locking behavior, you could:
- Use the real `_prerequisites_lock`, and
- Exercise the same registration path used in production (e.g., via the decorator) instead of mutating the dict directly.
If you intend only a smoke test, consider renaming or documenting the test accordingly so it doesn’t imply strong thread-safety guarantees.
Suggested implementation:
```python
class TestConcurrentAccessSmoke:
"""Smoke test for concurrent access to tool_prerequisites.
NOTE: This test intentionally does not validate the full locking protocol
in core.tool_filter_decorator; it only asserts that concurrent access does
not crash. Stronger thread-safety guarantees are covered elsewhere.
"""
```
1. If there are any references to `TestConcurrentAccess` (e.g., in comments or parametrized tests), rename them to `TestConcurrentAccessSmoke` to keep naming consistent.
2. If the test methods within this class also imply strong guarantees in their docstrings or names (e.g., `test_thread_safe_registration`), consider updating those docstrings/method names to reflect that this is a smoke test and not a full concurrency-proof of the locking protocol.
</issue_to_address>Help me be more useful! Please click 👍 or 👎 on each comment and I'll use the feedback to improve your reviews.
| private void ReregisterToolsAsync() | ||
| { | ||
| // Fire and forget - don't block UI | ||
| ThreadPool.QueueUserWorkItem(_ => | ||
| { | ||
| try | ||
| { | ||
| var transportManager = MCPServiceLocator.TransportManager; | ||
| var client = transportManager.GetClient(TransportMode.Http); | ||
| if (client != null && client.IsConnected) |
There was a problem hiding this comment.
suggestion (bug_risk): Avoid blocking on async ReregisterToolsAsync() using .Wait() inside ThreadPool work item.
Since this runs on a ThreadPool thread, synchronously waiting on ReregisterToolsAsync() can still cause deadlocks or thread pool starvation if the async method ever captures a context or awaits other blocking work. Instead, keep this path fully async by using an async delegate, e.g. ThreadPool.QueueUserWorkItem(async _ => { ... await client.ReregisterToolsAsync().ConfigureAwait(false); }); so you never block a worker thread.
Suggested implementation:
private void ReregisterToolsAsync()
{
// Fire and forget - don't block UI
ThreadPool.QueueUserWorkItem(async _ =>
{ var transportManager = MCPServiceLocator.TransportManager;
var client = transportManager.GetClient(TransportMode.Http);
if (client != null && client.IsConnected)
{
await client.ReregisterToolsAsync().ConfigureAwait(false);
}| @functools.wraps(func) | ||
| def _sync_wrapper(*args, **kwargs) -> Any: | ||
| # Check prerequisites before executing | ||
| # Note: For direct tool calls, we check here | ||
| # For tool list filtering, the middleware handles it | ||
| ctx = None | ||
| if args and hasattr(args[0], "get_state"): | ||
| ctx = args[0] | ||
| elif "ctx" in kwargs: | ||
| ctx = kwargs["ctx"] |
There was a problem hiding this comment.
suggestion (bug_risk): The sync wrapper’s event loop management is complex and may leak loops or conflict with existing ones.
In _sync_wrapper, the logic that obtains/creates an event loop and calls run_until_complete is fragile: when a loop is already running it just warns and skips the check, and when a new loop is created it’s never closed. If sync usage is required, either use a dedicated asyncio.run path (to get proper loop setup/teardown) or explicitly treat sync as a best-effort mode (e.g., skip the prereq check) rather than partially managing event loops in a way that may vary across runtimes.
Suggested implementation:
# Check prerequisites before executing.
# NOTE: sync invocation is treated as a *best-effort* mode.
# For direct tool calls from async code, prefer the async wrapper.
# For tool list filtering, the middleware handles prerequisite checks.The core of the requested change is in _sync_wrapper’s event-loop handling, which is below the snippet you shared. Because I can’t see the full function, you’ll need to replace its body directly. Here is a concrete implementation that addresses the comment by:
- Completely removing ad‑hoc loop creation /
run_until_complete. - Using
asyncio.runwhen there is no running loop (proper setup/teardown). - Treating sync as best‑effort when a loop is already running: we skip the async prerequisite check rather than trying to reuse or manipulate the existing loop.
You can replace the entire _sync_wrapper definition with:
@functools.wraps(func)
def _sync_wrapper(*args, **kwargs) -> Any:
"""
Synchronous wrapper for tools.
Sync usage is treated as best-effort: we avoid managing asyncio event
loops directly to prevent interference with existing runtimes.
- When there is no running event loop, we use `asyncio.run` to perform
prerequisite checks and then invoke the tool.
- When there is already a running loop (e.g. inside an async framework),
we skip the async prerequisite checks here; such cases should use the
async wrapper or middleware-level checks instead.
"""
import asyncio
import warnings
ctx = None
if args and hasattr(args[0], "get_state"):
ctx = args[0]
elif "ctx" in kwargs:
ctx = kwargs["ctx"]
# If there is already a running loop, we do *not* try to manage it here.
try:
running_loop = asyncio.get_running_loop()
except RuntimeError:
running_loop = None
if running_loop is not None:
# Best-effort mode: do not attempt to run async code from sync via
# run_until_complete on an already-running loop (which would error).
# Rely on the async path / middleware to enforce prerequisites.
warnings.warn(
"Tool called via synchronous wrapper while an event loop is already "
"running; skipping async prerequisite checks. Prefer using the async "
"wrapper in this context.",
RuntimeWarning,
stacklevel=2,
)
return func(*args, **kwargs)
# No running loop: we can safely perform async prerequisites via asyncio.run
async def _run_with_prereq() -> Any:
# This assumes you have an async prerequisite checker already used by
# the async wrapper. Replace the following call with the actual one,
# or no-op if checks are purely synchronous.
if ctx is not None:
# e.g. await _check_prerequisites(ctx, tool_name=func.__name__)
await _check_prerequisites(ctx, tool_name=func.__name__) # type: ignore[name-defined]
return func(*args, **kwargs)
return asyncio.run(_run_with_prereq())You will need to:
-
Ensure
asyncioandwarningsare imported at the top of the file if they are not already:import asyncio import warnings
-
Replace
await _check_prerequisites(...)with the actual async prerequisite check your async wrapper uses. If your current_sync_wrapperduplicated that logic inline, move that logic into a shared async helper (e.g.,_check_prerequisites(ctx, tool_name)). -
Remove any previous logic in
_sync_wrapperthat:- Calls
asyncio.get_event_loop()/asyncio.new_event_loop() - Calls
loop.run_until_complete(...) - Creates loops that are not later closed
- Calls
After these changes, the sync wrapper will no longer create or partially manage event loops itself, avoiding loop leaks and conflicts, while making it explicit that full prerequisite enforcement is guaranteed only on the async path.
| del tool_prerequisites[tool_name] | ||
|
|
||
|
|
||
| class TestConcurrentAccess: |
There was a problem hiding this comment.
suggestion (testing): The concurrent access test doesn’t really exercise the lock-based behavior and may give a false sense of thread-safety.
In TestConcurrentAccess, register_tools patches core.tool_filter_decorator._prerequisites_lock, so the real lock is never used and read_tools just iterates the dict without any synchronization. This only verifies that concurrent dict access doesn’t crash, not that your locking protocol in tool_filter_decorator is correct.
To actually validate the locking behavior, you could:
- Use the real
_prerequisites_lock, and - Exercise the same registration path used in production (e.g., via the decorator) instead of mutating the dict directly.
If you intend only a smoke test, consider renaming or documenting the test accordingly so it doesn’t imply strong thread-safety guarantees.
Suggested implementation:
class TestConcurrentAccessSmoke:
"""Smoke test for concurrent access to tool_prerequisites.
NOTE: This test intentionally does not validate the full locking protocol
in core.tool_filter_decorator; it only asserts that concurrent access does
not crash. Stronger thread-safety guarantees are covered elsewhere.
"""- If there are any references to
TestConcurrentAccess(e.g., in comments or parametrized tests), rename them toTestConcurrentAccessSmoketo keep naming consistent. - If the test methods within this class also imply strong guarantees in their docstrings or names (e.g.,
test_thread_safe_registration), consider updating those docstrings/method names to reflect that this is a smoke test and not a full concurrency-proof of the locking protocol.
There was a problem hiding this comment.
Pull request overview
This PR implements "State-Aware Tool Filtering" (#1), which filters available MCP tools based on the Unity Editor's current state (compilation, play mode, selection). The filtering prevents tools from appearing or executing when their prerequisites aren't met.
Changes:
- New
@prerequisite_checkdecorator andToolPrerequisiteclass for declaring per-tool prerequisites (no-compile, selection, play-mode, no-tests) EditorStateCacheextended to track and expose selection state and anadviceblock (blocking reasons, ready_for_tools) in the snapshotIMcpTransportClientinterface extended withReregisterToolsAsync()to push updated tool lists to the server after user toggle changes
Reviewed changes
Copilot reviewed 24 out of 25 changed files in this pull request and generated 9 comments.
Show a summary per file
| File | Description |
|---|---|
Server/src/core/tool_filter_decorator.py |
Core prerequisite decorator and registry |
Server/src/services/filter_middleware.py |
Tool list filtering middleware (not yet wired in) |
Server/src/services/resources/editor_state.py |
Adds EditorStateSelection model to editor state |
Server/src/services/tools/manage_scene.py |
Adds require_no_compile prerequisite |
Server/src/services/tools/manage_gameobject.py |
Adds require_paused_for_destructive prerequisite |
Server/src/services/tools/manage_components.py |
Adds require_selection prerequisite |
Server/src/services/tools/manage_asset.py |
Adds require_no_compile prerequisite |
Server/src/services/tools/manage_script.py |
Adds import for prerequisite_check (unused) |
Server/src/services/tools/__init__.py |
Re-exports prerequisite_check and tool_prerequisites |
MCPForUnity/Editor/Services/EditorStateCache.cs |
Adds selection tracking and advice block to snapshot; made public |
MCPForUnity/Editor/Services/Transport/IMcpTransportClient.cs |
Adds ReregisterToolsAsync() to interface |
MCPForUnity/Editor/Services/Transport/TransportManager.cs |
Makes GetClient public |
MCPForUnity/Editor/Services/Transport/Transports/WebSocketTransportClient.cs |
Implements ReregisterToolsAsync |
MCPForUnity/Editor/Services/Transport/Transports/StdioTransportClient.cs |
No-op ReregisterToolsAsync for stdio |
MCPForUnity/Editor/Windows/Components/Tools/McpToolsSection.cs |
Triggers reregistration on tool toggle |
MCPForUnity/Editor/Services/ToolDiscoveryService.cs |
Removes self-healing force-enable logic for built-in tools |
Server/tests/unit/test_tool_filter_decorator.py |
Unit tests for ToolPrerequisite.is_met() |
Server/tests/integration/test_tool_filtering_integration.py |
Integration tests for filtering flow |
TestProjects/UnityMCPTests/Assets/Tests/Editor/State/EditorStateSelectionTests.cs |
Unity editor tests for snapshot selection/advice |
TestProjects/UnityMCPTests/Assets/Tests/EditMode/Services/ToolDiscoveryServiceTests.cs |
Unity edit-mode tests for ToolDiscoveryService |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
| @mcp_for_unity_tool( | ||
| description="Manages components on GameObjects (add, remove, set_property). For reading component data, use the mcpforunity://scene/gameobject/{id}/components resource." | ||
| ) | ||
| @prerequisite_check(require_selection=True) |
There was a problem hiding this comment.
The manage_components tool accepts an explicit target parameter (GameObject by name/path/ID), meaning the user can operate on any GameObject regardless of what is currently selected in the Editor. Applying @prerequisite_check(require_selection=True) will block this tool whenever nothing is selected in the Unity Editor, even though the tool doesn't actually require a selection — it resolves the target via its target parameter. This would break a valid workflow like: "add a Rigidbody to the 'Enemy' GameObject" when the user has nothing selected. Consider removing this decorator or replacing it with a more appropriate constraint (e.g. require_no_compile=True).
| @prerequisite_check(require_selection=True) | |
| @prerequisite_check(require_no_compile=True) |
|
|
||
| from services.registry import mcp_for_unity_tool | ||
| from services.tools import get_unity_instance_from_context | ||
| from core.tool_filter_decorator import prerequisite_check |
There was a problem hiding this comment.
manage_script.py imports prerequisite_check but never uses the @prerequisite_check decorator anywhere in the file. This is an unused import that adds confusion — it suggests a decorator was intended to be applied to one or more script tools but was never added. Either apply the decorator to the appropriate functions (e.g. script-modifying tools should likely have require_no_compile=True) or remove the unused import.
| from core.tool_filter_decorator import prerequisite_check |
|
|
||
| // Trigger tool reregistration after bulk change | ||
| ReregisterToolsAsync(); |
There was a problem hiding this comment.
In SetAllToolsState, HandleToggleChange is called for each changed tool (line 278), and HandleToggleChange itself calls ReregisterToolsAsync() (line 238). Then, after the loop, SetAllToolsState calls ReregisterToolsAsync() again on line 284. This means a bulk toggle of N tools will fire N+1 reregistration requests — one per tool plus the final explicit call. The per-tool calls during the bulk operation are wasteful and redundant; only the final one is necessary. The HandleToggleChange calls from SetAllToolsState should either skip reregistration (e.g. via a parameter like triggerReregister = false) or the extra ReregisterToolsAsync() at line 284 should be removed.
| // Trigger tool reregistration after bulk change | |
| ReregisterToolsAsync(); |
| def test_concurrent_read_with_registration(self): | ||
| """Concurrent reads during registration should be thread-safe.""" | ||
| import threading | ||
|
|
||
| results = [] | ||
| errors = [] | ||
|
|
||
| def register_tools(): | ||
| try: | ||
| for i in range(5): | ||
| tool_name = f"concurrent_tool_{i}" | ||
| with patch("core.tool_filter_decorator._prerequisites_lock"): | ||
| tool_prerequisites[tool_name] = ToolPrerequisite( | ||
| require_selection=(i % 2 == 0) | ||
| ) | ||
| except Exception as e: | ||
| errors.append(e) | ||
|
|
||
| def read_tools(): | ||
| try: | ||
| for _ in range(10): | ||
| # Simulate reads | ||
| _ = list(tool_prerequisites.keys()) | ||
| except Exception as e: | ||
| errors.append(e) | ||
|
|
||
| # Start threads | ||
| threads = [ | ||
| threading.Thread(target=register_tools), | ||
| threading.Thread(target=read_tools), | ||
| threading.Thread(target=read_tools), | ||
| ] | ||
|
|
||
| for t in threads: | ||
| t.start() | ||
| for t in threads: | ||
| t.join() | ||
|
|
||
| # No errors should occur | ||
| assert len(errors) == 0 | ||
|
|
||
| # Clean up | ||
| with patch("core.tool_filter_decorator._prerequisites_lock"): | ||
| for i in range(5): | ||
| tool_prerequisites.pop(f"concurrent_tool_{i}", None) |
There was a problem hiding this comment.
The test_concurrent_read_with_registration test uses with patch("core.tool_filter_decorator._prerequisites_lock"): to mock the lock. patch replaces the lock with a MagicMock object, which means the with statement is a no-op (the lock is never actually acquired). As a result, the test runs the writer thread entirely without any lock, defeating its purpose of verifying thread safety. This also leaves a false impression that the code has been verified for concurrency. The test should either acquire the real lock (with _prerequisites_lock:) or use threading.Lock() properly. The same issue appears in the cleanup at lines 328-330.
| // Arrange - Wait for any pending compilation to finish | ||
| while (EditorApplication.isCompiling) | ||
| { | ||
| System.Threading.Thread.Sleep(100); |
There was a problem hiding this comment.
GetSnapshot_CompilationState_WhenNotCompiling uses System.Threading.Thread.Sleep(100) in a spin-loop to wait for compilation to finish. Unity Editor tests run on the main thread, and calling Thread.Sleep() on the main thread can block Unity's editor loop, preventing callbacks (like EditorApplication.update) from firing. This means Unity's internal compilation state may never progress while the test is sleeping. Consider using [UnityTest] with yield return null frames to yield control back to the editor, or simply skip the test if EditorApplication.isCompiling is true at start.
| // Arrange - Wait for any pending compilation to finish | |
| while (EditorApplication.isCompiling) | |
| { | |
| System.Threading.Thread.Sleep(100); | |
| // Skip the test if Unity is currently compiling to avoid blocking the editor loop. | |
| if (EditorApplication.isCompiling) | |
| { | |
| Assert.Ignore("Skipping GetSnapshot_CompilationState_WhenNotCompiling because Unity is currently compiling."); |
| Integration tests for the full decorator flow are in tests/integration/test_tool_filtering_integration.py | ||
| """ | ||
|
|
||
| import pytest |
There was a problem hiding this comment.
import pytest on line 7 is unused — no pytest.mark or other pytest APIs are referenced in this file. This should be removed to keep imports clean.
| import pytest |
| public void IsToolEnabled_ReturnsTrue_WhenNoPreferenceExistsAndToolIsBuiltIn() | ||
| { | ||
| // Arrange - Ensure no preference exists | ||
| string key = EditorPrefKeys.ToolEnabledPrefix + TestToolName; | ||
| if (EditorPrefs.HasKey(key)) | ||
| { | ||
| EditorPrefs.DeleteKey(key); | ||
| } | ||
|
|
||
| var service = new ToolDiscoveryService(); | ||
|
|
||
| // Act - For a non-existent tool, IsToolEnabled should return false | ||
| // (since metadata.AutoRegister defaults to false for non-existent tools) | ||
| bool result = service.IsToolEnabled(TestToolName); | ||
|
|
||
| // Assert - Non-existent tools return false (no metadata found) | ||
| Assert.IsFalse(result, "Non-existent tool should return false"); |
There was a problem hiding this comment.
The test method name IsToolEnabled_ReturnsTrue_WhenNoPreferenceExistsAndToolIsBuiltIn is misleading. The test body actually asserts IsFalse (not True), and the tool used (TestToolName = "test_tool_for_testing") is not a built-in tool — it's a non-existent test tool. The test name should reflect what it actually tests, e.g. IsToolEnabled_ReturnsFalse_WhenNoPreferenceExistsAndToolIsNotBuiltIn.
| # Check domain reload prerequisite | ||
| if self.require_no_compile: | ||
| if "domain_reload" in blocking_reasons: | ||
| return False, "domain_reload" |
There was a problem hiding this comment.
The is_met method contains a duplicate if self.require_no_compile: guard on line 64 (which should logically be the same block as line 59). Both the "compiling" and "domain_reload" checks guard with self.require_no_compile. While the behavior is currently correct (they still run sequentially), this creates a misleading code structure. The intent of the second block's comment "Check domain reload prerequisite" suggests it was meant to be a separate check, but there is no separate require_no_domain_reload flag. These two checks should be combined into a single if self.require_no_compile: block to avoid confusion and reduce the chance that a future change (e.g. renaming the flag or adding a new one) introduces a bug.
| var client = transportManager.GetClient(TransportMode.Http); | ||
| if (client != null && client.IsConnected) | ||
| { | ||
| client.ReregisterToolsAsync().Wait(); |
There was a problem hiding this comment.
Using .Wait() on an async Task inside a ThreadPool worker (line 252) can cause deadlocks if the async code uses a synchronization context that needs to post continuations back to the calling thread. While ThreadPool threads don't have a synchronization context attached (avoiding the classic Task.Wait() deadlock pattern), ReregisterToolsAsync ultimately calls GetEnabledToolsOnMainThreadAsync (in WebSocketTransportClient.SendRegisterToolsAsync), which requires dispatching work to Unity's main thread. If the main thread is busy, Wait() will block the ThreadPool thread indefinitely. Consider using ReregisterToolsAsync().ContinueWith(t => { if (t.IsFaulted) McpLog.Warn(...) }) or, better, making ReregisterToolsAsync a fire-and-forget at the call site.
| client.ReregisterToolsAsync().Wait(); | |
| client.ReregisterToolsAsync().ContinueWith(t => | |
| { | |
| if (t.IsFaulted) | |
| { | |
| var ex = t.Exception?.GetBaseException(); | |
| McpLog.Warn($"Failed to reregister tools: {ex?.Message ?? "Unknown error"}"); | |
| } | |
| }); |
Description
Type of Change
Save your change type
Changes Made
Testing/Screenshots/Recordings
Documentation Updates
tools/UPDATE_DOCS_PROMPT.md(recommended)tools/UPDATE_DOCS.mdRelated Issues
Additional Notes
Summary by Sourcery
Introduce state-aware tool filtering across Unity editor and MCP server, exposing richer editor state (selection and advice) and using it to control tool availability and registration.
New Features:
Enhancements:
Tests:
Summary by CodeRabbit
New Features
Tests