Skip to content

001 state aware tool filtering#864

Closed
whatevertogo wants to merge 7 commits intoCoplayDev:betafrom
whatevertogo:001-state-aware-tool-filtering
Closed

001 state aware tool filtering#864
whatevertogo wants to merge 7 commits intoCoplayDev:betafrom
whatevertogo:001-state-aware-tool-filtering

Conversation

@whatevertogo
Copy link
Contributor

@whatevertogo whatevertogo commented Mar 4, 2026

Description

Type of Change

Save your change type

  • Bug fix (non-breaking change that fixes an issue)
  • New feature (non-breaking change that adds functionality)
  • Breaking change (fix or feature that would cause existing functionality to change)
  • Documentation update
  • Refactoring (no functional changes)
  • Test update

Changes Made

Testing/Screenshots/Recordings

Documentation Updates

  • I have added/removed/modified tools or resources
  • If yes, I have updated all documentation files using:
    • The LLM prompt at tools/UPDATE_DOCS_PROMPT.md (recommended)
    • Manual updates following the guide at tools/UPDATE_DOCS.md

Related 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:

  • Expose Unity editor selection and advisory readiness data in the shared editor state schema for use by tools and server-side logic.
  • Add prerequisite-based tool filtering decorator and middleware on the server to hide tools based on compilation, selection, play mode, and test-running state.
  • Allow HTTP/WebSocket clients to trigger dynamic tool re-registration when tool enablement changes in the Unity editor.

Enhancements:

  • Make EditorStateCache public and extend transport metadata with tool-readiness advice derived from compilation, domain reload, and test state.
  • Simplify tool preference initialization logic so built-in and discovered tools rely on explicit enablement flags.
  • Expose transport clients from TransportManager for reuse and add a no-op re-registration hook for stdio transports.

Tests:

  • Add Unity edit-mode tests validating editor selection, compilation, play mode, and advice fields in the editor state snapshot.
  • Add Unity tests ensuring tool enablement is persisted via EditorPrefs and consistent across ToolDiscoveryService instances.
  • Add Python unit and integration tests covering tool prerequisite evaluation, async decorator behavior, and end-to-end state-based tool filtering.

Summary by CodeRabbit

  • New Features

    • Editor state now tracks selection information, including active GameObject and selection count.
    • Tools can now define prerequisites (e.g., compilation state, selection requirement, play mode restrictions) and are automatically shown/hidden based on editor readiness.
    • Added advisory information to editor state snapshots indicating tool availability and blocking reasons.
    • Reduced update throttling for faster editor state responsiveness.
  • Tests

    • Comprehensive test coverage for tool filtering and prerequisite logic.
    • Added editor state selection tracking tests.

whatevertogo and others added 7 commits February 5, 2026 00:33
- 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.
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>
Copilot AI review requested due to automatic review settings March 4, 2026 13:31
@sourcery-ai
Copy link
Contributor

sourcery-ai bot commented Mar 4, 2026

Reviewer's Guide

Implements 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 toggle

sequenceDiagram
    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
Loading

Flow diagram for tool list filtering based on editor state

flowchart 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"]
Loading

File-Level Changes

Change Details Files
Expose and extend EditorStateCache to track selection and provide tool-usage advice in the editor state snapshot.
  • Changed the editor state cache class to be publicly accessible and added selection tracking fields for active instance ID, selected object name, and selection count.
  • Hooked into Unity’s Selection.selectionChanged callback and update loop to detect selection changes and force state cache refreshes only when relevant state changes.
  • Extended the snapshot schema with editor.selection and advice fields, including ready_for_tools, blocking_reasons, and compile/test-related recommendations, and added a helper to build the advice object.
  • Updated GetSnapshot to optionally force a refresh and to rebuild the cached snapshot when needed, ensuring each call returns a clone.
MCPForUnity/Editor/Services/EditorStateCache.cs
Server/src/services/resources/editor_state.py
TestProjects/UnityMCPTests/Assets/Tests/Editor/State/EditorStateSelectionTests.cs
Add client transport support for dynamic tool re‑registration when tool enablement changes.
  • Exposed a public GetClient method on the transport manager to retrieve existing transport clients without creating them.
  • Implemented a ReregisterToolsAsync API on IMcpTransportClient and provided WebSocket and Stdio implementations, with WebSocket reusing the existing tool registration path and Stdio as a no-op.
  • Invoked asynchronous tool re-registration from the tools UI when a single tool’s toggle changes or when tools are bulk enabled/disabled, using a background thread to avoid blocking the UI and logging failures.
