mirror of
https://github.com/bytedance/deer-flow.git
synced 2026-04-25 11:18:22 +00:00
fix: validate bootstrap agent names before filesystem writes (#2274)
* fix: validate bootstrap agent names before filesystem writes * fix: tighten bootstrap agent-name validation
This commit is contained in:
parent
8e3591312a
commit
2176b2bbfc
@ -17,7 +17,7 @@ from deerflow.agents.middlewares.token_usage_middleware import TokenUsageMiddlew
|
||||
from deerflow.agents.middlewares.tool_error_handling_middleware import build_lead_runtime_middlewares
|
||||
from deerflow.agents.middlewares.view_image_middleware import ViewImageMiddleware
|
||||
from deerflow.agents.thread_state import ThreadState
|
||||
from deerflow.config.agents_config import load_agent_config
|
||||
from deerflow.config.agents_config import load_agent_config, validate_agent_name
|
||||
from deerflow.config.app_config import get_app_config
|
||||
from deerflow.config.memory_config import get_memory_config
|
||||
from deerflow.config.summarization_config import get_summarization_config
|
||||
@ -291,7 +291,7 @@ def make_lead_agent(config: RunnableConfig):
|
||||
subagent_enabled = cfg.get("subagent_enabled", False)
|
||||
max_concurrent_subagents = cfg.get("max_concurrent_subagents", 3)
|
||||
is_bootstrap = cfg.get("is_bootstrap", False)
|
||||
agent_name = cfg.get("agent_name")
|
||||
agent_name = validate_agent_name(cfg.get("agent_name"))
|
||||
|
||||
agent_config = load_agent_config(agent_name) if not is_bootstrap else None
|
||||
# Custom agent model from agent config (if any), or None to let _resolve_model_name pick the default
|
||||
|
||||
@ -15,6 +15,17 @@ SOUL_FILENAME = "SOUL.md"
|
||||
AGENT_NAME_PATTERN = re.compile(r"^[A-Za-z0-9-]+$")
|
||||
|
||||
|
||||
def validate_agent_name(name: str | None) -> str | None:
|
||||
"""Validate a custom agent name before using it in filesystem paths."""
|
||||
if name is None:
|
||||
return None
|
||||
if not isinstance(name, str):
|
||||
raise ValueError("Invalid agent name. Expected a string or None.")
|
||||
if not AGENT_NAME_PATTERN.fullmatch(name):
|
||||
raise ValueError(f"Invalid agent name '{name}'. Must match pattern: {AGENT_NAME_PATTERN.pattern}")
|
||||
return name
|
||||
|
||||
|
||||
class AgentConfig(BaseModel):
|
||||
"""Configuration for a custom agent."""
|
||||
|
||||
@ -46,8 +57,7 @@ def load_agent_config(name: str | None) -> AgentConfig | None:
|
||||
if name is None:
|
||||
return None
|
||||
|
||||
if not AGENT_NAME_PATTERN.match(name):
|
||||
raise ValueError(f"Invalid agent name '{name}'. Must match pattern: {AGENT_NAME_PATTERN.pattern}")
|
||||
name = validate_agent_name(name)
|
||||
agent_dir = get_paths().agent_dir(name)
|
||||
config_file = agent_dir / "config.yaml"
|
||||
|
||||
|
||||
@ -6,6 +6,7 @@ from langchain_core.tools import tool
|
||||
from langgraph.prebuilt import ToolRuntime
|
||||
from langgraph.types import Command
|
||||
|
||||
from deerflow.config.agents_config import validate_agent_name
|
||||
from deerflow.config.paths import get_paths
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
@ -25,8 +26,10 @@ def setup_agent(
|
||||
"""
|
||||
|
||||
agent_name: str | None = runtime.context.get("agent_name") if runtime.context else None
|
||||
agent_dir = None
|
||||
|
||||
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
|
||||
agent_dir.mkdir(parents=True, exist_ok=True)
|
||||
@ -55,7 +58,7 @@ def setup_agent(
|
||||
except Exception as e:
|
||||
import shutil
|
||||
|
||||
if agent_name and agent_dir.exists():
|
||||
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
|
||||
shutil.rmtree(agent_dir)
|
||||
logger.error(f"[agent_creator] Failed to create agent '{agent_name}': {e}", exc_info=True)
|
||||
|
||||
@ -113,6 +113,26 @@ def test_make_lead_agent_disables_thinking_when_model_does_not_support_it(monkey
|
||||
assert result["model"] is not None
|
||||
|
||||
|
||||
def test_make_lead_agent_rejects_invalid_bootstrap_agent_name(monkeypatch):
|
||||
app_config = _make_app_config([_make_model("safe-model", supports_thinking=False)])
|
||||
|
||||
monkeypatch.setattr(lead_agent_module, "get_app_config", lambda: app_config)
|
||||
|
||||
with pytest.raises(ValueError, match="Invalid agent name"):
|
||||
lead_agent_module.make_lead_agent(
|
||||
{
|
||||
"configurable": {
|
||||
"model_name": "safe-model",
|
||||
"thinking_enabled": False,
|
||||
"is_plan_mode": False,
|
||||
"subagent_enabled": False,
|
||||
"is_bootstrap": True,
|
||||
"agent_name": "../../../tmp/evil",
|
||||
}
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
def test_build_middlewares_uses_resolved_model_name_for_vision(monkeypatch):
|
||||
app_config = _make_app_config(
|
||||
[
|
||||
|
||||
40
backend/tests/test_setup_agent_tool.py
Normal file
40
backend/tests/test_setup_agent_tool.py
Normal file
@ -0,0 +1,40 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from pathlib import Path
|
||||
from types import SimpleNamespace
|
||||
|
||||
from deerflow.tools.builtins.setup_agent_tool import setup_agent
|
||||
|
||||
|
||||
class _DummyRuntime(SimpleNamespace):
|
||||
context: dict
|
||||
tool_call_id: str
|
||||
|
||||
|
||||
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"
|
||||
traversal_agent = f"../../../{outside_dir.name}/evil"
|
||||
runtime = _DummyRuntime(context={"agent_name": traversal_agent}, tool_call_id="tool-1")
|
||||
|
||||
result = setup_agent.func(soul="test soul", description="desc", runtime=runtime)
|
||||
|
||||
messages = result.update["messages"]
|
||||
assert len(messages) == 1
|
||||
assert "Invalid agent name" in messages[0].content
|
||||
assert not (tmp_path / "agents").exists()
|
||||
assert not (outside_dir / "evil" / "SOUL.md").exists()
|
||||
|
||||
|
||||
def test_setup_agent_rejects_absolute_agent_name_before_writing(tmp_path, monkeypatch):
|
||||
monkeypatch.setenv("DEER_FLOW_HOME", str(tmp_path))
|
||||
absolute_agent = str(tmp_path / "outside-agent")
|
||||
runtime = _DummyRuntime(context={"agent_name": absolute_agent}, tool_call_id="tool-2")
|
||||
|
||||
result = setup_agent.func(soul="test soul", description="desc", runtime=runtime)
|
||||
|
||||
messages = result.update["messages"]
|
||||
assert len(messages) == 1
|
||||
assert "Invalid agent name" in messages[0].content
|
||||
assert not (tmp_path / "agents").exists()
|
||||
assert not (Path(absolute_agent) / "SOUL.md").exists()
|
||||
Loading…
x
Reference in New Issue
Block a user