get_visible 此前比列表过滤更严格(仅 内置/本人/已发布),导致私密广场 智能体(转私密强制 published=false)和「我的」里已收藏但被取消发布的 智能体能在列表看到、点开问答却 404(前端连带报 X.map is not a function)。 - AgentRepository._visible_filter / MemoryAgentStore._is_fetchable: 单条读取放宽为所有列表口径的并集(+is_private 豁免、+本人收藏豁免) - 详情接口与 _enforce_agent_access 运行门禁:管理员回退 get_any, 对齐 ?admin=true 全量列表 - 普通未发布草稿对陌生人仍不可见,隐私不放开 - 本地新增 tests/test_agent_visibility.py 覆盖 SQL/Memory 双实现 (tests/ 目录在本仓库被 gitignore,不随提交) Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
323 lines
13 KiB
Python
323 lines
13 KiB
Python
"""In-memory custom agent store for memory persistence mode."""
|
|
|
|
from __future__ import annotations
|
|
|
|
from datetime import UTC, datetime
|
|
from typing import Any
|
|
|
|
from deerflow.config.agent_name_desensitize import name_matches
|
|
from deerflow.persistence.agents.base import AgentStore
|
|
|
|
|
|
def _now() -> str:
|
|
return datetime.now(UTC).isoformat()
|
|
|
|
|
|
def _copy(data: dict[str, Any]) -> dict[str, Any]:
|
|
copied = dict(data)
|
|
copied["builtin"] = copied.get("user_id") is None
|
|
return copied
|
|
|
|
|
|
class MemoryAgentStore(AgentStore):
|
|
def __init__(self) -> None:
|
|
self._agents: dict[str, dict[str, Any]] = {}
|
|
self._favorites: set[tuple[str, str]] = set()
|
|
self._pins: set[tuple[str, str]] = set()
|
|
# Listing cache mirrors (see AgentExtrasRow / AgentSkillRow):
|
|
self._skills: dict[str, list[str]] = {}
|
|
self._extras: dict[str, dict[str, Any]] = {}
|
|
|
|
@staticmethod
|
|
def _is_visible(agent: dict[str, Any], user_id: str) -> bool:
|
|
return agent.get("user_id") is None or agent.get("user_id") == user_id or bool(agent.get("published"))
|
|
|
|
def _is_fetchable(self, agent: dict[str, Any], user_id: str) -> bool:
|
|
"""Single-record visibility (``get_visible``) — the union of every
|
|
listing surface, so a list-visible agent always opens. Mirrors
|
|
``AgentRepository._visible_filter``: baseline visibility OR
|
|
private-square (password is the discovery gate, reads are not
|
|
re-gated) OR favorited (the "我的" listing keeps favorites even after
|
|
the owner unpublishes)."""
|
|
return self._is_visible(agent, user_id) or bool(agent.get("is_private")) or (user_id, agent["id"]) in self._favorites
|
|
|
|
async def list_visible(self, user_id: str) -> list[dict[str, Any]]:
|
|
rows = [_copy(agent) for agent in self._agents.values() if self._is_visible(agent, user_id)]
|
|
rows.sort(key=lambda agent: agent.get("created_at") or "")
|
|
return rows
|
|
|
|
async def get_visible(self, agent_id: str, user_id: str) -> dict[str, Any] | None:
|
|
agent = self._agents.get(agent_id)
|
|
if agent is None or not self._is_fetchable(agent, user_id):
|
|
return None
|
|
return _copy(agent)
|
|
|
|
async def get_any(self, agent_id: str) -> dict[str, Any] | None:
|
|
agent = self._agents.get(agent_id)
|
|
return _copy(agent) if agent is not None else None
|
|
|
|
async def get_owned(self, agent_id: str, user_id: str) -> dict[str, Any] | None:
|
|
agent = self._agents.get(agent_id)
|
|
if agent is None or agent.get("user_id") != user_id:
|
|
return None
|
|
return _copy(agent)
|
|
|
|
async def create(self, data: dict[str, Any]) -> dict[str, Any]:
|
|
now = _now()
|
|
row = {
|
|
"id": data["id"],
|
|
"name": data["name"],
|
|
"description": data.get("description") or "",
|
|
"user_id": data.get("user_id"),
|
|
"published": bool(data.get("published", False)),
|
|
"is_private": bool(data.get("is_private", False)),
|
|
"square_id": str(data.get("square_id") or ""),
|
|
"featured_order": None,
|
|
"created_at": now,
|
|
"updated_at": now,
|
|
}
|
|
self._agents[row["id"]] = row
|
|
return _copy(row)
|
|
|
|
async def ensure_builtin(self, data: dict[str, Any]) -> dict[str, Any]:
|
|
row = self._agents.get(data["id"])
|
|
if row is None:
|
|
row = {
|
|
"id": data["id"],
|
|
"name": data["name"],
|
|
"description": data.get("description") or "",
|
|
"user_id": None,
|
|
"published": bool(data.get("published", True)),
|
|
"featured_order": None,
|
|
"created_at": _now(),
|
|
"updated_at": _now(),
|
|
}
|
|
self._agents[row["id"]] = row
|
|
elif row.get("user_id") is None:
|
|
row["name"] = data["name"]
|
|
row["description"] = data.get("description") or ""
|
|
row["published"] = bool(data.get("published", True))
|
|
row["updated_at"] = _now()
|
|
return _copy(row)
|
|
|
|
async def update(self, agent_id: str, user_id: str, data: dict[str, Any]) -> dict[str, Any] | None:
|
|
row = self._agents.get(agent_id)
|
|
if row is None or row.get("user_id") != user_id:
|
|
return None
|
|
for key in ("name", "description", "published", "is_private", "square_id"):
|
|
if key in data:
|
|
row[key] = data[key]
|
|
row["updated_at"] = _now()
|
|
return _copy(row)
|
|
|
|
async def update_builtin(self, agent_id: str, data: dict[str, Any]) -> dict[str, Any] | None:
|
|
row = self._agents.get(agent_id)
|
|
if row is None or row.get("user_id") is not None:
|
|
return None
|
|
for key in ("name", "description", "published", "is_private", "square_id"):
|
|
if key in data:
|
|
row[key] = data[key]
|
|
row["updated_at"] = _now()
|
|
return _copy(row)
|
|
|
|
async def delete(self, agent_id: str, user_id: str) -> bool:
|
|
row = self._agents.get(agent_id)
|
|
if row is None or row.get("user_id") != user_id:
|
|
return False
|
|
self._agents.pop(agent_id, None)
|
|
return True
|
|
|
|
async def delete_builtin(self, agent_id: str) -> bool:
|
|
row = self._agents.get(agent_id)
|
|
if row is None or row.get("user_id") is not None:
|
|
return False
|
|
self._agents.pop(agent_id, None)
|
|
return True
|
|
|
|
async def count_owned(self, user_id: str) -> int:
|
|
return sum(1 for a in self._agents.values() if a.get("user_id") == user_id)
|
|
|
|
async def set_featured(self, agent_id: str, order: int | None) -> dict[str, Any] | None:
|
|
row = self._agents.get(agent_id)
|
|
if row is None:
|
|
return None
|
|
row["featured_order"] = order
|
|
row["updated_at"] = _now()
|
|
return _copy(row)
|
|
|
|
async def list_featured(self, user_id: str) -> list[dict[str, Any]]:
|
|
candidates = [a for a in self._agents.values() if a.get("featured_order") is not None and self._is_visible(a, user_id)]
|
|
candidates.sort(key=lambda a: (a.get("featured_order"), a.get("created_at") or ""))
|
|
return [_copy(a) for a in candidates]
|
|
|
|
async def list_by_ids(self, agent_ids: list[str], user_id: str) -> list[dict[str, Any]]:
|
|
order_map = {aid: idx for idx, aid in enumerate(agent_ids)}
|
|
rows = [self._agents[aid] for aid in agent_ids if aid in self._agents and self._is_visible(self._agents[aid], user_id)]
|
|
rows.sort(key=lambda a: order_map.get(a["id"], 0))
|
|
return [_copy(a) for a in rows]
|
|
|
|
async def list_private(self) -> list[dict[str, Any]]:
|
|
rows = [_copy(a) for a in self._agents.values() if bool(a.get("is_private"))]
|
|
rows.sort(key=lambda a: a.get("created_at") or "")
|
|
return rows
|
|
|
|
async def list_all(self) -> list[dict[str, Any]]:
|
|
rows = [_copy(a) for a in self._agents.values()]
|
|
rows.sort(key=lambda a: a.get("created_at") or "")
|
|
return rows
|
|
|
|
async def update_any(self, agent_id: str, data: dict[str, Any]) -> dict[str, Any] | None:
|
|
row = self._agents.get(agent_id)
|
|
if row is None:
|
|
return None
|
|
for key in ("name", "description", "published", "is_private", "square_id"):
|
|
if key in data:
|
|
row[key] = data[key]
|
|
row["updated_at"] = _now()
|
|
return _copy(row)
|
|
|
|
async def delete_any(self, agent_id: str) -> bool:
|
|
if agent_id not in self._agents:
|
|
return False
|
|
self._agents.pop(agent_id, None)
|
|
return True
|
|
|
|
# --- Paginated listing ---
|
|
|
|
def _matches_scope(self, agent: dict[str, Any], *, user_id: str, scope: str | None, admin: bool) -> bool:
|
|
if admin:
|
|
return True
|
|
if scope == "square":
|
|
return (agent.get("user_id") is None or bool(agent.get("published"))) and not bool(agent.get("is_private"))
|
|
if scope == "mine":
|
|
owned = agent.get("user_id") is not None and agent.get("user_id") == user_id
|
|
favorited = (user_id, agent["id"]) in self._favorites
|
|
return owned or favorited
|
|
if scope == "private":
|
|
return bool(agent.get("is_private"))
|
|
return self._is_visible(agent, user_id)
|
|
|
|
async def list_paginated(
|
|
self,
|
|
*,
|
|
user_id: str,
|
|
scope: str | None,
|
|
search: str | None,
|
|
tag_ids: list[str] | None,
|
|
page: int,
|
|
page_size: int,
|
|
admin: bool = False,
|
|
sort: str = "time_desc",
|
|
owner: str | None = None,
|
|
square_id: str | None = None,
|
|
square_is_default: bool = False,
|
|
) -> tuple[list[dict[str, Any]], int]:
|
|
# NOTE: the in-memory store has no tag-assignment data, so ``tag_ids``
|
|
# is ignored here. Tag-filter behavior is covered by the SQL repo tests.
|
|
owner_key = owner.strip() if owner and owner.strip() else None
|
|
square_key = square_id.strip() if square_id and square_id.strip() else None
|
|
|
|
def _square_ok(a: dict[str, Any]) -> bool:
|
|
if square_key is None:
|
|
return True
|
|
sid = str(a.get("square_id") or "")
|
|
return sid == square_key or (square_is_default and sid == "")
|
|
|
|
matched = [
|
|
a
|
|
for a in self._agents.values()
|
|
if self._matches_scope(a, user_id=user_id, scope=scope, admin=admin)
|
|
and (not search or not search.strip() or name_matches(a.get("name") or "", search))
|
|
and (owner_key is None or a.get("user_id") == owner_key)
|
|
and _square_ok(a)
|
|
]
|
|
if sort == "time_asc":
|
|
matched.sort(key=lambda a: (a.get("created_at") or "", a.get("id") or ""))
|
|
elif sort == "updated_desc":
|
|
matched.sort(
|
|
key=lambda a: (a.get("updated_at") or "", a.get("id") or ""),
|
|
reverse=True,
|
|
)
|
|
elif sort == "updated_asc":
|
|
matched.sort(key=lambda a: (a.get("updated_at") or "", a.get("id") or ""))
|
|
elif sort == "name_asc":
|
|
matched.sort(key=lambda a: ((a.get("name") or "").lower(), a.get("id") or ""))
|
|
elif sort == "name_desc":
|
|
matched.sort(
|
|
key=lambda a: ((a.get("name") or "").lower(), a.get("id") or ""),
|
|
reverse=True,
|
|
)
|
|
else:
|
|
matched.sort(
|
|
key=lambda a: (a.get("created_at") or "", a.get("id") or ""),
|
|
reverse=True,
|
|
)
|
|
pin_ids = {agent_id for uid, agent_id in self._pins if uid == user_id}
|
|
matched.sort(key=lambda a: (0 if a.get("id") in pin_ids else 1,))
|
|
total = len(matched)
|
|
start = max(page - 1, 0) * page_size
|
|
page_rows = matched[start : start + page_size]
|
|
return [_copy(a) for a in page_rows], total
|
|
|
|
# --- Listing cache (skills / tool_groups / model) ---
|
|
|
|
async def replace_skills(self, agent_id: str, skills: list[str] | None) -> None:
|
|
self._skills[agent_id] = [s for s in (skills or []) if isinstance(s, str)]
|
|
|
|
async def set_extras(self, agent_id: str, *, tool_groups: list[str] | None, model: str | None) -> None:
|
|
self._extras[agent_id] = {"tool_groups": tool_groups, "model": model}
|
|
|
|
async def get_extras_for(self, agent_ids: list[str]) -> dict[str, dict[str, Any]]:
|
|
out: dict[str, dict[str, Any]] = {}
|
|
for agent_id in agent_ids:
|
|
has_extras = agent_id in self._extras
|
|
has_skills = agent_id in self._skills
|
|
if not has_extras and not has_skills:
|
|
continue
|
|
extras = self._extras.get(agent_id, {})
|
|
out[agent_id] = {
|
|
"skills": list(self._skills.get(agent_id, [])),
|
|
"tool_groups": extras.get("tool_groups"),
|
|
"model": extras.get("model"),
|
|
}
|
|
return out
|
|
|
|
async def delete_extras_for_agent(self, agent_id: str) -> None:
|
|
self._skills.pop(agent_id, None)
|
|
self._extras.pop(agent_id, None)
|
|
|
|
# --- Favorites ("收藏") ---
|
|
|
|
async def add_favorite(self, user_id: str, agent_id: str) -> None:
|
|
self._favorites.add((user_id, agent_id))
|
|
|
|
async def remove_favorite(self, user_id: str, agent_id: str) -> bool:
|
|
if (user_id, agent_id) in self._favorites:
|
|
self._favorites.discard((user_id, agent_id))
|
|
return True
|
|
return False
|
|
|
|
async def list_favorite_ids(self, user_id: str) -> set[str]:
|
|
return {agent_id for (uid, agent_id) in self._favorites if uid == user_id}
|
|
|
|
async def delete_favorites_for_agent(self, agent_id: str) -> None:
|
|
self._favorites = {(uid, aid) for (uid, aid) in self._favorites if aid != agent_id}
|
|
|
|
# --- Pins ("置顶") ---
|
|
|
|
async def add_pin(self, user_id: str, agent_id: str) -> None:
|
|
self._pins.add((user_id, agent_id))
|
|
|
|
async def remove_pin(self, user_id: str, agent_id: str) -> bool:
|
|
key = (user_id, agent_id)
|
|
if key not in self._pins:
|
|
return False
|
|
self._pins.remove(key)
|
|
return True
|
|
|
|
async def list_pin_ids(self, user_id: str) -> set[str]:
|
|
return {agent_id for uid, agent_id in self._pins if uid == user_id}
|
|
|
|
async def delete_pins_for_agent(self, agent_id: str) -> None:
|
|
self._pins = {(uid, aid) for (uid, aid) in self._pins if aid != agent_id}
|