deerflow-code/offline-backend-20260512/backend/packages/harness/deerflow/persistence/agents/memory.py

"""In-memory custom agent store for memory persistence mode."""

from __future__ import annotations

from datetime import UTC, datetime
from typing import Any

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]] = {}

    @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"))

    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_visible(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)),
            "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"):
            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"):
            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]