diff --git a/UnityMcpBridge/Editor/Tools/ManageScript.cs b/UnityMcpBridge/Editor/Tools/ManageScript.cs index d92e6cb6..f22d4af9 100644 --- a/UnityMcpBridge/Editor/Tools/ManageScript.cs +++ b/UnityMcpBridge/Editor/Tools/ManageScript.cs @@ -90,6 +90,14 @@ private static bool TryResolveUnderAssets(string relDir, out string fullPathDir, relPathSafe = null; return false; } +#if NET6_0_OR_GREATER + if (!string.IsNullOrEmpty(di.LinkTarget)) + { + fullPathDir = null; + relPathSafe = null; + return false; + } +#endif } } catch { /* best effort; proceed */ } @@ -525,7 +533,11 @@ private static object ApplyTextEdits( int headerBoundary = 0; if (original.Length > 0 && original[0] == '\uFEFF') headerBoundary = 1; // skip BOM // Find first top-level using (very simple scan of start of file) - var mUsing = System.Text.RegularExpressions.Regex.Match(original, @"(?m)^(?:\uFEFF)?using\s+\w+", System.Text.RegularExpressions.RegexOptions.None); + var mUsing = System.Text.RegularExpressions.Regex.Match( + original, + @"(?m)^(?:\uFEFF)?(?:global\s+)?using(?:\s+static)?\b", + System.Text.RegularExpressions.RegexOptions.None + ); if (mUsing.Success) headerBoundary = Math.Min(Math.Max(headerBoundary, mUsing.Index), original.Length); foreach (var sp in spans) @@ -550,32 +562,34 @@ private static object ApplyTextEdits( { string methodName = mh.Groups[1].Value; // Find class span containing the edit - if (TryComputeClassSpan(original, name, null, out var clsStart, out var clsLen, out _)) + if (!TryComputeClassSpan(original, name, null, out var clsStart, out var clsLen, out _)) + { + FindEnclosingClassSpan(original, sp.start, out clsStart, out clsLen); + } + if (clsLen > 0 && + TryComputeMethodSpan(original, clsStart, clsLen, methodName, null, null, null, out var mStart, out var mLen, out _)) { - if (TryComputeMethodSpan(original, clsStart, clsLen, methodName, null, null, null, out var mStart, out var mLen, out _)) + // If the edit overlaps the method span significantly, treat as replace_method + if (sp.start <= mStart + 2 && sp.end >= mStart + 1) { - // If the edit overlaps the method span significantly, treat as replace_method - if (sp.start <= mStart + 2 && sp.end >= mStart + 1) + var methodOriginal = original.Substring(mStart, mLen); + int relStart = Math.Max(0, Math.Min(sp.start - mStart, methodOriginal.Length)); + int relEnd = Math.Max(relStart, Math.Min(sp.end - mStart, methodOriginal.Length)); + string replacementSnippet = methodOriginal + .Remove(relStart, relEnd - relStart) + .Insert(relStart, sp.text ?? string.Empty); + + var structEdits = new JArray(); + var op = new JObject { - var methodOriginal = original.Substring(mStart, mLen); - int relStart = Math.Max(0, Math.Min(sp.start - mStart, methodOriginal.Length)); - int relEnd = Math.Max(relStart, Math.Min(sp.end - mStart, methodOriginal.Length)); - string replacementSnippet = methodOriginal - .Remove(relStart, relEnd - relStart) - .Insert(relStart, sp.text ?? string.Empty); - - var structEdits = new JArray(); - var op = new JObject - { - ["mode"] = "replace_method", - ["className"] = name, - ["methodName"] = methodName, - ["replacement"] = replacementSnippet - }; - structEdits.Add(op); - // Reuse structured path - return EditScript(fullPath, relativePath, name, structEdits, new JObject{ ["refresh"] = "immediate", ["validate"] = "standard" }); - } + ["mode"] = "replace_method", + ["className"] = name, + ["methodName"] = methodName, + ["replacement"] = replacementSnippet + }; + structEdits.Add(op); + // Reuse structured path + return EditScript(fullPath, relativePath, name, structEdits, new JObject{ ["refresh"] = "immediate", ["validate"] = "standard" }); } } } @@ -737,7 +751,16 @@ private static bool CheckBalancedDelimiters(string text, out int line, out char char c = text[i]; char next = i + 1 < text.Length ? text[i + 1] : '\0'; - if (c == '\n') { line++; if (inSingle) inSingle = false; } + if (c == '\r') + { + // Treat CRLF as a single newline; skip the LF if present + if (next == '\n') { i++; } + line++; if (inSingle) inSingle = false; + } + else if (c == '\n') + { + line++; if (inSingle) inSingle = false; + } if (escape) { escape = false; continue; } @@ -1205,6 +1228,18 @@ private static bool ValidateClassSnippet(string snippet, string expectedName, ou #endif } + private static bool FindEnclosingClassSpan(string source, int index, out int start, out int length) + { + start = length = 0; + if (index < 0 || index > source.Length) return false; + var prefix = source.Substring(0, Math.Min(index, source.Length)); + var matches = Regex.Matches(prefix, @"(?m)\bclass\s+([A-Za-z_][A-Za-z0-9_]*)"); + if (matches.Count == 0) return false; + var m = matches[matches.Count - 1]; + var className = m.Groups[1].Value; + return TryComputeClassSpanBalanced(source, className, null, out start, out length, out _); + } + private static bool TryComputeClassSpan(string source, string className, string ns, out int start, out int length, out string why) { #if USE_ROSLYN diff --git a/UnityMcpBridge/UnityMcpServer~/src/port_discovery.py b/UnityMcpBridge/UnityMcpServer~/src/port_discovery.py index 070bde4d..828dd956 100644 --- a/UnityMcpBridge/UnityMcpServer~/src/port_discovery.py +++ b/UnityMcpBridge/UnityMcpServer~/src/port_discovery.py @@ -22,6 +22,21 @@ logger = logging.getLogger("unity-mcp-server") +FRAME_HEADER_SIZE = 8 +# Keep small; we're only looking for a tiny pong. 1 MiB is generous for probes. +MAX_FRAME_SIZE = 1 << 20 + + +# Module-level helper to avoid duplication and per-call redefinition +def _read_exact(sock: socket.socket, count: int) -> bytes: + buf = bytearray() + while len(buf) < count: + chunk = sock.recv(count - len(buf)) + if not chunk: + raise ConnectionError("Connection closed before reading expected bytes") + buf.extend(chunk) + return bytes(buf) + class PortDiscovery: """Handles port discovery from Unity Bridge registry""" REGISTRY_FILE = "unity-mcp-port.json" # legacy single-project file @@ -62,15 +77,6 @@ def _try_probe_unity_mcp(port: int) -> bool: pong are sent/received with an 8-byte big-endian length prefix. """ - def _read_exact(sock: socket.socket, count: int) -> bytes: - buf = bytearray() - while len(buf) < count: - chunk = sock.recv(count - len(buf)) - if not chunk: - raise ConnectionError("Connection closed before reading expected bytes") - buf.extend(chunk) - return bytes(buf) - try: with socket.create_connection(("127.0.0.1", port), PortDiscovery.CONNECT_TIMEOUT) as s: s.settimeout(PortDiscovery.CONNECT_TIMEOUT) @@ -81,12 +87,31 @@ def _read_exact(sock: socket.socket, count: int) -> bytes: if 'FRAMING=1' in text: header = struct.pack('>Q', len(payload)) s.sendall(header + payload) - resp_header = _read_exact(s, 8) + resp_header = _read_exact(s, FRAME_HEADER_SIZE) resp_len = struct.unpack('>Q', resp_header)[0] + # Defensive cap against unreasonable frame sizes + if resp_len > MAX_FRAME_SIZE: + return False data = _read_exact(s, resp_len) else: s.sendall(payload) - data = s.recv(512) + # Read a small bounded amount looking for pong + chunks = [] + total = 0 + data = b"" + while total < 1024: + try: + part = s.recv(512) + except socket.timeout: + break + if not part: + break + chunks.append(part) + total += len(part) + if b'"message":"pong"' in part: + break + if chunks: + data = b"".join(chunks) # Minimal validation: look for a success pong response if data and b'"message":"pong"' in data: return True