Skip to content

001 state aware tool filtering #864

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.

Sign up for GitHub

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

Closed
whatevertogo wants to merge 7 commits into from
Closed

001 state aware tool filtering #864

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
7 commits
Select commit Hold shift + click to select a range
db5e84c
Implement state-aware tool filtering middleware and associated tests
whatevertogo Feb 4, 2026
f403a46
feat: enhance tool filtering and state management with improved snaps…
whatevertogo Feb 4, 2026
254f530
fix: force refresh snapshot in selection state tests for accurate sta…
whatevertogo Feb 4, 2026
10d8d46
feat: implement state-aware tool filtering middleware and enhance edi…
whatevertogo Feb 5, 2026
78dc04c
Merge branch '001-state-aware-tool-filtering' of https://github.com/w…
whatevertogo Feb 5, 2026
4135bd2
Merge branch 'CoplayDev:beta' into 001-state-aware-tool-filtering
whatevertogo Feb 11, 2026
3341af3
fix: preserve tool enabled/disabled state across Unity restarts
whatevertogo Feb 11, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
109 changes: 103 additions & 6 deletions MCPForUnity/Editor/Services/EditorStateCache.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,7 @@
using UnityEditorInternal;
using UnityEditor.SceneManagement;
using UnityEngine;
using System.Collections.Generic;

namespace MCPForUnity.Editor.Services
{
Expand All @@ -15,7 +16,7 @@ namespace MCPForUnity.Editor.Services
/// Updated on the main thread via Editor callbacks and periodic update ticks.
/// </summary>
[InitializeOnLoad]
internal static class EditorStateCache
public static class EditorStateCache
{
private static readonly object LockObj = new();
private static long _sequence;
Expand All @@ -42,6 +43,11 @@ internal static class EditorStateCache
private static bool _lastTrackedTestsRunning;
private static string _lastTrackedActivityPhase;

// Selection state tracking for state-aware tool filtering
private static int _lastTrackedActiveInstanceID;
private static string _lastTrackedActiveGameObjectName;
private static int _lastTrackedSelectionCount;

private static JObject _cached;

private sealed class EditorStateSnapshot
Expand Down Expand Up @@ -75,6 +81,9 @@ private sealed class EditorStateSnapshot

[JsonProperty("transport")]
public EditorStateTransport Transport { get; set; }

[JsonProperty("advice")]
public EditorStateAdvice Advice { get; set; }
}

private sealed class EditorStateUnity
Expand Down Expand Up @@ -105,6 +114,24 @@ private sealed class EditorStateEditor

[JsonProperty("active_scene")]
public EditorStateActiveScene ActiveScene { get; set; }

[JsonProperty("selection")]
public EditorStateSelection Selection { get; set; }
}

private sealed class EditorStateSelection
{
[JsonProperty("has_selection")]
public bool HasSelection { get; set; }

[JsonProperty("active_instance_id")]
public int ActiveInstanceID { get; set; }

[JsonProperty("active_game_object_name")]
public string ActiveGameObjectName { get; set; }

[JsonProperty("selection_count")]
public int SelectionCount { get; set; }
}

private sealed class EditorStatePlayMode
Expand Down Expand Up @@ -230,6 +257,21 @@ private sealed class EditorStateLastRun
public object Counts { get; set; }
}

private sealed class EditorStateAdvice
{
[JsonProperty("ready_for_tools")]
public bool ReadyForTools { get; set; }

[JsonProperty("blocking_reasons")]
public string[] BlockingReasons { get; set; }

[JsonProperty("recommended_retry_after_ms")]
public long? RecommendedRetryAfterMs { get; set; }

[JsonProperty("recommended_next_action")]
public string RecommendedNextAction { get; set; }
}

