deer-flow/backend/packages/harness/deerflow/runtime/events/store/db.py

"""SQLAlchemy-backed RunEventStore implementation.

Persists events to the ``run_events`` table. Trace content is truncated
at ``max_trace_content`` bytes to avoid bloating the database.
"""

from __future__ import annotations

import asyncio
import json
import logging
import re
import weakref
from datetime import UTC, datetime
from typing import Any

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

from deerflow.persistence.models.run_event import RunEventRow
from deerflow.runtime.events.message_identity import message_identity
from deerflow.runtime.events.store.base import RunEventStore
from deerflow.runtime.user_context import AUTO, _AutoSentinel, get_current_user, resolve_user_id
from deerflow.utils.time import coerce_iso

logger = logging.getLogger(__name__)


class DbRunEventStore(RunEventStore):
    def __init__(self, session_factory: async_sessionmaker[AsyncSession], *, max_trace_content: int = 10240):
        self._sf = session_factory
        self._max_trace_content = max_trace_content
        # Per-thread asyncio locks serialize seq assignment for concurrent
        # in-process writers on the same thread. The DB-level FOR UPDATE /
        # advisory lock guards cross-process races; this guards the common
        # single-process case where two coroutines interleave between the
        # max(seq) read and the INSERT and would otherwise collide on seq.
        #
        # The weak registry preserves one lock generation while an admitted
        # holder/waiter still references it. A separate pin keeps the historical
        # one-lock-per-live-thread behavior until delete_by_thread() explicitly
        # retires that thread; after retirement, outstanding users alone keep
        # the generation alive until they drain.
        self._write_locks: weakref.WeakValueDictionary[str, asyncio.Lock] = weakref.WeakValueDictionary()
        self._write_lock_pins: dict[str, asyncio.Lock] = {}

    def _get_write_lock(self, thread_id: str) -> asyncio.Lock:
        """Return (creating if needed) the per-thread seq-assignment lock."""
        lock = self._write_locks.get(thread_id)
        if lock is None:
            lock = asyncio.Lock()
            self._write_locks[thread_id] = lock
        # A fresh caller after deletion makes the thread live again. Repin the
        # current generation so normal live-thread registry lifetime is stable.
        self._write_lock_pins[thread_id] = lock
        return lock

    @staticmethod
    def _row_to_dict(row: RunEventRow) -> dict:
        d = row.to_dict()
        d["metadata"] = d.pop("event_metadata", {})
        val = d.get("created_at")
        if isinstance(val, datetime):
            # SQLite drops tzinfo on read despite ``DateTime(timezone=True)``;
            # ``coerce_iso`` normalizes naive datetimes as UTC.
            d["created_at"] = coerce_iso(val)
        d.pop("id", None)
        # Restore structured content that was JSON-serialized on write.
        raw = d.get("content", "")
        metadata = d.get("metadata", {})
        if isinstance(raw, str) and (metadata.get("content_is_json") or metadata.get("content_is_dict")):
            try:
                d["content"] = json.loads(raw)
            except (json.JSONDecodeError, ValueError):
                # Content looked like JSON but failed to parse;
                # keep the raw string as-is.
                logger.debug("Failed to deserialize content as JSON for event seq=%s", d.get("seq"))
        return d

    def _truncate_trace(self, category: str, content: Any, metadata: dict | None) -> tuple[Any, dict]:
        if category == "trace":
            text = content if isinstance(content, str) else json.dumps(content, default=str, ensure_ascii=False)
            encoded = text.encode("utf-8")
            if len(encoded) > self._max_trace_content:
                # Truncate by bytes, then decode back (may cut a multi-byte char, so use errors="ignore")
                content = encoded[: self._max_trace_content].decode("utf-8", errors="ignore")
                metadata = {**(metadata or {}), "content_truncated": True, "original_byte_length": len(encoded)}
        return content, metadata or {}

    @staticmethod
    def _content_to_db(content: Any, metadata: dict | None) -> tuple[str, dict]:
        metadata = metadata or {}
        if isinstance(content, str):
            return content, metadata

        db_content = json.dumps(content, default=str, ensure_ascii=False)
        metadata = {**metadata, "content_is_json": True}
        if isinstance(content, dict):
            metadata["content_is_dict"] = True
        return db_content, metadata

    @staticmethod
    def _user_id_from_context() -> str | None:
        """Soft read of user_id from contextvar for write paths.

        Returns ``None`` (no filter / no stamp) if contextvar is unset,
        which is the expected case for background worker writes. HTTP
        request writes will have the contextvar set by auth middleware
        and get their user_id stamped automatically.

        Coerces ``user.id`` to ``str`` at the boundary: ``User.id`` is
        typed as ``UUID`` by the auth layer, but ``run_events.user_id``
        is ``VARCHAR(64)`` and aiosqlite cannot bind a raw UUID object
        to a VARCHAR column ("type 'UUID' is not supported") — the
        INSERT would silently roll back and the worker would hang.
        """
        user = get_current_user()
        return str(user.id) if user is not None else None

    #: Characters json.dumps escapes in the stored content (``ensure_ascii``
    #: is False, so non-ASCII survives verbatim and stays matchable).
    _LIKE_UNSAFE_ID = re.compile(r'["\\\x00-\x1f]')

    @classmethod
    def _prefilter_substrings(cls, wanted: set[str]) -> list[str] | None:
        """Return the raw ids to LIKE-match in ``content``, or ``None`` to full-scan.

        An identity is ``kind:raw_id`` and the raw id appears verbatim in the
        stored JSON string (``u1`` is a substring of a re-keyed ``u1__user``
        copy too), so a row not containing any wanted id cannot resolve any
        wanted identity. An id json.dumps would escape breaks that verbatim
        guarantee — one such id falls the whole set back to the full scan
        rather than silently missing it. LIKE wildcards are escaped, not
        rejected.
        """
        ids = []
        for identity in wanted:
            _kind, _sep, raw_id = identity.partition(":")
            if not raw_id or cls._LIKE_UNSAFE_ID.search(raw_id):
                return None
            ids.append(raw_id)
        return ids

    @staticmethod
    async def _acquire_thread_mutation_fence(session: AsyncSession, thread_id: str) -> None:
        """Take the cross-process thread mutation fence, if the dialect has one.

        PostgreSQL rejects ``SELECT max(...) FOR UPDATE`` because aggregate
        results are not lockable rows, so it serializes a thread's mutations with
        a transaction-level advisory lock keyed by ``thread_id``. This is the
        database half of the contract whose in-process half is
        ``_get_write_lock()``: every thread mutation — ``put``, ``put_batch``,
        ``put_if_absent`` and both deletions — takes this fence before touching
        rows, so an admitted writer can never land a row between a deletion's
        count and its commit.

        Dialects without a cross-process fence (SQLite) rely on the in-process
        per-thread lock alone, so this is a no-op there.
        """
        bind = session.get_bind()
        dialect_name = bind.dialect.name if bind is not None else ""

        if dialect_name == "postgresql":
            await session.execute(
                text("SELECT pg_advisory_xact_lock(hashtext(CAST(:thread_id AS text))::bigint)"),
                {"thread_id": thread_id},
            )

    @staticmethod
    async def _max_seq_for_thread(session: AsyncSession, thread_id: str) -> int | None:
        """Return the current max seq while serializing writers per thread.

        Takes the shared thread mutation fence before reading the aggregate, so
        the read is ordered against every other mutation of the same thread.
        Other dialects keep the existing row-locking statement.
        """
        await DbRunEventStore._acquire_thread_mutation_fence(session, thread_id)

        stmt = select(func.max(RunEventRow.seq)).where(RunEventRow.thread_id == thread_id)
        bind = session.get_bind()
        dialect_name = bind.dialect.name if bind is not None else ""

        if dialect_name == "postgresql":
            return await session.scalar(stmt)

        return await session.scalar(stmt.with_for_update())

    async def put(self, *, thread_id, run_id, event_type, category, content="", metadata=None, created_at=None):  # noqa: D401
        """Write a single event — low-frequency path only.

        This opens a dedicated transaction with a FOR UPDATE lock to
        assign a monotonic *seq*.  For high-throughput writes use
        :meth:`put_batch`, which acquires the lock once for the whole
        batch.  Currently the only caller is ``worker.run_agent`` for
        the initial ``human_message`` event (once per run).
        """
        content, metadata = self._truncate_trace(category, content, metadata)
        db_content, metadata = self._content_to_db(content, metadata)
        user_id = self._user_id_from_context()
        async with self._get_write_lock(thread_id):
            async with self._sf() as session:
                async with session.begin():
                    max_seq = await self._max_seq_for_thread(session, thread_id)
                    seq = (max_seq or 0) + 1
                    row = RunEventRow(
                        thread_id=thread_id,
                        run_id=run_id,
                        user_id=user_id,
                        event_type=event_type,
                        category=category,
                        content=db_content,
                        event_metadata=metadata,
                        seq=seq,
                        created_at=datetime.fromisoformat(created_at) if created_at else datetime.now(UTC),
                    )
                    session.add(row)
                return self._row_to_dict(row)

    async def put_batch(self, events):
        if not events:
            return []
        thread_ids = {e["thread_id"] for e in events}
        if len(thread_ids) > 1:
            raise ValueError(f"put_batch requires all events to belong to the same thread; got {thread_ids!r}")
        user_id = self._user_id_from_context()
        # All events belong to the same thread (validated above).
        thread_id = events[0]["thread_id"]
        async with self._get_write_lock(thread_id):
            async with self._sf() as session:
                async with session.begin():
                    max_seq = await self._max_seq_for_thread(session, thread_id)
                    seq = max_seq or 0
                    rows = []
                    for e in events:
                        seq += 1
                        content = e.get("content", "")
                        category = e.get("category", "trace")
                        metadata = e.get("metadata")
                        content, metadata = self._truncate_trace(category, content, metadata)
                        db_content, metadata = self._content_to_db(content, metadata)
                        row = RunEventRow(
                            thread_id=e["thread_id"],
                            run_id=e["run_id"],
                            user_id=e.get("user_id", user_id),
                            event_type=e["event_type"],
                            category=category,
                            content=db_content,
                            event_metadata=metadata,
                            seq=seq,
                            created_at=datetime.fromisoformat(e["created_at"]) if e.get("created_at") else datetime.now(UTC),
                        )
                        session.add(row)
                        rows.append(row)
                return [self._row_to_dict(r) for r in rows]

    async def put_if_absent(
        self,
        *,
        thread_id,
        run_id,
        event_type,
        category,
        content="",
        metadata=None,
        created_at=None,
    ):
        """Idempotently insert a run-scoped singleton event.

        ``_max_seq_for_thread`` takes the same PostgreSQL advisory lock used by
        every normal writer (and the in-process lock covers SQLite), so the
        existence check cannot race another ``put_if_absent`` or journal write.
        Terminal delivery receipts use this method on both the worker and
        recovery paths; ordinary event types remain append-only.
        """
        content, metadata = self._truncate_trace(category, content, metadata)
        db_content, metadata = self._content_to_db(content, metadata)
        user_id = self._user_id_from_context()
        async with self._get_write_lock(thread_id):
            async with self._sf() as session:
                async with session.begin():
                    max_seq = await self._max_seq_for_thread(session, thread_id)
                    stmt = (
                        select(RunEventRow)
                        .where(
                            RunEventRow.thread_id == thread_id,
                            RunEventRow.run_id == run_id,
                            RunEventRow.event_type == event_type,
                        )
                        .order_by(RunEventRow.seq.asc())
                        .limit(1)
                    )
                    existing = await session.scalar(stmt)
                    if existing is not None:
                        return self._row_to_dict(existing), False
                    row = RunEventRow(
                        thread_id=thread_id,
                        run_id=run_id,
                        user_id=user_id,
                        event_type=event_type,
                        category=category,
                        content=db_content,
                        event_metadata=metadata,
                        seq=(max_seq or 0) + 1,
                        created_at=datetime.fromisoformat(created_at) if created_at else datetime.now(UTC),
                    )
                    session.add(row)
                return self._row_to_dict(row), True

    async def list_messages(
        self,
        thread_id,
        *,
        limit=50,
        before_seq=None,
        after_seq=None,
        user_id: str | None | _AutoSentinel = AUTO,
    ):
        resolved_user_id = resolve_user_id(user_id, method_name="DbRunEventStore.list_messages")
        stmt = select(RunEventRow).where(RunEventRow.thread_id == thread_id, RunEventRow.category == "message")
        if resolved_user_id is not None:
            stmt = stmt.where(RunEventRow.user_id == resolved_user_id)
        if before_seq is not None:
            stmt = stmt.where(RunEventRow.seq < before_seq)
        if after_seq is not None:
            stmt = stmt.where(RunEventRow.seq > after_seq)

        if after_seq is not None:
            # Forward pagination: first `limit` records after cursor
            stmt = stmt.order_by(RunEventRow.seq.asc()).limit(limit)
            async with self._sf() as session:
                result = await session.execute(stmt)
                return [self._row_to_dict(r) for r in result.scalars()]
        else:
            # before_seq or default (latest): take last `limit` records, return ascending
            stmt = stmt.order_by(RunEventRow.seq.desc()).limit(limit)
            async with self._sf() as session:
                result = await session.execute(stmt)
                rows = list(result.scalars())
                return [self._row_to_dict(r) for r in reversed(rows)]

    async def list_events(
        self,
        thread_id,
        run_id,
        *,
        event_types=None,
        task_id=None,
        limit=500,
        after_seq=None,
        user_id: str | None | _AutoSentinel = AUTO,
    ):
        resolved_user_id = resolve_user_id(user_id, method_name="DbRunEventStore.list_events")
        stmt = select(RunEventRow).where(RunEventRow.thread_id == thread_id, RunEventRow.run_id == run_id)
        if resolved_user_id is not None:
            stmt = stmt.where(RunEventRow.user_id == resolved_user_id)
        if event_types:
            stmt = stmt.where(RunEventRow.event_type.in_(event_types))
        if task_id is not None:
            # Filter on metadata["task_id"] in SQL (before LIMIT) so cursor
            # pagination over a single subagent task stays correct (#3779). The
            # query is already scoped to (thread_id, run_id), so the JSON probe
            # only runs over this run's small candidate set; ``.as_string()``
            # renders to json_extract (SQLite) / ->> (Postgres).
            stmt = stmt.where(RunEventRow.event_metadata["task_id"].as_string() == task_id)
        if after_seq is not None:
            stmt = stmt.where(RunEventRow.seq > after_seq)
        stmt = stmt.order_by(RunEventRow.seq.asc()).limit(limit)
        async with self._sf() as session:
            result = await session.execute(stmt)
            return [self._row_to_dict(r) for r in result.scalars()]

    async def list_messages_by_run(
        self,
        thread_id,
        run_id,
        *,
        limit=50,
        before_seq=None,
        after_seq=None,
        user_id: str | None | _AutoSentinel = AUTO,
    ):
        resolved_user_id = resolve_user_id(user_id, method_name="DbRunEventStore.list_messages_by_run")
        stmt = select(RunEventRow).where(
            RunEventRow.thread_id == thread_id,
            RunEventRow.run_id == run_id,
            RunEventRow.category == "message",
        )
        if resolved_user_id is not None:
            stmt = stmt.where(RunEventRow.user_id == resolved_user_id)
        if before_seq is not None:
            stmt = stmt.where(RunEventRow.seq < before_seq)
        if after_seq is not None:
            stmt = stmt.where(RunEventRow.seq > after_seq)

        if after_seq is not None:
            stmt = stmt.order_by(RunEventRow.seq.asc()).limit(limit)
            async with self._sf() as session:
                result = await session.execute(stmt)
                return [self._row_to_dict(r) for r in result.scalars()]
        else:
            stmt = stmt.order_by(RunEventRow.seq.desc()).limit(limit)
            async with self._sf() as session:
                result = await session.execute(stmt)
                rows = list(result.scalars())
                return [self._row_to_dict(r) for r in reversed(rows)]

    async def get_last_visible_ai_seq_by_run(
        self,
        thread_id,
        run_ids,
        *,
        user_id: str | None | _AutoSentinel = AUTO,
    ):
        if not run_ids:
            return {}
        resolved_user_id = resolve_user_id(user_id, method_name="DbRunEventStore.get_last_visible_ai_seq_by_run")
        caller = RunEventRow.event_metadata["caller"].as_string()
        # RunJournal canonically persists AI message rows as
        # ``llm.ai.response``; ``ai_message`` remains for legacy compatibility.
        stmt = (
            select(RunEventRow.run_id, func.max(RunEventRow.seq))
            .where(
                RunEventRow.thread_id == thread_id,
                RunEventRow.run_id.in_(run_ids),
                RunEventRow.category == "message",
                RunEventRow.event_type.in_(("llm.ai.response", "ai_message")),
                ~func.coalesce(caller, "").like("middleware:%"),
            )
            .group_by(RunEventRow.run_id)
        )
        if resolved_user_id is not None:
            stmt = stmt.where(RunEventRow.user_id == resolved_user_id)
        async with self._sf() as session:
            result = await session.execute(stmt)
            return {run_id: seq for run_id, seq in result if isinstance(seq, int)}

    async def count_messages(
        self,
        thread_id,
        *,
        user_id: str | None | _AutoSentinel = AUTO,
    ):
        resolved_user_id = resolve_user_id(user_id, method_name="DbRunEventStore.count_messages")
        stmt = select(func.count()).select_from(RunEventRow).where(RunEventRow.thread_id == thread_id, RunEventRow.category == "message")
        if resolved_user_id is not None:
            stmt = stmt.where(RunEventRow.user_id == resolved_user_id)
        async with self._sf() as session:
            return await session.scalar(stmt) or 0

    async def get_message_seqs(
        self,
        thread_id,
        identities,
        *,
        user_id: str | None | _AutoSentinel = AUTO,
    ):
        wanted = set(identities)
        if not wanted:
            return {}
        resolved_user_id = resolve_user_id(user_id, method_name="DbRunEventStore.get_message_seqs")
        # ``content`` is a TEXT column holding a JSON *string* (see
        # ``_content_to_db``), not a JSON column, so the identity fields cannot
        # be projected in SQL — matching rows are decoded here instead. The
        # ``content`` column carries full tool outputs, and a wanted identity
        # absent from the feed (a message still streaming) defeats the early
        # exit below — so without a prefilter a `/state`/`/history` read of a
        # long thread pays a full fetch-and-decode of every message row. The
        # LIKE prefilter keeps that cost in SQL: only rows containing a wanted
        # id as a raw substring are fetched (false positives are re-checked by
        # ``message_identity``; ids the prefilter cannot express fall back to
        # the full scan).
        stmt = select(RunEventRow.seq, RunEventRow.content).where(RunEventRow.thread_id == thread_id, RunEventRow.category == "message").order_by(RunEventRow.seq)
        if resolved_user_id is not None:
            stmt = stmt.where(RunEventRow.user_id == resolved_user_id)
        prefilter_ids = self._prefilter_substrings(wanted)
        if prefilter_ids is not None:
            stmt = stmt.where(or_(*[RunEventRow.content.like(f"%{i.replace('%', '\\%').replace('_', '\\_')}%", escape="\\") for i in prefilter_ids]))

        found: dict[str, int] = {}
        async with self._sf() as session:
            result = await session.execute(stmt)
            for seq, raw in result:
                # Plain-text content (never a message dict) is skipped without
                # paying for a failed JSON parse.
                if not isinstance(raw, str) or not raw.startswith("{"):
                    continue
                try:
                    content = json.loads(raw)
                except (json.JSONDecodeError, ValueError):
                    continue
                if not isinstance(content, dict):
                    continue
                identity = message_identity(content)
                # Earliest seq wins: a message re-persisted later keeps the
                # position it first occupied in the feed.
                if identity in wanted and identity not in found:
                    found[identity] = seq
                    # Later rows can only be re-persisted copies that already
                    # lose that tiebreak, so the scan (and its JSON decoding)
                    # ends with the last wanted seq instead of the thread's
                    # full message count.
                    if len(found) == len(wanted):
                        break
        return found

    async def delete_by_thread(
        self,
        thread_id,
        *,
        user_id: str | None | _AutoSentinel = AUTO,
    ):
        """Delete every event of *thread_id* inside the thread mutation fence.

        Deletion takes the same critical section as the writers — the in-process
        per-thread lock plus, on PostgreSQL, the transaction advisory lock — so a
        writer admitted before this call can no longer land a row between the
        count below and the commit, which would resurrect a deleted thread. The
        JSONL store serializes deletion the same way (``_run_mutation``).
        """
        resolved_user_id = resolve_user_id(user_id, method_name="DbRunEventStore.delete_by_thread")
        async with self._get_write_lock(thread_id):
            async with self._sf() as session:
                async with session.begin():
                    await self._acquire_thread_mutation_fence(session, thread_id)
                    count_conditions = [RunEventRow.thread_id == thread_id]
                    if resolved_user_id is not None:
                        count_conditions.append(RunEventRow.user_id == resolved_user_id)
                    count_stmt = select(func.count()).select_from(RunEventRow).where(*count_conditions)
                    count = await session.scalar(count_stmt) or 0
                    if count > 0:
                        await session.execute(delete(RunEventRow).where(*count_conditions))
            # Retire the live-thread pin, but never remove the weak registry
            # entry directly. asyncio.Lock.release() clears ``locked()`` before
            # a queued waiter resumes, so an unlocked check can observe the
            # handoff window and split one thread onto two lock generations.
            # Holders/waiters keep the old generation alive until they drain; a
            # later caller therefore resolves that same lock instead of racing
            # it with a fresh one.
            self._write_lock_pins.pop(thread_id, None)
        return count

    async def delete_by_run(
        self,
        thread_id,
        run_id,
        *,
        user_id: str | None | _AutoSentinel = AUTO,
    ):
        """Delete one run's events inside the thread mutation fence.

        Shares ``delete_by_thread``'s critical section; deleting a single run
        leaves the thread alive, so the write-lock pin is deliberately kept.
        """
        resolved_user_id = resolve_user_id(user_id, method_name="DbRunEventStore.delete_by_run")
        async with self._get_write_lock(thread_id):
            async with self._sf() as session:
                async with session.begin():
                    await self._acquire_thread_mutation_fence(session, thread_id)
                    count_conditions = [RunEventRow.thread_id == thread_id, RunEventRow.run_id == run_id]
                    if resolved_user_id is not None:
                        count_conditions.append(RunEventRow.user_id == resolved_user_id)
                    count_stmt = select(func.count()).select_from(RunEventRow).where(*count_conditions)
                    count = await session.scalar(count_stmt) or 0
                    if count > 0:
                        await session.execute(delete(RunEventRow).where(*count_conditions))
        return count