Skip to content

Commit f2c57ca

Browse files
msanatanmsanatan
authored
Add testing and move menu items to resources (#316)
* deps: add tomli>=2.3.0 dependency to UnityMcpServer package * feat: dynamically fetch package version from pyproject.toml for telemetry * Add pydantic * feat: add resource registry for MCP resource auto-discovery * feat: add telemetry decorator for tracking MCP resource usage * feat: add auto-discovery and registration system for MCP resources * feat: add resource registration to MCP server initialization * feat: add MCPResponse model class for standardized API responses * refactor: replace Debug.Log calls with McpLog wrapper for consistent logging * feat: add test discovery endpoints for Unity Test Framework integration We haven't connected them as yet, still thinking about how to do this neatly * Fix server setup * refactor: reduce log verbosity by changing individual resource/tool registration logs to debug level * chore: bump mcp[cli] dependency from 1.15.0 to 1.17.0 * refactor: remove Context parameter and add uri keyword argument in resource decorator The Context parameter doesn't work on our version of FastMCP * chore: upgrade Python base image to 3.13 and simplify Dockerfile setup * fix: apply telemetry decorator before mcp.tool to ensure proper wrapping order * fix: swap order of telemetry and resource decorators to properly wrap handlers * fix: update log prefixes for consistency in logging methods * Fix compile errors * feat: extend command registry to support both tools and resources * Run get tests as a coroutine because it doesn't return results immediately This works but it spams logs like crazy, maybe there's a better/simpler way * refactor: migrate from coroutines to async/await for test retrieval and command execution * feat: add optional error field to MCPResponse model * Increased timeout because loading tests can take some time * Make message optional so error responses that only have success and error don't cause Pydantic errors * Set max_retries to 5 This connection module needs a lookover. The retries should be an exponential backoff and we could structure why it's failing so much * Use pydantic model to structure the error output * fix: initialize data field in GetTestsResponse to avoid potential errors * Don't return path parameter * feat: add Unity test runner execution with structured results and Python bindings * refactor: simplify GetTests by removing mode filtering and related parsing logic * refactor: move test runner functionality into dedicated service interface * feat: add resource retrieval telemetry tracking with new record type and helper function * fix: convert tool functions to async and await ctx.info calls * refactor: reorganize menu item functionality into separate execute and get commands An MCP resource for retrieval, and a simple command to execute. Because it's a resource, it's easier for the user to see what's in the menu items * refactor: rename manage_menu_item to execute_menu_item and update tool examples to use async/await We'll eventually put a section for resources * Revert "fix: convert tool functions to async and await ctx.info calls" This reverts commit 012ea6b. * fix: replace tomllib with tomli for Python 3.10 compatibility in telemetry module * Remove confusing comment * refactor: improve error handling and simplify test retrieval logic in GetTests commands * No cache by default * docs: remove redundant comment for HandleCommand method in ExecuteMenuItem
1 parent 5ad5b4a commit f2c57ca

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

54 files changed

+2052
-546
lines changed

‎MCPForUnity/Editor/Helpers/McpLog.cs‎

Lines changed: 6 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,9 @@ namespace MCPForUnity.Editor.Helpers
55
{
66
internal static class McpLog
77
{
8-
private const string Prefix = "<b><color=#2EA3FF>MCP-FOR-UNITY</color></b>:";
8+
private const string LogPrefix = "<b><color=#2EA3FF>MCP-FOR-UNITY</color></b>:";
9+
private const string WarnPrefix = "<b><color=#cc7a00>MCP-FOR-UNITY</color></b>:";
10+
private const string ErrorPrefix = "<b><color=#cc3333>MCP-FOR-UNITY</color></b>:";
911

1012
private static bool IsDebugEnabled()
1113
{
@@ -15,17 +17,17 @@ private static bool IsDebugEnabled()
1517
public static void Info(string message, bool always = true)
1618
{
1719
if (!always && !IsDebugEnabled()) return;
18-
Debug.Log($"{Prefix} {message}");
20+
Debug.Log($"{LogPrefix} {message}");
1921
}
2022

2123
public static void Warn(string message)
2224
{
23-
Debug.LogWarning($"<color=#cc7a00>{Prefix} {message}</color>");
25+
Debug.LogWarning($"{WarnPrefix} {message}");
2426
}
2527

2628
public static void Error(string message)
2729
{
28-
Debug.LogError($"<color=#cc3333>{Prefix} {message}</color>");
30+
Debug.LogError($"{ErrorPrefix} {message}");
2931
}
3032
}
3133
}

‎MCPForUnity/Editor/MCPForUnityBridge.cs‎

Lines changed: 88 additions & 42 deletions
Original file line numberDiff line numberDiff line change
@@ -14,11 +14,30 @@
1414
using MCPForUnity.Editor.Helpers;
1515
using MCPForUnity.Editor.Models;
1616
using MCPForUnity.Editor.Tools;
17-
using MCPForUnity.Editor.Tools.MenuItems;
1817
using MCPForUnity.Editor.Tools.Prefabs;
1918

2019
namespace MCPForUnity.Editor
2120
{
21+
22+
/// <summary>
23+
/// Outbound message structure for the writer thread
24+
/// </summary>
25+
class Outbound
26+
{
27+
public byte[] Payload;
28+
public string Tag;
29+
public int? ReqId;
30+
}
31+
32+
/// <summary>
33+
/// Queued command structure for main thread processing
34+
/// </summary>
35+
class QueuedCommand
36+
{
37+
public string CommandJson;
38+
public TaskCompletionSource<string> Tcs;
39+
public bool IsExecuting;
40+
}
2241
[InitializeOnLoad]
2342
public static partial class MCPForUnityBridge
2443
{
@@ -28,13 +47,6 @@ public static partial class MCPForUnityBridge
2847
private static readonly object startStopLock = new();
2948
private static readonly object clientsLock = new();
3049
private static readonly System.Collections.Generic.HashSet<TcpClient> activeClients = new();
31-
// Single-writer outbox for framed responses
32-
private class Outbound
33-
{
34-
public byte[] Payload;
35-
public string Tag;
36-
public int? ReqId;
37-
}
3850
private static readonly BlockingCollection<Outbound> _outbox = new(new ConcurrentQueue<Outbound>());
3951
private static CancellationTokenSource cts;
4052
private static Task listenerTask;
@@ -45,10 +57,7 @@ private class Outbound
4557
private static double nextStartAt = 0.0f;
4658
private static double nextHeartbeatAt = 0.0f;
4759
private static int heartbeatSeq = 0;
48-
private static Dictionary<
49-
string,
50-
(string commandJson, TaskCompletionSource<string> tcs)
51-
> commandQueue = new();
60+
private static Dictionary<string, QueuedCommand> commandQueue = new();
5261
private static int mainThreadId;
5362
private static int currentUnityPort = 6400; // Dynamic port, starts with default
5463
private static bool isAutoConnectMode = false;
@@ -96,7 +105,7 @@ public static void StartAutoConnect()
96105
}
97106
catch (Exception ex)
98107
{
99-
Debug.LogError($"Auto-connect failed: {ex.Message}");
108+
McpLog.Error($"Auto-connect failed: {ex.Message}");
100109

101110
// Record telemetry for connection failure
102111
TelemetryHelper.RecordBridgeConnection(false, ex.Message);
@@ -297,7 +306,7 @@ public static void Start()
297306
{
298307
if (IsDebugEnabled())
299308
{
300-
Debug.Log($"<b><color=#2EA3FF>MCP-FOR-UNITY</color></b>: MCPForUnityBridge already running on port {currentUnityPort}");
309+
McpLog.Info($"MCPForUnityBridge already running on port {currentUnityPort}");
301310
}
302311
return;
303312
}
@@ -383,7 +392,7 @@ public static void Start()
383392
isAutoConnectMode = false;
384393
string platform = Application.platform.ToString();
385394
string serverVer = ReadInstalledServerVersionSafe();
386-
Debug.Log($"<b><color=#2EA3FF>MCP-FOR-UNITY</color></b>: MCPForUnityBridge started on port {currentUnityPort}. (OS={platform}, server={serverVer})");
395+
McpLog.Info($"MCPForUnityBridge started on port {currentUnityPort}. (OS={platform}, server={serverVer})");
387396
// Start background listener with cooperative cancellation
388397
cts = new CancellationTokenSource();
389398
listenerTask = Task.Run(() => ListenerLoopAsync(cts.Token));
@@ -403,7 +412,7 @@ public static void Start()
403412
}
404413
catch (SocketException ex)
405414
{
406-
Debug.LogError($"Failed to start TCP listener: {ex.Message}");
415+
McpLog.Error($"Failed to start TCP listener: {ex.Message}");
407416
}
408417
}
409418
}
@@ -437,7 +446,7 @@ public static void Stop()
437446
}
438447
catch (Exception ex)
439448
{
440-
Debug.LogError($"Error stopping MCPForUnityBridge: {ex.Message}");
449+
McpLog.Error($"Error stopping MCPForUnityBridge: {ex.Message}");
441450
}
442451
}
443452

