deer-flow/backend/packages/harness/deerflow/persistence/thread_meta/memory.py

"""In-memory ThreadMetaStore backed by LangGraph BaseStore.

Used when database.backend=memory. Delegates to the LangGraph Store's
``("threads",)`` namespace — the same namespace used by the Gateway
router for thread records.
"""

from __future__ import annotations

import uuid
from typing import Any

from langgraph.store.base import BaseStore

from deerflow.persistence.json_compat import json_value_matches
from deerflow.persistence.thread_meta.base import PROJECT_FILTER_UNSET, THREAD_ARCHIVED_METADATA_KEY, THREAD_PINNED_METADATA_KEY, ThreadMetaStore, ThreadOwnershipConflictError, _ProjectFilterUnset
from deerflow.runtime.keyed_lock import AsyncKeyedLockTable
from deerflow.runtime.user_context import AUTO, _AutoSentinel, resolve_user_id
from deerflow.utils.time import coerce_iso, now_iso

THREADS_NS: tuple[str, ...] = ("threads",)
SEARCH_PAGE_SIZE = 500


class MemoryThreadMetaStore(ThreadMetaStore):
    def __init__(self, store: BaseStore) -> None:
        self._store = store
        self._thread_locks = AsyncKeyedLockTable[str]()

    async def _get_owned_record(
        self,
        thread_id: str,
        user_id: str | None | _AutoSentinel,
        method_name: str,
    ) -> dict | None:
        """Fetch a record and verify ownership. Returns a mutable copy, or None."""
        resolved = resolve_user_id(user_id, method_name=method_name)
        item = await self._store.aget(THREADS_NS, thread_id)
        if item is None:
            return None
        record = dict(item.value)
        if resolved is not None and record.get("user_id") != resolved:
            return None
        return record

    async def create(
        self,
        thread_id: str,
        *,
        assistant_id: str | None = None,
        user_id: str | None | _AutoSentinel = AUTO,
        display_name: str | None = None,
        metadata: dict | None = None,
        project_id: str | None = None,
    ) -> dict:
        # Memory mode has no projects backend in Phase 1. Fail closed exactly
        # like the SQL store does for a missing/foreign/archived project: a
        # create carrying ``project_id`` raises ``ProjectNotAssignableError``
        # (the router maps it to 404) instead of silently persisting an
        # unassigned thread that a run would then proceed under. Mirrors
        # ``set_project`` below, which already reports rejection.
        if project_id is not None:
            from deerflow.persistence.projects import ProjectNotAssignableError

            raise ProjectNotAssignableError(project_id)
        resolved_user_id = resolve_user_id(user_id, method_name="MemoryThreadMetaStore.create")
        async with self._thread_locks.hold(thread_id):
            existing = await self._store.aget(THREADS_NS, thread_id)
            if existing is not None and resolved_user_id is not None and existing.value.get("user_id") != resolved_user_id:
                raise ThreadOwnershipConflictError(thread_id)
            now = now_iso()
            record: dict[str, Any] = {
                "thread_id": thread_id,
                "incarnation": (existing.value.get("incarnation") if existing is not None else None) or uuid.uuid4().hex,
                "assistant_id": assistant_id,
                "user_id": resolved_user_id,
                "display_name": display_name,
                "status": "idle",
                "metadata": metadata or {},
                "values": {},
                "created_at": now,
                "updated_at": now,
            }
            await self._store.aput(THREADS_NS, thread_id, record)
            return record

    async def claim_unowned(self, thread_id: str, owner: str) -> bool:
        async with self._thread_locks.hold(thread_id):
            item = await self._store.aget(THREADS_NS, thread_id)
            if item is None or item.value.get("user_id") is not None:
                return False
            record = dict(item.value)
            record["user_id"] = owner
            await self._store.aput(THREADS_NS, thread_id, record)
            return True

    async def set_project(self, thread_id: str, project_id: str | None, *, user_id: str | None | _AutoSentinel = AUTO) -> bool:
        # Memory mode has no projects backend in Phase 1: membership moves
        # are unsupported and always report rejection.
        return False

    async def get(self, thread_id: str, *, user_id: str | None | _AutoSentinel = AUTO) -> dict | None:
        return await self._get_owned_record(thread_id, user_id, "MemoryThreadMetaStore.get")

    async def search(
        self,
        *,
        metadata: dict[str, Any] | None = None,
        status: str | None = None,
        archived: bool | None = None,
        project_id: str | None | _ProjectFilterUnset = PROJECT_FILTER_UNSET,
        limit: int = 100,
        offset: int = 0,
        user_id: str | None | _AutoSentinel = AUTO,
    ) -> list[dict[str, Any]]:
        """Search threads by materializing matches, then sorting in Python.

        The memory backend loads all matching rows in chunks before slicing so
        it can mirror SQL's pinned-first ordering. Use the SQL store for
        scalable paginated I/O.
        """
        resolved_user_id = resolve_user_id(user_id, method_name="MemoryThreadMetaStore.search")
        if not isinstance(project_id, _ProjectFilterUnset) and project_id is not None:
            # Memory mode has no projects backend: no thread can be a member.
            return []
        filter_dict: dict[str, Any] = {}
        if status:
            filter_dict["status"] = status
        if resolved_user_id is not None:
            filter_dict["user_id"] = resolved_user_id

        items = []
        search_offset = 0
        while True:
            page = await self._store.asearch(
                THREADS_NS,
                filter=filter_dict or None,
                limit=SEARCH_PAGE_SIZE,
                offset=search_offset,
            )
            if not page:
                break
            items.extend(page)
            if len(page) < SEARCH_PAGE_SIZE:
                break
            search_offset += len(page)

        records = [self._item_to_dict(item) for item in items]
        if metadata:
            records = [record for record in records if isinstance(record.get("metadata"), dict) and all(json_value_matches(record["metadata"], key, value) for key, value in metadata.items())]
        if archived is not None:
            records = [record for record in records if ((record.get("metadata") or {}).get(THREAD_ARCHIVED_METADATA_KEY) is True) == archived]
        records.sort(key=self._sort_key, reverse=True)
        return records[offset : offset + limit]

    async def check_access(self, thread_id: str, user_id: str, *, require_existing: bool = False) -> bool:
        item = await self._store.aget(THREADS_NS, thread_id)
        if item is None:
            return not require_existing
        record_user_id = item.value.get("user_id")
        if record_user_id is None:
            return True
        return record_user_id == user_id

    async def update_display_name(
        self,
        thread_id: str,
        display_name: str,
        *,
        remove_metadata_keys: tuple[str, ...] = (),
        user_id: str | None | _AutoSentinel = AUTO,
    ) -> None:
        async with self._thread_locks.hold(thread_id):
            record = await self._get_owned_record(thread_id, user_id, "MemoryThreadMetaStore.update_display_name")
            if record is None:
                return
            record["display_name"] = display_name
            metadata = dict(record.get("metadata") or {})
            for key in remove_metadata_keys:
                metadata.pop(key, None)
            record["metadata"] = metadata
            record["updated_at"] = now_iso()
            await self._store.aput(THREADS_NS, thread_id, record)

    async def update_status(self, thread_id: str, status: str, *, user_id: str | None | _AutoSentinel = AUTO) -> None:
        async with self._thread_locks.hold(thread_id):
            record = await self._get_owned_record(thread_id, user_id, "MemoryThreadMetaStore.update_status")
            if record is None:
                return
            record["status"] = status
            record["updated_at"] = now_iso()
            await self._store.aput(THREADS_NS, thread_id, record)

    async def update_metadata(self, thread_id: str, metadata: dict, *, touch: bool = True, user_id: str | None | _AutoSentinel = AUTO) -> None:
        async with self._thread_locks.hold(thread_id):
            record = await self._get_owned_record(thread_id, user_id, "MemoryThreadMetaStore.update_metadata")
            if record is None:
                return
            merged = dict(record.get("metadata") or {})
            merged.update(metadata)
            record["metadata"] = merged
            if touch:
                record["updated_at"] = now_iso()
            await self._store.aput(THREADS_NS, thread_id, record)

    async def update_owner(self, thread_id: str, owner_user_id: str, *, user_id: str | None | _AutoSentinel = AUTO) -> None:
        async with self._thread_locks.hold(thread_id):
            record = await self._get_owned_record(thread_id, user_id, "MemoryThreadMetaStore.update_owner")
            if record is None:
                return
            record["user_id"] = owner_user_id
            record["updated_at"] = now_iso()
            await self._store.aput(THREADS_NS, thread_id, record)

    async def delete(self, thread_id: str, *, user_id: str | None | _AutoSentinel = AUTO) -> None:
        async with self._thread_locks.hold(thread_id):
            record = await self._get_owned_record(thread_id, user_id, "MemoryThreadMetaStore.delete")
            if record is None:
                return
            await self._store.adelete(THREADS_NS, thread_id)

    @staticmethod
    def _item_to_dict(item) -> dict[str, Any]:
        """Convert a Store SearchItem to the dict format expected by callers."""
        val = item.value
        return {
            "thread_id": item.key,
            "incarnation": val.get("incarnation"),
            "assistant_id": val.get("assistant_id"),
            "user_id": val.get("user_id"),
            "display_name": val.get("display_name"),
            "status": val.get("status", "idle"),
            "metadata": val.get("metadata", {}),
            # ``coerce_iso`` heals legacy unix-second values written by
            # earlier Gateway versions that called ``str(time.time())``.
            "created_at": coerce_iso(val.get("created_at", "")),
            "updated_at": coerce_iso(val.get("updated_at", "")),
        }

    @staticmethod
    def _sort_key(record: dict[str, Any]) -> tuple[bool, str, str]:
        metadata = record.get("metadata")
        pinned = isinstance(metadata, dict) and metadata.get(THREAD_PINNED_METADATA_KEY) is True
        return (pinned, str(record.get("updated_at") or ""), str(record.get("thread_id") or ""))