From e05b75fa63dc2b784b3b646d4da167eeed188913 Mon Sep 17 00:00:00 2001 From: Claude Date: Mon, 15 Jun 2026 05:23:15 +0000 Subject: [PATCH] fix(transport): guard liveness watchdog against long-running commands MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The receive loop awaits HandleExecuteAsync inline, so while a command runs it stops reading inbound frames and _lastInboundUtcTicks goes stale. A command longer than the liveness timeout (~37.5s) — test runs, asset imports, batch ops — would make the watchdog falsely abort the socket and kill the in-flight command. Track an in-flight command counter and skip the watchdog while a command is executing; refresh liveness on completion (before decrementing) so the gap until the next ReceiveAsync drains buffered server pings can't race into a false trip. A genuinely dead socket still surfaces via the result-send failure or the next faulted ReceiveAsync. Extract the timeout derivation and trip decision into pure static helpers and cover them with unit tests. --- .../Transports/WebSocketTransportClient.cs | 37 +++++- .../Services/WebSocketTransportClientTests.cs | 115 ++++++++++++++++++ 2 files changed, 148 insertions(+), 4 deletions(-) diff --git a/MCPForUnity/Editor/Services/Transport/Transports/WebSocketTransportClient.cs b/MCPForUnity/Editor/Services/Transport/Transports/WebSocketTransportClient.cs index a4b83e7fe..9a754fd00 100644 --- a/MCPForUnity/Editor/Services/Transport/Transports/WebSocketTransportClient.cs +++ b/MCPForUnity/Editor/Services/Transport/Transports/WebSocketTransportClient.cs @@ -63,6 +63,7 @@ public class WebSocketTransportClient : IMcpTransportClient, IDisposable private TimeSpan _socketKeepAliveInterval = DefaultKeepAliveInterval; private TimeSpan _inboundLivenessTimeout = DefaultInboundLivenessTimeout; private long _lastInboundUtcTicks; + private int _commandsInFlight; private volatile bool _isConnected; private int _isReconnectingFlag; private TransportState _state = TransportState.Disconnected(TransportDisplayName, "Transport not started"); @@ -528,9 +529,23 @@ private void ApplyWelcome(JObject payload) _socketKeepAliveInterval = TimeSpan.FromSeconds(safeSeconds); } - // Allow ~2.5 missed server keep-alives before declaring the link dead. - double livenessSeconds = Math.Max(30.0, _keepAliveInterval.TotalSeconds * 2.5); - _inboundLivenessTimeout = TimeSpan.FromSeconds(livenessSeconds); + _inboundLivenessTimeout = ComputeInboundLivenessTimeout(_keepAliveInterval); + } + + // Allow ~2.5 missed server keep-alives before declaring the link dead, with a 30s + // floor so a small server cadence can't produce a trigger-happy timeout. + private static TimeSpan ComputeInboundLivenessTimeout(TimeSpan keepAliveInterval) + { + double livenessSeconds = Math.Max(30.0, keepAliveInterval.TotalSeconds * 2.5); + return TimeSpan.FromSeconds(livenessSeconds); + } + + // The watchdog only trips on genuine link silence. While a command is in flight the + // receive loop is intentionally parked in HandleExecuteAsync rather than reading, so + // the resulting inbound silence is expected and must not be mistaken for a dead socket. + private static bool ShouldTripLivenessWatchdog(TimeSpan sinceInbound, TimeSpan livenessTimeout, int commandsInFlight) + { + return commandsInFlight == 0 && sinceInbound > livenessTimeout; } private async Task HandleRegisteredAsync(JObject payload, CancellationToken token) @@ -642,6 +657,10 @@ private async Task HandleExecuteAsync(JObject payload, CancellationToken token) }; string responseJson; + // Dispatch parks the receive loop here while the command runs, so flag it + // in-flight to keep the liveness watchdog from mistaking that expected + // silence for a dead socket on long-running commands (tests, imports, etc.). + Interlocked.Increment(ref _commandsInFlight); try { using var timeoutCts = CancellationTokenSource.CreateLinkedTokenSource(token); @@ -664,6 +683,14 @@ private async Task HandleExecuteAsync(JObject payload, CancellationToken token) error = ex.Message }); } + finally + { + // Refresh liveness before clearing the flag so the watchdog doesn't trip in + // the gap before the next ReceiveAsync drains the pings the server buffered + // during execution. Stamp first, then decrement. + Interlocked.Exchange(ref _lastInboundUtcTicks, DateTime.UtcNow.Ticks); + Interlocked.Decrement(ref _commandsInFlight); + } JToken resultToken; try @@ -707,8 +734,10 @@ private async Task KeepAliveLoopAsync(CancellationToken token) // TCP that ReceiveAsync never faults on) — abort and reconnect rather than // zombie. Abort() faults the parked ReceiveAsync; the reconnect is guarded // against double-firing by _isReconnectingFlag in HandleSocketClosureAsync. + // A command in flight is skipped: the receive loop is parked on dispatch, + // not on a dead socket, so its silence is expected. TimeSpan sinceInbound = TimeSpan.FromTicks(DateTime.UtcNow.Ticks - Interlocked.Read(ref _lastInboundUtcTicks)); - if (sinceInbound > _inboundLivenessTimeout) + if (ShouldTripLivenessWatchdog(sinceInbound, _inboundLivenessTimeout, Volatile.Read(ref _commandsInFlight))) { McpLog.Warn($"[WebSocket] No server traffic for {sinceInbound.TotalSeconds:0}s (liveness timeout {_inboundLivenessTimeout.TotalSeconds:0}s); aborting and reconnecting."); try { _socket.Abort(); } catch { } diff --git a/TestProjects/UnityMCPTests/Assets/Tests/EditMode/Services/WebSocketTransportClientTests.cs b/TestProjects/UnityMCPTests/Assets/Tests/EditMode/Services/WebSocketTransportClientTests.cs index 7e7797ec8..1eb85ccdb 100644 --- a/TestProjects/UnityMCPTests/Assets/Tests/EditMode/Services/WebSocketTransportClientTests.cs +++ b/TestProjects/UnityMCPTests/Assets/Tests/EditMode/Services/WebSocketTransportClientTests.cs @@ -82,6 +82,121 @@ public void BuildConnectionCandidateUris_LocalhostFallbacks_PreserveSchemePortPa } } + [Test] + public void ComputeInboundLivenessTimeout_ScalesKeepAliveByTwoAndAHalf() + { + TimeSpan result = InvokeComputeInboundLivenessTimeout(TimeSpan.FromSeconds(15)); + + // 15s * 2.5 = 37.5s, comfortably above the 30s floor. + Assert.AreEqual(37.5, result.TotalSeconds, 0.0001); + } + + [Test] + public void ComputeInboundLivenessTimeout_AppliesThirtySecondFloor() + { + // 5s * 2.5 = 12.5s, which the floor lifts to 30s. + TimeSpan small = InvokeComputeInboundLivenessTimeout(TimeSpan.FromSeconds(5)); + Assert.AreEqual(30.0, small.TotalSeconds, 0.0001); + + // A zero/degenerate cadence must still produce the floor, never zero. + TimeSpan zero = InvokeComputeInboundLivenessTimeout(TimeSpan.Zero); + Assert.AreEqual(30.0, zero.TotalSeconds, 0.0001); + } + + [Test] + public void ComputeInboundLivenessTimeout_LargeCadenceScalesAboveFloor() + { + // 60s * 2.5 = 150s, well above the floor. + TimeSpan result = InvokeComputeInboundLivenessTimeout(TimeSpan.FromSeconds(60)); + Assert.AreEqual(150.0, result.TotalSeconds, 0.0001); + } + + [Test] + public void ShouldTripLivenessWatchdog_SilencePastTimeoutWithNoCommand_Trips() + { + bool trip = InvokeShouldTripLivenessWatchdog( + TimeSpan.FromSeconds(50), TimeSpan.FromSeconds(38), commandsInFlight: 0); + + Assert.IsTrue(trip); + } + + [Test] + public void ShouldTripLivenessWatchdog_CommandInFlight_DoesNotTrip() + { + // The core false-trip guard: an in-flight command parks the receive loop, so the + // inbound silence is expected and must not be treated as a dead socket. + bool oneInFlight = InvokeShouldTripLivenessWatchdog( + TimeSpan.FromSeconds(50), TimeSpan.FromSeconds(38), commandsInFlight: 1); + Assert.IsFalse(oneInFlight); + + bool manyInFlight = InvokeShouldTripLivenessWatchdog( + TimeSpan.FromSeconds(5000), TimeSpan.FromSeconds(38), commandsInFlight: 3); + Assert.IsFalse(manyInFlight); + } + + [Test] + public void ShouldTripLivenessWatchdog_WithinTimeout_DoesNotTrip() + { + bool underTimeout = InvokeShouldTripLivenessWatchdog( + TimeSpan.FromSeconds(20), TimeSpan.FromSeconds(38), commandsInFlight: 0); + Assert.IsFalse(underTimeout); + + // Exactly at the timeout is not past it (strictly-greater comparison). + bool atTimeout = InvokeShouldTripLivenessWatchdog( + TimeSpan.FromSeconds(38), TimeSpan.FromSeconds(38), commandsInFlight: 0); + Assert.IsFalse(atTimeout); + } + + private static TimeSpan InvokeComputeInboundLivenessTimeout(TimeSpan keepAliveInterval) + { + MethodInfo method = ResolveStaticMethod("ComputeInboundLivenessTimeout", typeof(TimeSpan)); + if (method == null) + { + Assert.Fail("Expected private static ComputeInboundLivenessTimeout(TimeSpan) to exist."); + } + object result = method.Invoke(null, new object[] { keepAliveInterval }); + Assert.IsInstanceOf(result); + return (TimeSpan)result; + } + + private static bool InvokeShouldTripLivenessWatchdog(TimeSpan sinceInbound, TimeSpan livenessTimeout, int commandsInFlight) + { + MethodInfo method = ResolveStaticMethod( + "ShouldTripLivenessWatchdog", typeof(TimeSpan), typeof(TimeSpan), typeof(int)); + if (method == null) + { + Assert.Fail("Expected private static ShouldTripLivenessWatchdog(TimeSpan, TimeSpan, int) to exist."); + } + object result = method.Invoke(null, new object[] { sinceInbound, livenessTimeout, commandsInFlight }); + Assert.IsInstanceOf(result); + return (bool)result; + } + + private static MethodInfo ResolveStaticMethod(string name, params Type[] parameterTypes) + { + const BindingFlags flags = BindingFlags.NonPublic | BindingFlags.Public | BindingFlags.Static; + + MethodInfo direct = typeof(WebSocketTransportClient).GetMethod(name, flags, binder: null, types: parameterTypes, modifiers: null); + if (direct != null) + { + return direct; + } + + // Fallback across loaded assemblies, mirroring candidate-builder resolution for + // environments where multiple copies of the type may be loaded. + foreach (Assembly assembly in AppDomain.CurrentDomain.GetAssemblies()) + { + Type candidateType = assembly.GetType(WebSocketTransportClientTypeName); + MethodInfo method = candidateType?.GetMethod(name, flags, binder: null, types: parameterTypes, modifiers: null); + if (method != null) + { + return method; + } + } + + return null; + } + private static List InvokeBuildConnectionCandidateUris(Uri endpoint) { if (BuildConnectionCandidateUrisMethod == null)