525 lines
19 KiB
Python
525 lines
19 KiB
Python
from __future__ import annotations
|
|
|
|
import json
|
|
from uuid import uuid4
|
|
|
|
import pytest
|
|
import pytest_asyncio
|
|
import yaml
|
|
from sqlalchemy import select
|
|
from sqlalchemy.ext.asyncio import async_sessionmaker, create_async_engine
|
|
|
|
from app.gateway.routers.agents import (
|
|
_backfill_listing_cache,
|
|
_load_config_extras,
|
|
_write_agent_config,
|
|
)
|
|
from app.gateway.services import build_run_config
|
|
from deerflow.config import paths as paths_module
|
|
from deerflow.config.agents_config import AgentConfig
|
|
from deerflow.persistence.agents import AgentRepository, MemoryAgentStore
|
|
from deerflow.persistence.agents.model import AgentFavoriteRow, AgentRow
|
|
from deerflow.persistence.base import Base
|
|
from deerflow.persistence.positions.memory import MemoryPositionStore
|
|
from deerflow.persistence.roundtable_chains.model import RoundtableChainRow
|
|
from deerflow.persistence.tags.memory import MemoryTagStore
|
|
from deerflow.services.position_sync import PositionSyncService
|
|
|
|
|
|
@pytest.fixture
|
|
def isolated_paths(tmp_path, monkeypatch):
|
|
"""Redirect get_paths() at a fresh temp directory for the duration of one test."""
|
|
monkeypatch.setenv("DEER_FLOW_HOME", str(tmp_path))
|
|
monkeypatch.setattr(paths_module, "_paths", None)
|
|
yield tmp_path
|
|
monkeypatch.setattr(paths_module, "_paths", None)
|
|
|
|
|
|
@pytest_asyncio.fixture
|
|
async def sql_session_factory(tmp_path):
|
|
engine = create_async_engine(f"sqlite+aiosqlite:///{tmp_path / 'agent-tests.db'}")
|
|
async with engine.begin() as conn:
|
|
# Importing PositionSyncService above registers the agent / skill /
|
|
# scheduled-task rows it materializes into.
|
|
await conn.run_sync(Base.metadata.create_all)
|
|
sf = async_sessionmaker(engine, expire_on_commit=False)
|
|
try:
|
|
yield sf
|
|
finally:
|
|
await engine.dispose()
|
|
|
|
|
|
async def _add_sql_agent(sf, *, agent_id: str, user_id: str | None) -> None:
|
|
async with sf() as session:
|
|
session.add(
|
|
AgentRow(
|
|
id=agent_id,
|
|
name=agent_id,
|
|
description="",
|
|
user_id=user_id,
|
|
published=True,
|
|
is_private=False,
|
|
square_id="",
|
|
)
|
|
)
|
|
await session.commit()
|
|
|
|
|
|
async def _add_sql_favorite(sf, *, user_id: str, agent_id: str, origin: str) -> None:
|
|
async with sf() as session:
|
|
session.add(
|
|
AgentFavoriteRow(
|
|
id=uuid4().hex,
|
|
user_id=user_id,
|
|
agent_id=agent_id,
|
|
origin=origin,
|
|
position_id="default" if origin == "position" else None,
|
|
)
|
|
)
|
|
await session.commit()
|
|
|
|
|
|
async def _add_sql_chain(sf, *, chain_id: str, agent_ids: list[str]) -> None:
|
|
async with sf() as session:
|
|
session.add(
|
|
RoundtableChainRow(
|
|
id=chain_id,
|
|
user_id="chain-owner",
|
|
title=chain_id,
|
|
description="",
|
|
seats=json.dumps(
|
|
[{"agent_id": agent_id, "name": agent_id, "description": ""} for agent_id in agent_ids],
|
|
ensure_ascii=False,
|
|
),
|
|
)
|
|
)
|
|
await session.commit()
|
|
|
|
|
|
async def _sql_favorite_rows(sf, user_id: str) -> set[tuple[str, str]]:
|
|
async with sf() as session:
|
|
rows = (
|
|
await session.execute(
|
|
select(AgentFavoriteRow.agent_id, AgentFavoriteRow.origin).where(
|
|
AgentFavoriteRow.user_id == user_id
|
|
)
|
|
)
|
|
).all()
|
|
return {(agent_id, origin) for agent_id, origin in rows}
|
|
|
|
|
|
async def _default_positions() -> MemoryPositionStore:
|
|
positions = MemoryPositionStore()
|
|
await positions.create_position(
|
|
{
|
|
"id": "default",
|
|
"name": "Default",
|
|
"description": "",
|
|
"is_default": True,
|
|
}
|
|
)
|
|
return positions
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_agent_store_allows_duplicate_chinese_names_and_filters_visibility() -> None:
|
|
store = MemoryAgentStore()
|
|
chinese_name = "\u4e2d\u6587\u52a9\u624b"
|
|
|
|
await store.ensure_builtin({"id": "builtin-agent", "name": chinese_name, "description": "built in"})
|
|
await store.create({"id": "private-a", "name": chinese_name, "description": "private", "user_id": "user-a"})
|
|
await store.create({"id": "published-b", "name": chinese_name, "description": "published", "user_id": "user-b", "published": True})
|
|
|
|
user_a_ids = {agent["id"] for agent in await store.list_visible("user-a")}
|
|
assert user_a_ids == {"builtin-agent", "private-a", "published-b"}
|
|
|
|
user_c_ids = {agent["id"] for agent in await store.list_visible("user-c")}
|
|
assert user_c_ids == {"builtin-agent", "published-b"}
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_position_sync_default_does_not_collect_builtin_agents(sql_session_factory) -> None:
|
|
await _add_sql_agent(sql_session_factory, agent_id="builtin-auto", user_id=None)
|
|
await _add_sql_agent(sql_session_factory, agent_id="builtin-manual", user_id=None)
|
|
await _add_sql_favorite(
|
|
sql_session_factory,
|
|
user_id="user-a",
|
|
agent_id="builtin-auto",
|
|
origin="position",
|
|
)
|
|
await _add_sql_favorite(
|
|
sql_session_factory,
|
|
user_id="user-a",
|
|
agent_id="builtin-manual",
|
|
origin="user",
|
|
)
|
|
|
|
sync = PositionSyncService(
|
|
sql_session_factory,
|
|
await _default_positions(),
|
|
MemoryTagStore(),
|
|
)
|
|
await sync.sync_user("user-a")
|
|
|
|
assert await _sql_favorite_rows(sql_session_factory, "user-a") == {
|
|
("builtin-manual", "user")
|
|
}
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_position_sync_tagged_agents_are_not_auto_favorited(sql_session_factory) -> None:
|
|
await _add_sql_agent(sql_session_factory, agent_id="builtin-tagged", user_id=None)
|
|
await _add_sql_agent(sql_session_factory, agent_id="custom-tagged", user_id="owner")
|
|
await _add_sql_agent(sql_session_factory, agent_id="chain-tagged", user_id="owner")
|
|
await _add_sql_chain(sql_session_factory, chain_id="chain-1", agent_ids=["chain-tagged"])
|
|
|
|
positions = await _default_positions()
|
|
tags = MemoryTagStore()
|
|
await tags.create_tag({"id": "agent-tag", "name": "Agent tag", "tag_type": "agent"})
|
|
await tags.set_target_tags("agent", "builtin-tagged", ["agent-tag"])
|
|
await tags.set_target_tags("agent", "custom-tagged", ["agent-tag"])
|
|
await tags.set_target_tags("agent", "chain-tagged", ["agent-tag"])
|
|
await positions.set_tag_bindings("default", {"agent": ["agent-tag"]})
|
|
|
|
sync = PositionSyncService(sql_session_factory, positions, tags)
|
|
await sync.sync_user("user-a")
|
|
|
|
assert await _sql_favorite_rows(sql_session_factory, "user-a") == {
|
|
("custom-tagged", "position")
|
|
}
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_mine_scope_ignores_auto_agent_favorites_but_keeps_manual_builtin(sql_session_factory) -> None:
|
|
await _add_sql_agent(sql_session_factory, agent_id="builtin-auto", user_id=None)
|
|
await _add_sql_agent(sql_session_factory, agent_id="builtin-manual", user_id=None)
|
|
await _add_sql_agent(sql_session_factory, agent_id="custom-auto", user_id="owner")
|
|
await _add_sql_agent(sql_session_factory, agent_id="chain-auto", user_id="owner")
|
|
await _add_sql_chain(sql_session_factory, chain_id="chain-1", agent_ids=["chain-auto"])
|
|
await _add_sql_favorite(
|
|
sql_session_factory,
|
|
user_id="user-a",
|
|
agent_id="builtin-auto",
|
|
origin="position",
|
|
)
|
|
await _add_sql_favorite(
|
|
sql_session_factory,
|
|
user_id="user-a",
|
|
agent_id="builtin-manual",
|
|
origin="user",
|
|
)
|
|
await _add_sql_favorite(
|
|
sql_session_factory,
|
|
user_id="user-a",
|
|
agent_id="custom-auto",
|
|
origin="position",
|
|
)
|
|
await _add_sql_favorite(
|
|
sql_session_factory,
|
|
user_id="user-a",
|
|
agent_id="chain-auto",
|
|
origin="position",
|
|
)
|
|
|
|
store = AgentRepository(sql_session_factory)
|
|
rows, total = await store.list_paginated(
|
|
user_id="user-a",
|
|
scope="mine",
|
|
search=None,
|
|
tag_ids=None,
|
|
page=1,
|
|
page_size=20,
|
|
)
|
|
|
|
assert total == 2
|
|
assert {row["id"] for row in rows} == {"builtin-manual", "custom-auto"}
|
|
assert await store.list_favorite_ids("user-a") == {"builtin-manual", "custom-auto"}
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_manual_favorite_upgrades_existing_auto_builtin_favorite(sql_session_factory) -> None:
|
|
await _add_sql_agent(sql_session_factory, agent_id="builtin-agent", user_id=None)
|
|
await _add_sql_favorite(
|
|
sql_session_factory,
|
|
user_id="user-a",
|
|
agent_id="builtin-agent",
|
|
origin="position",
|
|
)
|
|
|
|
store = AgentRepository(sql_session_factory)
|
|
assert await store.list_favorite_ids("user-a") == set()
|
|
|
|
await store.add_favorite("user-a", "builtin-agent")
|
|
|
|
assert await _sql_favorite_rows(sql_session_factory, "user-a") == {
|
|
("builtin-agent", "user")
|
|
}
|
|
assert await store.list_favorite_ids("user-a") == {"builtin-agent"}
|
|
rows, total = await store.list_paginated(
|
|
user_id="user-a",
|
|
scope="mine",
|
|
search=None,
|
|
tag_ids=None,
|
|
page=1,
|
|
page_size=20,
|
|
)
|
|
assert total == 1
|
|
assert [row["id"] for row in rows] == ["builtin-agent"]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_agent_store_only_owner_can_update_or_delete() -> None:
|
|
store = MemoryAgentStore()
|
|
|
|
await store.ensure_builtin({"id": "builtin-agent", "name": "Built in"})
|
|
await store.create({"id": "private-a", "name": "Owned", "user_id": "user-a"})
|
|
|
|
assert await store.update("builtin-agent", "user-a", {"name": "Nope"}) is None
|
|
assert await store.update("private-a", "user-b", {"name": "Nope"}) is None
|
|
|
|
updated = await store.update("private-a", "user-a", {"name": "Renamed"})
|
|
assert updated is not None
|
|
assert updated["name"] == "Renamed"
|
|
|
|
assert await store.delete("private-a", "user-b") is False
|
|
assert await store.delete("private-a", "user-a") is True
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_delete_builtin_only_targets_builtins() -> None:
|
|
store = MemoryAgentStore()
|
|
|
|
await store.ensure_builtin({"id": "builtin-agent", "name": "Built in"})
|
|
await store.create({"id": "owned-agent", "name": "Owned", "user_id": "user-a"})
|
|
|
|
# delete_builtin refuses to touch a user-owned agent
|
|
assert await store.delete_builtin("owned-agent") is False
|
|
assert await store.get_any("owned-agent") is not None
|
|
|
|
# delete_builtin succeeds on a built-in
|
|
assert await store.delete_builtin("builtin-agent") is True
|
|
assert await store.get_any("builtin-agent") is None
|
|
|
|
# And no-ops on a missing id
|
|
assert await store.delete_builtin("does-not-exist") is False
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_update_builtin_only_targets_builtins() -> None:
|
|
store = MemoryAgentStore()
|
|
|
|
await store.ensure_builtin({"id": "builtin-agent", "name": "Built in"})
|
|
await store.create({"id": "owned-agent", "name": "Owned", "user_id": "user-a"})
|
|
|
|
# update_builtin succeeds on a built-in agent
|
|
updated = await store.update_builtin("builtin-agent", {"name": "Renamed by admin"})
|
|
assert updated is not None
|
|
assert updated["name"] == "Renamed by admin"
|
|
|
|
# update_builtin refuses to touch a user-owned agent (returns None)
|
|
assert await store.update_builtin("owned-agent", {"name": "Hijacked"}) is None
|
|
|
|
# And on a missing agent
|
|
assert await store.update_builtin("does-not-exist", {"name": "Ghost"}) is None
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_get_editable_or_raise_admin_can_edit_builtin() -> None:
|
|
"""Route helper: admin gets edit access to any agent; non-admin still 403."""
|
|
from fastapi import HTTPException
|
|
|
|
from app.gateway.routers.agents import _get_editable_or_raise
|
|
|
|
store = MemoryAgentStore()
|
|
await store.ensure_builtin({"id": "builtin-agent", "name": "Built in"})
|
|
await store.create({"id": "owned-agent", "name": "Owned", "user_id": "user-a"})
|
|
|
|
# Owner edits their own agent → returns ('owner', record)
|
|
record, mode = await _get_editable_or_raise(
|
|
store, "owned-agent", "user-a", is_admin=False
|
|
)
|
|
assert mode == "owner" and record["id"] == "owned-agent"
|
|
|
|
# Admin edits a built-in agent → returns ('admin', record)
|
|
record, mode = await _get_editable_or_raise(
|
|
store, "builtin-agent", "user-a", is_admin=True
|
|
)
|
|
assert mode == "admin" and record["id"] == "builtin-agent"
|
|
|
|
# Non-admin tries to edit a built-in → 403
|
|
with pytest.raises(HTTPException) as exc:
|
|
await _get_editable_or_raise(store, "builtin-agent", "user-a", is_admin=False)
|
|
assert exc.value.status_code == 403
|
|
|
|
# Admin edits someone else's custom agent → returns ('admin', record).
|
|
# Behavior changed: admins manage the full agent catalog (built-in + any
|
|
# user's agents), not just built-ins. Non-admins are still owner-only.
|
|
await store.create(
|
|
{"id": "published-other", "name": "Pub", "user_id": "user-x", "published": True}
|
|
)
|
|
record, mode = await _get_editable_or_raise(
|
|
store, "published-other", "user-a", is_admin=True
|
|
)
|
|
assert mode == "admin" and record["id"] == "published-other"
|
|
|
|
|
|
def test_build_run_config_injects_custom_assistant_as_agent_id() -> None:
|
|
config = build_run_config("thread-1", None, None, assistant_id="agent-123")
|
|
|
|
assert config["configurable"]["thread_id"] == "thread-1"
|
|
assert config["configurable"]["agent_id"] == "agent-123"
|
|
assert "agent_name" not in config["configurable"]
|
|
|
|
|
|
def test_write_agent_config_persists_skills_into_config_yaml(isolated_paths) -> None:
|
|
agent_id = "agent-001"
|
|
_write_agent_config(
|
|
agent_id,
|
|
name="中文助手",
|
|
description="desc",
|
|
skills=["s1", "s2"],
|
|
model="gpt-x",
|
|
tool_groups=["tg-a"],
|
|
)
|
|
|
|
config_file = paths_module.get_paths().agent_dir(agent_id) / "config.yaml"
|
|
raw = yaml.safe_load(config_file.read_text(encoding="utf-8"))
|
|
|
|
assert raw["id"] == agent_id
|
|
assert raw["name"] == "中文助手"
|
|
assert raw["description"] == "desc"
|
|
assert raw["skills"] == ["s1", "s2"]
|
|
assert raw["model"] == "gpt-x"
|
|
assert raw["tool_groups"] == ["tg-a"]
|
|
|
|
|
|
def test_load_config_extras_returns_values_written_to_yaml(isolated_paths) -> None:
|
|
agent_id = "agent-002"
|
|
_write_agent_config(
|
|
agent_id,
|
|
name="agent-002",
|
|
description="",
|
|
skills=["s1"],
|
|
model=None,
|
|
tool_groups=None,
|
|
)
|
|
|
|
extras = _load_config_extras(agent_id)
|
|
assert extras == {"skills": ["s1"], "model": None, "tool_groups": None, "common_questions": None, "template_name": None, "seat_skill_directive": None}
|
|
|
|
|
|
def test_load_config_extras_missing_directory_returns_none_triple(isolated_paths) -> None:
|
|
extras = _load_config_extras("does-not-exist")
|
|
assert extras == {"skills": None, "model": None, "tool_groups": None, "common_questions": None, "template_name": None, "seat_skill_directive": None}
|
|
|
|
|
|
def test_write_agent_config_persists_common_questions(isolated_paths) -> None:
|
|
agent_id = "agent-cq-1"
|
|
# Structured {title, prompt} questions round-trip as clean dict mappings.
|
|
questions = [
|
|
{"title": "能做什么", "prompt": "这个智能体能做什么？"},
|
|
{"title": "上传文件", "prompt": "如何上传文件？"},
|
|
]
|
|
_write_agent_config(agent_id, name="助手", description="", common_questions=questions)
|
|
|
|
config_file = paths_module.get_paths().agent_dir(agent_id) / "config.yaml"
|
|
raw = yaml.safe_load(config_file.read_text(encoding="utf-8"))
|
|
assert raw["common_questions"] == questions
|
|
|
|
extras = _load_config_extras(agent_id)
|
|
assert [q.model_dump() for q in extras["common_questions"]] == questions
|
|
|
|
|
|
def test_write_agent_config_coerces_legacy_string_questions(isolated_paths) -> None:
|
|
"""Legacy ``list[str]`` configs coerce to {title, prompt} with title=content."""
|
|
agent_id = "agent-cq-legacy"
|
|
_write_agent_config(
|
|
agent_id, name="助手", description="", common_questions=["如何上传文件？"]
|
|
)
|
|
|
|
config_file = paths_module.get_paths().agent_dir(agent_id) / "config.yaml"
|
|
raw = yaml.safe_load(config_file.read_text(encoding="utf-8"))
|
|
assert raw["common_questions"] == [
|
|
{"title": "如何上传文件？", "prompt": "如何上传文件？"}
|
|
]
|
|
|
|
|
|
def test_write_agent_config_omits_common_questions_when_none(isolated_paths) -> None:
|
|
agent_id = "agent-cq-2"
|
|
_write_agent_config(agent_id, name="x", description="")
|
|
|
|
config_file = paths_module.get_paths().agent_dir(agent_id) / "config.yaml"
|
|
raw = yaml.safe_load(config_file.read_text(encoding="utf-8"))
|
|
assert "common_questions" not in raw
|
|
assert _load_config_extras(agent_id)["common_questions"] is None
|
|
|
|
|
|
def test_write_agent_config_explicit_empty_common_questions_is_persisted(isolated_paths) -> None:
|
|
agent_id = "agent-cq-3"
|
|
_write_agent_config(agent_id, name="x", description="", common_questions=[])
|
|
|
|
assert _load_config_extras(agent_id)["common_questions"] == []
|
|
|
|
|
|
def test_write_agent_config_explicit_empty_skills_is_persisted(isolated_paths) -> None:
|
|
agent_id = "agent-003"
|
|
_write_agent_config(agent_id, name="x", description="", skills=[])
|
|
|
|
extras = _load_config_extras(agent_id)
|
|
assert extras["skills"] == []
|
|
|
|
|
|
def test_write_agent_config_omits_skills_when_none(isolated_paths) -> None:
|
|
agent_id = "agent-004"
|
|
_write_agent_config(agent_id, name="x", description="")
|
|
|
|
config_file = paths_module.get_paths().agent_dir(agent_id) / "config.yaml"
|
|
raw = yaml.safe_load(config_file.read_text(encoding="utf-8"))
|
|
|
|
assert "skills" not in raw
|
|
assert "model" not in raw
|
|
assert "tool_groups" not in raw
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_backfill_listing_cache_mirrors_config_values() -> None:
|
|
store = MemoryAgentStore()
|
|
cfg = AgentConfig(id="a1", name="Agent", skills=["s1", "s2"], tool_groups=["g1"], model="m1")
|
|
await _backfill_listing_cache(store, [("a1", cfg)])
|
|
|
|
extras = await store.get_extras_for(["a1"])
|
|
assert extras["a1"]["skills"] == ["s1", "s2"]
|
|
assert extras["a1"]["tool_groups"] == ["g1"]
|
|
assert extras["a1"]["model"] == "m1"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_backfill_listing_cache_is_idempotent_and_skips_unchanged(monkeypatch) -> None:
|
|
store = MemoryAgentStore()
|
|
cfg = AgentConfig(id="a1", name="Agent", skills=["s1"], tool_groups=None, model=None)
|
|
await _backfill_listing_cache(store, [("a1", cfg)])
|
|
|
|
# On a second run with identical config, no re-write should occur.
|
|
calls = {"replace": 0, "set": 0}
|
|
orig_replace = store.replace_skills
|
|
orig_set = store.set_extras
|
|
|
|
async def _replace(agent_id, skills):
|
|
calls["replace"] += 1
|
|
await orig_replace(agent_id, skills)
|
|
|
|
async def _set(agent_id, *, tool_groups, model):
|
|
calls["set"] += 1
|
|
await orig_set(agent_id, tool_groups=tool_groups, model=model)
|
|
|
|
monkeypatch.setattr(store, "replace_skills", _replace)
|
|
monkeypatch.setattr(store, "set_extras", _set)
|
|
|
|
await _backfill_listing_cache(store, [("a1", cfg)])
|
|
assert calls == {"replace": 0, "set": 0}
|
|
|
|
# Changing the config triggers a re-write.
|
|
cfg2 = AgentConfig(id="a1", name="Agent", skills=["s1", "s2"], tool_groups=None, model=None)
|
|
await _backfill_listing_cache(store, [("a1", cfg2)])
|
|
assert calls["replace"] == 1
|