Merge pull request #8 from dsarno/codex/streamline-protocol-framing-in-c#-and-python

Author: dsarno

UnityMcpBridge/Editor/UnityMcpBridge.cs

@@ -431,15 +431,7 @@ private static async Task HandleClientAsync(TcpClient client)
                         if (true)
                         {
                             // Enforced framed mode for this connection
-                            byte[] header = await ReadExactAsync(stream, 8, FrameIOTimeoutMs);
-                            ulong payloadLen = ReadUInt64BigEndian(header);
-                            if (payloadLen == 0UL || payloadLen > MaxFrameBytes)
-                            {
-                                throw new System.IO.IOException($"Invalid framed length: {payloadLen}");
-                            }
-                            int payloadLenInt = checked((int)payloadLen);
-                            byte[] payload = await ReadExactAsync(stream, payloadLenInt, FrameIOTimeoutMs);
-                            commandText = System.Text.Encoding.UTF8.GetString(payload);
+                            commandText = await ReadFrameAsUtf8Async(stream, FrameIOTimeoutMs);
                         }
 
                         try
@@ -459,16 +451,7 @@ private static async Task HandleClientAsync(TcpClient client)
                                 /*lang=json,strict*/
                                 "{\"status\":\"success\",\"result\":{\"message\":\"pong\"}}"
                             );
-                            if ((ulong)pingResponseBytes.Length > MaxFrameBytes)
-                            {
-                                throw new System.IO.IOException($"Frame too large: {pingResponseBytes.Length}");
-                            }
-                            {
-                                byte[] outHeader = new byte[8];
-                                WriteUInt64BigEndian(outHeader, (ulong)pingResponseBytes.Length);
-                                await stream.WriteAsync(outHeader, 0, outHeader.Length);
-                            }
-                            await stream.WriteAsync(pingResponseBytes, 0, pingResponseBytes.Length);
+                            await WriteFrameAsync(stream, pingResponseBytes);
                             continue;
                         }
 
@@ -479,16 +462,7 @@ private static async Task HandleClientAsync(TcpClient client)
 
                         string response = await tcs.Task;
                         byte[] responseBytes = System.Text.Encoding.UTF8.GetBytes(response);
-                        if ((ulong)responseBytes.Length > MaxFrameBytes)
-                        {
-                            throw new System.IO.IOException($"Frame too large: {responseBytes.Length}");
-                        }
-                        {
-                            byte[] outHeader = new byte[8];
-                            WriteUInt64BigEndian(outHeader, (ulong)responseBytes.Length);
-                            await stream.WriteAsync(outHeader, 0, outHeader.Length);
-                        }
-                        await stream.WriteAsync(responseBytes, 0, responseBytes.Length);
+                        await WriteFrameAsync(stream, responseBytes);
                     }
                     catch (Exception ex)
                     {
@@ -499,22 +473,6 @@ private static async Task HandleClientAsync(TcpClient client)
             }
         }
 
