Skip to content

Commit 17cd543

Browse files
msanatanmsanatan
authored
Fix stdio reloads (#402)
* First pass at MCP client refactor * Restore original text instructions Well most of them, I modified a few * Move configurators to their own folder It's less clusterd * Remvoe override for Windsurf because we no longer need to use it * Add Antigravity configs Works like Windsurf, but it sucks ass * Add some docs for properties * Add comprehensive MCP client configurators documentation * Add missing imports (#7) * Handle Linux paths when unregistering CLI commands * Construct a JSON error in a much more secure fashion * Fix stdio auto-reconnect after domain reloads We mirror what we've done with the HTTP/websocket connection We also ensure the states from the stdio/HTTP connections are handled separately. Things now work as expected * Fix ActiveMode to return resolved transport mode instead of preferred mode The ActiveMode property now calls ResolvePreferredMode() to return the actual active transport mode rather than just the preferred mode setting. * Minor improvements for stdio bridge - Consolidated the !useHttp && isRunning checks into a single shouldResume flag. - Wrapped the fire-and-forget StopAsync in a continuation that logs faults (matching the HTTP handler pattern). - Wrapped StartAsync in a continuation that logs failures and only triggers the health check on success. * Refactor TransportManager to use switch expressions and improve error handling - Replace if-else chains with switch expressions for better readability and exhaustiveness checking - Add GetClient() helper method to centralize client retrieval logic - Wrap StopAsync in try-catch to log failures when stopping a failed transport - Use client.TransportName instead of mode.ToString() for consistent naming in error messages
1 parent f94cb24 commit 17cd543

File tree

7 files changed

+233
-107
lines changed

7 files changed

+233
-107
lines changed

‎MCPForUnity/Editor/Constants/EditorPrefKeys.cs‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,7 @@ internal static class EditorPrefKeys
1111
internal const string ValidationLevel = "MCPForUnity.ValidationLevel";
1212
internal const string UnitySocketPort = "MCPForUnity.UnitySocketPort";
1313
internal const string ResumeHttpAfterReload = "MCPForUnity.ResumeHttpAfterReload";
14+
internal const string ResumeStdioAfterReload = "MCPForUnity.ResumeStdioAfterReload";
1415

1516
internal const string UvxPathOverride = "MCPForUnity.UvxPath";
1617
internal const string ClaudeCliPathOverride = "MCPForUnity.ClaudeCliPath";

‎MCPForUnity/Editor/Services/BridgeControlService.cs‎

Lines changed: 19 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -49,13 +49,21 @@ private static BridgeVerificationResult BuildVerificationResult(TransportState s
4949
};
5050
}
5151

52-
public bool IsRunning => _transportManager.GetState().IsConnected;
52+
public bool IsRunning
53+
{
54+
get
55+
{
56+
var mode = ResolvePreferredMode();
57+
return _transportManager.IsRunning(mode);
58+
}
59+
}
5360