private sealed class EditorStateTransport
{
[JsonProperty("unity_bridge_connected")]
Expand All @@ -249,6 +291,7 @@ static EditorStateCache()

EditorApplication.update += OnUpdate;
EditorApplication.playModeStateChanged += _ => ForceUpdate("playmode");
Selection.selectionChanged += () => ForceUpdate("selection");

AssemblyReloadEvents.beforeAssemblyReload += () =>
{
Expand Down Expand Up @@ -296,6 +339,11 @@ private static void OnUpdate()
bool isUpdating = EditorApplication.isUpdating;
bool testsRunning = TestRunStatus.IsRunning;

// Selection state reading for state-aware tool filtering
int activeInstanceID = Selection.activeInstanceID;
string activeGameObjectName = Selection.activeGameObject?.name ?? string.Empty;
int selectionCount = Selection.count;

var activityPhase = "idle";
if (testsRunning)
{
Expand Down Expand Up @@ -326,7 +374,10 @@ private static void OnUpdate()
|| _lastTrackedIsPaused != isPaused
|| _lastTrackedIsUpdating != isUpdating
|| _lastTrackedTestsRunning != testsRunning
|| _lastTrackedActivityPhase != activityPhase;
|| _lastTrackedActivityPhase != activityPhase
|| _lastTrackedActiveInstanceID != activeInstanceID
|| _lastTrackedActiveGameObjectName != activeGameObjectName
|| _lastTrackedSelectionCount != selectionCount;

if (!hasChanges)
{
Expand All @@ -344,6 +395,9 @@ private static void OnUpdate()
_lastTrackedIsUpdating = isUpdating;
_lastTrackedTestsRunning = testsRunning;
_lastTrackedActivityPhase = activityPhase;
_lastTrackedActiveInstanceID = activeInstanceID;
_lastTrackedActiveGameObjectName = activeGameObjectName;
_lastTrackedSelectionCount = selectionCount;

_lastUpdateTimeSinceStartup = now;
ForceUpdate("tick");
Expand Down Expand Up @@ -404,6 +458,11 @@ private static JObject BuildSnapshot(string reason)
activityPhase = "playmode_transition";
}

// Read current selection state directly for snapshot
int currentActiveInstanceID = Selection.activeInstanceID;
string currentActiveGameObjectName = Selection.activeGameObject?.name ?? string.Empty;
int currentSelectionCount = Selection.count;

var snapshot = new EditorStateSnapshot
{
SchemaVersion = "unity-mcp/editor_state@2",
Expand Down Expand Up @@ -431,6 +490,13 @@ private static JObject BuildSnapshot(string reason)
Path = scenePath,
Guid = sceneGuid,
Name = scene.name ?? string.Empty
},
Selection = new EditorStateSelection
{
HasSelection = currentSelectionCount > 0,
ActiveInstanceID = currentActiveInstanceID,
ActiveGameObjectName = currentActiveGameObjectName,
SelectionCount = currentSelectionCount
}
},
Activity = new EditorStateActivity
Expand Down Expand Up @@ -482,20 +548,51 @@ private static JObject BuildSnapshot(string reason)
{
UnityBridgeConnected = null,
LastMessageUnixMs = null
}
},
Advice = BuildEditorStateAdvice(isCompiling, testsRunning)
};

return JObject.FromObject(snapshot);
}

public static JObject GetSnapshot()
private static EditorStateAdvice BuildEditorStateAdvice(bool isCompiling, bool testsRunning)
{
var blockingReasons = new List<string>();

if (isCompiling)
{
blockingReasons.Add("compiling");
}

if (_domainReloadPending)
{
blockingReasons.Add("domain_reload");
}

if (testsRunning)
{
blockingReasons.Add("tests_running");
}

bool readyForTools = blockingReasons.Count == 0;

return new EditorStateAdvice
{
ReadyForTools = readyForTools,
BlockingReasons = blockingReasons.ToArray(),
RecommendedRetryAfterMs = isCompiling ? 1000 : null,
RecommendedNextAction = isCompiling ? "wait_for_compile" : null
};
}

public static JObject GetSnapshot(bool forceRefresh = false)
{
lock (LockObj)
{
// Defensive: if something went wrong early, rebuild once.
if (_cached == null)
if (_cached == null || forceRefresh)
{
_cached = BuildSnapshot("rebuild");
_cached = BuildSnapshot(forceRefresh ? "get_snapshot_force" : "get_snapshot_rebuild");
}

// Always return a fresh clone to prevent mutation bugs.
Expand Down
10 changes: 0 additions & 10 deletions MCPForUnity/Editor/Services/ToolDiscoveryService.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -226,16 +226,6 @@ private void EnsurePreferenceInitialized(ToolMetadata metadata)
{
bool defaultValue = metadata.AutoRegister || metadata.IsBuiltIn;
EditorPrefs.SetBool(key, defaultValue);
return;
}

if (metadata.IsBuiltIn && !metadata.AutoRegister)
{
bool currentValue = EditorPrefs.GetBool(key, metadata.AutoRegister);
if (currentValue == metadata.AutoRegister)
{
EditorPrefs.SetBool(key, true);
}
}
}

Expand Down
1 change: 1 addition & 0 deletions MCPForUnity/Editor/Services/Transport/IMcpTransportClient.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -14,5 +14,6 @@ public interface IMcpTransportClient
Task<bool> StartAsync();
Task StopAsync();
Task<bool> VerifyAsync();
Task ReregisterToolsAsync();
}
}
24 changes: 14 additions & 10 deletions MCPForUnity/Editor/Services/Transport/TransportManager.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -42,16 +42,6 @@ private IMcpTransportClient GetOrCreateClient(TransportMode mode)
};
}

