deer-flow/backend/packages/harness/deerflow/runtime/stream_bridge/memory.py

"""In-memory stream bridge backed by an in-process event log."""

from __future__ import annotations

import asyncio
import logging
import re
import time
from collections.abc import AsyncIterator
from copy import deepcopy
from dataclasses import dataclass, field
from typing import Any

from .base import DEFAULT_HEARTBEAT_INTERVAL_SECONDS, END_SENTINEL, HEARTBEAT_SENTINEL, StreamBridge, StreamEvent, StreamGap, StreamItem

logger = logging.getLogger(__name__)
_MEMORY_STREAM_ID_RE = re.compile(r"[0-9]+-([0-9]+)")


@dataclass
class _RunStream:
    events: list[StreamEvent] = field(default_factory=list)
    condition: asyncio.Condition = field(default_factory=asyncio.Condition)
    ended: bool = False
    start_offset: int = 0
    last_evicted_event_id: str | None = None
    retired: bool = False


class MemoryStreamBridge(StreamBridge):
    """Per-run in-memory event log implementation.

    Events are retained for a bounded time window per run so late subscribers
    and reconnecting clients can replay buffered events from ``Last-Event-ID``.
    """

    def __init__(self, *, queue_maxsize: int = 256, heartbeat_interval: float = DEFAULT_HEARTBEAT_INTERVAL_SECONDS) -> None:
        super().__init__(heartbeat_interval=heartbeat_interval)
        self._maxsize = max(1, queue_maxsize)
        self._streams: dict[str, _RunStream] = {}
        self._counters: dict[str, int] = {}

    # -- helpers ---------------------------------------------------------------

    def _get_or_create_stream(self, run_id: str) -> _RunStream:
        if run_id not in self._streams:
            self._streams[run_id] = _RunStream()
            self._counters[run_id] = 0
        return self._streams[run_id]

    def _next_id(self, run_id: str) -> str:
        self._counters[run_id] = self._counters.get(run_id, 0) + 1
        ts = int(time.time() * 1000)
        seq = self._counters[run_id] - 1
        return f"{ts}-{seq}"

    @staticmethod
    def _parse_event_seq(event_id: str) -> int | None:
        """Extract the per-run sequence number from a ``{ts}-{seq}`` event id.

        ``seq`` (assigned by :meth:`_next_id`) increases by one per published
        event, so it equals the event's absolute offset within the run. Returns
        ``None`` for ids that do not match the expected format.
        """
        match = _MEMORY_STREAM_ID_RE.fullmatch(event_id)
        if match is None:
            return None
        try:
            return int(match.group(1))
        except ValueError:
            # Numeric-looking reconnect ids can exceed Python's digit limit.
            return None

    @staticmethod
    def _make_gap(stream: _RunStream, requested_event_id: str | None) -> StreamGap:
        return StreamGap(
            requested_event_id=requested_event_id,
            earliest_available_event_id=stream.events[0].id if stream.events else None,
            latest_available_event_id=stream.events[-1].id if stream.events else None,
        )

    def _resolve_start_offset(self, stream: _RunStream, last_event_id: str | None) -> int | StreamGap:
        # Cleanup can win the Condition after a new reader captures this generation.
        if stream.retired:
            return StreamGap(last_event_id, None, None)
        if last_event_id is None:
            return stream.start_offset
        if last_event_id == stream.last_evicted_event_id:
            return stream.start_offset

        # Event ids embed a per-run, monotonically increasing ``seq`` that equals
        # the event's absolute offset, so locate the event by arithmetic in O(1)
        # rather than scanning the retained buffer. Retained ids are verified at
        # the computed index. The exact last evicted id is a known predecessor
        # whose successors are all retained. Other ids below the watermark cannot
        # be verified and take the conservative gap path. Unknown ids at or above the
        # watermark keep the legacy replay-from-earliest behavior.
        seq = self._parse_event_seq(last_event_id)
        if seq is not None:
            if seq < stream.start_offset:
                return self._make_gap(stream, last_event_id)
            local_index = seq - stream.start_offset
            if 0 <= local_index < len(stream.events) and stream.events[local_index].id == last_event_id:
                return stream.start_offset + local_index + 1

        if stream.events:
            logger.warning(
                "last_event_id=%s not found in retained buffer; replaying from earliest retained event",
                last_event_id,
            )
        return stream.start_offset

    async def stream_exists(self, run_id: str) -> bool:
        """Return whether the in-process event log still has data for *run_id*."""
        return run_id in self._streams

    # -- StreamBridge API ------------------------------------------------------

    async def publish(self, run_id: str, event: str, data: Any) -> None:
        data = deepcopy(data)
        stream = self._get_or_create_stream(run_id)
        entry = StreamEvent(id=self._next_id(run_id), event=event, data=data)
        async with stream.condition:
            stream.events.append(entry)
            if len(stream.events) > self._maxsize:
                overflow = len(stream.events) - self._maxsize
                stream.last_evicted_event_id = stream.events[overflow - 1].id
                del stream.events[:overflow]
                stream.start_offset += overflow
            stream.condition.notify_all()

    async def publish_end(self, run_id: str) -> None:
        stream = self._get_or_create_stream(run_id)
        async with stream.condition:
            stream.ended = True
            stream.condition.notify_all()

    async def subscribe(
        self,
        run_id: str,
        *,
        last_event_id: str | None = None,
        heartbeat_interval: float | None = None,
    ) -> AsyncIterator[StreamItem]:
        heartbeat_interval = self._resolve_heartbeat_interval(heartbeat_interval)
        stream = self._get_or_create_stream(run_id)
        async with stream.condition:
            start = self._resolve_start_offset(stream, last_event_id)
            if isinstance(start, StreamGap):
                gap = start
                next_offset = stream.start_offset
            else:
                gap = None
                next_offset = start

        if gap is not None:
            yield gap
            return

        cursor_event_id = last_event_id

        while True:
            async with stream.condition:
                if stream.retired:
                    entry: StreamItem = END_SENTINEL if stream.ended and next_offset >= stream.start_offset else StreamGap(cursor_event_id, None, None)
                    should_stop = True
                elif next_offset < stream.start_offset:
                    logger.warning(
                        "subscriber for run %s fell behind retained buffer at offset %s",
                        run_id,
                        next_offset,
                    )
                    entry = self._make_gap(stream, cursor_event_id)
                    should_stop = True
                else:
                    should_stop = False

                    local_index = next_offset - stream.start_offset
                    if 0 <= local_index < len(stream.events):
                        entry = stream.events[local_index]
                        next_offset += 1
                        cursor_event_id = entry.id
                    elif stream.ended:
                        entry = END_SENTINEL
                    else:
                        try:
                            await asyncio.wait_for(stream.condition.wait(), timeout=heartbeat_interval)
                        except TimeoutError:
                            entry = HEARTBEAT_SENTINEL
                        else:
                            continue

            if entry is END_SENTINEL:
                yield END_SENTINEL
                return
            # Consumers identify sentinels with `is`; never deep-copy them.
            yield deepcopy(entry) if isinstance(entry, StreamEvent) and entry is not HEARTBEAT_SENTINEL else entry
            if should_stop:
                return

    async def cleanup(self, run_id: str, *, delay: float = 0) -> None:
        if delay > 0:
            await asyncio.sleep(delay)
        stream = self._streams.get(run_id)
        if stream is None:
            return
        async with stream.condition:
            # Established generators retain this object after registry removal.
            # Retire that generation, release its payloads, and wake readers so
            # they can reload durable state rather than waiting on an orphan.
            stream.retired = True
            stream.start_offset += len(stream.events)
            stream.events.clear()
            if self._streams.get(run_id) is stream:
                self._streams.pop(run_id, None)
                self._counters.pop(run_id, None)
            stream.condition.notify_all()

    async def close(self) -> None:
        for run_id in tuple(self._streams):
            await self.cleanup(run_id)