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

"""In-memory ``SkillStore`` implementation for tests."""

from __future__ import annotations

import copy
from datetime import UTC, datetime
from typing import Any

from deerflow.persistence.skills.base import SkillStore


def _now() -> str:
    return datetime.now(UTC).isoformat()


def _copy(record: dict[str, Any] | None) -> dict[str, Any] | None:
    return copy.deepcopy(record) if record is not None else None


def _with_derived(record: dict[str, Any]) -> dict[str, Any]:
    out = dict(record)
    out["builtin"] = out.get("owner_user_id") is None
    out.setdefault("always_on", False)
    return out


class MemorySkillStore(SkillStore):
    def __init__(self) -> None:
        self._skills: dict[str, dict[str, Any]] = {}
        # user_id -> {skill_name}
        self._favorites: dict[str, set[str]] = {}

    @staticmethod
    def _is_visible(record: dict[str, Any], user_id: str) -> bool:
        owner = record.get("owner_user_id")
        if owner is None:
            return True
        if owner == user_id:
            return True
        return bool(record.get("published"))

    async def list_visible(self, user_id: str) -> list[dict[str, Any]]:
        return [
            _with_derived(_copy(record))  # type: ignore[arg-type]
            for record in self._skills.values()
            if self._is_visible(record, user_id)
        ]

    async def get_visible(self, name: str, user_id: str) -> dict[str, Any] | None:
        record = self._skills.get(name)
        if record is None or not self._is_visible(record, user_id):
            return None
        return _with_derived(_copy(record))  # type: ignore[arg-type]

    async def get_any(self, name: str) -> dict[str, Any] | None:
        record = self._skills.get(name)
        if record is None:
            return None
        return _with_derived(_copy(record))  # type: ignore[arg-type]

    async def get_owned(self, name: str, user_id: str) -> dict[str, Any] | None:
        record = self._skills.get(name)
        if record is None or record.get("owner_user_id") != user_id:
            return None
        return _with_derived(_copy(record))  # type: ignore[arg-type]

    async def create(self, data: dict[str, Any]) -> dict[str, Any]:
        name = data["name"]
        if name in self._skills:
            raise ValueError(f"Skill '{name}' already exists")
        record = {
            "name": name,
            "owner_user_id": data.get("owner_user_id"),
            "published": bool(data.get("published", False)),
            "detail": data.get("detail"),
            "name_zh": data.get("name_zh"),
            "always_on": bool(data.get("always_on", False)),
            "square_id": str(data.get("square_id") or ""),
            "created_at": _now(),
            "updated_at": _now(),
        }
        self._skills[name] = record
        return _with_derived(_copy(record))  # type: ignore[arg-type]

    async def ensure_legacy(self, data: dict[str, Any]) -> dict[str, Any]:
        name = data["name"]
        record = self._skills.get(name)
        if record is None:
            record = {
                "name": name,
                "owner_user_id": data.get("owner_user_id"),
                "published": bool(data.get("published", True)),
                "detail": data.get("detail"),
                "name_zh": data.get("name_zh"),
                "created_at": _now(),
                "updated_at": _now(),
            }
            self._skills[name] = record
        elif record.get("owner_user_id") is None and data.get("owner_user_id") is not None:
            record["owner_user_id"] = data["owner_user_id"]
            record["updated_at"] = _now()
        return _with_derived(_copy(record))  # type: ignore[arg-type]

    async def update(self, name: str, user_id: str, data: dict[str, Any]) -> dict[str, Any] | None:
        record = self._skills.get(name)
        if record is None or record.get("owner_user_id") != user_id:
            return None
        if "published" in data:
            record["published"] = bool(data["published"])
        if "detail" in data:
            record["detail"] = data["detail"]
        if "name_zh" in data:
            record["name_zh"] = data["name_zh"]
        if "always_on" in data:
            record["always_on"] = bool(data["always_on"])
        if "square_id" in data:
            record["square_id"] = str(data["square_id"] or "")
        record["updated_at"] = _now()
        return _with_derived(_copy(record))  # type: ignore[arg-type]

    async def update_admin(self, name: str, data: dict[str, Any]) -> dict[str, Any] | None:
        record = self._skills.get(name)
        if record is None:
            return None
        if "published" in data:
            record["published"] = bool(data["published"])
        if "owner_user_id" in data:
            record["owner_user_id"] = data["owner_user_id"]
        if "detail" in data:
            record["detail"] = data["detail"]
        if "name_zh" in data:
            record["name_zh"] = data["name_zh"]
        if "always_on" in data:
            record["always_on"] = bool(data["always_on"])
        if "square_id" in data:
            record["square_id"] = str(data["square_id"] or "")
        record["updated_at"] = _now()
        return _with_derived(_copy(record))  # type: ignore[arg-type]

    async def delete(self, name: str, user_id: str) -> bool:
        record = self._skills.get(name)
        if record is None or record.get("owner_user_id") != user_id:
            return False
        del self._skills[name]
        return True

    async def delete_admin(self, name: str) -> bool:
        if name not in self._skills:
            return False
        del self._skills[name]
        return True

    async def count_owned(self, user_id: str) -> int:
        return sum(1 for s in self._skills.values() if s.get("owner_user_id") == user_id)

    async def list_all(self) -> list[dict[str, Any]]:
        return [
            _with_derived(_copy(record))  # type: ignore[arg-type]
            for record in self._skills.values()
        ]

    async def list_published_names(self) -> set[str]:
        return {name for name, record in self._skills.items() if record.get("published")}

    async def list_always_on_names(self) -> set[str]:
        return {name for name, record in self._skills.items() if record.get("always_on")}

    async def list_favorite_skill_names(self, user_id: str) -> set[str]:
        return set(self._favorites.get(user_id, set()))

    # test helper — not part of the abstract API
    def add_favorite(self, user_id: str, skill_name: str) -> None:
        self._favorites.setdefault(user_id, set()).add(skill_name)