-        private static async System.Threading.Tasks.Task<byte[]> ReadExactAsync(NetworkStream stream, int count)
-        {
-            byte[] data = new byte[count];
-            int offset = 0;
-            while (offset < count)
-            {
-                int r = await stream.ReadAsync(data, offset, count - offset);
-                if (r == 0)
-                {
-                    throw new System.IO.IOException("Connection closed before reading expected bytes");
-                }
-                offset += r;
-            }
-            return data;
-        }
-
         // Timeout-aware exact read helper; avoids indefinite stalls
         private static async System.Threading.Tasks.Task<byte[]> ReadExactAsync(NetworkStream stream, int count, int timeoutMs)
         {
@@ -538,6 +496,35 @@ private static async System.Threading.Tasks.Task<byte[]> ReadExactAsync(NetworkS
             return data;
         }
 
+        private static async System.Threading.Tasks.Task WriteFrameAsync(NetworkStream stream, byte[] payload)
+        {
+            if ((ulong)payload.LongLength > MaxFrameBytes)
+            {
+                throw new System.IO.IOException($"Frame too large: {payload.LongLength}");
+            }
+            byte[] header = new byte[8];
+            WriteUInt64BigEndian(header, (ulong)payload.LongLength);
+            await stream.WriteAsync(header, 0, header.Length);
+            await stream.WriteAsync(payload, 0, payload.Length);
+        }
+
+        private static async System.Threading.Tasks.Task<string> ReadFrameAsUtf8Async(NetworkStream stream, int timeoutMs)
+        {
+            byte[] header = await ReadExactAsync(stream, 8, timeoutMs);
+            ulong payloadLen = ReadUInt64BigEndian(header);
+            if (payloadLen == 0UL || payloadLen > MaxFrameBytes)
+            {
+                throw new System.IO.IOException($"Invalid framed length: {payloadLen}");
+            }
+            if (payloadLen > int.MaxValue)
+            {
+                throw new System.IO.IOException("Frame too large for buffer");
+            }
+            int count = (int)payloadLen;
+            byte[] payload = await ReadExactAsync(stream, count, timeoutMs);
+            return System.Text.Encoding.UTF8.GetString(payload);
+        }
+
         private static ulong ReadUInt64BigEndian(byte[] buffer)
         {
             if (buffer == null || buffer.Length < 8) return 0UL;

UnityMcpBridge/UnityMcpServer~/src/server.py

@@ -9,6 +9,8 @@
 from tools import register_all_tools
 from unity_connection import get_unity_connection, UnityConnection
 from pathlib import Path
+import os
+import hashlib
 
 # Configure logging: strictly stderr/file only (never stdout)
 stderr_handler = logging.StreamHandler(stream=sys.stderr)
@@ -98,52 +100,44 @@ def asset_creation_strategy() -> str:
 class _:
     pass
 
-import os
-import hashlib
-
-def _unity_assets_root() -> str:
-    # Heuristic: from the Unity project root (one level up from Library/ProjectSettings), 'Assets'
-    # Here, assume server runs from repo; let clients pass absolute paths under project too.
-    return None
+PROJECT_ROOT = Path(os.environ.get("UNITY_PROJECT_ROOT", Path.cwd())).resolve()
+ASSETS_ROOT = (PROJECT_ROOT / "Assets").resolve()
 
-def _safe_path(uri: str) -> str | None:
-    # URIs: unity://path/Assets/... or file:///absolute
+def _resolve_safe_path_from_uri(uri: str) -> Path | None:
+    raw: str | None = None
     if uri.startswith("unity://path/"):
-        p = uri[len("unity://path/"):]
-        return p
-    if uri.startswith("file://"):
-        return uri[len("file://"):]
-    # Minimal tolerance for plain Assets/... paths
-    if uri.startswith("Assets/"):
-        return uri
-    return None
+        raw = uri[len("unity://path/"):]
+    elif uri.startswith("file://"):
+        raw = uri[len("file://"):]
+    elif uri.startswith("Assets/"):
+        raw = uri
+    if raw is None:
+        return None
+    p = (PROJECT_ROOT / raw).resolve()
+    try:
+        p.relative_to(PROJECT_ROOT)
+    except ValueError:
+        return None
+    return p
 
 @mcp.resource.list()
 def list_resources(ctx: Context) -> list[dict]:
-    # Lightweight: expose only C# under Assets by default
     assets = []
     try:
-        root = os.getcwd()
-        for base, _, files in os.walk(os.path.join(root, "Assets")):
-            for f in files:
-                if f.endswith(".cs"):
-                    rel = os.path.relpath(os.path.join(base, f), root).replace("\\", "/")
-                    assets.append({
-                        "uri": f"unity://path/{rel}",
-                        "name": os.path.basename(rel)
-                    })
+        for p in ASSETS_ROOT.rglob("*.cs"):
+            rel = p.relative_to(PROJECT_ROOT).as_posix()
+            assets.append({"uri": f"unity://path/{rel}", "name": p.name})
     except Exception:
         pass
     return assets
 
 @mcp.resource.read()
 def read_resource(ctx: Context, uri: str) -> dict:
-    path = _safe_path(uri)
-    if not path or not os.path.exists(path):
+    p = _resolve_safe_path_from_uri(uri)
+    if not p or not p.exists():
         return {"mimeType": "text/plain", "text": f"Resource not found: {uri}"}
     try:
-        with open(path, "r", encoding="utf-8") as f:
-            text = f.read()
+        text = p.read_text(encoding="utf-8")
         sha = hashlib.sha256(text.encode("utf-8")).hexdigest()
         return {"mimeType": "text/plain", "text": text, "metadata": {"sha256": sha}}
     except Exception as e:

UnityMcpBridge/UnityMcpServer~/src/tools/__init__.py

@@ -1,3 +1,4 @@
+import logging
 from .manage_script_edits import register_manage_script_edits_tools
 from .manage_script import register_manage_script_tools
 from .manage_scene import register_manage_scene_tools
@@ -8,9 +9,11 @@
 from .read_console import register_read_console_tools
 from .execute_menu_item import register_execute_menu_item_tools
 
+logger = logging.getLogger("unity-mcp-server")
+
 def register_all_tools(mcp):
     """Register all refactored tools with the MCP server."""
-    # Note: Do not print to stdout; Claude treats stdout as MCP JSON. Use logging.
+    logger.info("Registering Unity MCP Server refactored tools...")
     # Prefer the surgical edits tool so LLMs discover it first
     register_manage_script_edits_tools(mcp)
     register_manage_script_tools(mcp)
@@ -21,4 +24,4 @@ def register_all_tools(mcp):
     register_manage_shader_tools(mcp)
     register_read_console_tools(mcp)
     register_execute_menu_item_tools(mcp)
-    # Do not print to stdout here either.
+    logger.info("Unity MCP Server tool registration complete.")

UnityMcpBridge/UnityMcpServer~/src/tools/manage_script.py

@@ -60,7 +60,7 @@ def create_script(
             "namespace": namespace,
             "scriptType": script_type,
         }
-        if contents is not None:
+        if contents:
             params["encodedContents"] = base64.b64encode(contents.encode("utf-8")).decode("utf-8")
             params["contentsEncoded"] = True
         params = {k: v for k, v in params.items() if v is not None}
@@ -107,7 +107,7 @@ def manage_script(
         - Edits should use apply_text_edits.
 
         Args:
-            action: Operation ('create', 'read', 'update', 'delete').
+            action: Operation ('create', 'read', 'delete').
             name: Script name (no .cs extension).
             path: Asset path (default: "Assets/").
             contents: C# code for 'create'/'update'.
@@ -132,8 +132,8 @@ def manage_script(
             }
 
             # Base64 encode the contents if they exist to avoid JSON escaping issues
-            if contents is not None:
-                if action in ['create', 'update']:
+            if contents:
+                if action == 'create':
                     params["encodedContents"] = base64.b64encode(contents.encode('utf-8')).decode('utf-8')
                     params["contentsEncoded"] = True
                 else:
@@ -143,22 +143,22 @@ def manage_script(
 
             response = send_command_with_retry("manage_script", params)
 
-            if isinstance(response, dict) and response.get("success"):
-                if response.get("data", {}).get("contentsEncoded"):
-                    decoded_contents = base64.b64decode(response["data"]["encodedContents"]).decode('utf-8')
-                    response["data"]["contents"] = decoded_contents
-                    del response["data"]["encodedContents"]
-                    del response["data"]["contentsEncoded"]
-
-                return {
-                    "success": True,
-                    "message": response.get("message", "Operation successful."),
-                    "data": response.get("data"),
-                }
-            return response if isinstance(response, dict) else {
-                "success": False,
-                "message": str(response),
-            }
+            if isinstance(response, dict):
+                if response.get("success"):
+                    if response.get("data", {}).get("contentsEncoded"):
+                        decoded_contents = base64.b64decode(response["data"]["encodedContents"]).decode('utf-8')
+                        response["data"]["contents"] = decoded_contents
+                        del response["data"]["encodedContents"]
+                        del response["data"]["contentsEncoded"]
+
+                    return {
+                        "success": True,
+                        "message": response.get("message", "Operation successful."),
+                        "data": response.get("data"),
+                    }
+                return response
+
+            return {"success": False, "message": str(response)}
 
         except Exception as e:
             return {

UnityMcpBridge/UnityMcpServer~/src/tools/manage_script_edits.py

@@ -51,10 +51,11 @@ def _apply_edits_locally(original_text: str, edits: List[Dict[str, Any]]) -> str
             end_line = int(edit.get("endLine", start_line))
             replacement = edit.get("text", "")
             lines = text.splitlines(keepends=True)
-            if start_line < 1 or end_line < start_line or end_line > len(lines):
+            max_end = len(lines) + 1
+            if start_line < 1 or end_line < start_line or end_line > max_end:
                 raise RuntimeError("replace_range out of bounds")
             a = start_line - 1
-            b = end_line
+            b = min(end_line, len(lines))
             rep = replacement
             if rep and not rep.endswith("\n"):
                 rep += "\n"
@@ -88,7 +89,8 @@ def script_apply_edits(
         script_type: str = "MonoBehaviour",
         namespace: str = "",
     ) -> Dict[str, Any]:
-        # If the edits request structured class/method ops, route directly to Unity's 'edit' action
+        # If the edits request structured class/method ops, route directly to Unity's 'edit' action.
+        # These bypass local text validation/encoding since Unity performs the semantic changes.
         for e in edits or []:
             op = (e.get("op") or e.get("operation") or e.get("type") or e.get("mode") or "").strip().lower()
             if op in ("replace_class", "delete_class", "replace_method", "delete_method", "insert_method"):

UnityMcpBridge/UnityMcpServer~/src/unity_connection.py

@@ -49,12 +49,23 @@ def connect(self) -> bool:
                     self.use_framing = True
                     logger.debug('Unity MCP handshake received: FRAMING=1 (strict)')
                 else:
+                    try:
+                        msg = b'Unity MCP requires FRAMING=1'
+                        header = struct.pack('>Q', len(msg))
+                        self.sock.sendall(header + msg)
+                    except Exception:
+                        pass
                     raise ConnectionError(f'Unity MCP requires FRAMING=1, got: {text!r}')
             finally:
                 self.sock.settimeout(config.connection_timeout)
             return True
         except Exception as e:
             logger.error(f"Failed to connect to Unity: {str(e)}")
+            try:
+                if self.sock:
+                    self.sock.close()
+            except Exception:
+                pass
             self.sock = None
             return False
 
@@ -83,7 +94,7 @@ def receive_full_response(self, sock, buffer_size=config.buffer_size) -> bytes:
             try:
                 header = self._read_exact(sock, 8)
                 payload_len = struct.unpack('>Q', header)[0]
-                if payload_len == 0 or payload_len > (64 * 1024 * 1024):
+                if payload_len > (64 * 1024 * 1024):
                     raise Exception(f"Invalid framed length: {payload_len}")
                 payload = self._read_exact(sock, payload_len)
                 logger.info(f"Received framed response ({len(payload)} bytes)")