private IMcpTransportClient GetClient(TransportMode mode)
{
return mode switch
{
TransportMode.Http => _httpClient,
TransportMode.Stdio => _stdioClient,
_ => throw new ArgumentOutOfRangeException(nameof(mode), mode, "Unsupported transport mode"),
};
}

public async Task<bool> StartAsync(TransportMode mode)
{
IMcpTransportClient client = GetOrCreateClient(mode);
Expand Down Expand Up @@ -128,6 +118,20 @@ public TransportState GetState(TransportMode mode)

public bool IsRunning(TransportMode mode) => GetState(mode).IsConnected;

/// <summary>
/// Gets the active transport client for the specified mode.
/// Returns null if the client hasn't been created yet.
/// </summary>
public IMcpTransportClient GetClient(TransportMode mode)
{
return mode switch
{
TransportMode.Http => _httpClient,
TransportMode.Stdio => _stdioClient,
_ => throw new ArgumentOutOfRangeException(nameof(mode), mode, "Unsupported transport mode"),
};
}

private void UpdateState(TransportMode mode, TransportState state)
{
switch (mode)
Expand Down
7 changes: 7 additions & 0 deletions MCPForUnity/Editor/Services/Transport/Transports/StdioTransportClient.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -46,5 +46,12 @@ public Task<bool> VerifyAsync()
return Task.FromResult(running);
}

public Task ReregisterToolsAsync()
{
// Stdio transport doesn't support dynamic tool reregistration
// Tools are registered at server startup
return Task.CompletedTask;
}

}
}
23 changes: 23 additions & 0 deletions MCPForUnity/Editor/Services/Transport/Transports/WebSocketTransportClient.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -506,6 +506,29 @@ private async Task SendRegisterToolsAsync(CancellationToken token)
McpLog.Info($"[WebSocket] Sent {tools.Count} tools registration", false);
}

public async Task ReregisterToolsAsync()
{
if (!IsConnected || _lifecycleCts == null)
{
McpLog.Warn("[WebSocket] Cannot reregister tools: not connected");
return;
}

try
{
await SendRegisterToolsAsync(_lifecycleCts.Token).ConfigureAwait(false);
McpLog.Info("[WebSocket] Tool reregistration completed", false);
}
catch (System.OperationCanceledException)
{
McpLog.Warn("[WebSocket] Tool reregistration cancelled");
}
catch (System.Exception ex)
{
McpLog.Error($"[WebSocket] Tool reregistration failed: {ex.Message}");
}
}

private async Task HandleExecuteAsync(JObject payload, CancellationToken token)
{
string commandId = payload.Value<string>("id");
Expand Down
29 changes: 29 additions & 0 deletions MCPForUnity/Editor/Windows/Components/Tools/McpToolsSection.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,9 +1,11 @@
using System;
using System.Collections.Generic;
using System.Linq;
using System.Threading;
using MCPForUnity.Editor.Constants;
using MCPForUnity.Editor.Helpers;
using MCPForUnity.Editor.Services;
using MCPForUnity.Editor.Services.Transport;
using MCPForUnity.Editor.Tools;
using UnityEditor;
using UnityEngine.UIElements;
Expand Down Expand Up @@ -231,6 +233,30 @@ private void HandleToggleChange(ToolMetadata tool, bool enabled, bool updateSumm
{
UpdateSummary();
}

// Trigger tool reregistration with connected MCP server
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)
Comment on lines +241 to +250
Copy link
Contributor

@sourcery-ai sourcery-ai bot Mar 4, 2026

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

{
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.
}
}
catch (Exception ex)
{
McpLog.Warn($"Failed to reregister tools: {ex.Message}");
}
});
}

private void SetAllToolsState(bool enabled)
Expand All @@ -253,6 +279,9 @@ private void SetAllToolsState(bool enabled)
}

UpdateSummary();

// Trigger tool reregistration after bulk change
ReregisterToolsAsync();
Comment on lines +282 to +284
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.
}

private void UpdateSummary()
Expand Down
Loading
Loading