Skip to content

Commit 01f5aeb

Browse files
msanatanmsanatan
committed
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
1 parent 63496c5 commit 01f5aeb

File tree

7 files changed

+195
-108
lines changed

7 files changed

+195
-108
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 => _preferredMode;
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: 90 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,90 @@
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+
if (!useHttp && isRunning)
30+
{
31+
EditorPrefs.SetBool(EditorPrefKeys.ResumeStdioAfterReload, true);
32+
}
33+
else
34+
{
35+
EditorPrefs.DeleteKey(EditorPrefKeys.ResumeStdioAfterReload);
36+
}
37+
38+
if (!useHttp && isRunning)
39+
{
40+
// Stop only the stdio bridge; leave HTTP untouched if it is running concurrently.
41+
MCPServiceLocator.TransportManager.StopAsync(TransportMode.Stdio);
42+
}
43+
}
44+
catch (Exception ex)
45+
{
46+
McpLog.Warn($"Failed to persist stdio reload flag: {ex.Message}");
47+
}
48+
}
49+
50+
private static void OnAfterAssemblyReload()
51+
{
52+
bool resume = false;
53+
try
54+
{
55+
resume = EditorPrefs.GetBool(EditorPrefKeys.ResumeStdioAfterReload, false);
56+
bool useHttp = EditorPrefs.GetBool(EditorPrefKeys.UseHttpTransport, true);
57+
resume = resume && !useHttp;
58+
if (resume)
59+
{
60+
EditorPrefs.DeleteKey(EditorPrefKeys.ResumeStdioAfterReload);
61+
}
62+
}
63+
catch (Exception ex)
64+
{
65+
McpLog.Warn($"Failed to read stdio reload flag: {ex.Message}");
66+
}
67+
68+
if (!resume)
69+
{
70+
return;
71+
}
72+
73+
// Restart via TransportManager so state stays in sync; if it fails (port busy), rely on UI to retry.
74+
TryStartBridgeImmediate();
75+
}
76+
77+
private static void TryStartBridgeImmediate()
78+
{
79+
try
80+
{
81+
MCPServiceLocator.TransportManager.StartAsync(TransportMode.Stdio);
82+
MCPForUnity.Editor.Windows.MCPForUnityEditorWindow.RequestHealthVerification();
83+
}
84+
catch (Exception ex)
85+
{
86+
McpLog.Warn($"Failed to resume stdio bridge after reload: {ex.Message}");
87+
}
88+
}
89+
}
90+
}

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

‎MCPForUnity/Editor/Services/Transport/TransportManager.cs‎

Lines changed: 65 additions & 41 deletions
Original file line numberDiff line numberDiff line change
@@ -10,8 +10,10 @@ namespace MCPForUnity.Editor.Services.Transport
1010
/// </summary>
1111
public class TransportManager
1212
{
13-
private IMcpTransportClient _active;
14-
private TransportMode? _activeMode;
13+
private IMcpTransportClient _httpClient;
14+
private IMcpTransportClient _stdioClient;
15+
private TransportState _httpState = TransportState.Disconnected("http");
16+
private TransportState _stdioState = TransportState.Disconnected("stdio");
1517
private Func<IMcpTransportClient> _webSocketFactory;
1618
private Func<IMcpTransportClient> _stdioFactory;
1719

@@ -22,8 +24,8 @@ public TransportManager()
2224
() => new StdioTransportClient());
2325
}
2426

25-
public IMcpTransportClient ActiveTransport => _active;
26-
public TransportMode? ActiveMode => _activeMode;
27+
public IMcpTransportClient ActiveTransport => null; // Deprecated single-transport accessor
28+
public TransportMode? ActiveMode => null; // Deprecated single-transport accessor
2729