@@ -465,7 +474,7 @@ public static void Stop()
465474
try { AssemblyReloadEvents.afterAssemblyReload -= OnAfterAssemblyReload; } catch { }
466475
try { EditorApplication.quitting -= Stop; } catch { }
467476

468-
if (IsDebugEnabled()) Debug.Log("<b><color=#2EA3FF>MCP-FOR-UNITY</color></b>: MCPForUnityBridge stopped.");
477+
if (IsDebugEnabled()) McpLog.Info("MCPForUnityBridge stopped.");
469478
}
470479

471480
private static async Task ListenerLoopAsync(CancellationToken token)
@@ -504,7 +513,7 @@ private static async Task ListenerLoopAsync(CancellationToken token)
504513
{
505514
if (isRunning && !token.IsCancellationRequested)
506515
{
507-
if (IsDebugEnabled()) Debug.LogError($"Listener error: {ex.Message}");
516+
if (IsDebugEnabled()) McpLog.Error($"Listener error: {ex.Message}");
508517
}
509518
}
510519
}
@@ -524,7 +533,7 @@ private static async Task HandleClientAsync(TcpClient client, CancellationToken
524533
if (IsDebugEnabled())
525534
{
526535
var ep = client.Client?.RemoteEndPoint?.ToString() ?? "unknown";
527-
Debug.Log($"<b><color=#2EA3FF>UNITY-MCP</color></b>: Client connected {ep}");
536+
McpLog.Info($"Client connected {ep}");
528537
}
529538
}
530539
catch { }
@@ -544,11 +553,11 @@ private static async Task HandleClientAsync(TcpClient client, CancellationToken
544553
#else
545554
await stream.WriteAsync(handshakeBytes, 0, handshakeBytes.Length, cts.Token).ConfigureAwait(false);
546555
#endif
547-
if (IsDebugEnabled()) MCPForUnity.Editor.Helpers.McpLog.Info("Sent handshake FRAMING=1 (strict)", always: false);
556+
if (IsDebugEnabled()) McpLog.Info("Sent handshake FRAMING=1 (strict)", always: false);
548557
}
549558
catch (Exception ex)
550559
{
551-
if (IsDebugEnabled()) MCPForUnity.Editor.Helpers.McpLog.Warn($"Handshake failed: {ex.Message}");
560+
if (IsDebugEnabled()) McpLog.Warn($"Handshake failed: {ex.Message}");
552561
return; // abort this client
553562
}
554563

