From fc94e90f6caed2a0198af9836314dc79111c9548 Mon Sep 17 00:00:00 2001 From: imhaoran <117565557+thresh111@users.noreply.github.com> Date: Mon, 20 Apr 2026 20:17:30 +0800 Subject: [PATCH] =?UTF-8?q?fix(setup-agent):=20prevent=20data=20loss=20whe?= =?UTF-8?q?n=20setup=20fails=20on=20existing=20agen=E2=80=A6=20(#2254)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * 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 --- .../tools/builtins/setup_agent_tool.py | 6 +- backend/tests/test_setup_agent_tool.py | 87 +++++++++++++++++++ 2 files changed, 91 insertions(+), 2 deletions(-) diff --git a/backend/packages/harness/deerflow/tools/builtins/setup_agent_tool.py b/backend/packages/harness/deerflow/tools/builtins/setup_agent_tool.py index 922ad7b68..a42f8bbef 100644 --- a/backend/packages/harness/deerflow/tools/builtins/setup_agent_tool.py +++ b/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)]}) diff --git a/backend/tests/test_setup_agent_tool.py b/backend/tests/test_setup_agent_tool.py index 72ac03fb5..482fe1358 100644 --- a/backend/tests/test_setup_agent_tool.py +++ b/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()