5461
public int CurrentPort
5562
{
5663
get
5764
{
58-
var state = _transportManager.GetState();
65+
var mode = ResolvePreferredMode();
66+
var state = _transportManager.GetState(mode);
5967
if (state.Port.HasValue)
6068
{
6169
return state.Port.Value;
@@ -67,7 +75,7 @@ public int CurrentPort
6775
}
6876

6977
public bool IsAutoConnectMode => StdioBridgeHost.IsAutoConnectMode();
70-
public TransportMode? ActiveMode => _transportManager.ActiveMode;
78+
public TransportMode? ActiveMode => ResolvePreferredMode();
7179

7280
public async Task<bool> StartAsync()
7381
{
@@ -92,7 +100,8 @@ public async Task StopAsync()
92100
{
93101
try
94102
{
95-
await _transportManager.StopAsync();
103+
var mode = ResolvePreferredMode();
104+
await _transportManager.StopAsync(mode);
96105
}
97106
catch (Exception ex)
98107
{
@@ -102,17 +111,17 @@ public async Task StopAsync()
102111

103112
public async Task<BridgeVerificationResult> VerifyAsync()
104113
{
105-
var mode = _transportManager.ActiveMode ?? ResolvePreferredMode();
106-
bool pingSucceeded = await _transportManager.VerifyAsync();
107-
var state = _transportManager.GetState();
114+
var mode = ResolvePreferredMode();
115+
bool pingSucceeded = await _transportManager.VerifyAsync(mode);
116+
var state = _transportManager.GetState(mode);
108117
return BuildVerificationResult(state, mode, pingSucceeded);
109118
}
110119

111120
public BridgeVerificationResult Verify(int port)
112121
{
113-
var mode = _transportManager.ActiveMode ?? ResolvePreferredMode();
114-
bool pingSucceeded = _transportManager.VerifyAsync().GetAwaiter().GetResult();
115-
var state = _transportManager.GetState();
122+
var mode = ResolvePreferredMode();
123+
bool pingSucceeded = _transportManager.VerifyAsync(mode).GetAwaiter().GetResult();
124+
var state = _transportManager.GetState(mode);
116125

117126
if (mode == TransportMode.Stdio)
118127
{

‎MCPForUnity/Editor/Services/HttpBridgeReloadHandler.cs‎

Lines changed: 9 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -24,8 +24,8 @@ private static void OnBeforeAssemblyReload()
2424
{
2525
try
2626
{
27-
var bridge = MCPServiceLocator.Bridge;
28-
bool shouldResume = bridge.IsRunning && bridge.ActiveMode == TransportMode.Http;
27+
var transport = MCPServiceLocator.TransportManager;
28+
bool shouldResume = transport.IsRunning(TransportMode.Http);
2929

3030
if (shouldResume)
3131
{
@@ -36,9 +36,9 @@ private static void OnBeforeAssemblyReload()
3636
EditorPrefs.DeleteKey(EditorPrefKeys.ResumeHttpAfterReload);
3737
}
3838

39-
if (bridge.IsRunning)
39+
if (shouldResume)
4040
{
41-
var stopTask = bridge.StopAsync();
41+
var stopTask = transport.StopAsync(TransportMode.Http);
4242
stopTask.ContinueWith(t =>
4343
{
4444
if (t.IsFaulted && t.Exception != null)
@@ -59,7 +59,9 @@ private static void OnAfterAssemblyReload()
5959
bool resume = false;
6060
try
6161
{
62-
resume = EditorPrefs.GetBool(EditorPrefKeys.ResumeHttpAfterReload, false);
62+
// Only resume HTTP if it is still the selected transport.
63+
bool useHttp = EditorPrefs.GetBool(EditorPrefKeys.UseHttpTransport, true);
64+
resume = useHttp && EditorPrefs.GetBool(EditorPrefKeys.ResumeHttpAfterReload, false);
6365
if (resume)
6466
{
6567
EditorPrefs.DeleteKey(EditorPrefKeys.ResumeHttpAfterReload);
@@ -90,7 +92,7 @@ private static void OnAfterAssemblyReload()
9092
{
9193
try
9294
{
93-
var startTask = MCPServiceLocator.Bridge.StartAsync();
95+
var startTask = MCPServiceLocator.TransportManager.StartAsync(TransportMode.Http);
9496
startTask.ContinueWith(t =>
9597
{
9698
if (t.IsFaulted)
@@ -123,7 +125,7 @@ private static void OnAfterAssemblyReload()
123125
{
124126
try
125127
{
126-
bool started = await MCPServiceLocator.Bridge.StartAsync();
128+
bool started = await MCPServiceLocator.TransportManager.StartAsync(TransportMode.Http);
127129
if (!started)
128130
{
129131
McpLog.Warn("Failed to resume HTTP MCP bridge after domain reload");

‎MCPForUnity/Editor/Services/StdioBridgeReloadHandler.cs‎

Lines changed: 104 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,104 @@
1+
using System;
2+
using UnityEditor;
3+
using MCPForUnity.Editor.Constants;
4+
using MCPForUnity.Editor.Helpers;
5+
using MCPForUnity.Editor.Services.Transport;
6+
using MCPForUnity.Editor.Services.Transport.Transports;
7+
8+
namespace MCPForUnity.Editor.Services
9+
{
10+
/// <summary>
11+
/// Ensures the legacy stdio bridge resumes after domain reloads, mirroring the HTTP handler.
12+
/// </summary>
13+
[InitializeOnLoad]
14+
internal static class StdioBridgeReloadHandler
15+
{
16+
static StdioBridgeReloadHandler()
17+
{
18+
AssemblyReloadEvents.beforeAssemblyReload += OnBeforeAssemblyReload;
19+
AssemblyReloadEvents.afterAssemblyReload += OnAfterAssemblyReload;
20+
}
21+
22+
private static void OnBeforeAssemblyReload()
23+
{
24+
try
25+
{
26+
// Only persist resume intent when stdio is the active transport and the bridge is running.
27+
bool useHttp = EditorPrefs.GetBool(EditorPrefKeys.UseHttpTransport, true);
28+
bool isRunning = MCPServiceLocator.TransportManager.IsRunning(TransportMode.Stdio);
29+
bool shouldResume = !useHttp && isRunning;
30+
31+
if (shouldResume)
32+
{
33+
EditorPrefs.SetBool(EditorPrefKeys.ResumeStdioAfterReload, true);
34+
35+
// Stop only the stdio bridge; leave HTTP untouched if it is running concurrently.
36+
var stopTask = MCPServiceLocator.TransportManager.StopAsync(TransportMode.Stdio);
37+
stopTask.ContinueWith(t =>
38+
{
39+
if (t.IsFaulted && t.Exception != null)
40+
{
41+
McpLog.Warn($"Error stopping stdio bridge before reload: {t.Exception.GetBaseException()?.Message}");
42+
}
43+
}, System.Threading.Tasks.TaskScheduler.Default);
44+
}
45+
else
46+
{
47+
EditorPrefs.DeleteKey(EditorPrefKeys.ResumeStdioAfterReload);
48+
}
49+
}
50+
catch (Exception ex)
51+
{
52+
McpLog.Warn($"Failed to persist stdio reload flag: {ex.Message}");
53+
}
54+
}
55+
56+
private static void OnAfterAssemblyReload()
57+
{
58+
bool resume = false;
59+
try
60+
{
61+
resume = EditorPrefs.GetBool(EditorPrefKeys.ResumeStdioAfterReload, false);
62+
bool useHttp = EditorPrefs.GetBool(EditorPrefKeys.UseHttpTransport, true);
63+
resume = resume && !useHttp;
64+
if (resume)
65+
{
66+
EditorPrefs.DeleteKey(EditorPrefKeys.ResumeStdioAfterReload);
67+
}
68+
}
69+
catch (Exception ex)
70+
{
71+
McpLog.Warn($"Failed to read stdio reload flag: {ex.Message}");
72+
}
73+
74+
if (!resume)
75+
{
76+
return;
77+
}
78+
79+
// Restart via TransportManager so state stays in sync; if it fails (port busy), rely on UI to retry.
80+
TryStartBridgeImmediate();
81+
}
82+
83+
private static void TryStartBridgeImmediate()
84+
{
85+
var startTask = MCPServiceLocator.TransportManager.StartAsync(TransportMode.Stdio);
86+
startTask.ContinueWith(t =>
87+
{
88+
if (t.IsFaulted)
89+
{
90+
var baseEx = t.Exception?.GetBaseException();
91+
McpLog.Warn($"Failed to resume stdio bridge after reload: {baseEx?.Message}");
92+
return;
93+
}
94+
if (!t.Result)
95+
{
96+
McpLog.Warn("Failed to resume stdio bridge after domain reload");
97+
return;
98+
}
99+
100+
MCPForUnity.Editor.Windows.MCPForUnityEditorWindow.RequestHealthVerification();
101+
}, System.Threading.Tasks.TaskScheduler.Default);
102+
}
103+
}
104+
}

‎MCPForUnity/Editor/Services/StdioBridgeReloadHandler.cs.meta‎

Lines changed: 11 additions & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

0 commit comments

Comments
 (0)