Skip to content

Improve port probing and script edit tooling #19

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 1 commit into from
Aug 18, 2025
Merged

Improve port probing and script edit tooling #19

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
85 changes: 60 additions & 25 deletions UnityMcpBridge/Editor/Tools/ManageScript.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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 */ }
Expand Down Expand Up @@ -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)
Expand All @@ -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);
Comment on lines +565 to +567
Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

logic: The fallback logic here changes behavior - if TryComputeClassSpan fails, it now tries FindEnclosingClassSpan instead of failing immediately. This could mask legitimate errors where the specified class doesn't exist.

}
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" });
}
}
}
Expand Down Expand Up @@ -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; }

Expand Down Expand Up @@ -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
Expand Down
47 changes: 36 additions & 11 deletions UnityMcpBridge/UnityMcpServer~/src/port_discovery.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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)
Expand All @@ -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
Expand Down