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