fix(transport): guard liveness watchdog against long-running commands + tests

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 { }

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<TimeSpan>(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<bool>(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<Uri> InvokeBuildConnectionCandidateUris(Uri endpoint)
         {
             if (BuildConnectionCandidateUrisMethod == null)