@@ -564,7 +573,7 @@ private static async Task HandleClientAsync(TcpClient client, CancellationToken
564573
if (IsDebugEnabled())
565574
{
566575
var preview = commandText.Length > 120 ? commandText.Substring(0, 120) + "…" : commandText;
567-
MCPForUnity.Editor.Helpers.McpLog.Info($"recv framed: {preview}", always: false);
576+
McpLog.Info($"recv framed: {preview}", always: false);
568577
}
569578
}
570579
catch { }
@@ -585,7 +594,12 @@ private static async Task HandleClientAsync(TcpClient client, CancellationToken
585594

586595
lock (lockObj)
587596
{
588-
commandQueue[commandId] = (commandText, tcs);
597+
commandQueue[commandId] = new QueuedCommand
598+
{
599+
CommandJson = commandText,
600+
Tcs = tcs,
601+
IsExecuting = false
602+
};
589603
}
590604

591605
// Wait for the handler to produce a response, but do not block indefinitely
@@ -623,7 +637,7 @@ private static async Task HandleClientAsync(TcpClient client, CancellationToken
623637

624638
if (IsDebugEnabled())
625639
{
626-
try { MCPForUnity.Editor.Helpers.McpLog.Info("[MCP] sending framed response", always: false); } catch { }
640+
try { McpLog.Info("[MCP] sending framed response", always: false); } catch { }
627641
}
628642
// Crash-proof and self-reporting writer logs (direct write to this client's stream)
629643
long seq = System.Threading.Interlocked.Increment(ref _ioSeq);
@@ -662,11 +676,11 @@ private static async Task HandleClientAsync(TcpClient client, CancellationToken
662676
|| ex is System.IO.IOException;
663677
if (isBenign)
664678
{
665-
if (IsDebugEnabled()) MCPForUnity.Editor.Helpers.McpLog.Info($"Client handler: {msg}", always: false);
679+
if (IsDebugEnabled()) McpLog.Info($"Client handler: {msg}", always: false);
666680
}
667681
else
668682
{
669-
MCPForUnity.Editor.Helpers.McpLog.Error($"Client handler error: {msg}");
683+
McpLog.Error($"Client handler error: {msg}");
670684
}
671685
break;
672686
}
@@ -817,19 +831,25 @@ private static void ProcessCommands()
817831
}
818832

819833
// Snapshot under lock, then process outside to reduce contention
820-
List<(string id, string text, TaskCompletionSource<string> tcs)> work;
834+
List<(string id, QueuedCommand command)> work;
821835
lock (lockObj)
822836
{
823-
work = commandQueue
824-
.Select(kvp => (kvp.Key, kvp.Value.commandJson, kvp.Value.tcs))
825-
.ToList();
837+
work = new List<(string, QueuedCommand)>(commandQueue.Count);
838+
foreach (var kvp in commandQueue)
839+
{
840+
var queued = kvp.Value;
841+
if (queued.IsExecuting) continue;
842+
queued.IsExecuting = true;
843+
work.Add((kvp.Key, queued));
844+
}
826845
}
827846

828847
foreach (var item in work)
829848
{
830849
string id = item.id;
831-
string commandText = item.text;
832-
TaskCompletionSource<string> tcs = item.tcs;
850+
QueuedCommand queuedCommand = item.command;
851+
string commandText = queuedCommand.CommandJson;
852+
TaskCompletionSource<string> tcs = queuedCommand.Tcs;
833853

834854
try
835855
{
@@ -894,13 +914,41 @@ private static void ProcessCommands()
894914
}
895915
else
896916
{
897-
string responseJson = ExecuteCommand(command);
898-
tcs.SetResult(responseJson);
917+
// Use JObject for parameters as handlers expect this
918+
JObject paramsObject = command.@params ?? new JObject();
919+
920+
// Execute command (may be sync or async)
921+
object result = CommandRegistry.ExecuteCommand(command.type, paramsObject, tcs);
922+
923+
// If result is null, it means async execution - TCS will be completed by the awaited task
924+
// In this case, DON'T remove from queue yet, DON'T complete TCS
925+
if (result == null)
926+
{
927+
// Async command - the task continuation will complete the TCS
928+
// Setup cleanup when TCS completes - schedule on next frame to avoid race conditions
929+
string asyncCommandId = id;
930+
_ = tcs.Task.ContinueWith(_ =>
931+
{
932+
// Use EditorApplication.delayCall to schedule cleanup on main thread, next frame
933+
EditorApplication.delayCall += () =>
934+
{
935+
lock (lockObj)
936+
{
937+
commandQueue.Remove(asyncCommandId);
938+
}
939+
};
940+
});
941+
continue; // Skip the queue removal below
942+
}
943+
944+
// Synchronous result - complete TCS now
945+
var response = new { status = "success", result };
946+
tcs.SetResult(JsonConvert.SerializeObject(response));
899947
}
900948
}
901949
catch (Exception ex)
902950
{
903-
Debug.LogError($"Error processing command: {ex.Message}\n{ex.StackTrace}");
951+
McpLog.Error($"Error processing command: {ex.Message}\n{ex.StackTrace}");
904952

905953
var response = new
906954
{
@@ -915,7 +963,7 @@ private static void ProcessCommands()
915963
tcs.SetResult(responseJson);
916964
}
917965

918-
// Remove quickly under lock
966+
// Remove from queue (only for sync commands - async ones skip with 'continue' above)
919967
lock (lockObj) { commandQueue.Remove(id); }
920968
}
921969
}
@@ -1051,9 +1099,7 @@ private static string ExecuteCommand(Command command)
10511099
catch (Exception ex)
10521100
{
10531101
// Log the detailed error in Unity for debugging
1054-
Debug.LogError(
1055-
$"Error executing command '{command?.type ?? "Unknown"}': {ex.Message}\n{ex.StackTrace}"
1056-
);
1102+
McpLog.Error($"Error executing command '{command?.type ?? "Unknown"}': {ex.Message}\n{ex.StackTrace}");
10571103

10581104
// Standard error response format
10591105
var response = new
@@ -1074,11 +1120,11 @@ private static object HandleManageScene(JObject paramsObject)
10741120
{
10751121
try
10761122
{
1077-
if (IsDebugEnabled()) Debug.Log("[MCP] manage_scene: dispatching to main thread");
1123+
if (IsDebugEnabled()) McpLog.Info("[MCP] manage_scene: dispatching to main thread");
10781124
var sw = System.Diagnostics.Stopwatch.StartNew();
10791125
var r = InvokeOnMainThreadWithTimeout(() => ManageScene.HandleCommand(paramsObject), FrameIOTimeoutMs);
10801126
sw.Stop();
1081-
if (IsDebugEnabled()) Debug.Log($"[MCP] manage_scene: completed in {sw.ElapsedMilliseconds} ms");
1127+
if (IsDebugEnabled()) McpLog.Info($"[MCP] manage_scene: completed in {sw.ElapsedMilliseconds} ms");
10821128
return r ?? Response.Error("manage_scene returned null (timeout or error)");
10831129
}
10841130
catch (Exception ex)

‎TestProjects/UnityMCPTests/Assets/Tests/EditMode/Tools/MenuItems.meta‎ renamed to ‎MCPForUnity/Editor/Resources.meta‎

Lines changed: 1 addition & 1 deletion
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎MCPForUnity/Editor/Resources/McpForUnityResourceAttribute.cs‎

Lines changed: 37 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,37 @@
1+
using System;
2+
3+
namespace MCPForUnity.Editor.Resources
4+
{
5+
/// <summary>
6+
/// Marks a class as an MCP resource handler for auto-discovery.
7+
/// The class must have a public static HandleCommand(JObject) method.
8+
/// </summary>
9+
[AttributeUsage(AttributeTargets.Class, AllowMultiple = false)]
10+
public class McpForUnityResourceAttribute : Attribute
11+
{
12+
/// <summary>
13+
/// The resource name used to route requests to this resource.
14+
/// If not specified, defaults to the PascalCase class name converted to snake_case.
15+
/// </summary>
16+
public string ResourceName { get; }
17+
18+
/// <summary>
19+
/// Create an MCP resource attribute with auto-generated resource name.
20+
/// The resource name will be derived from the class name (PascalCase → snake_case).
21+
/// Example: ManageAsset → manage_asset
22+
/// </summary>
23+
public McpForUnityResourceAttribute()
24+
{
25+
ResourceName = null; // Will be auto-generated
26+
}
27+
28+
/// <summary>
29+
/// Create an MCP resource attribute with explicit resource name.
30+
/// </summary>
31+
/// <param name="resourceName">The resource name (e.g., "manage_asset")</param>
32+
public McpForUnityResourceAttribute(string resourceName)
33+
{
34+
ResourceName = resourceName;
35+
}
36+
}
37+
}

‎TestProjects/UnityMCPTests/Assets/Tests/EditMode/Tools/MenuItems/ManageMenuItemTests.cs.meta‎ renamed to ‎MCPForUnity/Editor/Resources/McpForUnityResourceAttribute.cs.meta‎

Lines changed: 1 addition & 1 deletion
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎MCPForUnity/Editor/Tools/MenuItems.meta‎ renamed to ‎MCPForUnity/Editor/Resources/MenuItems.meta‎

Lines changed: 1 addition & 1 deletion
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

0 commit comments

Comments
 (0)