MCPForUnity/Editor/Windows/Components/Tools/McpToolsSection.cs
MCPForUnity/Editor/Services/Transport/TransportManager.cs
MCPForUnity/Editor/Services/Transport/Transports/WebSocketTransportClient.cs
MCPForUnity/Editor/Services/Transport/Transports/StdioTransportClient.cs
MCPForUnity/Editor/Services/Transport/IMcpTransportClient.cs
Simplify tool enablement preference initialization and add tests around ToolDiscoveryService behavior.
  • Removed special-case logic that forced built-in tools to true when their AutoRegister flag was false, leaving initial values to come purely from metadata or existing prefs.
  • Added edit-mode tests that verify SetToolEnabled writes to EditorPrefs and that enabled state is persisted and read consistently across ToolDiscoveryService instances.
  • Ensured test setup/teardown cleans up EditorPrefs for the test tool key.
MCPForUnity/Editor/Services/ToolDiscoveryService.cs
TestProjects/UnityMCPTests/Assets/Tests/EditMode/Services/ToolDiscoveryServiceTests.cs
TestProjects/UnityMCPTests/Assets/Tests/EditMode/Services/ToolDiscoveryServiceTests.cs.meta
Introduce server-side tool prerequisite/decorator framework and middleware for state-aware tool filtering.
  • Added a ToolPrerequisite model and prerequisite_check decorator that register per-tool prerequisite flags (compile, selection, play mode, tests) into a global registry and enforce them at execution time by querying editor_state and returning MCPResponse errors when unmet.
  • Implemented is_met logic to interpret editor_state.advice, editor.selection, play_mode, and tests sections, including fail-open behavior when state is missing or unknown.
  • Created a filter_middleware helper that uses the global prerequisite registry and current editor state to filter the list of tools that should be exposed to the LLM, with fail-safe behavior on errors.
  • Introduced a FilterResult helper class (currently used in tests) for structured reporting of filtering decisions.
Server/src/core/tool_filter_decorator.py
Server/src/services/filter_middleware.py
Apply prerequisite decorators to specific tools so their availability respects compilation, selection, and play mode state.
  • Marked asset and scene management tools to require_no_compile so they are hidden/blocked while Unity is compiling.
  • Marked manage_components to require_selection so it requires an active selection.
  • Marked manage_gameobject to require_paused_for_destructive so destructive operations are only allowed when not actively playing or when play mode is paused.
Server/src/services/tools/manage_asset.py
Server/src/services/tools/manage_components.py
Server/src/services/tools/manage_gameobject.py
Server/src/services/tools/manage_scene.py
Add unit and integration test coverage for tool filtering logic and decorator behavior on the server side.
  • Created integration tests that exercise tool filtering end-to-end, mocking editor_state responses to validate hiding/showing tools based on compilation, selection, and play mode, plus fail-safe behavior when state is unavailable.
  • Verified that the async prerequisite_check wrapper behaves correctly for async tools, both when prerequisites are met and when they fail, and that it reuses the registered ToolPrerequisite instance.
  • Added unit tests for ToolPrerequisite.is_met with various combinations of editor states and requirements, as well as thread-safety and registration behavior of the global tool_prerequisites dictionary.
Server/tests/integration/test_tool_filtering_integration.py
Server/tests/unit/test_tool_filter_decorator.py

Tips and commands

Interacting with Sourcery

  • Trigger a new review: Comment @sourcery-ai review on the pull request.
  • Continue discussions: Reply directly to Sourcery's review comments.
  • Generate a GitHub issue from a review comment: Ask Sourcery to create an
    issue from a review comment by replying to it. You can also reply to a
    review comment with @sourcery-ai issue to create an issue from it.
  • Generate a pull request title: Write @sourcery-ai anywhere in the pull
    request title to generate a title at any time. You can also comment
    @sourcery-ai title on the pull request to (re-)generate the title at any time.
  • Generate a pull request summary: Write @sourcery-ai summary anywhere in
    the pull request body to generate a PR summary at any time exactly where you
    want it. You can also comment @sourcery-ai summary on the pull request to
    (re-)generate the summary at any time.
  • Generate reviewer's guide: Comment @sourcery-ai guide on the pull
    request to (re-)generate the reviewer's guide at any time.
  • Resolve all Sourcery comments: Comment @sourcery-ai resolve on the
    pull request to resolve all Sourcery comments. Useful if you've already
    addressed all the comments and don't want to see them anymore.
  • Dismiss all Sourcery reviews: Comment @sourcery-ai dismiss on the pull
    request to dismiss all existing Sourcery reviews. Especially useful if you
    want to start fresh with a new review - don't forget to comment
    @sourcery-ai review to trigger a new review!

Customizing Your Experience

Access your dashboard to:

  • Enable or disable review features such as the Sourcery-generated pull request
    summary, the reviewer's guide, and others.
  • Change the review language.
  • Add, remove or edit custom review instructions.
  • Adjust other review settings.

Getting Help

@whatevertogo
Copy link
Contributor Author