2830
public void Configure(
2931
Func<IMcpTransportClient> webSocketFactory,
@@ -33,68 +35,90 @@ public void Configure(
3335
_stdioFactory = stdioFactory ?? throw new ArgumentNullException(nameof(stdioFactory));
3436
}
3537

36-
public async Task<bool> StartAsync(TransportMode mode)
38+
private IMcpTransportClient GetOrCreateClient(TransportMode mode)
3739
{
38-
await StopAsync();
39-
40-
IMcpTransportClient next = mode switch
40+
return mode switch
4141
{
42-
TransportMode.Stdio => _stdioFactory(),
43-
TransportMode.Http => _webSocketFactory(),
44-
_ => throw new ArgumentOutOfRangeException(nameof(mode), mode, "Unsupported transport mode")
45-
} ?? throw new InvalidOperationException($"Factory returned null for transport mode {mode}");
42+
TransportMode.Http => _httpClient ??= _webSocketFactory(),
43+
TransportMode.Stdio => _stdioClient ??= _stdioFactory(),
44+
_ => throw new ArgumentOutOfRangeException(nameof(mode), mode, "Unsupported transport mode"),
45+
};
46+
}
47+
48+
public async Task<bool> StartAsync(TransportMode mode)
49+
{
50+
IMcpTransportClient client = GetOrCreateClient(mode);
4651

47-
bool started = await next.StartAsync();
52+
bool started = await client.StartAsync();
4853
if (!started)
4954
{
50-
await next.StopAsync();
51-
_active = null;
52-
_activeMode = null;
55+
await client.StopAsync();
56+
UpdateState(mode, TransportState.Disconnected(mode.ToString().ToLowerInvariant(), "Failed to start"));
5357
return false;
5458
}
5559

56-
_active = next;
57-
_activeMode = mode;
60+
UpdateState(mode, client.State ?? TransportState.Connected(client.TransportName));
5861
return true;
5962
}
6063

61-
public async Task StopAsync()
64+
public async Task StopAsync(TransportMode? mode = null)
6265
{
63-
if (_active != null)
66+
async Task StopClient(IMcpTransportClient client, TransportMode clientMode)
67+
{
68+
if (client == null) return;
69+
try { await client.StopAsync(); }
70+
catch (Exception ex) { McpLog.Warn($"Error while stopping transport {client.TransportName}: {ex.Message}"); }
71+
finally { UpdateState(clientMode, TransportState.Disconnected(client.TransportName)); }
72+
}
73+
74+
if (mode == null)
75+
{
76+
await StopClient(_httpClient, TransportMode.Http);
77+
await StopClient(_stdioClient, TransportMode.Stdio);
78+
return;
79+
}
80+
81+
if (mode == TransportMode.Http)
82+
{
83+
await StopClient(_httpClient, TransportMode.Http);
84+
}
85+
else
6486
{
65-
try
66-
{
67-
await _active.StopAsync();
68-
}
69-
catch (Exception ex)
70-
{
71-
McpLog.Warn($"Error while stopping transport {_active.TransportName}: {ex.Message}");
72-
}
73-
finally
74-
{
75-
_active = null;
76-
_activeMode = null;
77-
}
87+
await StopClient(_stdioClient, TransportMode.Stdio);
7888
}
7989
}
8090

81-
public async Task<bool> VerifyAsync()
91+
public async Task<bool> VerifyAsync(TransportMode mode)
8292
{
83-
if (_active == null)
93+
IMcpTransportClient client = mode == TransportMode.Http ? _httpClient : _stdioClient;
94+
if (client == null)
8495
{
8596
return false;
8697
}
87-
return await _active.VerifyAsync();
98+
99+
bool ok = await client.VerifyAsync();
100+
var state = client.State ?? TransportState.Disconnected(client.TransportName, "No state reported");
101+
UpdateState(mode, state);
102+
return ok;
88103
}
89104

90-
public TransportState GetState()
105+
public TransportState GetState(TransportMode mode)
91106
{
92-
if (_active == null)
107+
return mode == TransportMode.Http ? _httpState : _stdioState;
108+
}
109+
110+
public bool IsRunning(TransportMode mode) => GetState(mode).IsConnected;
111+
112+
private void UpdateState(TransportMode mode, TransportState state)
113+
{
114+
if (mode == TransportMode.Http)
93115
{
94-
return TransportState.Disconnected(_activeMode?.ToString()?.ToLowerInvariant() ?? "unknown", "Transport not started");
116+
_httpState = state;
117+
}
118+
else
119+
{
120+
_stdioState = state;
95121
}
96-
97-
return _active.State ?? TransportState.Disconnected(_active.TransportName, "No state reported");
98122
}
99123
}
100124

0 commit comments

Comments
 (0)