diff --git a/UnityMcpBridge/Editor/Tools/ManageScript.cs b/UnityMcpBridge/Editor/Tools/ManageScript.cs index d2df4584..0d2fae60 100644 --- a/UnityMcpBridge/Editor/Tools/ManageScript.cs +++ b/UnityMcpBridge/Editor/Tools/ManageScript.cs @@ -8,10 +8,12 @@ using UnityEngine; using UnityMcpBridge.Editor.Helpers; using System.Threading; +using System.Security.Cryptography; #if USE_ROSLYN using Microsoft.CodeAnalysis; using Microsoft.CodeAnalysis.CSharp; +using Microsoft.CodeAnalysis.Formatting; #endif #if UNITY_EDITOR @@ -193,12 +195,12 @@ public static object HandleCommand(JObject @params) case "apply_text_edits": { var edits = @params["edits"] as JArray; - string precondition = @params["precondition_sha256"]?.ToString(); // optional, currently ignored here - return ApplyTextEdits(fullPath, relativePath, name, edits); + string precondition = @params["precondition_sha256"]?.ToString(); + return ApplyTextEdits(fullPath, relativePath, name, edits, precondition); } case "validate": { - string level = @params["level"]?.ToString()?.ToLowerInvariant() ?? "standard"; + string level = @params["level"]?.ToString()?.ToLowerInvariant() ?? "basic"; var chosen = level switch { "basic" => ValidationLevel.Basic, @@ -209,13 +211,19 @@ public static object HandleCommand(JObject @params) try { fileText = File.ReadAllText(fullPath); } catch (Exception ex) { return Response.Error($"Failed to read script: {ex.Message}"); } - bool ok = ValidateScriptSyntax(fileText, chosen, out string[] diags); - var result = new + bool ok = ValidateScriptSyntax(fileText, chosen, out string[] diagsRaw); + var diags = (diagsRaw ?? Array.Empty()).Select(s => { - isValid = ok, - diagnostics = diags ?? Array.Empty() - }; - return ok ? Response.Success("Validation completed.", result) : Response.Error("Validation failed.", result); + var m = Regex.Match(s, @"^(ERROR|WARNING|INFO): (.*?)(?: \(Line (\d+)\))?$"); + string severity = m.Success ? m.Groups[1].Value.ToLowerInvariant() : "info"; + string message = m.Success ? m.Groups[2].Value : s; + int lineNum = m.Success && int.TryParse(m.Groups[3].Value, out var l) ? l : 0; + return new { line = lineNum, col = 0, severity, message }; + }).ToArray(); + + var result = new { diagnostics = diags }; + return ok ? Response.Success("Validation completed.", result) + : Response.Error("Validation failed.", result); } case "edit": return Response.Error("Deprecated: use apply_text_edits. Structured 'edit' mode has been retired in favor of simple text edits."); @@ -299,9 +307,10 @@ string namespaceName try { File.Delete(tmp); } catch { } } + var uri = $"unity://path/{relativePath}"; var ok = Response.Success( $"Script '{name}.cs' created successfully at '{relativePath}'.", - new { path = relativePath, scheduledRefresh = true } + new { uri, scheduledRefresh = true } ); // Schedule heavy work AFTER replying @@ -423,11 +432,14 @@ string contents /// /// Apply simple text edits specified by line/column ranges. Applies transactionally and validates result. /// + private const int MaxEditPayloadBytes = 15 * 1024; + private static object ApplyTextEdits( string fullPath, string relativePath, string name, - JArray edits) + JArray edits, + string preconditionSha256) { if (!File.Exists(fullPath)) return Response.Error($"Script not found at '{relativePath}'."); @@ -438,8 +450,15 @@ private static object ApplyTextEdits( try { original = File.ReadAllText(fullPath); } catch (Exception ex) { return Response.Error($"Failed to read script: {ex.Message}"); } + string currentSha = ComputeSha256(original); + if (!string.IsNullOrEmpty(preconditionSha256) && !preconditionSha256.Equals(currentSha, StringComparison.OrdinalIgnoreCase)) + { + return Response.Error("stale_file", new { status = "stale_file", expected_sha256 = preconditionSha256, current_sha256 = currentSha }); + } + // Convert edits to absolute index ranges var spans = new List<(int start, int end, string text)>(); + int totalBytes = 0; foreach (var e in edits) { try @@ -457,6 +476,7 @@ private static object ApplyTextEdits( if (eidx < sidx) (sidx, eidx) = (eidx, sidx); spans.Add((sidx, eidx, newText)); + totalBytes += System.Text.Encoding.UTF8.GetByteCount(newText); } catch (Exception ex) { @@ -464,6 +484,11 @@ private static object ApplyTextEdits( } } + if (totalBytes > MaxEditPayloadBytes) + { + return Response.Error("too_large", new { status = "too_large", limitBytes = MaxEditPayloadBytes, hint = "split into smaller edits" }); + } + // Ensure non-overlap and apply from back to front spans = spans.OrderByDescending(t => t.start).ToList(); for (int i = 1; i < spans.Count; i++) @@ -478,10 +503,40 @@ private static object ApplyTextEdits( working = working.Remove(sp.start, sp.end - sp.start).Insert(sp.start, sp.text ?? string.Empty); } - // Validate result - var level = GetValidationLevelFromGUI(); - if (!ValidateScriptSyntax(working, level, out var errors)) - return Response.Error("Script validation failed:\n" + string.Join("\n", errors ?? Array.Empty())); + if (!CheckBalancedDelimiters(working, out int line, out char expected)) + { + int startLine = Math.Max(1, line - 5); + int endLine = line + 5; + string hint = $"unbalanced_braces at line {line}. Call resources/read for lines {startLine}-{endLine} and resend a smaller apply_text_edits that restores balance."; + return Response.Error(hint, new { status = "unbalanced_braces", line, expected = expected.ToString() }); + } + +#if USE_ROSLYN + var tree = CSharpSyntaxTree.ParseText(working); + var diagnostics = tree.GetDiagnostics().Where(d => d.Severity == DiagnosticSeverity.Error).Take(3) + .Select(d => new { + line = d.Location.GetLineSpan().StartLinePosition.Line + 1, + col = d.Location.GetLineSpan().StartLinePosition.Character + 1, + code = d.Id, + message = d.GetMessage() + }).ToArray(); + if (diagnostics.Length > 0) + { + return Response.Error("syntax_error", new { status = "syntax_error", diagnostics }); + } + + // Optional formatting + try + { + var root = tree.GetRoot(); + var workspace = new AdhocWorkspace(); + root = Microsoft.CodeAnalysis.Formatting.Formatter.Format(root, workspace); + working = root.ToFullString(); + } + catch { } +#endif + + string newSha = ComputeSha256(working); // Atomic write and schedule refresh try @@ -495,7 +550,17 @@ private static object ApplyTextEdits( catch (IOException) { File.Copy(tmp, fullPath, true); try { File.Delete(tmp); } catch { } } ManageScriptRefreshHelpers.ScheduleScriptRefresh(relativePath); - return Response.Success($"Applied {spans.Count} text edit(s) to '{relativePath}'.", new { path = relativePath, editsApplied = spans.Count, scheduledRefresh = true }); + return Response.Success( + $"Applied {spans.Count} text edit(s) to '{relativePath}'.", + new + { + applied = spans.Count, + unchanged = 0, + sha256 = newSha, + uri = $"unity://path/{relativePath}", + scheduledRefresh = true + } + ); } catch (Exception ex) { @@ -522,6 +587,84 @@ private static bool TryIndexFromLineCol(string text, int line1, int col1, out in index = -1; return false; } + private static string ComputeSha256(string contents) + { + using (var sha = SHA256.Create()) + { + var bytes = System.Text.Encoding.UTF8.GetBytes(contents); + var hash = sha.ComputeHash(bytes); + return BitConverter.ToString(hash).Replace("-", string.Empty).ToLowerInvariant(); + } + } + + private static bool CheckBalancedDelimiters(string text, out int line, out char expected) + { + var braceStack = new Stack(); + var parenStack = new Stack(); + var bracketStack = new Stack(); + bool inString = false, inChar = false, inSingle = false, inMulti = false, escape = false; + line = 1; expected = '\0'; + + for (int i = 0; i < text.Length; i++) + { + char c = text[i]; + char next = i + 1 < text.Length ? text[i + 1] : '\0'; + + if (c == '\n') { line++; if (inSingle) inSingle = false; } + + if (escape) { escape = false; continue; } + + if (inString) + { + if (c == '\\') { escape = true; } + else if (c == '"') inString = false; + continue; + } + if (inChar) + { + if (c == '\\') { escape = true; } + else if (c == '\'') inChar = false; + continue; + } + if (inSingle) continue; + if (inMulti) + { + if (c == '*' && next == '/') { inMulti = false; i++; } + continue; + } + + if (c == '"') { inString = true; continue; } + if (c == '\'') { inChar = true; continue; } + if (c == '/' && next == '/') { inSingle = true; i++; continue; } + if (c == '/' && next == '*') { inMulti = true; i++; continue; } + + switch (c) + { + case '{': braceStack.Push(line); break; + case '}': + if (braceStack.Count == 0) { expected = '{'; return false; } + braceStack.Pop(); + break; + case '(': parenStack.Push(line); break; + case ')': + if (parenStack.Count == 0) { expected = '('; return false; } + parenStack.Pop(); + break; + case '[': bracketStack.Push(line); break; + case ']': + if (bracketStack.Count == 0) { expected = '['; return false; } + bracketStack.Pop(); + break; + } + } + + if (braceStack.Count > 0) { line = braceStack.Peek(); expected = '}'; return false; } + if (parenStack.Count > 0) { line = parenStack.Peek(); expected = ')'; return false; } + if (bracketStack.Count > 0) { line = bracketStack.Peek(); expected = ']'; return false; } + + return true; + } + private static object DeleteScript(string fullPath, string relativePath) { if (!File.Exists(fullPath)) @@ -537,7 +680,8 @@ private static object DeleteScript(string fullPath, string relativePath) { AssetDatabase.Refresh(); return Response.Success( - $"Script '{Path.GetFileName(relativePath)}' moved to trash successfully." + $"Script '{Path.GetFileName(relativePath)}' moved to trash successfully.", + new { deleted = true } ); } else diff --git a/UnityMcpBridge/UnityMcpServer~/src/tools/manage_asset.py b/UnityMcpBridge/UnityMcpServer~/src/tools/manage_asset.py index 19ac0c2e..ccafb047 100644 --- a/UnityMcpBridge/UnityMcpServer~/src/tools/manage_asset.py +++ b/UnityMcpBridge/UnityMcpServer~/src/tools/manage_asset.py @@ -76,4 +76,4 @@ async def manage_asset( # Use centralized async retry helper to avoid blocking the event loop result = await async_send_command_with_retry("manage_asset", params_dict, loop=loop) # Return the result obtained from Unity - return result if isinstance(result, dict) else {"success": False, "message": str(result)} \ No newline at end of file + return result if isinstance(result, dict) else {"success": False, "message": str(result)} diff --git a/UnityMcpBridge/UnityMcpServer~/src/tools/manage_script.py b/UnityMcpBridge/UnityMcpServer~/src/tools/manage_script.py index af44a446..f7836da3 100644 --- a/UnityMcpBridge/UnityMcpServer~/src/tools/manage_script.py +++ b/UnityMcpBridge/UnityMcpServer~/src/tools/manage_script.py @@ -1,14 +1,95 @@ from mcp.server.fastmcp import FastMCP, Context -from typing import Dict, Any +from typing import Dict, Any, List from unity_connection import get_unity_connection, send_command_with_retry from config import config import time import os import base64 + def register_manage_script_tools(mcp: FastMCP): """Register all script management tools with the MCP server.""" + def _split_uri(uri: str) -> tuple[str, str]: + if uri.startswith("unity://path/"): + path = uri[len("unity://path/") :] + elif uri.startswith("file://"): + path = uri[len("file://") :] + else: + path = uri + path = path.replace("\\", "/") + name = os.path.splitext(os.path.basename(path))[0] + directory = os.path.dirname(path) + return name, directory + + @mcp.tool() + def apply_text_edits( + ctx: Context, + uri: str, + edits: List[Dict[str, Any]], + precondition_sha256: str | None = None, + ) -> Dict[str, Any]: + """Apply small text edits to a C# script identified by URI.""" + name, directory = _split_uri(uri) + params = { + "action": "apply_text_edits", + "name": name, + "path": directory, + "edits": edits, + "precondition_sha256": precondition_sha256, + } + params = {k: v for k, v in params.items() if v is not None} + resp = send_command_with_retry("manage_script", params) + return resp if isinstance(resp, dict) else {"success": False, "message": str(resp)} + + @mcp.tool() + def create_script( + ctx: Context, + path: str, + contents: str = "", + script_type: str | None = None, + namespace: str | None = None, + ) -> Dict[str, Any]: + """Create a new C# script at the given path.""" + name = os.path.splitext(os.path.basename(path))[0] + directory = os.path.dirname(path) + params: Dict[str, Any] = { + "action": "create", + "name": name, + "path": directory, + "namespace": namespace, + "scriptType": script_type, + } + if contents is not None: + 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} + resp = send_command_with_retry("manage_script", params) + return resp if isinstance(resp, dict) else {"success": False, "message": str(resp)} + + @mcp.tool() + def delete_script(ctx: Context, uri: str) -> Dict[str, Any]: + """Delete a C# script by URI.""" + name, directory = _split_uri(uri) + params = {"action": "delete", "name": name, "path": directory} + resp = send_command_with_retry("manage_script", params) + return resp if isinstance(resp, dict) else {"success": False, "message": str(resp)} + + @mcp.tool() + def validate_script( + ctx: Context, uri: str, level: str = "basic" + ) -> Dict[str, Any]: + """Validate a C# script and return diagnostics.""" + name, directory = _split_uri(uri) + params = { + "action": "validate", + "name": name, + "path": directory, + "level": level, + } + resp = send_command_with_retry("manage_script", params) + return resp if isinstance(resp, dict) else {"success": False, "message": str(resp)} + @mcp.tool() def manage_script( ctx: Context, @@ -17,12 +98,13 @@ def manage_script( path: str, contents: str, script_type: str, - namespace: str + namespace: str, ) -> Dict[str, Any]: - """Manage C# scripts in Unity. + """Compatibility router for legacy script operations. IMPORTANT: - - This router is minimized. Use resources/read for file content and 'script_apply_edits' for changes. + - Direct file reads should use resources/read. + - Edits should use apply_text_edits. Args: action: Operation ('create', 'read', 'update', 'delete'). @@ -38,7 +120,7 @@ def manage_script( try: # Deprecate full-file update path entirely if action == 'update': - return {"success": False, "message": "Deprecated: use script_apply_edits (line/col edits) or resources/read + small edits."} + return {"success": False, "message": "Deprecated: use apply_text_edits or resources/read + small edits."} # Prepare parameters for Unity params = { @@ -46,36 +128,40 @@ def manage_script( "name": name, "path": path, "namespace": namespace, - "scriptType": script_type + "scriptType": script_type, } - + # Base64 encode the contents if they exist to avoid JSON escaping issues if contents is not None: if action in ['create', 'update']: - # Encode content for safer transmission params["encodedContents"] = base64.b64encode(contents.encode('utf-8')).decode('utf-8') params["contentsEncoded"] = True else: params["contents"] = contents - - # Remove None values so they don't get sent as null + params = {k: v for k, v in params.items() if v is not None} - # Send command via centralized retry helper response = send_command_with_retry("manage_script", params) - - # Process response from Unity + if isinstance(response, dict) and response.get("success"): - # If the response contains base64 encoded content, decode it 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)} + + 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), + } except Exception as e: - # Handle Python-side errors (e.g., connection issues) - return {"success": False, "message": f"Python error managing script: {str(e)}"} \ No newline at end of file + return { + "success": False, + "message": f"Python error managing script: {str(e)}", + } diff --git a/test_unity_socket_framing.py b/test_unity_socket_framing.py index b0e179c9..c24064a1 100644 --- a/test_unity_socket_framing.py +++ b/test_unity_socket_framing.py @@ -3,7 +3,10 @@ HOST = "127.0.0.1" PORT = 6400 -SIZE_MB = int(sys.argv[1]) if len(sys.argv) > 1 else 5 # e.g., 5 or 10 +try: + SIZE_MB = int(sys.argv[1]) +except (IndexError, ValueError): + SIZE_MB = 5 # e.g., 5 or 10 FILL = "R" def recv_exact(sock, n): diff --git a/tests/test_script_tools.py b/tests/test_script_tools.py new file mode 100644 index 00000000..9b953a1a --- /dev/null +++ b/tests/test_script_tools.py @@ -0,0 +1,123 @@ +import sys +import pathlib +import importlib.util +import types +import pytest + +# add server src to path and load modules without triggering package imports +ROOT = pathlib.Path(__file__).resolve().parents[1] +SRC = ROOT / "UnityMcpBridge" / "UnityMcpServer~" / "src" +sys.path.insert(0, str(SRC)) + +# stub mcp.server.fastmcp to satisfy imports without full dependency +mcp_pkg = types.ModuleType("mcp") +server_pkg = types.ModuleType("mcp.server") +fastmcp_pkg = types.ModuleType("mcp.server.fastmcp") + +class _Dummy: + pass + +fastmcp_pkg.FastMCP = _Dummy +fastmcp_pkg.Context = _Dummy +server_pkg.fastmcp = fastmcp_pkg +mcp_pkg.server = server_pkg +sys.modules.setdefault("mcp", mcp_pkg) +sys.modules.setdefault("mcp.server", server_pkg) +sys.modules.setdefault("mcp.server.fastmcp", fastmcp_pkg) + +def load_module(path, name): + spec = importlib.util.spec_from_file_location(name, path) + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + +manage_script_module = load_module(SRC / "tools" / "manage_script.py", "manage_script_module") +manage_asset_module = load_module(SRC / "tools" / "manage_asset.py", "manage_asset_module") + + +class DummyMCP: + def __init__(self): + self.tools = {} + + def tool(self): + def decorator(func): + self.tools[func.__name__] = func + return func + return decorator + +def setup_manage_script(): + mcp = DummyMCP() + manage_script_module.register_manage_script_tools(mcp) + return mcp.tools + +def setup_manage_asset(): + mcp = DummyMCP() + manage_asset_module.register_manage_asset_tools(mcp) + return mcp.tools + +def test_apply_text_edits_long_file(monkeypatch): + tools = setup_manage_script() + apply_edits = tools["apply_text_edits"] + captured = {} + + def fake_send(cmd, params): + captured["cmd"] = cmd + captured["params"] = params + return {"success": True} + + monkeypatch.setattr(manage_script_module, "send_command_with_retry", fake_send) + + edit = {"startLine": 1005, "startCol": 0, "endLine": 1005, "endCol": 5, "newText": "Hello"} + resp = apply_edits(None, "unity://path/Assets/Scripts/LongFile.cs", [edit]) + assert captured["cmd"] == "manage_script" + assert captured["params"]["action"] == "apply_text_edits" + assert captured["params"]["edits"][0]["startLine"] == 1005 + assert resp["success"] is True + +def test_sequential_edits_use_precondition(monkeypatch): + tools = setup_manage_script() + apply_edits = tools["apply_text_edits"] + calls = [] + + def fake_send(cmd, params): + calls.append(params) + return {"success": True, "sha256": f"hash{len(calls)}"} + + monkeypatch.setattr(manage_script_module, "send_command_with_retry", fake_send) + + edit1 = {"startLine": 1, "startCol": 0, "endLine": 1, "endCol": 0, "newText": "//header\n"} + resp1 = apply_edits(None, "unity://path/Assets/Scripts/File.cs", [edit1]) + edit2 = {"startLine": 2, "startCol": 0, "endLine": 2, "endCol": 0, "newText": "//second\n"} + resp2 = apply_edits(None, "unity://path/Assets/Scripts/File.cs", [edit2], precondition_sha256=resp1["sha256"]) + + assert calls[1]["precondition_sha256"] == resp1["sha256"] + assert resp2["sha256"] == "hash2" + +def test_manage_asset_prefab_modify_request(monkeypatch): + tools = setup_manage_asset() + manage_asset = tools["manage_asset"] + captured = {} + + async def fake_async(cmd, params, loop=None): + captured["cmd"] = cmd + captured["params"] = params + return {"success": True} + + monkeypatch.setattr(manage_asset_module, "async_send_command_with_retry", fake_async) + monkeypatch.setattr(manage_asset_module, "get_unity_connection", lambda: object()) + + async def run(): + resp = await manage_asset( + None, + action="modify", + path="Assets/Prefabs/Player.prefab", + properties={"hp": 100}, + ) + assert captured["cmd"] == "manage_asset" + assert captured["params"]["action"] == "modify" + assert captured["params"]["path"] == "Assets/Prefabs/Player.prefab" + assert captured["params"]["properties"] == {"hp": 100} + assert resp["success"] is True + + import asyncio + asyncio.run(run())