deerflow-code/offline-backend-20260512/backend/packages/harness/deerflow/persistence/thread_meta/sql.py

"""SQLAlchemy-backed thread metadata repository."""

from __future__ import annotations

from datetime import UTC, datetime
from typing import Any

from sqlalchemy import func, or_, select, update
from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker

from deerflow.persistence.json_compat import JsonPredicate, supports_json_filter
from deerflow.persistence.task_deeplink_threads.sql import (
    delete_task_deeplink_index,
    sync_task_deeplink_index,
    update_task_deeplink_card,
)
from deerflow.persistence.thread_meta.base import EXCLUDED_SYSTEM_THREAD_TYPES, ThreadMetaStore, is_excluded_system_thread, is_shared_access_thread
from deerflow.persistence.thread_meta.model import ThreadMetaRow
from deerflow.runtime.user_context import AUTO, _AutoSentinel, resolve_user_id

_INDEXED_METADATA_COLUMNS = {
    "taskId": ThreadMetaRow.task_id,
    "agent_id": ThreadMetaRow.agent_id,
    "actionId": ThreadMetaRow.action_id,
}

_INDEXED_METADATA_ATTRIBUTES = {
    "taskId": "task_id",
    "agent_id": "agent_id",
    "actionId": "action_id",
}


def _indexed_metadata_values(metadata: dict[str, Any] | None) -> dict[str, str]:
    """Return metadata predicates that have a maintained SQL projection.

    Task deep-link callers use string ids.  Restrict the SQL shortcut to
    strings so legacy callers that intentionally store another JSON type keep
    the exact Python-side equality semantics of the generic metadata filter.
    """
    return {
        key: value
        for key, value in (metadata or {}).items()
        if key in _INDEXED_METADATA_COLUMNS and isinstance(value, str)
    }


def _indexed_thread_metadata(metadata: dict[str, Any] | None) -> dict[str, str | None]:
    """Build ORM column values from authoritative thread metadata.

    Always include every projection so a metadata update that changes a value
    from a string to another JSON type cannot leave a stale indexed value.
    """
    values = _indexed_metadata_values(metadata)
    return {
        attribute: values.get(metadata_key)
        for metadata_key, attribute in _INDEXED_METADATA_ATTRIBUTES.items()
    }