@codex review it

@coderabbitai
Copy link
Contributor

coderabbitai bot commented Mar 4, 2026

Caution

Review failed

Pull request was closed or merged during review

📝 Walkthrough

Walkthrough

This 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

Cohort / File(s) Summary
EditorStateCache & Selection State
MCPForUnity/Editor/Services/EditorStateCache.cs
Made EditorStateCache public; added selection state (HasSelection, ActiveInstanceID, ActiveGameObjectName, SelectionCount); introduced EditorStateAdvice with ReadyForTools, BlockingReasons, and RecommendedRetryAfterMs; expanded GetSnapshot with forceRefresh parameter; subscribed to Selection.selectionChanged; reduced MinUpdateIntervalSeconds to 1.0 second.
Transport & Reregistration
MCPForUnity/Editor/Services/Transport/IMcpTransportClient.cs, MCPForUnity/Editor/Services/Transport/TransportManager.cs, MCPForUnity/Editor/Services/Transport/Transports/StdioTransportClient.cs, MCPForUnity/Editor/Services/Transport/Transports/WebSocketTransportClient.cs, MCPForUnity/Editor/Windows/Components/Tools/McpToolsSection.cs
Added ReregisterToolsAsync method to IMcpTransportClient interface; made TransportManager.GetClient public with XML doc; implemented ReregisterToolsAsync in StdioTransportClient (no-op) and WebSocketTransportClient (performs re-registration with error handling); wired tool reregistration in McpToolsSection after individual toggles and bulk state updates.
ToolDiscoveryService
MCPForUnity/Editor/Services/ToolDiscoveryService.cs
Removed conditional post-initialization adjustment logic in EnsurePreferenceInitialized that forced values under specific conditions.
Server-Side Prerequisite Infrastructure
Server/src/core/tool_filter_decorator.py, Server/src/services/filter_middleware.py, Server/src/services/resources/editor_state.py, Server/src/services/tools/__init__.py
Introduced ToolPrerequisite class with flags for compile, selection, destructive, and test constraints; implemented prerequisite_check decorator (sync/async) with runtime checks; added filter_middleware with get_tools_matching_state for visibility filtering; expanded EditorStateEditor with EditorStateSelection (has_selection, active_instance_id, active_game_object_name, selection_count); exposed decorators in tools/init.py.
Tool Prerequisite Decorators
Server/src/services/tools/manage_asset.py, Server/src/services/tools/manage_components.py, Server/src/services/tools/manage_gameobject.py, Server/src/services/tools/manage_scene.py, Server/src/services/tools/manage_script.py
Applied @prerequisite_check decorators with appropriate flags: manage_asset (require_no_compile), manage_components (require_selection), manage_gameobject (require_paused_for_destructive), manage_scene (require_no_compile), manage_script (import only).
Test Suite - Unit Tests
Server/tests/unit/test_tool_filter_decorator.py
Comprehensive unit tests for ToolPrerequisite across individual and combined flags (compilation, selection, play mode, tests), edge cases (missing state sections, fail-open behavior), decorator registration, and concurrent access safety.
Test Suite - Integration Tests
Server/tests/integration/test_tool_filtering_integration.py
Integration test suites covering compilation filtering, selection filtering, play mode filtering, fail-safe behavior on state errors, FilterResult serialization, and async decorator wrapping with prerequisite evaluation.
Test Suite - C# Tests
TestProjects/UnityMCPTests/Assets/Tests/EditMode/Services/ToolDiscoveryServiceTests.cs
New test class validating ToolDiscoveryService SetToolEnabled, IsToolEnabled, and persistence across instances; includes setup/teardown for EditorPrefs cleanup.
Test Suite - Editor State Tests
TestProjects/UnityMCPTests/Assets/Tests/Editor/State/EditorStateSelectionTests.cs
Test suite for EditorStateCache snapshot behavior covering selection presence/absence, active instance ID, compilation state, blocking reasons, play mode, required fields, and snapshot cloning on each call.
Test Metadata
TestProjects/UnityMCPTests/Assets/Tests/EditMode/Services/ToolDiscoveryServiceTests.cs.meta, TestProjects/UnityMCPTests/Assets/Tests/Editor.meta, TestProjects/UnityMCPTests/Assets/Tests/Editor/State.meta, TestProjects/UnityMCPTests/Assets/Tests/Editor/State/EditorStateSelectionTests.cs.meta
Unity asset metadata files for test organization; no executable logic 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

codex

Suggested reviewers

  • dsarno
  • msanatan
  • justinpbarnett

Poem

🐰 Tools now know when they're ready to play,
Checking the state of the editor each day,
No compiles mid-way, selections must be,
Prerequisites gate what the user will see! 🔐

🚥 Pre-merge checks | ❌ 3

