diff --git a/README-DEV.md b/README-DEV.md index debcffc7..1df2a0cd 100644 --- a/README-DEV.md +++ b/README-DEV.md @@ -66,6 +66,59 @@ To find it reliably: Note: In recent builds, the Python server sources are also bundled inside the package under `UnityMcpServer~/src`. This is handy for local testing or pointing MCP clients directly at the packaged server. +## MCP Bridge Stress Test + +An on-demand stress utility exercises the MCP bridge with multiple concurrent clients while triggering real script reloads via immediate script edits (no menu calls required). + +### Script +- `tools/stress_mcp.py` + +### What it does +- Starts N TCP clients against the Unity MCP bridge (default port auto-discovered from `~/.unity-mcp/unity-mcp-status-*.json`). +- Sends lightweight framed `ping` keepalives to maintain concurrency. +- In parallel, appends a unique marker comment to a target C# file using `manage_script.apply_text_edits` with: + - `options.refresh = "immediate"` to force an import/compile immediately (triggers domain reload), and + - `precondition_sha256` computed from the current file contents to avoid drift. +- Uses EOF insertion to avoid header/`using`-guard edits. + +### Usage (local) +```bash +# Recommended: use the included large script in the test project +python3 tools/stress_mcp.py \ + --duration 60 \ + --clients 8 \ + --unity-file "TestProjects/UnityMCPTests/Assets/Scripts/LongUnityScriptClaudeTest.cs" +``` + +Flags: +- `--project` Unity project path (auto-detected to the included test project by default) +- `--unity-file` C# file to edit (defaults to the long test script) +- `--clients` number of concurrent clients (default 10) +- `--duration` seconds to run (default 60) + +### Expected outcome +- No Unity Editor crashes during reload churn +- Immediate reloads after each applied edit (no `Assets/Refresh` menu calls) +- Some transient disconnects or a few failed calls may occur during domain reload; the tool retries and continues +- JSON summary printed at the end, e.g.: + - `{"port": 6400, "stats": {"pings": 28566, "applies": 69, "disconnects": 0, "errors": 0}}` + +### Notes and troubleshooting +- Immediate vs debounced: + - The tool sets `options.refresh = "immediate"` so changes compile instantly. If you only need churn (not per-edit confirmation), switch to debounced to reduce mid-reload failures. +- Precondition required: + - `apply_text_edits` requires `precondition_sha256` on larger files. The tool reads the file first to compute the SHA. +- Edit location: + - To avoid header guards or complex ranges, the tool appends a one-line marker at EOF each cycle. +- Read API: + - The bridge currently supports `manage_script.read` for file reads. You may see a deprecation warning; it's harmless for this internal tool. +- Transient failures: + - Occasional `apply_errors` often indicate the connection reloaded mid-reply. Edits still typically apply; the loop continues on the next iteration. + +### CI guidance +- Keep this out of default PR CI due to Unity/editor requirements and runtime variability. +- Optionally run it as a manual workflow or nightly job on a Unity-capable runner. + ## CI Test Workflow (GitHub Actions) We provide a CI job to run a Natural Language Editing mini-suite against the Unity test project. It spins up a headless Unity container and connects via the MCP bridge. diff --git a/TestProjects/UnityMCPTests/Assets/Tests/EditMode/Tools/ManageGameObjectTests.cs b/TestProjects/UnityMCPTests/Assets/Tests/EditMode/Tools/ManageGameObjectTests.cs index db2b525d..34138999 100644 --- a/TestProjects/UnityMCPTests/Assets/Tests/EditMode/Tools/ManageGameObjectTests.cs +++ b/TestProjects/UnityMCPTests/Assets/Tests/EditMode/Tools/ManageGameObjectTests.cs @@ -261,6 +261,31 @@ public void SetComponentProperties_CollectsAllFailuresAndAppliesValidOnes() // The collect-and-continue behavior means we should get an error response // that contains info about the failed properties, but valid ones were still applied // This proves the collect-and-continue behavior is working + + // Harden: verify structured error response with failures list contains both invalid fields + var successProp = result.GetType().GetProperty("success"); + Assert.IsNotNull(successProp, "Result should expose 'success' property"); + Assert.IsFalse((bool)successProp.GetValue(result), "Result.success should be false for partial failure"); + + var dataProp = result.GetType().GetProperty("data"); + Assert.IsNotNull(dataProp, "Result should include 'data' with errors"); + var dataVal = dataProp.GetValue(result); + Assert.IsNotNull(dataVal, "Result.data should not be null"); + var errorsProp = dataVal.GetType().GetProperty("errors"); + Assert.IsNotNull(errorsProp, "Result.data should include 'errors' list"); + var errorsEnum = errorsProp.GetValue(dataVal) as System.Collections.IEnumerable; + Assert.IsNotNull(errorsEnum, "errors should be enumerable"); + + bool foundRotatoin = false; + bool foundInvalidProp = false; + foreach (var err in errorsEnum) + { + string s = err?.ToString() ?? string.Empty; + if (s.Contains("rotatoin")) foundRotatoin = true; + if (s.Contains("invalidProp")) foundInvalidProp = true; + } + Assert.IsTrue(foundRotatoin, "errors should mention the misspelled 'rotatoin' property"); + Assert.IsTrue(foundInvalidProp, "errors should mention the 'invalidProp' property"); } [Test] @@ -307,6 +332,28 @@ public void SetComponentProperties_ContinuesAfterException() // The key test: processing continued after the exception and set useGravity // This proves the collect-and-continue behavior works even with exceptions + + // Harden: verify structured error response contains velocity failure + var successProp2 = result.GetType().GetProperty("success"); + Assert.IsNotNull(successProp2, "Result should expose 'success' property"); + Assert.IsFalse((bool)successProp2.GetValue(result), "Result.success should be false when an exception occurs for a property"); + + var dataProp2 = result.GetType().GetProperty("data"); + Assert.IsNotNull(dataProp2, "Result should include 'data' with errors"); + var dataVal2 = dataProp2.GetValue(result); + Assert.IsNotNull(dataVal2, "Result.data should not be null"); + var errorsProp2 = dataVal2.GetType().GetProperty("errors"); + Assert.IsNotNull(errorsProp2, "Result.data should include 'errors' list"); + var errorsEnum2 = errorsProp2.GetValue(dataVal2) as System.Collections.IEnumerable; + Assert.IsNotNull(errorsEnum2, "errors should be enumerable"); + + bool foundVelocityError = false; + foreach (var err in errorsEnum2) + { + string s = err?.ToString() ?? string.Empty; + if (s.Contains("velocity")) { foundVelocityError = true; break; } + } + Assert.IsTrue(foundVelocityError, "errors should include a message referencing 'velocity'"); } } } \ No newline at end of file diff --git a/UnityMcpBridge/Editor/MCPForUnityBridge.cs b/UnityMcpBridge/Editor/MCPForUnityBridge.cs index 1a175847..0fadce31 100644 --- a/UnityMcpBridge/Editor/MCPForUnityBridge.cs +++ b/UnityMcpBridge/Editor/MCPForUnityBridge.cs @@ -23,6 +23,11 @@ public static partial class MCPForUnityBridge private static bool isRunning = false; private static readonly object lockObj = new(); private static readonly object startStopLock = new(); + private static readonly object clientsLock = new(); + private static readonly System.Collections.Generic.HashSet activeClients = new(); + private static CancellationTokenSource cts; + private static Task listenerTask; + private static int processingCommands = 0; private static bool initScheduled = false; private static bool ensureUpdateHooked = false; private static bool isStarting = false; @@ -193,9 +198,15 @@ private static void EnsureStartedOnEditorIdle() } isStarting = true; - // Attempt start; if it succeeds, remove the hook to avoid overhead - Start(); - isStarting = false; + try + { + // Attempt start; if it succeeds, remove the hook to avoid overhead + Start(); + } + finally + { + isStarting = false; + } if (isRunning) { EditorApplication.update -= EnsureStartedOnEditorIdle; @@ -319,8 +330,17 @@ public static void Start() string platform = Application.platform.ToString(); string serverVer = ReadInstalledServerVersionSafe(); Debug.Log($"MCP-FOR-UNITY: MCPForUnityBridge started on port {currentUnityPort}. (OS={platform}, server={serverVer})"); - Task.Run(ListenerLoop); + // Start background listener with cooperative cancellation + cts = new CancellationTokenSource(); + listenerTask = Task.Run(() => ListenerLoopAsync(cts.Token)); EditorApplication.update += ProcessCommands; + // Ensure lifecycle events are (re)subscribed in case Stop() removed them earlier in-domain + try { AssemblyReloadEvents.beforeAssemblyReload -= OnBeforeAssemblyReload; } catch { } + try { AssemblyReloadEvents.beforeAssemblyReload += OnBeforeAssemblyReload; } catch { } + try { AssemblyReloadEvents.afterAssemblyReload -= OnAfterAssemblyReload; } catch { } + try { AssemblyReloadEvents.afterAssemblyReload += OnAfterAssemblyReload; } catch { } + try { EditorApplication.quitting -= Stop; } catch { } + try { EditorApplication.quitting += Stop; } catch { } // Write initial heartbeat immediately heartbeatSeq++; WriteHeartbeat(false, "ready"); @@ -335,6 +355,7 @@ public static void Start() public static void Stop() { + Task toWait = null; lock (startStopLock) { if (!isRunning) @@ -346,23 +367,55 @@ public static void Stop() { // Mark as stopping early to avoid accept logging during disposal isRunning = false; - // Mark heartbeat one last time before stopping - WriteHeartbeat(false, "stopped"); - listener?.Stop(); + + // Quiesce background listener quickly + var cancel = cts; + cts = null; + try { cancel?.Cancel(); } catch { } + + try { listener?.Stop(); } catch { } listener = null; - EditorApplication.update -= ProcessCommands; - if (IsDebugEnabled()) Debug.Log("MCP-FOR-UNITY: MCPForUnityBridge stopped."); + + // Capture background task to wait briefly outside the lock + toWait = listenerTask; + listenerTask = null; } catch (Exception ex) { Debug.LogError($"Error stopping MCPForUnityBridge: {ex.Message}"); } } + + // Proactively close all active client sockets to unblock any pending reads + TcpClient[] toClose; + lock (clientsLock) + { + toClose = activeClients.ToArray(); + activeClients.Clear(); + } + foreach (var c in toClose) + { + try { c.Close(); } catch { } + } + + // Give the background loop a short window to exit without blocking the editor + if (toWait != null) + { + try { toWait.Wait(100); } catch { } + } + + // Now unhook editor events safely + try { EditorApplication.update -= ProcessCommands; } catch { } + try { AssemblyReloadEvents.beforeAssemblyReload -= OnBeforeAssemblyReload; } catch { } + try { AssemblyReloadEvents.afterAssemblyReload -= OnAfterAssemblyReload; } catch { } + try { EditorApplication.quitting -= Stop; } catch { } + + if (IsDebugEnabled()) Debug.Log("MCP-FOR-UNITY: MCPForUnityBridge stopped."); } - private static async Task ListenerLoop() + private static async Task ListenerLoopAsync(CancellationToken token) { - while (isRunning) + while (isRunning && !token.IsCancellationRequested) { try { @@ -378,19 +431,23 @@ private static async Task ListenerLoop() client.ReceiveTimeout = 60000; // 60 seconds // Fire and forget each client connection - _ = HandleClientAsync(client); + _ = Task.Run(() => HandleClientAsync(client, token), token); } catch (ObjectDisposedException) { // Listener was disposed during stop/reload; exit quietly - if (!isRunning) + if (!isRunning || token.IsCancellationRequested) { break; } } + catch (OperationCanceledException) + { + break; + } catch (Exception ex) { - if (isRunning) + if (isRunning && !token.IsCancellationRequested) { if (IsDebugEnabled()) Debug.LogError($"Listener error: {ex.Message}"); } @@ -398,11 +455,14 @@ private static async Task ListenerLoop() } } - private static async Task HandleClientAsync(TcpClient client) + private static async Task HandleClientAsync(TcpClient client, CancellationToken token) { using (client) using (NetworkStream stream = client.GetStream()) { + lock (clientsLock) { activeClients.Add(client); } + try + { // Framed I/O only; legacy mode removed try { @@ -437,12 +497,12 @@ private static async Task HandleClientAsync(TcpClient client) return; // abort this client } - while (isRunning) + while (isRunning && !token.IsCancellationRequested) { try { // Strict framed mode only: enforced framed I/O for this connection - string commandText = await ReadFrameAsUtf8Async(stream, FrameIOTimeoutMs); + string commandText = await ReadFrameAsUtf8Async(stream, FrameIOTimeoutMs, token).ConfigureAwait(false); try { @@ -454,7 +514,7 @@ private static async Task HandleClientAsync(TcpClient client) } catch { } string commandId = Guid.NewGuid().ToString(); - TaskCompletionSource tcs = new(); + var tcs = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously); // Special handling for ping command to avoid JSON parsing if (commandText.Trim() == "ping") @@ -473,7 +533,7 @@ private static async Task HandleClientAsync(TcpClient client) commandQueue[commandId] = (commandText, tcs); } - string response = await tcs.Task; + string response = await tcs.Task.ConfigureAwait(false); byte[] responseBytes = System.Text.Encoding.UTF8.GetBytes(response); await WriteFrameAsync(stream, responseBytes); } @@ -496,6 +556,11 @@ private static async Task HandleClientAsync(TcpClient client) break; } } + } + finally + { + lock (clientsLock) { activeClients.Remove(client); } + } } } @@ -574,9 +639,9 @@ private static async System.Threading.Tasks.Task WriteFrameAsync(NetworkStream s #endif } - private static async System.Threading.Tasks.Task ReadFrameAsUtf8Async(NetworkStream stream, int timeoutMs) + private static async System.Threading.Tasks.Task ReadFrameAsUtf8Async(NetworkStream stream, int timeoutMs, CancellationToken cancel) { - byte[] header = await ReadExactAsync(stream, 8, timeoutMs); + byte[] header = await ReadExactAsync(stream, 8, timeoutMs, cancel).ConfigureAwait(false); ulong payloadLen = ReadUInt64BigEndian(header); if (payloadLen > MaxFrameBytes) { @@ -589,7 +654,7 @@ private static async System.Threading.Tasks.Task ReadFrameAsUtf8Async(Ne throw new System.IO.IOException("Frame too large for buffer"); } int count = (int)payloadLen; - byte[] payload = await ReadExactAsync(stream, count, timeoutMs); + byte[] payload = await ReadExactAsync(stream, count, timeoutMs, cancel).ConfigureAwait(false); return System.Text.Encoding.UTF8.GetString(payload); } @@ -624,6 +689,10 @@ private static void WriteUInt64BigEndian(byte[] dest, ulong value) private static void ProcessCommands() { + if (!isRunning) return; + if (Interlocked.Exchange(ref processingCommands, 1) == 1) return; // reentrancy guard + try + { // Heartbeat without holding the queue lock double now = EditorApplication.timeSinceStartup; if (now >= nextHeartbeatAt) @@ -734,6 +803,11 @@ private static void ProcessCommands() // Remove quickly under lock lock (lockObj) { commandQueue.Remove(id); } } + } + finally + { + Interlocked.Exchange(ref processingCommands, 0); + } } // Helper method to check if a string is valid JSON @@ -865,8 +939,7 @@ private static void OnBeforeAssemblyReload() { // Stop cleanly before reload so sockets close and clients see 'reloading' try { Stop(); } catch { } - WriteHeartbeat(true, "reloading"); - LogBreadcrumb("Reload"); + // Avoid file I/O or heavy work here } private static void OnAfterAssemblyReload() diff --git a/UnityMcpBridge/Editor/Tools/ExecuteMenuItem.cs b/UnityMcpBridge/Editor/Tools/ExecuteMenuItem.cs index 5adb476d..306cf8d3 100644 --- a/UnityMcpBridge/Editor/Tools/ExecuteMenuItem.cs +++ b/UnityMcpBridge/Editor/Tools/ExecuteMenuItem.cs @@ -27,7 +27,7 @@ public static class ExecuteMenuItem /// public static object HandleCommand(JObject @params) { - string action = @params["action"]?.ToString().ToLower() ?? "execute"; // Default action + string action = (@params["action"]?.ToString())?.ToLowerInvariant() ?? "execute"; // Default action try { @@ -96,14 +96,15 @@ private static object ExecuteItem(JObject @params) try { - // Trace incoming execute requests - Debug.Log($"[ExecuteMenuItem] Request to execute menu: '{menuPath}'"); + // Trace incoming execute requests (debug-gated) + McpLog.Info($"[ExecuteMenuItem] Request to execute menu: '{menuPath}'", always: false); // Execute synchronously. This code runs on the Editor main thread in our bridge path. bool executed = EditorApplication.ExecuteMenuItem(menuPath); if (executed) { - Debug.Log($"[ExecuteMenuItem] Executed successfully: '{menuPath}'"); + // Success trace (debug-gated) + McpLog.Info($"[ExecuteMenuItem] Executed successfully: '{menuPath}'", always: false); return Response.Success( $"Executed menu item: '{menuPath}'", new { executed = true, menuPath } diff --git a/UnityMcpBridge/Editor/Tools/ManageGameObject.cs b/UnityMcpBridge/Editor/Tools/ManageGameObject.cs index 0e19382b..c3357ed9 100644 --- a/UnityMcpBridge/Editor/Tools/ManageGameObject.cs +++ b/UnityMcpBridge/Editor/Tools/ManageGameObject.cs @@ -814,9 +814,34 @@ string searchMethod // Return component errors if any occurred (after processing all components) if (componentErrors.Count > 0) { + // Aggregate flattened error strings to make tests/API assertions simpler + var aggregatedErrors = new System.Collections.Generic.List(); + foreach (var errorObj in componentErrors) + { + try + { + var dataProp = errorObj?.GetType().GetProperty("data"); + var dataVal = dataProp?.GetValue(errorObj); + if (dataVal != null) + { + var errorsProp = dataVal.GetType().GetProperty("errors"); + var errorsEnum = errorsProp?.GetValue(dataVal) as System.Collections.IEnumerable; + if (errorsEnum != null) + { + foreach (var item in errorsEnum) + { + var s = item?.ToString(); + if (!string.IsNullOrEmpty(s)) aggregatedErrors.Add(s); + } + } + } + } + catch { } + } + return Response.Error( $"One or more component property operations failed on '{targetGo.name}'.", - new { componentErrors = componentErrors } + new { componentErrors = componentErrors, errors = aggregatedErrors } ); } diff --git a/tools/stress_mcp.py b/tools/stress_mcp.py new file mode 100644 index 00000000..bd14c35a --- /dev/null +++ b/tools/stress_mcp.py @@ -0,0 +1,333 @@ +#!/usr/bin/env python3 +import asyncio +import argparse +import json +import os +import struct +import time +from pathlib import Path +import random +import sys + + +TIMEOUT = float(os.environ.get("MCP_STRESS_TIMEOUT", "2.0")) +DEBUG = os.environ.get("MCP_STRESS_DEBUG", "").lower() in ("1", "true", "yes") + + +def dlog(*args): + if DEBUG: + print(*args, file=sys.stderr) + + +def find_status_files() -> list[Path]: + home = Path.home() + status_dir = Path(os.environ.get("UNITY_MCP_STATUS_DIR", home / ".unity-mcp")) + if not status_dir.exists(): + return [] + return sorted(status_dir.glob("unity-mcp-status-*.json"), key=lambda p: p.stat().st_mtime, reverse=True) + + +def discover_port(project_path: str | None) -> int: + # Default bridge port if nothing found + default_port = 6400 + files = find_status_files() + for f in files: + try: + data = json.loads(f.read_text()) + port = int(data.get("unity_port", 0) or 0) + proj = data.get("project_path") or "" + if project_path: + # Match status for the given project if possible + if proj and project_path in proj: + if 0 < port < 65536: + return port + else: + if 0 < port < 65536: + return port + except Exception: + pass + return default_port + + +async def read_exact(reader: asyncio.StreamReader, n: int) -> bytes: + buf = b"" + while len(buf) < n: + chunk = await reader.read(n - len(buf)) + if not chunk: + raise ConnectionError("Connection closed while reading") + buf += chunk + return buf + + +async def read_frame(reader: asyncio.StreamReader) -> bytes: + header = await read_exact(reader, 8) + (length,) = struct.unpack(">Q", header) + if length <= 0 or length > (64 * 1024 * 1024): + raise ValueError(f"Invalid frame length: {length}") + return await read_exact(reader, length) + + +async def write_frame(writer: asyncio.StreamWriter, payload: bytes) -> None: + header = struct.pack(">Q", len(payload)) + writer.write(header) + writer.write(payload) + await asyncio.wait_for(writer.drain(), timeout=TIMEOUT) + + +async def do_handshake(reader: asyncio.StreamReader) -> None: + # Server sends a single line handshake: "WELCOME UNITY-MCP 1 FRAMING=1\n" + line = await reader.readline() + if not line or b"WELCOME UNITY-MCP" not in line: + raise ConnectionError(f"Unexpected handshake from server: {line!r}") + + +def make_ping_frame() -> bytes: + return b"ping" + + +def make_execute_menu_item(menu_path: str) -> bytes: + # Retained for manual debugging; not used in normal stress runs + payload = {"type": "execute_menu_item", "params": {"action": "execute", "menu_path": menu_path}} + return json.dumps(payload).encode("utf-8") + + +async def client_loop(idx: int, host: str, port: int, stop_time: float, stats: dict): + reconnect_delay = 0.2 + while time.time() < stop_time: + writer = None + try: + # slight stagger to prevent burst synchronization across clients + await asyncio.sleep(0.003 * (idx % 11)) + reader, writer = await asyncio.wait_for(asyncio.open_connection(host, port), timeout=TIMEOUT) + await asyncio.wait_for(do_handshake(reader), timeout=TIMEOUT) + # Send a quick ping first + await write_frame(writer, make_ping_frame()) + _ = await asyncio.wait_for(read_frame(reader), timeout=TIMEOUT) # ignore content + + # Main activity loop (keep-alive + light load). Edit spam handled by reload_churn_task. + while time.time() < stop_time: + # Ping-only; edits are sent via reload_churn_task to avoid console spam + await write_frame(writer, make_ping_frame()) + _ = await asyncio.wait_for(read_frame(reader), timeout=TIMEOUT) + stats["pings"] += 1 + await asyncio.sleep(0.02 + random.uniform(-0.003, 0.003)) + + except (ConnectionError, OSError, asyncio.IncompleteReadError, asyncio.TimeoutError): + stats["disconnects"] += 1 + dlog(f"[client {idx}] disconnect/backoff {reconnect_delay}s") + await asyncio.sleep(reconnect_delay) + reconnect_delay = min(reconnect_delay * 1.5, 2.0) + continue + except Exception: + stats["errors"] += 1 + dlog(f"[client {idx}] unexpected error") + await asyncio.sleep(0.2) + continue + finally: + if writer is not None: + try: + writer.close() + await writer.wait_closed() + except Exception: + pass + + +async def reload_churn_task(project_path: str, stop_time: float, unity_file: str | None, host: str, port: int, stats: dict, storm_count: int = 1): + # Use script edit tool to touch a C# file, which triggers compilation reliably + path = Path(unity_file) if unity_file else None + seq = 0 + proj_root = Path(project_path).resolve() if project_path else None + # Build candidate list for storm mode + candidates: list[Path] = [] + if proj_root: + try: + for p in (proj_root / "Assets").rglob("*.cs"): + candidates.append(p.resolve()) + except Exception: + candidates = [] + if path and path.exists(): + rp = path.resolve() + if rp not in candidates: + candidates.append(rp) + while time.time() < stop_time: + try: + if path and path.exists(): + # Determine files to touch this cycle + targets: list[Path] + if storm_count and storm_count > 1 and candidates: + k = min(max(1, storm_count), len(candidates)) + targets = random.sample(candidates, k) + else: + targets = [path] + + for tpath in targets: + # Build a tiny ApplyTextEdits request that toggles a trailing comment + relative = None + try: + # Derive Unity-relative path under Assets/ (cross-platform) + resolved = tpath.resolve() + parts = list(resolved.parts) + if "Assets" in parts: + i = parts.index("Assets") + relative = Path(*parts[i:]).as_posix() + elif proj_root and str(resolved).startswith(str(proj_root)): + rel = resolved.relative_to(proj_root) + parts2 = list(rel.parts) + if "Assets" in parts2: + i2 = parts2.index("Assets") + relative = Path(*parts2[i2:]).as_posix() + except Exception: + relative = None + + if relative: + # Derive name and directory for ManageScript and compute precondition SHA + EOF position + name_base = Path(relative).stem + dir_path = str(Path(relative).parent).replace('\\', '/') + + # 1) Read current contents via manage_script.read to compute SHA and true EOF location + contents = None + read_success = False + for attempt in range(3): + writer = None + try: + reader, writer = await asyncio.wait_for(asyncio.open_connection(host, port), timeout=TIMEOUT) + await asyncio.wait_for(do_handshake(reader), timeout=TIMEOUT) + read_payload = { + "type": "manage_script", + "params": { + "action": "read", + "name": name_base, + "path": dir_path + } + } + await write_frame(writer, json.dumps(read_payload).encode("utf-8")) + resp = await asyncio.wait_for(read_frame(reader), timeout=TIMEOUT) + + read_obj = json.loads(resp.decode("utf-8", errors="ignore")) + result = read_obj.get("result", read_obj) if isinstance(read_obj, dict) else {} + if result.get("success"): + data_obj = result.get("data", {}) + contents = data_obj.get("contents") or "" + read_success = True + break + except Exception: + # retry with backoff + await asyncio.sleep(0.2 * (2 ** attempt) + random.uniform(0.0, 0.1)) + finally: + if 'writer' in locals() and writer is not None: + try: + writer.close() + await writer.wait_closed() + except Exception: + pass + + if not read_success or contents is None: + stats["apply_errors"] = stats.get("apply_errors", 0) + 1 + await asyncio.sleep(0.5) + continue + + # Compute SHA and EOF insertion point + import hashlib + sha = hashlib.sha256(contents.encode("utf-8")).hexdigest() + lines = contents.splitlines(keepends=True) + # Insert at true EOF (safe against header guards) + end_line = len(lines) + 1 # 1-based exclusive end + end_col = 1 + + # Build a unique marker append; ensure it begins with a newline if needed + marker = f"// MCP_STRESS seq={seq} time={int(time.time())}" + seq += 1 + insert_text = ("\n" if not contents.endswith("\n") else "") + marker + "\n" + + # 2) Apply text edits with immediate refresh and precondition + apply_payload = { + "type": "manage_script", + "params": { + "action": "apply_text_edits", + "name": name_base, + "path": dir_path, + "edits": [ + { + "startLine": end_line, + "startCol": end_col, + "endLine": end_line, + "endCol": end_col, + "newText": insert_text + } + ], + "precondition_sha256": sha, + "options": {"refresh": "immediate", "validate": "standard"} + } + } + + apply_success = False + for attempt in range(3): + writer = None + try: + reader, writer = await asyncio.wait_for(asyncio.open_connection(host, port), timeout=TIMEOUT) + await asyncio.wait_for(do_handshake(reader), timeout=TIMEOUT) + await write_frame(writer, json.dumps(apply_payload).encode("utf-8")) + resp = await asyncio.wait_for(read_frame(reader), timeout=TIMEOUT) + try: + data = json.loads(resp.decode("utf-8", errors="ignore")) + result = data.get("result", data) if isinstance(data, dict) else {} + ok = bool(result.get("success", False)) + if ok: + stats["applies"] = stats.get("applies", 0) + 1 + apply_success = True + break + except Exception: + # fall through to retry + pass + except Exception: + # retry with backoff + await asyncio.sleep(0.2 * (2 ** attempt) + random.uniform(0.0, 0.1)) + finally: + if 'writer' in locals() and writer is not None: + try: + writer.close() + await writer.wait_closed() + except Exception: + pass + if not apply_success: + stats["apply_errors"] = stats.get("apply_errors", 0) + 1 + + except Exception: + pass + await asyncio.sleep(1.0) + + +async def main(): + ap = argparse.ArgumentParser(description="Stress test the Unity MCP bridge with concurrent clients and reload churn") + ap.add_argument("--host", default="127.0.0.1") + ap.add_argument("--project", default=str(Path(__file__).resolve().parents[1] / "TestProjects" / "UnityMCPTests")) + ap.add_argument("--unity-file", default=str(Path(__file__).resolve().parents[1] / "TestProjects" / "UnityMCPTests" / "Assets" / "Scripts" / "LongUnityScriptClaudeTest.cs")) + ap.add_argument("--clients", type=int, default=10) + ap.add_argument("--duration", type=int, default=60) + ap.add_argument("--storm-count", type=int, default=1, help="Number of scripts to touch each cycle") + args = ap.parse_args() + + port = discover_port(args.project) + stop_time = time.time() + max(10, args.duration) + + stats = {"pings": 0, "menus": 0, "mods": 0, "disconnects": 0, "errors": 0} + tasks = [] + + # Spawn clients + for i in range(max(1, args.clients)): + tasks.append(asyncio.create_task(client_loop(i, args.host, port, stop_time, stats))) + + # Spawn reload churn task + tasks.append(asyncio.create_task(reload_churn_task(args.project, stop_time, args.unity_file, args.host, port, stats, storm_count=args.storm_count))) + + await asyncio.gather(*tasks, return_exceptions=True) + print(json.dumps({"port": port, "stats": stats}, indent=2)) + + +if __name__ == "__main__": + try: + asyncio.run(main()) + except KeyboardInterrupt: + pass + +