beadhub 0.1.0__py3-none-any.whl

beadhub/events.py

@@ -0,0 +1,704 @@
+"""Event publishing and streaming via Redis pub/sub.
+
+This module provides the infrastructure for real-time event streaming:
+- Event types for messages, escalations, and beads
+- EventBus for publishing events to Redis pub/sub channels
+- Helpers for SSE streaming
+"""
+
+from __future__ import annotations
+
+import asyncio
+import json
+import logging
+import time
+from dataclasses import asdict, dataclass, field
+from datetime import datetime, timedelta, timezone
+from enum import Enum
+from typing import Any, AsyncIterator, Awaitable, Callable, Optional
+
+from redis.asyncio import Redis
+from redis.asyncio.client import PubSub
+from redis.exceptions import ResponseError
+
+logger = logging.getLogger(__name__)
+
+
+class EventCategory(str, Enum):
+    """Categories of events that can be streamed."""
+
+    RESERVATION = "reservation"
+    MESSAGE = "message"
+    ESCALATION = "escalation"
+    BEAD = "bead"
+
+
+@dataclass
+class Event:
+    """Base class for all events."""
+
+    workspace_id: str
+    type: str = ""
+    timestamp: str = field(default_factory=lambda: datetime.now(timezone.utc).isoformat())
+    project_slug: str | None = None
+
+    def to_dict(self) -> dict[str, Any]:
+        return asdict(self)
+
+    def to_json(self) -> str:
+        return json.dumps(self.to_dict())
+
+    @property
+    def category(self) -> EventCategory:
+        """Extract category from event type (e.g., 'message.delivered' -> 'message')."""
+        return EventCategory(self.type.split(".")[0])
+
+
+@dataclass
+class ReservationAcquiredEvent(Event):
+    """Event emitted when reservations are acquired."""
+
+    type: str = field(default="reservation.acquired", init=False)
+    paths: list[str] = field(default_factory=list)
+    alias: str = ""
+    ttl_seconds: int = 0
+    bead_id: str | None = None
+    reason: str | None = None
+    exclusive: bool = True
+
+
+@dataclass
+class ReservationReleasedEvent(Event):
+    """Event emitted when reservations are released."""
+
+    type: str = field(default="reservation.released", init=False)
+    paths: list[str] = field(default_factory=list)
+    alias: str = ""
+
+
+@dataclass
+class ReservationRenewedEvent(Event):
+    """Event emitted when reservation TTLs are extended."""
+
+    type: str = field(default="reservation.renewed", init=False)
+    paths: list[str] = field(default_factory=list)
+    alias: str = ""
+    ttl_seconds: int = 0
+
+
+@dataclass
+class MessageDeliveredEvent(Event):
+    """Event emitted when a message is delivered to a workspace inbox."""
+
+    type: str = field(default="message.delivered", init=False)
+    message_id: str = ""
+    from_workspace: str = ""
+    from_alias: str = ""
+    subject: str = ""
+    priority: str = "normal"
+
+
+@dataclass
+class MessageAcknowledgedEvent(Event):
+    """Event emitted when a message is acknowledged."""
+
+    type: str = field(default="message.acknowledged", init=False)
+    message_id: str = ""
+
+
+@dataclass
+class EscalationCreatedEvent(Event):
+    """Event emitted when an escalation is created."""
+
+    type: str = field(default="escalation.created", init=False)
+    escalation_id: str = ""
+    alias: str = ""
+    subject: str = ""
+
+
+@dataclass
+class EscalationRespondedEvent(Event):
+    """Event emitted when an escalation receives a response."""
+
+    type: str = field(default="escalation.responded", init=False)
+    escalation_id: str = ""
+    response: str = ""
+
+
+@dataclass
+class BeadStatusChangedEvent(Event):
+    """Event emitted when a bead's status changes."""
+
+    type: str = field(default="bead.status_changed", init=False)
+    project_id: str = ""
+    bead_id: str = ""
+    repo: str = ""
+    old_status: str = ""
+    new_status: str = ""
+
+
+# =============================================================================
+# Chat Events (for real-time agent chat sessions)
+# =============================================================================
+
+
+@dataclass
+class ChatEvent:
+    """Base class for chat session events.
+
+    Chat events are scoped to a session, not a workspace.
+    """
+
+    session_id: str
+    type: str = ""
+    timestamp: str = field(default_factory=lambda: datetime.now(timezone.utc).isoformat())
+
+    def to_dict(self) -> dict[str, Any]:
+        return asdict(self)
+
+    def to_json(self) -> str:
+        return json.dumps(self.to_dict())
+
+
+@dataclass
+class ChatMessageEvent(ChatEvent):
+    """Event emitted when a message is sent in a chat session.
+
+    Sessions are persistent and do not require explicit join/leave state.
+    """
+
+    type: str = field(default="message", init=False)
+    message_id: str = ""
+    from_agent: str = ""
+    body: str = ""
+    sender_leaving: bool = False  # True when sender left the conversation
+    hang_on: bool = False  # True when sender requests more time
+    extends_wait_seconds: int = 0  # How long to extend wait (for hang_on)
+
+
+@dataclass
+class ChatReadReceiptEvent(ChatEvent):
+    """Event emitted when a participant reads messages in a chat session.
+
+    Notifies the original sender that their message was read, allowing their
+    CLI to extend the wait timeout.
+    """
+
+    type: str = field(default="read_receipt", init=False)
+    reader: str = ""  # workspace_id of the reader
+    reader_alias: str = ""  # alias of the reader
+    up_to_message_id: str = ""  # messages read up to this ID
+    extends_wait_seconds: int = 0  # How long to extend the sender's wait
+
+
+def _channel_name(workspace_id: str) -> str:
+    """Generate Redis channel name for a workspace."""
+    return f"events:{workspace_id}"
+
+
+def _chat_channel_name(session_id: str) -> str:
+    """Generate Redis channel name for a chat session."""
+    return f"chat:{session_id}"
+
+
+async def publish_event(redis: Redis, event: Event) -> int:
+    """Publish an event to the workspace's Redis pub/sub channel.
+
+    Args:
+        redis: Redis client
+        event: Event to publish
+
+    Returns:
+        Number of subscribers that received the message
+    """
+    channel = _channel_name(event.workspace_id)
+    message = event.to_json()
+    count = await redis.publish(channel, message)
+    logger.debug(f"Published {event.type} to {channel}, {count} subscribers")
+    return count
+
+
+async def stream_events(
+    redis: Redis,
+    workspace_id: str,
+    event_types: Optional[set[str]] = None,
+    keepalive_seconds: int = 30,
+) -> AsyncIterator[str]:
+    """Stream events for a workspace as SSE-formatted strings.
+
+    Args:
+        redis: Redis client
+        workspace_id: Workspace to stream events for
+        event_types: Optional set of event categories to filter (e.g., {'message', 'bead'})
+                     If None, all events are streamed.
+        keepalive_seconds: Seconds between keepalive comments
+
+    Yields:
+        SSE-formatted event strings (e.g., "data: {...}\\n\\n")
+    """
+    async for event in stream_events_multi(redis, [workspace_id], event_types, keepalive_seconds):
+        yield event
+
+
+async def stream_events_multi(
+    redis: Redis,
+    workspace_ids: list[str],
+    event_types: Optional[set[str]] = None,
+    keepalive_seconds: int = 30,
+    check_disconnected: Optional[Callable[[], Awaitable[bool]]] = None,
+) -> AsyncIterator[str]:
+    """Stream events for multiple workspaces as SSE-formatted strings.
+
+    Args:
+        redis: Redis client
+        workspace_ids: List of workspace IDs to stream events for
+        event_types: Optional set of event categories to filter (e.g., {'message', 'bead'})
+                     If None, all events are streamed.
+        keepalive_seconds: Seconds between keepalive comments
+        check_disconnected: Optional async callback to check if client has disconnected.
+                           When provided and returns True, the stream ends cleanly.
+
+    Yields:
+        SSE-formatted event strings (e.g., "data: {...}\\n\\n")
+    """
+    channels = [_channel_name(ws_id) for ws_id in workspace_ids]
+
+    # Empty workspace list: send keepalives for a limited time.
+    # This handles new projects with no workspaces yet while preventing
+    # resource leaks if disconnect detection fails.
+    if not channels:
+        max_duration_seconds = 5 * 60  # 5 minutes
+        max_keepalives = max_duration_seconds // keepalive_seconds
+        keepalive_count = 0
+
+        while keepalive_count < max_keepalives:
+            # Check for client disconnect
+            if check_disconnected and await check_disconnected():
+                logger.debug("Client disconnected (empty workspace list)")
+                return
+            await asyncio.sleep(keepalive_seconds)
+            yield ": keepalive\n\n"
+            keepalive_count += 1
+
+        logger.debug("Empty workspace stream reached max duration, closing")
+        return
+
+    pubsub: PubSub = redis.pubsub()
+
+    try:
+        await pubsub.subscribe(*channels)
+        logger.debug(f"Subscribed to {len(channels)} channels")
+
+        last_keepalive = asyncio.get_event_loop().time()
+
+        while True:
+            # Check for messages with timeout for keepalive
+            try:
+                message = await asyncio.wait_for(
+                    pubsub.get_message(ignore_subscribe_messages=True, timeout=1.0),
+                    timeout=keepalive_seconds,
+                )
+            except asyncio.TimeoutError:
+                message = None
+
+            # Check for client disconnect
+            if check_disconnected and await check_disconnected():
+                logger.debug(f"Client disconnected, ending stream for {len(channels)} channels")
+                return
+
+            current_time = asyncio.get_event_loop().time()
+
+            if message is not None and message["type"] == "message":
+                data = message["data"]
+                if isinstance(data, bytes):
+                    data = data.decode("utf-8")
+
+                # Parse event to check category filter
+                try:
+                    event_data = json.loads(data)
+                    event_category = event_data.get("type", "").split(".")[0]
+
+                    # Apply filter if specified
+                    if event_types is None or event_category in event_types:
+                        yield f"data: {data}\n\n"
+                        last_keepalive = current_time
+                except json.JSONDecodeError:
+                    logger.warning(f"Invalid JSON in event: {data}")
+                    continue
+
+            # Send keepalive comment if needed
+            if current_time - last_keepalive >= keepalive_seconds:
+                yield ": keepalive\n\n"
+                last_keepalive = current_time
+
+    except asyncio.CancelledError:
+        logger.debug(f"Stream cancelled for {len(channels)} channels")
+        raise
+    finally:
+        await pubsub.unsubscribe(*channels)
+        await pubsub.aclose()
+        logger.debug(f"Unsubscribed from {len(channels)} channels")
+
+
+async def publish_chat_event(redis: Redis, event: ChatEvent) -> int:
+    """Publish a chat event to the session's Redis pub/sub channel.
+
+    Args:
+        redis: Redis client
+        event: Chat event to publish
+
+    Returns:
+        Number of subscribers that received the message
+    """
+    channel = _chat_channel_name(event.session_id)
+    message = event.to_json()
+    count = await redis.publish(channel, message)
+    logger.debug(f"Published chat {event.type} to {channel}, {count} subscribers")
+    return count
+
+
+def _chat_waiting_key(session_id: str) -> str:
+    """Redis key for tracking workspaces waiting in a chat session."""
+    return f"chat:waiting:{session_id}"
+
+
+def _chat_deadline_key(session_id: str) -> str:
+    """Redis key for tracking workspace deadlines in a chat session.
+
+    Stores deadline timestamps so we can show time remaining to other participants.
+    """
+    return f"chat:deadline:{session_id}"
+
+
+async def _ensure_waiting_key_zset(redis: Redis, waiting_key: str, keep_ttl: bool = True) -> None:
+    """Ensure the chat waiting key uses a ZSET (member -> last_seen timestamp).
+
+    If the key exists with an unexpected Redis type, ZSET operations will raise
+    WRONGTYPE. This function replaces the key with a ZSET (best-effort).
+
+    This function is best-effort and safe to call on hot paths.
+    """
+    try:
+        key_type = await redis.type(waiting_key)
+    except Exception:
+        logger.warning(
+            "Failed to read Redis key type for waiting_key=%s", waiting_key, exc_info=True
+        )
+        return
+
+    # redis-py returns bytes for TYPE.
+    if isinstance(key_type, bytes):
+        key_type = key_type.decode("utf-8", errors="ignore")
+
+    if key_type in ("none", ""):
+        return
+    if key_type == "zset":
+        return
+
+    ttl = None
+    if keep_ttl:
+        try:
+            ttl = await redis.ttl(waiting_key)
+        except Exception:
+            logger.warning(
+                "Failed to read Redis TTL for waiting_key=%s", waiting_key, exc_info=True
+            )
+            ttl = None
+
+    if key_type == "set":
+        try:
+            members = await redis.smembers(waiting_key)
+        except Exception:
+            logger.warning(
+                "Failed to read Redis set members for waiting_key=%s", waiting_key, exc_info=True
+            )
+            members = set()
+
+        try:
+            await redis.delete(waiting_key)
+        except Exception:
+            logger.warning(
+                "Failed to delete Redis key waiting_key=%s for type migration",
+                waiting_key,
+                exc_info=True,
+            )
+            return
+
+        mapping: dict[str, float] = {}
+        now = time.time()
+        for m in members:
+            if isinstance(m, bytes):
+                m = m.decode("utf-8", errors="ignore")
+            mapping[str(m)] = now
+
+        if mapping:
+            try:
+                await redis.zadd(waiting_key, mapping)
+            except Exception:
+                logger.warning(
+                    "Failed to write Redis zset waiting_key=%s for type migration",
+                    waiting_key,
+                    exc_info=True,
+                )
+                return
+
+        if ttl is not None and ttl > 0:
+            try:
+                await redis.expire(waiting_key, ttl)
+            except Exception:
+                logger.warning(
+                    "Failed to restore Redis TTL for waiting_key=%s after type migration",
+                    waiting_key,
+                    exc_info=True,
+                )
+        return
+
+    # Unknown type; safest is to delete it to avoid crashing endpoints.
+    try:
+        await redis.delete(waiting_key)
+    except Exception:
+        logger.warning(
+            "Failed to delete Redis key waiting_key=%s with unknown type",
+            waiting_key,
+            exc_info=True,
+        )
+
+
+async def is_workspace_waiting(
+    redis: Redis,
+    session_id: str,
+    workspace_id: str,
+    max_age_seconds: int = 90,
+) -> bool:
+    """Check if a workspace is currently waiting (connected to SSE) in a session.
+
+    We track a per-workspace "last seen" timestamp in a Redis sorted set.
+    This avoids false positives when the Redis key TTL is refreshed by other
+    connected workspaces (stale members can no longer stick indefinitely).
+    """
+    waiting_key = _chat_waiting_key(session_id)
+    await _ensure_waiting_key_zset(redis, waiting_key)
+    try:
+        score = await redis.zscore(waiting_key, workspace_id)
+    except ResponseError as e:
+        if "WRONGTYPE" in str(e):
+            await _ensure_waiting_key_zset(redis, waiting_key)
+            score = await redis.zscore(waiting_key, workspace_id)
+        else:
+            raise
+    if score is None:
+        return False
+    now = time.time()
+    if float(score) < now - max_age_seconds:
+        await redis.zrem(waiting_key, workspace_id)
+        # Best-effort: clean up stale deadline too.
+        try:
+            await redis.hdel(_chat_deadline_key(session_id), workspace_id)
+        except Exception:
+            logger.warning(
+                "Failed to delete stale chat deadline session_id=%s workspace_id=%s",
+                session_id,
+                workspace_id,
+                exc_info=True,
+            )
+        return False
+    return True
+
+
+async def get_workspace_deadline(
+    redis: Redis,
+    session_id: str,
+    workspace_id: str,
+) -> Optional[datetime]:
+    """Get the deadline for a workspace waiting in a chat session.
+
+    Returns the deadline timestamp if the workspace has one set, or None if no
+    deadline is stored.
+
+    Args:
+        redis: Redis client
+        session_id: Chat session ID
+        workspace_id: Workspace ID to check
+
+    Returns:
+        Deadline as datetime (UTC) or None if not set
+    """
+    deadline_key = _chat_deadline_key(session_id)
+    try:
+        deadline_str = await redis.hget(deadline_key, workspace_id)
+        if deadline_str:
+            if isinstance(deadline_str, bytes):
+                deadline_str = deadline_str.decode("utf-8")
+            return datetime.fromisoformat(deadline_str)
+    except Exception:
+        logger.warning(
+            "Failed to get workspace deadline session_id=%s workspace_id=%s",
+            session_id,
+            workspace_id,
+            exc_info=True,
+        )
+    return None
+
+
+async def extend_workspace_deadline(
+    redis: Redis,
+    session_id: str,
+    workspace_id: str,
+    extends_seconds: int,
+) -> Optional[datetime]:
+    """Extend a workspace's stored wait deadline for a chat session.
+
+    This only updates Redis metadata used for time-remaining displays. The client
+    must still extend its own local wait timeout to avoid disconnecting.
+
+    Returns the new deadline when one exists and was updated, otherwise None.
+    """
+    if extends_seconds <= 0:
+        return None
+
+    deadline_key = _chat_deadline_key(session_id)
+    current = await get_workspace_deadline(redis, session_id, workspace_id)
+    if current is None:
+        return None
+
+    if current.tzinfo is None:
+        current = current.replace(tzinfo=timezone.utc)
+
+    new_deadline = current + timedelta(seconds=extends_seconds)
+
+    try:
+        await redis.hset(deadline_key, workspace_id, new_deadline.isoformat())
+        # Keep the deadline hash alive while the waiting zset is alive.
+        waiting_ttl = await redis.ttl(_chat_waiting_key(session_id))
+        if waiting_ttl and waiting_ttl > 0:
+            await redis.expire(deadline_key, waiting_ttl)
+    except Exception:
+        logger.warning(
+            "Failed to extend workspace deadline session_id=%s workspace_id=%s",
+            session_id,
+            workspace_id,
+            exc_info=True,
+        )
+        return None
+
+    return new_deadline
+
+
+async def stream_chat_events(
+    redis: Redis,
+    session_id: str,
+    workspace_id: str,
+    keepalive_seconds: int = 30,
+    check_disconnected: Optional[Callable[[], Awaitable[bool]]] = None,
+    deadline: Optional[datetime] = None,
+) -> AsyncIterator[str]:
+    """Stream chat events for a session as SSE-formatted strings.
+
+    Tracks workspace connection state in Redis to enable "waiting" detection.
+    When workspace connects, it's added to the waiting set.
+    When workspace disconnects (timeout, reply received, crash), it's removed.
+
+    Args:
+        redis: Redis client
+        session_id: Chat session to stream events for
+        workspace_id: Workspace ID of the connecting client
+        keepalive_seconds: Seconds between keepalive comments
+        check_disconnected: Optional async callback to check if client has disconnected.
+                           When provided and returns True, the stream ends cleanly.
+        deadline: Optional deadline timestamp for this workspace's wait.
+                  If provided, stored in Redis so other participants can see time remaining.
+
+    Yields:
+        SSE-formatted event strings with event type (e.g., "event: message\\ndata: {...}\\n\\n")
+    """
+    channel = _chat_channel_name(session_id)
+    waiting_key = _chat_waiting_key(session_id)
+    deadline_key = _chat_deadline_key(session_id)
+    pubsub: PubSub = redis.pubsub()
+
+    # TTL for waiting set - 3x keepalive to handle missed refreshes
+    waiting_ttl = keepalive_seconds * 3
+
+    try:
+        await _ensure_waiting_key_zset(redis, waiting_key)
+
+        # Track that this workspace is waiting for a reply.
+        # Use a per-workspace heartbeat timestamp so stale entries decay even
+        # if other workspaces keep the key alive.
+        await redis.zadd(waiting_key, {workspace_id: time.time()})
+        await redis.expire(waiting_key, waiting_ttl)
+
+        # Store deadline if provided
+        if deadline:
+            await redis.hset(deadline_key, workspace_id, deadline.isoformat())
+            await redis.expire(deadline_key, waiting_ttl)
+
+        logger.debug(f"Workspace {workspace_id} now waiting in session {session_id}")
+
+        await pubsub.subscribe(channel)
+        logger.debug(f"Subscribed to chat channel {channel}")
+
+        last_keepalive = asyncio.get_event_loop().time()
+
+        while True:
+            try:
+                message = await asyncio.wait_for(
+                    pubsub.get_message(ignore_subscribe_messages=True, timeout=1.0),
+                    timeout=keepalive_seconds,
+                )
+            except asyncio.TimeoutError:
+                message = None
+
+            # Check for client disconnect
+            if check_disconnected and await check_disconnected():
+                logger.debug(f"Client disconnected, ending chat stream for {channel}")
+                return
+
+            current_time = asyncio.get_event_loop().time()
+
+            if message is not None and message["type"] == "message":
+                data = message["data"]
+                if isinstance(data, bytes):
+                    data = data.decode("utf-8")
+
+                try:
+                    event_data = json.loads(data)
+                    event_type = event_data.get("type", "message")
+                    # SSE format with event type for client-side event handling
+                    yield f"event: {event_type}\ndata: {data}\n\n"
+                    last_keepalive = current_time
+                except json.JSONDecodeError:
+                    logger.warning(f"Invalid JSON in chat event: {data}")
+                    continue
+
+            if current_time - last_keepalive >= keepalive_seconds:
+                yield ": keepalive\n\n"
+                last_keepalive = current_time
+                # Refresh heartbeat + TTL to prevent expiration while connected
+                await _ensure_waiting_key_zset(redis, waiting_key)
+                await redis.zadd(waiting_key, {workspace_id: time.time()})
+                await redis.expire(waiting_key, waiting_ttl)
+                # Also refresh deadline TTL (key may exist even if this stream didn't set it).
+                try:
+                    await redis.expire(deadline_key, waiting_ttl)
+                except Exception:
+                    logger.warning(
+                        "Failed to refresh chat deadline TTL session_id=%s workspace_id=%s",
+                        session_id,
+                        workspace_id,
+                        exc_info=True,
+                    )
+
+    except asyncio.CancelledError:
+        logger.debug(f"Chat stream cancelled for {channel}")
+        raise
+    finally:
+        # Remove from waiting set - workspace is no longer waiting
+        await redis.zrem(waiting_key, workspace_id)
+        # Also remove deadline
+        await redis.hdel(deadline_key, workspace_id)
+        logger.debug(f"Workspace {workspace_id} no longer waiting in session {session_id}")
+
+        await pubsub.unsubscribe(channel)
+        await pubsub.aclose()
+        logger.debug(f"Unsubscribed from chat channel {channel}")