❌ Failed checks (2 warnings, 1 inconclusive)

Check name Status Explanation Resolution
Description check ⚠️ Warning The PR description is largely incomplete, retaining template placeholders without substantive content filled in for required sections. Complete the description by selecting a change type, listing specific changes made, providing testing details, confirming documentation updates, and linking any related issues.
Docstring Coverage ⚠️ Warning Docstring coverage is 38.95% which is insufficient. The required threshold is 80.00%. Write docstrings for the functions missing them to satisfy the coverage threshold.
Title check ❓ Inconclusive The title 'state aware tool filtering' is partially related to the changeset, referring to a real aspect of the change, but lacks specificity and clarity. Consider a more descriptive title that clearly indicates the scope, such as 'Add state-aware tool filtering for editor and server' or 'Implement prerequisite-based tool visibility control'.

✏️ Tip: You can configure your own custom pre-merge checks in the settings.

✨ Finishing Touches
🧪 Generate unit tests (beta)
  • Create PR with unit tests
  • Post copyable unit tests in a comment

Tip

Try Coding Plans. Let us write the prompt for your AI agent so you can ship faster (with fewer bugs).
Share your feedback on Discord.


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.

❤️ Share

Comment @coderabbitai help to get the list of available commands and usage tips.

@chatgpt-codex-connector
Copy link

To use Codex here, create a Codex account and connect to github.

@whatevertogo
Copy link
Contributor Author

@codex review it

@chatgpt-codex-connector
Copy link

To use Codex here, create a Codex account and connect to github.

Copy link
Contributor

@sourcery-ai sourcery-ai bot left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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 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.
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 wrappers 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 its 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 cant see the full function, youll need to replace its body directly. Here is a concrete implementation that addresses the comment by:

1. Completely removing adhoc loop creation / `run_until_complete`.
2. Using `asyncio.run` when there is **no** running loop (proper setup/teardown).
3. Treating sync as **besteffort** 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 doesnt 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 doesnt 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 doesnt 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>

Sourcery is free for open source - if you like our reviews please consider sharing them ✨
Help me be more useful! Please click 👍 or 👎 on each comment and I'll use the feedback to improve your reviews.

Comment on lines +241 to +250
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)
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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);
                    }

Comment on lines +138 to +147
@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"]
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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:

  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:

@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:

    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.

del tool_prerequisites[tool_name]


class TestConcurrentAccess:
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.
    """
  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.

Copy link
Contributor

Copilot AI left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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_check decorator and ToolPrerequisite class for declaring per-tool prerequisites (no-compile, selection, play-mode, no-tests)
  • EditorStateCache extended to track and expose selection state and an advice block (blocking reasons, ready_for_tools) in the snapshot
  • IMcpTransportClient interface extended with ReregisterToolsAsync() 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)
Copy link

Copilot AI Mar 4, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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).

Suggested change
@prerequisite_check(require_selection=True)
@prerequisite_check(require_no_compile=True)

Copilot uses AI. Check for mistakes.

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
Copy link

Copilot AI Mar 4, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Suggested change
from core.tool_filter_decorator import prerequisite_check

Copilot uses AI. Check for mistakes.
Comment on lines +282 to +284

// Trigger tool reregistration after bulk change
ReregisterToolsAsync();
Copy link

Copilot AI Mar 4, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Suggested change
// Trigger tool reregistration after bulk change
ReregisterToolsAsync();

Copilot uses AI. Check for mistakes.
Comment on lines +286 to +330
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)
Copy link

Copilot AI Mar 4, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Copilot uses AI. Check for mistakes.
Comment on lines +188 to +191
// Arrange - Wait for any pending compilation to finish
while (EditorApplication.isCompiling)
{
System.Threading.Thread.Sleep(100);
Copy link

Copilot AI Mar 4, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Suggested change
// 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.");

Copilot uses AI. Check for mistakes.
Integration tests for the full decorator flow are in tests/integration/test_tool_filtering_integration.py
"""

import pytest
Copy link

Copilot AI Mar 4, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Suggested change
import pytest

Copilot uses AI. Check for mistakes.
Comment on lines +51 to +67
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");
Copy link

Copilot AI Mar 4, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Copilot uses AI. Check for mistakes.
Comment on lines +63 to +66
# Check domain reload prerequisite
if self.require_no_compile:
if "domain_reload" in blocking_reasons:
return False, "domain_reload"
Copy link

Copilot AI Mar 4, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Copilot uses AI. Check for mistakes.
var client = transportManager.GetClient(TransportMode.Http);
if (client != null && client.IsConnected)
{
client.ReregisterToolsAsync().Wait();
Copy link

Copilot AI Mar 4, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Suggested change
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"}");
}
});

Copilot uses AI. Check for mistakes.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants