Skip to content

Commit ca4fd79

Browse files
dsarnodsarno
committed
Enhance MCP tools validation and error handling
1 parent 764e625 commit ca4fd79

File tree

2 files changed

+159
-157
lines changed

2 files changed

+159
-157
lines changed

Diff for: ‎Editor/Tools/NotifyMessageTool.cs

+35-13
Original file line numberDiff line numberDiff line change
@@ -1,15 +1,27 @@
1+
using System;
12
using System.Threading.Tasks;
23
using McpUnity.Unity;
34
using UnityEngine;
5+
using UnityEditor;
46
using Newtonsoft.Json.Linq;
57

68
namespace McpUnity.Tools
79
{
810
/// <summary>
9-
/// Tool for sending notification messages to the Unity console
11+
/// Tool for displaying messages in the Unity console
1012
/// </summary>
1113
public class NotifyMessageTool : McpToolBase
1214
{
15+
/// <summary>
16+
/// Supported message types for Unity console output
17+
/// </summary>
18+
private enum MessageType
19+
{
20+
Info,
21+
Warning,
22+
Error
23+
}
24+
1325
public NotifyMessageTool()
1426
{
1527
Name = "notify_message";
@@ -25,35 +37,45 @@ public override JObject Execute(JObject parameters)
2537
// Extract parameters
2638
string message = parameters["message"]?.ToObject<string>();
2739
string type = parameters["type"]?.ToObject<string>()?.ToLower() ?? "info";
28-
40+
41+
// Validate message
2942
if (string.IsNullOrEmpty(message))
3043
{
3144
return McpUnitySocketHandler.CreateErrorResponse(
32-
"Required parameter 'message' not provided",
45+
"Required parameter 'message' not provided or is empty",
46+
"validation_error"
47+
);
48+
}
49+
50+
// Parse and validate message type
51+
if (!Enum.TryParse<MessageType>(type, true, out var messageType))
52+
{
53+
return McpUnitySocketHandler.CreateErrorResponse(
54+
$"Invalid message type '{type}'. Valid types are: info, warning, error",
3355
"validation_error"
3456
);
3557
}
36-
58+
3759
// Log the message based on type
38-
switch (type)
60+
switch (messageType)
3961
{
40-
case "error":
41-
Debug.LogError($"[MCP]: {message}");
62+
case MessageType.Warning:
63+
Debug.LogWarning($"[MCP Unity] {message}");
4264
break;
43-
case "warning":
44-
Debug.LogWarning($"[MCP]: {message}");
65+
case MessageType.Error:
66+
Debug.LogError($"[MCP Unity] {message}");
4567
break;
4668
default:
47-
Debug.Log($"[MCP]: {message}");
69+
Debug.Log($"[MCP Unity] {message}");
4870
break;
4971
}
50-
72+
5173
// Create the response
5274
return new JObject
5375
{
5476
["success"] = true,
55-
["type"] = "text",
56-
["message"] = $"Message displayed: {message}"
77+
["message"] = $"Message displayed: {message}",
78+
["type"] = "text"
5779
};
5880
}
5981
}

Diff for: ‎Editor/Tools/RunTestsTool.cs

+124-144
Original file line numberDiff line numberDiff line change
@@ -1,12 +1,11 @@
11
using System;
2-
using System.Threading;
2+
using System.Collections.Generic;
33
using System.Threading.Tasks;
44
using McpUnity.Unity;
5+
using UnityEditor;
6+
using UnityEditor.TestTools.TestRunner.Api;
57
using UnityEngine;
68
using Newtonsoft.Json.Linq;
7-
using System.Collections.Generic;
8-
using UnityEditor.TestTools.TestRunner.Api;
9-
using McpUnity.Services;
109

1110
namespace McpUnity.Tools
1211
{
@@ -15,170 +14,151 @@ namespace McpUnity.Tools
1514
/// </summary>
1615
public class RunTestsTool : McpToolBase, ICallbacks
1716
{
18-
private readonly ITestRunnerService _testRunnerService;
19-
20-
private bool _isRunning = false;
21-
private TaskCompletionSource<JObject> _testRunCompletionSource;
22-
private List<TestResult> _testResults = new List<TestResult>();
23-
24-
// Structure to store test results
25-
private class TestResult
17+
/// <summary>
18+
/// Supported test modes for Unity Test Runner
19+
/// </summary>
20+
private enum TestMode
2621
{
27-
public string Name { get; set; }
28-
public string FullName { get; set; }
29-
public string ResultState { get; set; }
30-
public string Message { get; set; }
31-
public double Duration { get; set; }
32-
public bool Passed => ResultState == "Passed";
22+
EditMode,
23+
PlayMode
3324
}
34-
35-
public RunTestsTool(ITestRunnerService testRunnerService)
25+
26+
private TaskCompletionSource<JObject> _testCompletionSource;
27+
private int _testCount;
28+
private int _passCount;
29+
private int _failCount;
30+
private List<JObject> _testResults;
31+
32+
public RunTestsTool()
3633
{
3734
Name = "run_tests";
38-
Description = "Runs tests using Unity's Test Runner";
35+
Description = "Runs Unity's Test Runner tests";
3936
IsAsync = true;
40-
41-
_testRunnerService = testRunnerService;
42-
43-
// Register callbacks with the TestRunnerApi
44-
_testRunnerService.TestRunnerApi.RegisterCallbacks(this);
4537
}
46-
47-
/// <summary>
48-
/// Executes the RunTests tool asynchronously on the main thread.
49-
/// </summary>
50-
/// <param name="parameters">Tool parameters, including optional 'testMode' and 'testFilter'.</param>
51-
/// <param name="tcs">TaskCompletionSource to set the result or exception.</param>
52-
public override void ExecuteAsync(JObject parameters, TaskCompletionSource<JObject> tcs)
38+
39+
public override JObject Execute(JObject parameters)
5340
{
54-
// Check if tests are already running
55-
if (_isRunning)
41+
throw new NotSupportedException("This tool only supports async execution. Please use ExecuteAsync instead.");
42+
}
43+
44+
public override async Task<JObject> ExecuteAsync(JObject parameters)
45+
{
46+
// Extract parameters
47+
string testMode = parameters["testMode"]?.ToObject<string>() ?? "EditMode";
48+
string testFilter = parameters["testFilter"]?.ToObject<string>();
49+
50+
// Validate test mode
51+
if (!Enum.TryParse<TestMode>(testMode, true, out var mode))
5652
{
57-
tcs.SetResult(McpUnitySocketHandler.CreateErrorResponse(
58-
"Tests are already running. Please wait for them to complete.",
59-
"test_runner_busy"
60-
));
61-
return;
53+
return McpUnitySocketHandler.CreateErrorResponse(
54+
$"Invalid test mode '{testMode}'. Valid modes are: EditMode, PlayMode",
55+
"validation_error"
56+
);
6257
}
63-
64-
// Extract parameters
65-
string testModeStr = parameters["testMode"]?.ToObject<string>() ?? "editmode";
66-
string testFilter = parameters["testFilter"]?.ToObject<string>() ?? "";
67-
68-
// Parse test mode
69-
TestMode testMode;
70-
switch (testModeStr.ToLowerInvariant())
58+
59+
// Initialize test tracking
60+
_testCount = 0;
61+
_passCount = 0;
62+
_failCount = 0;
63+
_testResults = new List<JObject>();
64+
_testCompletionSource = new TaskCompletionSource<JObject>();
65+
66+
try
7167
{
72-
case "playmode":
73-
testMode = TestMode.PlayMode;
74-
break;
75-
case "editmode":
76-
testMode = TestMode.EditMode;
77-
break;
78-
default:
79-
testMode = TestMode.EditMode;
80-
break;
68+
var testRunnerApi = ScriptableObject.CreateInstance<TestRunnerApi>();
69+
testRunnerApi.RegisterCallbacks(this);
70+
71+
var filter = new Filter
72+
{
73+
testMode = mode == TestMode.EditMode ? TestMode.EditMode : TestMode.PlayMode
74+
};
75+
76+
if (!string.IsNullOrEmpty(testFilter))
77+
{
78+
filter.testNames = new[] { testFilter };
79+
}
80+
81+
testRunnerApi.Execute(new ExecutionSettings(filter));
82+
83+
// Wait for test completion or timeout after 5 minutes
84+
var timeoutTask = Task.Delay(TimeSpan.FromMinutes(5));
85+
var completedTask = await Task.WhenAny(_testCompletionSource.Task, timeoutTask);
86+
87+
if (completedTask == timeoutTask)
88+
{
89+
return McpUnitySocketHandler.CreateErrorResponse(
90+
"Failed to run tests: Request timed out",
91+
"timeout_error"
92+
);
93+
}
94+
95+
return await _testCompletionSource.Task;
8196
}
82-
83-
// Log the execution
84-
Debug.Log($"[MCP Unity] Running tests: Mode={testMode}, Filter={testFilter}");
85-
86-
// Reset state
87-
_isRunning = true;
88-
_testResults.Clear();
89-
_testRunCompletionSource = tcs;
90-
91-
// Execute tests using the TestRunnerService
92-
_testRunnerService.ExecuteTests(
93-
testMode,
94-
testFilter,
95-
tcs
96-
);
97-
}
98-
99-
#region ICallbacks Implementation
100-
101-
// Called when a test run starts
102-
public void RunStarted(ITestAdaptor testsToRun)
103-
{
104-
Debug.Log($"[MCP Unity] Test run started: {testsToRun.Name}");
105-
}
106-
107-
// Called when a test runs
108-
public void TestStarted(ITestAdaptor test)
109-
{
110-
// Nothing to do here
111-
}
112-
113-
// Called when a test finishes
114-
public void TestFinished(ITestResultAdaptor result)
115-
{
116-
_testResults.Add(new TestResult
97+
catch (Exception ex)
11798
{
118-
Name = result.Test.Name,
119-
FullName = result.FullName,
120-
ResultState = result.ResultState,
121-
Message = result.Message,
122-
Duration = result.Duration
123-
});
124-
125-
Debug.Log($"[MCP Unity] Test finished: {result.Test.Name} - {result.ResultState}");
99+
return McpUnitySocketHandler.CreateErrorResponse(
100+
$"Failed to run tests: {ex.Message}",
101+
"execution_error"
102+
);
103+
}
126104
}
127-
128-
// Called when a test run completes
129-
public void RunFinished(ITestResultAdaptor result)
105+
106+
public void RunStarted(ITestAdapterRef testsToRun)
130107
{
131-
Debug.Log($"[MCP Unity] Test run completed: {result.Test.Name} - {result.ResultState}");
132-
133-
_isRunning = false;
108+
_testCount = testsToRun.TestCaseCount;
134109

135-
// Create test results summary
136-
var summary = new JObject
110+
if (_testCount == 0)
137111
{
138-
["testCount"] = _testResults.Count,
139-
["passCount"] = _testResults.FindAll(r => r.Passed).Count,
140-
["duration"] = result.Duration,
141-
["success"] = result.ResultState == "Passed",
142-
["status"] = "completed",
143-
["message"] = $"Test run completed: {result.Test.Name} - {result.ResultState}"
144-
};
145-
146-
// Add test results array
147-
var resultArray = new JArray();
148-
foreach (var testResult in _testResults)
149-
{
150-
resultArray.Add(new JObject
112+
_testCompletionSource.TrySetResult(new JObject
151113
{
152-
["name"] = testResult.Name,
153-
["fullName"] = testResult.FullName,
154-
["result"] = testResult.ResultState,
155-
["message"] = testResult.Message,
156-
["duration"] = testResult.Duration
114+
["success"] = false,
115+
["message"] = "No tests found matching the specified criteria",
116+
["testCount"] = 0,
117+
["passCount"] = 0,
118+
["failCount"] = 0,
119+
["results"] = new JArray()
157120
});
158121
}
159-
summary["results"] = resultArray;
160-
161-
// Set the test run completion result
162-
try
122+
}
123+
124+
public void RunFinished(ITestResultAdapterRef testResults)
125+
{
126+
var response = new JObject
163127
{
164-
_testRunCompletionSource.SetResult(new JObject
165-
{
166-
["success"] = true,
167-
["type"] = "text",
168-
["message"] = summary["message"].Value<string>()
169-
});
170-
}
171-
catch (Exception ex)
128+
["success"] = _failCount == 0,
129+
["message"] = $"Tests completed: {_passCount} passed, {_failCount} failed",
130+
["testCount"] = _testCount,
131+
["passCount"] = _passCount,
132+
["failCount"] = _failCount,
133+
["results"] = JArray.FromObject(_testResults)
134+
};
135+
136+
_testCompletionSource.TrySetResult(response);
137+
}
138+
139+
public void TestStarted(ITestAdapterRef test)
140+
{
141+
// Optional: Add test started tracking if needed
142+
}
143+
144+
public void TestFinished(ITestResultAdapterRef result)
145+
{
146+
if (result.TestStatus == TestStatus.Passed)
172147
{
173-
Debug.LogError($"[MCP Unity] Failed to set test results: {ex.Message}");
174-
_testRunCompletionSource.TrySetException(ex);
148+
_passCount++;
175149
}
176-
finally
150+
else if (result.TestStatus == TestStatus.Failed)
177151
{
178-
_testRunCompletionSource = null;
152+
_failCount++;
179153
}
154+
155+
_testResults.Add(new JObject
156+
{
157+
["name"] = result.Name,
158+
["status"] = result.TestStatus.ToString(),
159+
["message"] = result.Message,
160+
["duration"] = result.Duration
161+
});
180162
}
181-
182-
#endregion
183163
}
184164
}

0 commit comments

Comments
 (0)