Merge pull request #990 from zaferdace/feat/save-prefab-stage

feat(manage_editor): add save_prefab_stage action

Author: Scriptwonder

MCPForUnity/Editor/Tools/ManageEditor.cs

@@ -140,6 +140,8 @@ public static object HandleCommand(JObject @params)
                 // Prefab Stage
                 case "open_prefab_stage":
                     return OpenPrefabStage(prefabPath);
+                case "save_prefab_stage":
+                    return SavePrefabStage();
                 case "close_prefab_stage":
                     return ClosePrefabStage();
 
@@ -180,7 +182,7 @@ public static object HandleCommand(JObject @params)
 
                 default:
                     return new ErrorResponse(
-                        $"Unknown action: '{action}'. Supported actions: play, pause, stop, set_active_tool, add_tag, remove_tag, add_layer, remove_layer, open_prefab_stage, close_prefab_stage, deploy_package, restore_package, undo, redo. Use MCP resources for reading editor state, project info, tags, layers, selection, windows, prefab stage, and active tool."
+                        $"Unknown action: '{action}'. Supported actions: play, pause, stop, set_active_tool, add_tag, remove_tag, add_layer, remove_layer, open_prefab_stage, save_prefab_stage, close_prefab_stage, deploy_package, restore_package, undo, redo. Use MCP resources for reading editor state, project info, tags, layers, selection, windows, prefab stage, and active tool."
                     );
             }
         }
@@ -460,6 +462,32 @@ private static object OpenPrefabStage(string requestedPath)
             }
         }
 