class ThreadMetaRepository(ThreadMetaStore):
    def __init__(self, session_factory: async_sessionmaker[AsyncSession]) -> None:
        self._sf = session_factory

    @staticmethod
    def _row_to_dict(row: ThreadMetaRow) -> dict[str, Any]:
        d = row.to_dict()
        d["metadata"] = d.pop("metadata_json", {})
        for key in ("created_at", "updated_at"):
            val = d.get(key)
            if isinstance(val, datetime):
                d[key] = val.isoformat()
        return d

    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,
    ) -> dict:
        # Auto-resolve user_id from contextvar when AUTO; explicit None
        # creates an orphan row (used by migration scripts).
        resolved_user_id = resolve_user_id(user_id, method_name="ThreadMetaRepository.create")
        now = datetime.now(UTC)
        row = ThreadMetaRow(
            thread_id=thread_id,
            assistant_id=assistant_id,
            user_id=resolved_user_id,
            **_indexed_thread_metadata(metadata),
            display_name=display_name,
            status="idle",
            metadata_json=metadata or {},
            created_at=now,
            updated_at=now,
        )
        async with self._sf() as session:
            session.add(row)
            await sync_task_deeplink_index(session, row)
            await session.commit()
            # Deliberately no session.refresh(): every column above is assigned
            # explicitly and the session factory uses expire_on_commit=False, so
            # the in-memory row is already complete. The re-SELECT would run on a
            # different pooled connection than the INSERT, which a MySQL
            # read/write-splitting endpoint can route to a lagging replica --
            # the row comes back empty and SQLAlchemy raises
            # "Could not refresh instance". Concurrency makes that far more
            # likely, so a single-user deployment never sees it.
            return self._row_to_dict(row)

    async def get(
        self,
        thread_id: str,
        *,
        user_id: str | None | _AutoSentinel = AUTO,
    ) -> dict | None:
        resolved_user_id = resolve_user_id(user_id, method_name="ThreadMetaRepository.get")
        async with self._sf() as session:
            row = await session.get(ThreadMetaRow, thread_id)
            if row is None:
                return None
            # Enforce owner filter unless explicitly bypassed (user_id=None).
            if resolved_user_id is not None and row.user_id != resolved_user_id:
                return None
            return self._row_to_dict(row)

    async def check_access(self, thread_id: str, user_id: str, *, require_existing: bool = False) -> bool:
        """Check if ``user_id`` has access to ``thread_id``.

        Two modes — one row, two distinct semantics depending on what
        the caller is about to do:

        - ``require_existing=False`` (default, permissive):
          Returns True for: row missing (untracked legacy thread),
          ``row.user_id`` is None (shared / pre-auth data),
          or ``row.user_id == user_id``. Use for **read-style**
          decorators where treating an untracked thread as accessible
          preserves backward-compat.

        - ``require_existing=True`` (strict):
          Returns True **only** when the row exists AND
          (``row.user_id == user_id`` OR ``row.user_id is None``).
          Use for **destructive / mutating** decorators (DELETE, PATCH,
          state-update) so a thread that has *already been deleted*
          cannot be re-targeted by any caller — closing the
          delete-idempotence cross-user gap where the row vanishing
          made every other user appear to "own" it.
        """
        async with self._sf() as session:
            row = await session.get(ThreadMetaRow, thread_id)
            if row is None:
                return not require_existing
            if row.user_id:
                from deerflow.runtime.thread_paths import remember_thread_owner

                remember_thread_owner(thread_id, row.user_id)
            md = row.metadata_json or {}
            # 全局共享、不区分用户（任何登录用户都可读写/续聊）：
            #  ① 任务深链对话（metadata.taskId）；
            #  ② 圆桌 / 任务分析等系统编排线程（thread_type 见 SHARED_ACCESS_THREAD_TYPES）
            #     —— loopback 派活签发的 JWT 可能与建 thread 时 HTTP 上下文用户不是同一人，
            #     若按属主拦截会把刚建好的总控 thread 打成 404 not accessible。
            if is_shared_access_thread(md):
                return True
            if row.user_id is None:
                return True
            return row.user_id == user_id

    @staticmethod
    def _query_conditions(*, dialect, metadata, status, user_id, exclude_system, query):
        """Share exact search/count predicates without changing owner resolution.

        Complex metadata and unsupported dialects retain their exact Python
        filter. No rejected key/value is silently ignored. Server numeric JSON
        comparisons retain that path: DECIMAL and DOUBLE cannot reproduce all
        Python float/int comparisons (including large numeric task ids).
        """
        conditions = []
        if user_id is not None:
            conditions.append(ThreadMetaRow.user_id == user_id)
        if status:
            conditions.append(ThreadMetaRow.status == status)
        if query and query.strip():
            conditions.append(func.lower(ThreadMetaRow.display_name).contains(query.strip().lower()))
        # Retain the master branch's indexed task/agent/action narrowing.
        # JSON predicates below still enforce exact case/type semantics.
        for key, value in _indexed_metadata_values(metadata).items():
            conditions.append(_INDEXED_METADATA_COLUMNS[key] == value)
        supported = dialect in {"sqlite", "postgresql", "mysql"}
        residual = {}
        for key, value in (metadata or {}).items():
            numeric_server = dialect in {"mysql", "postgresql"} and isinstance(value, (bool, int, float))
            if supported and supports_json_filter(key, value) and not numeric_server:
                conditions.append(JsonPredicate(ThreadMetaRow.metadata_json, key, value=value))
            else:
                residual[key] = value
        if exclude_system and supported:
            column = ThreadMetaRow.metadata_json
            conditions.append(~or_(
                JsonPredicate(column, "system", "truthy"),
                JsonPredicate(column, "task_analysis_seat", "nonblank"),
                *(JsonPredicate(column, "thread_type", value=value) for value in sorted(EXCLUDED_SYSTEM_THREAD_TYPES)),
            ))
        return conditions, residual, exclude_system and not supported

    async def search(
        self,
        *,
        metadata: dict | None = None,
        status: str | None = None,
        limit: int = 100,
        offset: int = 0,
        user_id: str | None | _AutoSentinel = AUTO,
        exclude_system: bool = False,
        query: str | None = None,
    ) -> list[dict]:
        resolved_user_id = resolve_user_id(user_id, method_name="ThreadMetaRepository.search")
        # Preserve task-scoped sharing, still applying EVERY metadata filter
        # (taskId + agent_id + actionId) before the pagination window.
        if metadata and metadata.get("taskId"):
            resolved_user_id = None
        async with self._sf() as session:
            conditions, residual, filter_system = self._query_conditions(
                dialect=session.get_bind().dialect.name,
                metadata=metadata, status=status, user_id=resolved_user_id,
                exclude_system=exclude_system, query=query,
            )
            stmt = select(ThreadMetaRow).where(*conditions).order_by(
                ThreadMetaRow.updated_at.desc(), ThreadMetaRow.thread_id.desc(),
            )
            if not residual and not filter_system:
                result = await session.execute(stmt.limit(limit).offset(offset))
                return [self._row_to_dict(row) for row in result.scalars()]

            # Compatibility path for nested values, unusual literal keys,
            # numeric server filters, or the separate openGauss dialect.
            needed = offset + limit
            batch_size = max(500, needed)
            matched = []
            scan_offset = 0
            while len(matched) < needed:
                result = await session.execute(stmt.limit(batch_size).offset(scan_offset))
                rows = result.scalars().all()
                if not rows:
                    break
                for row in rows:
                    record = self._row_to_dict(row)
                    md = record.get("metadata") or {}
                    if all(md.get(key) == value for key, value in residual.items()) and not (
                        filter_system and is_excluded_system_thread(md)
                    ):
                        matched.append(record)
                if len(rows) < batch_size:
                    break
                scan_offset += len(rows)
            return matched[offset:needed]

    async def count(
        self,
        *,
        metadata: dict | None = None,
        status: str | None = None,
        user_id: str | None | _AutoSentinel = AUTO,
        exclude_system: bool = False,
        query: str | None = None,
    ) -> int:
        resolved_user_id = resolve_user_id(user_id, method_name="ThreadMetaRepository.count")
        if metadata and metadata.get("taskId"):
            resolved_user_id = None
        async with self._sf() as session:
            conditions, residual, filter_system = self._query_conditions(
                dialect=session.get_bind().dialect.name,
                metadata=metadata, status=status, user_id=resolved_user_id,
                exclude_system=exclude_system, query=query,
            )
            if not residual and not filter_system:
                stmt = select(func.count()).select_from(ThreadMetaRow).where(*conditions)
                return int((await session.execute(stmt)).scalar_one())

            stmt = select(ThreadMetaRow.metadata_json).where(*conditions).order_by(ThreadMetaRow.thread_id)
            total = 0
            scan_offset = 0
            batch_size = 1000
            while True:
                rows = (await session.execute(stmt.limit(batch_size).offset(scan_offset))).scalars().all()
                if not rows:
                    break
                for md in rows:
                    md = md or {}
                    if all(md.get(key) == value for key, value in residual.items()) and not (
                        filter_system and is_excluded_system_thread(md)
                    ):
                        total += 1
                if len(rows) < batch_size:
                    break
                scan_offset += len(rows)
            return total

    async def _check_ownership(self, session: AsyncSession, thread_id: str, resolved_user_id: str | None) -> bool:
        """Return True if the row exists and is owned (or filter bypassed)."""
        if resolved_user_id is None:
            return True  # explicit bypass
        row = await session.get(ThreadMetaRow, thread_id)
        if row is None:
            return False
        # 任务深链对话（metadata.taskId）全局共享：任何用户都可更新标题/状态/updated_at（续聊置顶）。
        if (row.metadata_json or {}).get("taskId"):
            return True
        return row.user_id == resolved_user_id

    async def update_display_name(
        self,
        thread_id: str,
        display_name: str,
        *,
        user_id: str | None | _AutoSentinel = AUTO,
    ) -> None:
        """Update the display_name (title) for a thread."""
        resolved_user_id = resolve_user_id(user_id, method_name="ThreadMetaRepository.update_display_name")
        async with self._sf() as session:
            if not await self._check_ownership(session, thread_id, resolved_user_id):
                return
            now = datetime.now(UTC)
            await session.execute(update(ThreadMetaRow).where(ThreadMetaRow.thread_id == thread_id).values(display_name=display_name, updated_at=now))
            await update_task_deeplink_card(session, thread_id, display_name=display_name, updated_at=now)
            await session.commit()

    async def update_status(
        self,
        thread_id: str,
        status: str,
        *,
        user_id: str | None | _AutoSentinel = AUTO,
    ) -> None:
        resolved_user_id = resolve_user_id(user_id, method_name="ThreadMetaRepository.update_status")
        async with self._sf() as session:
            if not await self._check_ownership(session, thread_id, resolved_user_id):
                return
            now = datetime.now(UTC)
            await session.execute(update(ThreadMetaRow).where(ThreadMetaRow.thread_id == thread_id).values(status=status, updated_at=now))
            await update_task_deeplink_card(session, thread_id, status=status, updated_at=now)
            await session.commit()

    async def update_assistant_id(
        self,
        thread_id: str,
        assistant_id: str,
        *,
        user_id: str | None | _AutoSentinel = AUTO,
    ) -> None:
        resolved_user_id = resolve_user_id(user_id, method_name="ThreadMetaRepository.update_assistant_id")
        async with self._sf() as session:
            if not await self._check_ownership(session, thread_id, resolved_user_id):
                return
            now = datetime.now(UTC)
            await session.execute(update(ThreadMetaRow).where(ThreadMetaRow.thread_id == thread_id).values(assistant_id=assistant_id, updated_at=now))
            await update_task_deeplink_card(session, thread_id, updated_at=now)
            await session.commit()

    async def update_metadata(
        self,
        thread_id: str,
        metadata: dict,
        *,
        user_id: str | None | _AutoSentinel = AUTO,
    ) -> None:
        """Merge ``metadata`` into ``metadata_json``.

        Read-modify-write inside a single session/transaction so concurrent
        callers see consistent state. No-op if the row does not exist or
        the user_id check fails.
        """
        resolved_user_id = resolve_user_id(user_id, method_name="ThreadMetaRepository.update_metadata")
        async with self._sf() as session:
            row = await session.get(ThreadMetaRow, thread_id)
            if row is None:
                return
            # 任务深链对话（metadata.taskId）全局共享：跳过 owner 校验。
            task_shared = bool((row.metadata_json or {}).get("taskId"))
            if resolved_user_id is not None and not task_shared and row.user_id != resolved_user_id:
                return
            merged = dict(row.metadata_json or {})
            merged.update(metadata)
            row.metadata_json = merged
            for key, value in _indexed_thread_metadata(merged).items():
                setattr(row, key, value)
            row.updated_at = datetime.now(UTC)
            await sync_task_deeplink_index(session, row)
            await session.commit()

    async def delete(
        self,
        thread_id: str,
        *,
        user_id: str | None | _AutoSentinel = AUTO,
    ) -> None:
        resolved_user_id = resolve_user_id(user_id, method_name="ThreadMetaRepository.delete")
        async with self._sf() as session:
            row = await session.get(ThreadMetaRow, thread_id)
            if row is None:
                return
            # 任务深链对话（metadata.taskId）全局共享：任何用户都可删除。
            task_shared = bool((row.metadata_json or {}).get("taskId"))
            if resolved_user_id is not None and not task_shared and row.user_id != resolved_user_id:
                return
            await delete_task_deeplink_index(session, thread_id)
            await session.delete(row)
            await session.commit()