fix(setup-agent): prevent data loss when setup fails on existing agen… (#2254)

* fix(setup-agent): prevent data loss when setup fails on existing agent directory

Record whether the agent directory pre-existed before mkdir, and only
run shutil.rmtree cleanup when the directory was newly created during
this call. Previously, any failure would delete the entire directory
including pre-existing SOUL.md and config.yaml.

* fix: address PR review — init variables before try, remove unused result

* style: fix ruff I001 import block formatting in test file

* style: add missing blank lines between top-level definitions in test file

backend/packages/harness/deerflow/tools/builtins/setup_agent_tool.py

@@ -27,11 +27,13 @@ def setup_agent(
 
     agent_name: str | None = runtime.context.get("agent_name") if runtime.context else None
     agent_dir = None
+    is_new_dir = False
 
     try:
         agent_name = validate_agent_name(agent_name)
         paths = get_paths()
         agent_dir = paths.agent_dir(agent_name) if agent_name else paths.base_dir
+        is_new_dir = not agent_dir.exists()
         agent_dir.mkdir(parents=True, exist_ok=True)
 
         if agent_name:
@@ -58,8 +60,8 @@ def setup_agent(
     except Exception as e:
         import shutil
 
-        if agent_name and agent_dir is not None and agent_dir.exists():
-            # Cleanup the custom agent directory only if it was created but an error occurred during setup
+        if agent_name and is_new_dir and agent_dir is not None and agent_dir.exists():
+            # Cleanup the custom agent directory only if it was newly created during this call
             shutil.rmtree(agent_dir)
         logger.error(f"[agent_creator] Failed to create agent '{agent_name}': {e}", exc_info=True)
         return Command(update={"messages": [ToolMessage(content=f"Error: {e}", tool_call_id=runtime.tool_call_id)]})

backend/tests/test_setup_agent_tool.py

@@ -1,16 +1,48 @@
+"""Tests for setup_agent tool — validates agent name security and data loss prevention."""
+
 from __future__ import annotations
 
 from pathlib import Path
 from types import SimpleNamespace
+from unittest.mock import MagicMock, patch
 
 from deerflow.tools.builtins.setup_agent_tool import setup_agent
 
+# --- Helpers ---
+
 
 class _DummyRuntime(SimpleNamespace):
     context: dict
     tool_call_id: str
 
 
+def _make_runtime(agent_name: str | None = "test-agent") -> MagicMock:
+    runtime = MagicMock()
+    runtime.context = {"agent_name": agent_name}
+    runtime.tool_call_id = "call_1"
+    return runtime
+
+
+def _make_paths_mock(tmp_path: Path):
+    paths = MagicMock()
+    paths.base_dir = tmp_path
+    paths.agent_dir = lambda name: tmp_path / "agents" / name
+    return paths
+
+
+def _call_setup_agent(tmp_path: Path, soul: str, description: str, agent_name: str = "test-agent"):
+    """Call the underlying setup_agent function directly, bypassing langchain tool wrapper."""
+    with patch("deerflow.tools.builtins.setup_agent_tool.get_paths", return_value=_make_paths_mock(tmp_path)):
+        return setup_agent.func(
+            soul=soul,
+            description=description,
+            runtime=_make_runtime(agent_name),
+        )
+
+
+# --- Agent name validation tests ---
+
+
 def test_setup_agent_rejects_invalid_agent_name_before_writing(tmp_path, monkeypatch):
     monkeypatch.setenv("DEER_FLOW_HOME", str(tmp_path))
     outside_dir = tmp_path.parent / "outside-target"
@@ -38,3 +70,58 @@ def test_setup_agent_rejects_absolute_agent_name_before_writing(tmp_path, monkey
     assert "Invalid agent name" in messages[0].content
     assert not (tmp_path / "agents").exists()
     assert not (Path(absolute_agent) / "SOUL.md").exists()
+
+
+# --- Data loss prevention tests ---
+
+
+class TestSetupAgentNoDataLoss:
+    """Ensure shutil.rmtree only removes directories created during the current call."""
+
+    def test_existing_agent_dir_preserved_on_failure(self, tmp_path: Path):
+        """If the agent directory already exists and setup fails,
+        the directory and its contents must NOT be deleted."""
+        agent_dir = tmp_path / "agents" / "test-agent"
+        agent_dir.mkdir(parents=True)
+        old_soul = agent_dir / "SOUL.md"
+        old_soul.write_text("original soul content")
+
+        with patch("deerflow.tools.builtins.setup_agent_tool.get_paths", return_value=_make_paths_mock(tmp_path)):
+            # Force soul_file.write_text to raise after directory already exists
+            with patch.object(Path, "write_text", side_effect=OSError("disk full")):
+                setup_agent.func(
+                    soul="new soul",
+                    description="desc",
+                    runtime=_make_runtime(),
+                )
+
+        # Directory must still exist
+        assert agent_dir.exists(), "Pre-existing agent directory was deleted on failure"
+        # Original SOUL.md should still be on disk (not deleted by rmtree)
+        assert old_soul.exists(), "Pre-existing SOUL.md was deleted on failure"
+
+    def test_new_agent_dir_cleaned_up_on_failure(self, tmp_path: Path):
+        """If the agent directory is newly created and setup fails,
+        the directory should be cleaned up."""
+        agent_dir = tmp_path / "agents" / "test-agent"
+        assert not agent_dir.exists()
+
+        with patch("deerflow.tools.builtins.setup_agent_tool.get_paths", return_value=_make_paths_mock(tmp_path)):
+            with patch("yaml.dump", side_effect=OSError("write error")):
+                setup_agent.func(
+                    soul="new soul",
+                    description="desc",
+                    runtime=_make_runtime(),
+                )
+
+        # Newly created directory should be cleaned up
+        assert not agent_dir.exists(), "Newly created agent directory was not cleaned up on failure"
+
+    def test_successful_setup_creates_files(self, tmp_path: Path):
+        """Happy path: setup_agent creates config.yaml and SOUL.md."""
+        _call_setup_agent(tmp_path, soul="# My Agent", description="A test agent")
+
+        agent_dir = tmp_path / "agents" / "test-agent"
+        assert agent_dir.exists()
+        assert (agent_dir / "SOUL.md").read_text() == "# My Agent"
+        assert (agent_dir / "config.yaml").exists()