+        private static object SavePrefabStage()
+        {
+            try
+            {
+                var prefabStage = PrefabStageUtility.GetCurrentPrefabStage();
+                if (prefabStage == null)
+                {
+                    return new ErrorResponse("Not currently in prefab editing mode. Open a prefab stage first with open_prefab_stage.");
+                }
+
+                string prefabPath = prefabStage.assetPath;
+                EditorSceneManager.MarkSceneDirty(prefabStage.scene);
+                bool saved = EditorSceneManager.SaveScene(prefabStage.scene);
+                if (!saved)
+                {
+                    return new ErrorResponse($"Failed to save prefab stage for '{prefabPath}'. The file may be read-only or the disk may be full.");
+                }
+
+                return new SuccessResponse($"Saved prefab stage changes for '{prefabPath}'.", new { prefabPath, saved });
+            }
+            catch (Exception e)
+            {
+                return new ErrorResponse($"Error saving prefab stage: {e.Message}");
+            }
+        }
+
         private static object ClosePrefabStage()
         {
             try

Server/src/cli/commands/prefab.py

@@ -69,32 +69,24 @@ def close_stage(save: bool):
 
 
 @prefab.command("save")
-@click.option(
-    "--force", "-f",
-    is_flag=True,
-    help="Force save even if no changes detected. Useful for automated workflows."
-)
 @handle_unity_errors
-def save_stage(force: bool):
+def save_stage():
     """Save the currently open prefab stage.
 
     \b
     Examples:
         unity-mcp prefab save
-        unity-mcp prefab save --force
     """
     config = get_config()
 
     params: dict[str, Any] = {
-        "action": "save_open_stage",
+        "action": "save_prefab_stage",
     }
-    if force:
-        params["force"] = True
 
-    result = run_command("manage_prefabs", params, config)
+    result = run_command("manage_editor", params, config)
     click.echo(format_output(result, config.format))
     if result.get("success"):
-        print_success("Saved prefab")
+        print_success("Saved prefab stage")
 
 
 @prefab.command("info")

Server/src/services/resources/prefab.py

@@ -57,7 +57,7 @@ async def get_prefab_api_docs(_ctx: Context) -> MCPResponse:
         "workflow": [
             "1. Use manage_asset action=search filterType=Prefab to find prefabs",
             "2. Use the asset path to access detailed data via resources below",
-            "3. Use manage_editor action=open_prefab_stage / close_prefab_stage for prefab editing UI transitions"
+            "3. Use manage_editor action=open_prefab_stage / save_prefab_stage / close_prefab_stage for prefab editing UI transitions"
         ],
         "path_encoding": {
             "note": "Prefab paths must be URL-encoded when used in resource URIs",
@@ -80,7 +80,7 @@ async def get_prefab_api_docs(_ctx: Context) -> MCPResponse:
             }
         },
         "related_tools": {
-            "manage_editor": "Open/close prefab stages in the Unity Editor UI",
+            "manage_editor": "Open/save/close prefab stages in the Unity Editor UI",
             "manage_prefabs": "Headless prefab inspection and modification without opening prefab stages",
             "manage_asset": "Search for prefab assets, get asset info",
             "manage_gameobject": "Modify GameObjects in open prefab stage",

Server/src/services/tools/manage_editor.py

@@ -10,14 +10,14 @@
 from transport.legacy.unity_connection import async_send_command_with_retry
 
 @mcp_for_unity_tool(
-    description="Controls and queries the Unity editor's state and settings. Read-only actions: telemetry_status, telemetry_ping. Modifying actions: play, pause, stop, set_active_tool, add_tag, remove_tag, add_layer, remove_layer, open_prefab_stage, close_prefab_stage, deploy_package, restore_package, undo, redo. open_prefab_stage opens a prefab asset in Unity's prefab editing mode. deploy_package copies the configured MCPForUnity source folder into the project's installed package location (triggers recompile, no confirmation dialog). restore_package reverts to the pre-deployment backup. undo/redo perform Unity editor undo/redo and return the affected group name.",
+    description="Controls and queries the Unity editor's state and settings. Read-only actions: telemetry_status, telemetry_ping. Modifying actions: play, pause, stop, set_active_tool, add_tag, remove_tag, add_layer, remove_layer, open_prefab_stage, save_prefab_stage, close_prefab_stage, deploy_package, restore_package, undo, redo. open_prefab_stage opens a prefab asset in Unity's prefab editing mode. save_prefab_stage saves changes in the currently open prefab stage back to the prefab asset. deploy_package copies the configured MCPForUnity source folder into the project's installed package location (triggers recompile, no confirmation dialog). restore_package reverts to the pre-deployment backup. undo/redo perform Unity editor undo/redo and return the affected group name.",
     annotations=ToolAnnotations(
         title="Manage Editor",
     ),
 )
 async def manage_editor(
     ctx: Context,
-    action: Annotated[Literal["telemetry_status", "telemetry_ping", "play", "pause", "stop", "set_active_tool", "add_tag", "remove_tag", "add_layer", "remove_layer", "open_prefab_stage", "close_prefab_stage", "deploy_package", "restore_package", "undo", "redo"], "Get and update the Unity Editor state. open_prefab_stage opens a prefab asset in prefab editing mode; close_prefab_stage exits prefab editing mode and returns to the main scene stage. deploy_package copies the configured MCPForUnity source into the project's package location (triggers recompile). restore_package reverts the last deployment from backup. undo/redo perform editor undo/redo."],
+    action: Annotated[Literal["telemetry_status", "telemetry_ping", "play", "pause", "stop", "set_active_tool", "add_tag", "remove_tag", "add_layer", "remove_layer", "open_prefab_stage", "save_prefab_stage", "close_prefab_stage", "deploy_package", "restore_package", "undo", "redo"], "Get and update the Unity Editor state. open_prefab_stage opens a prefab asset in prefab editing mode; save_prefab_stage saves changes in the open prefab stage back to the asset; close_prefab_stage exits prefab editing mode and returns to the main scene stage. deploy_package copies the configured MCPForUnity source into the project's package location (triggers recompile). restore_package reverts the last deployment from backup. undo/redo perform editor undo/redo."],
     tool_name: Annotated[str,
                          "Tool name when setting active tool"] | None = None,
     tag_name: Annotated[str,

Server/tests/test_manage_editor.py

@@ -55,7 +55,7 @@ def test_redo_forwards_to_unity(mock_unity):
 UNITY_FORWARDED_ACTIONS = [
     "play", "pause", "stop", "set_active_tool",
     "add_tag", "remove_tag", "add_layer", "remove_layer",
-    "open_prefab_stage", "close_prefab_stage", "deploy_package", "restore_package",
+    "open_prefab_stage", "save_prefab_stage", "close_prefab_stage", "deploy_package", "restore_package",
     "undo", "redo",
 ]
 
@@ -159,3 +159,35 @@ def test_open_prefab_stage_rejects_conflicting_path_inputs(mock_unity):
     )
     assert result["success"] is False
     assert "Provide only one of prefab_path or path" in result.get("message", "")
+
+
+# ── save_prefab_stage ────────────────────────────────────────────────
+
+
+def test_manage_editor_description_mentions_save_prefab_stage():
+    """The tool description should advertise the save_prefab_stage action."""
+    editor_tool = next(
+        (t for t in get_registered_tools() if t["name"] == "manage_editor"), None
+    )
+    assert editor_tool is not None
+    desc = editor_tool.get("description") or editor_tool.get("kwargs", {}).get("description", "")
+    assert "save_prefab_stage" in desc
+
+
+def test_save_prefab_stage_forwards_to_unity(mock_unity):
+    """save_prefab_stage should forward to Unity without extra parameters."""
+    result = asyncio.run(manage_editor(SimpleNamespace(), action="save_prefab_stage"))
+    assert result["success"] is True
+    assert mock_unity["params"]["action"] == "save_prefab_stage"
+    assert mock_unity["tool_name"] == "manage_editor"
+
+
+def test_save_prefab_stage_omits_none_params(mock_unity):
+    """save_prefab_stage should not include toolName, tagName, layerName, or path params."""
+    asyncio.run(manage_editor(SimpleNamespace(), action="save_prefab_stage"))
+    params = mock_unity["params"]
+    assert "toolName" not in params
+    assert "tagName" not in params
+    assert "layerName" not in params
+    assert "prefabPath" not in params
+    assert "path" not in params

unity-mcp-skill/references/tools-reference.md

@@ -730,6 +730,7 @@ manage_editor(action="add_layer", layer_name="Projectiles")
 manage_editor(action="remove_layer", layer_name="OldLayer")
 
 manage_editor(action="open_prefab_stage", prefab_path="Assets/Prefabs/Enemy.prefab")
+manage_editor(action="save_prefab_stage")   # Save changes in the open prefab stage
 manage_editor(action="close_prefab_stage")  # Exit prefab editing mode back to main scene
 
 # Package deployment (no confirmation dialog — designed for LLM-driven iteration)