kubiya-control-plane-api 0.9.15__py3-none-any.whl

control_plane_api/app/routers/executions/streaming/deduplication.py

@@ -0,0 +1,352 @@
+"""
+Message deduplication for streaming execution data.
+
+This module provides a bounded-memory message deduplication system using LRU caching
+to prevent duplicate messages from being sent during streaming execution sessions.
+
+Key Features:
+- Two-level deduplication: message ID + content signature
+- LRU cache with bounded memory (max 1000 entries)
+- Backward compatibility with old message ID formats
+- Content-based deduplication for assistant messages within 5-second window
+- Thread-safe operation
+
+Test Strategy:
+- Unit test deduplication with various message types (user, assistant, system)
+- Test LRU eviction at boundary (1001st entry)
+- Test backward compatibility with old message ID formats (timestamp-based vs turn-based)
+- Test content signature collision handling with near-duplicate messages
+- Test edge cases: empty content, None message_id, missing timestamps
+- Test thread safety with concurrent message processing
+"""
+
+import hashlib
+import logging
+from collections import OrderedDict
+from datetime import datetime
+from typing import Any, Dict, Optional
+
+logger = logging.getLogger(__name__)
+
+
+class MessageDeduplicator:
+    """
+    Handles message deduplication with bounded memory using LRU cache.
+
+    This class provides two-level deduplication:
+    1. Message ID deduplication - prevents exact duplicate messages
+    2. Content signature deduplication - prevents near-duplicate assistant messages
+       within a 5-second window (handles retry/regeneration scenarios)
+
+    The deduplicator also normalizes old message ID formats for backward compatibility:
+    - New format: {execution_id}_{role}_{turn_number} (deterministic, turn-based)
+    - Old formats: {execution_id}_{role}_{timestamp_micros} or {execution_id}_{role}_{idx}
+
+    Memory is bounded using LRU eviction at 1000 entries to prevent memory leaks
+    in long-running streaming sessions.
+    """
+
+    # LRU cache size - limits memory usage to ~1000 message IDs + signatures
+    MAX_CACHE_SIZE = 1000
+
+    # Time window for content signature deduplication (seconds)
+    CONTENT_DEDUP_WINDOW = 5.0
+
+    # Length of content to use for signature (characters)
+    CONTENT_SIGNATURE_LENGTH = 200
+
+    def __init__(self, max_size: int = MAX_CACHE_SIZE):
+        """
+        Initialize the message deduplicator.
+
+        Args:
+            max_size: Maximum number of entries in LRU cache (default: 1000)
+        """
+        self.max_size = max_size
+
+        # LRU cache for sent message IDs
+        # OrderedDict maintains insertion order, move_to_end() implements LRU
+        self._sent_ids: OrderedDict[str, bool] = OrderedDict()
+
+        # Content signature cache for assistant messages
+        # Maps content_signature -> (message_data, timestamp)
+        self._content_cache: OrderedDict[str, tuple[Dict[str, Any], datetime]] = OrderedDict()
+
+        # Statistics for monitoring
+        self._stats = {
+            "messages_checked": 0,
+            "duplicates_by_id": 0,
+            "duplicates_by_content": 0,
+            "evictions": 0,
+            "normalized_ids": 0,
+        }
+
+    def is_sent(self, message: Dict[str, Any]) -> bool:
+        """
+        Check if a message has already been sent (deduplicate).
+
+        This method performs two-level deduplication:
+        1. Check if message ID is in sent cache
+        2. For assistant messages, check content signature within time window
+
+        Args:
+            message: Message dictionary with keys: message_id, role, content, timestamp
+
+        Returns:
+            True if message is a duplicate and should be skipped, False otherwise
+        """
+        self._stats["messages_checked"] += 1
+
+        message_id = message.get("message_id")
+        role = message.get("role")
+
+        # Normalize old message ID formats for backward compatibility
+        if message_id:
+            message_id = self._normalize_message_id(message_id, message)
+            message["message_id"] = message_id
+
+        # Level 1: Check message ID deduplication
+        if message_id and message_id in self._sent_ids:
+            self._stats["duplicates_by_id"] += 1
+            logger.debug(
+                "duplicate_message_id_detected",
+                message_id=message_id,
+                role=role,
+                content_preview=(message.get("content", "") or "")[:50]
+            )
+            # Move to end to mark as recently used (LRU)
+            self._sent_ids.move_to_end(message_id)
+            return True
+
+        # Level 2: Check content signature deduplication (assistant messages only)
+        if role == "assistant":
+            content = message.get("content", "") or ""
+            timestamp_str = message.get("timestamp", "")
+
+            if content and timestamp_str:
+                content_sig = self._content_signature(content)
+
+                if content_sig in self._content_cache:
+                    prev_msg, prev_timestamp = self._content_cache[content_sig]
+
+                    # Check if messages are within deduplication window
+                    try:
+                        current_timestamp = self._parse_timestamp(timestamp_str)
+                        time_diff = abs((current_timestamp - prev_timestamp).total_seconds())
+
+                        if time_diff <= self.CONTENT_DEDUP_WINDOW:
+                            self._stats["duplicates_by_content"] += 1
+                            logger.debug(
+                                "duplicate_content_signature_detected",
+                                message_id=message_id,
+                                content_signature=content_sig[:16],
+                                time_diff=time_diff,
+                                prev_message_id=prev_msg.get("message_id")
+                            )
+                            # Move to end (LRU)
+                            self._content_cache.move_to_end(content_sig)
+                            return True
+                    except Exception as e:
+                        # If timestamp parsing fails, don't skip the message
+                        logger.warning(
+                            "failed_to_parse_timestamp_for_content_dedup",
+                            timestamp=timestamp_str,
+                            error=str(e)
+                        )
+
+        return False
+
+    def mark_sent(self, message: Dict[str, Any]) -> None:
+        """
+        Mark a message as sent (add to deduplication cache).
+
+        This method:
+        1. Adds message ID to sent cache
+        2. For assistant messages, adds content signature to cache
+        3. Enforces LRU eviction when cache exceeds max size
+
+        Args:
+            message: Message dictionary with keys: message_id, role, content, timestamp
+        """
+        message_id = message.get("message_id")
+        role = message.get("role")
+
+        # Add message ID to sent cache
+        if message_id:
+            self._sent_ids[message_id] = True
+            self._evict_if_needed(self._sent_ids)
+
+        # Add content signature for assistant messages
+        if role == "assistant":
+            content = message.get("content", "") or ""
+            timestamp_str = message.get("timestamp", "")
+
+            if content and timestamp_str:
+                try:
+                    content_sig = self._content_signature(content)
+                    timestamp = self._parse_timestamp(timestamp_str)
+                    self._content_cache[content_sig] = (message, timestamp)
+                    self._evict_if_needed(self._content_cache)
+                except Exception as e:
+                    logger.warning(
+                        "failed_to_cache_content_signature",
+                        message_id=message_id,
+                        error=str(e)
+                    )
+
+    def _content_signature(self, content: str) -> str:
+        """
+        Generate a content signature for deduplication.
+
+        The signature is based on:
+        - First 200 characters of normalized content
+        - Lowercase, stripped of whitespace
+        - MD5 hashed for fixed-length signature
+
+        Args:
+            content: Message content string
+
+        Returns:
+            MD5 hash of normalized content (32-character hex string)
+        """
+        if not content:
+            return ""
+
+        # Normalize: strip, lowercase, first 200 chars
+        normalized = content.strip().lower()[:self.CONTENT_SIGNATURE_LENGTH]
+
+        # Hash for fixed-length signature
+        return hashlib.md5(normalized.encode()).hexdigest()
+
+    def _normalize_message_id(self, message_id: str, message: Dict[str, Any]) -> str:
+        """
+        Normalize old message ID formats for backward compatibility.
+
+        Message ID formats:
+        - New (turn-based): {execution_id}_{role}_{turn_number}
+          Example: "exec123_assistant_5"
+        - Old (timestamp-based): {execution_id}_{role}_{timestamp_micros}
+          Example: "exec123_assistant_1234567890123456"
+        - Old (index-based): {execution_id}_{role}_{idx}
+          Example: "exec123_assistant_42" (ambiguous with turn-based)
+
+        Detection heuristic:
+        - If last part is < 10000, assume turn-based (new format) - keep as-is
+        - If last part is >= 10000, assume timestamp-based (old format) - use content hash
+        - If can't parse, use content hash
+
+        Args:
+            message_id: Original message ID
+            message: Full message dictionary for content hash fallback
+
+        Returns:
+            Normalized message ID (may be unchanged if already in new format)
+        """
+        if not message_id:
+            return message_id
+
+        parts = message_id.split("_")
+
+        # Check if format is: {execution_id}_{role}_{number}
+        if len(parts) >= 3 and parts[-2] in ["user", "assistant", "system"]:
+            try:
+                last_part = int(parts[-1])
+
+                # Turn numbers are small (1-100), timestamps are huge (1e15)
+                if last_part < 10000:
+                    # New format (turn-based) - keep as-is
+                    return message_id
+
+                # Old format (timestamp-based) - normalize to content hash
+                self._stats["normalized_ids"] += 1
+
+            except (ValueError, IndexError):
+                # Can't parse as number - might be hash or other format
+                # Use content hash for stability
+                self._stats["normalized_ids"] += 1
+
+        # Generate stable ID based on content hash
+        content = message.get("content", "") or ""
+        role = message.get("role", "unknown")
+        execution_id = parts[0] if parts else "unknown"
+
+        content_hash = hashlib.md5(content.encode()).hexdigest()[:8]
+        normalized_id = f"{execution_id}_{role}_{content_hash}"
+
+        logger.debug(
+            "normalized_old_message_id_format",
+            old_id=message_id,
+            new_id=normalized_id,
+            role=role
+        )
+
+        return normalized_id
+
+    def _parse_timestamp(self, timestamp_str: str) -> datetime:
+        """
+        Parse ISO format timestamp string.
+
+        Handles both with and without 'Z' suffix.
+
+        Args:
+            timestamp_str: ISO format timestamp (e.g., "2024-01-15T10:30:00Z")
+
+        Returns:
+            datetime object
+
+        Raises:
+            ValueError: If timestamp cannot be parsed
+        """
+        # Handle 'Z' suffix for UTC timestamps
+        normalized = timestamp_str.replace('Z', '+00:00')
+        return datetime.fromisoformat(normalized)
+
+    def _evict_if_needed(self, cache: OrderedDict) -> None:
+        """
+        Evict oldest entry from cache if it exceeds max size (LRU eviction).
+
+        Args:
+            cache: OrderedDict cache to check and evict from
+        """
+        if len(cache) > self.max_size:
+            # Remove oldest entry (first item in OrderedDict)
+            cache.popitem(last=False)
+            self._stats["evictions"] += 1
+
+            logger.debug(
+                "lru_cache_eviction",
+                cache_size=len(cache),
+                max_size=self.max_size,
+                total_evictions=self._stats["evictions"]
+            )
+
+    def get_stats(self) -> Dict[str, int]:
+        """
+        Get deduplication statistics.
+
+        Returns:
+            Dictionary with statistics:
+            - messages_checked: Total messages checked
+            - duplicates_by_id: Duplicates found by message ID
+            - duplicates_by_content: Duplicates found by content signature
+            - evictions: LRU cache evictions performed
+            - normalized_ids: Old message IDs normalized
+        """
+        return self._stats.copy()
+
+    def reset(self) -> None:
+        """
+        Reset the deduplicator (clear all caches and statistics).
+
+        Useful for testing or starting a new session.
+        """
+        self._sent_ids.clear()
+        self._content_cache.clear()
+        self._stats = {
+            "messages_checked": 0,
+            "duplicates_by_id": 0,
+            "duplicates_by_content": 0,
+            "evictions": 0,
+            "normalized_ids": 0,
+        }
+        logger.debug("message_deduplicator_reset")

control_plane_api/app/routers/executions/streaming/event_buffer.py

@@ -0,0 +1,353 @@
+"""
+EventBuffer class for SSE gap recovery support.
+
+Implements a ring buffer pattern using collections.deque for bounded memory,
+supporting event replay and gap detection for Server-Sent Events (SSE) streams.
+"""
+
+import time
+import json
+from collections import deque
+from typing import Dict, Any, List, Optional, Tuple
+from structlog import get_logger
+
+logger = get_logger(__name__)
+
+
+class EventBuffer:
+    """
+    Ring buffer for SSE event storage with gap recovery support.
+
+    Uses a bounded deque to store recent SSE events, enabling:
+    - Event replay on reconnection (Last-Event-ID pattern)
+    - Gap detection between client and server state
+    - Automatic memory management with size-based eviction
+
+    Events are stored as tuples: (event_id, event_type, data, size)
+    Event IDs follow format: {execution_id}_{counter}_{timestamp_micros}
+
+    Limits:
+    - Max 100 events (ring buffer evicts oldest)
+    - Max 100KB total buffer size (evicts oldest until under limit)
+
+    Thread Safety:
+    - Uses deque which is thread-safe for append/popleft operations
+    - Size tracking uses atomic operations
+
+    Test Strategy:
+    - Unit test buffer capacity limits (100 events)
+    - Unit test size limits (100KB)
+    - Test replay from various positions
+    - Test gap detection accuracy
+    - Test edge cases (empty buffer, no gaps, large gaps)
+    - Test event ID parsing with various formats
+    """
+
+    MAX_BUFFER_EVENTS = 100
+    MAX_BUFFER_SIZE = 100 * 1024  # 100KB
+
+    def __init__(self, execution_id: str):
+        """
+        Initialize EventBuffer for a specific execution.
+
+        Args:
+            execution_id: The execution ID this buffer is associated with
+        """
+        self.execution_id = execution_id
+        self.buffer: deque = deque(maxlen=self.MAX_BUFFER_EVENTS)
+        self._current_size = 0
+
+    def add_event(self, event_id: str, event_type: str, data: str) -> None:
+        """
+        Add event to buffer with automatic eviction.
+
+        Events are stored as (event_id, event_type, data, size) tuples.
+        If buffer exceeds MAX_BUFFER_SIZE, oldest events are evicted.
+        The deque's maxlen ensures automatic eviction at MAX_BUFFER_EVENTS.
+
+        Args:
+            event_id: Unique event identifier (format: {execution_id}_{counter}_{timestamp})
+            event_type: Type of event (e.g., 'status', 'tool_call', 'message')
+            data: JSON string of event data
+        """
+        event_size = len(data)
+        self._current_size += event_size
+
+        # Add event to buffer (deque automatically evicts oldest if at maxlen)
+        # Note: If eviction happens, we need to track it for size accounting
+        old_len = len(self.buffer)
+        self.buffer.append((event_id, event_type, data, event_size))
+
+        # If deque evicted an event (length didn't increase), subtract its size
+        if len(self.buffer) == old_len and old_len == self.MAX_BUFFER_EVENTS:
+            # Deque evicted the oldest event, need to account for it
+            # We can't access the evicted event, so we need to recalculate
+            self._recalculate_size()
+
+        # Remove old events if buffer exceeds size limit
+        while self._current_size > self.MAX_BUFFER_SIZE and len(self.buffer) > 1:
+            _, _, _, old_size = self.buffer.popleft()
+            self._current_size -= old_size
+
+        logger.debug(
+            "event_buffered",
+            execution_id=self.execution_id,
+            event_id=event_id,
+            event_type=event_type,
+            buffer_count=len(self.buffer),
+            buffer_size=self._current_size
+        )
+
+    def _recalculate_size(self) -> None:
+        """Recalculate buffer size from scratch."""
+        self._current_size = sum(size for _, _, _, size in self.buffer)
+
+    def replay_from_id(self, last_event_id: str) -> List[Tuple[str, str, str]]:
+        """
+        Replay events after the given event ID.
+
+        Parses the last_event_id to extract the counter, then returns all
+        buffered events with higher counters in chronological order.
+
+        Args:
+            last_event_id: Last event ID received by client
+
+        Returns:
+            List of (event_id, event_type, data) tuples to replay
+        """
+        last_counter = self._parse_event_id(last_event_id)
+
+        if last_counter is None or not self.buffer:
+            return []
+
+        replay_events: List[Tuple[str, str, str]] = []
+
+        for buf_event_id, buf_event_type, buf_data, _ in self.buffer:
+            buf_counter = self._parse_event_id(buf_event_id)
+
+            if buf_counter is not None and buf_counter > last_counter:
+                replay_events.append((buf_event_id, buf_event_type, buf_data))
+
+        if replay_events:
+            logger.info(
+                "replaying_buffered_events",
+                execution_id=self.execution_id,
+                last_counter=last_counter,
+                replay_count=len(replay_events)
+            )
+
+        return replay_events
+
+    def detect_gap(self, last_event_id: str, current_event_id: str) -> Optional[Dict[str, Any]]:
+        """
+        Detect gap between last and current event IDs.
+
+        Compares sequence counters to identify missing events. Returns gap
+        information if events are missing from the buffer.
+
+        Args:
+            last_event_id: Last event ID received by client
+            current_event_id: Current event ID being processed
+
+        Returns:
+            Gap information dict if gap detected, None otherwise:
+            {
+                "gap_detected": True,
+                "last_counter": int,
+                "current_counter": int,
+                "missing_count": int,
+                "reason": str
+            }
+        """
+        last_counter = self._parse_event_id(last_event_id)
+        current_counter = self._parse_event_id(current_event_id)
+
+        if last_counter is None or current_counter is None:
+            return None
+
+        # Check if there's a gap
+        expected_counter = last_counter + 1
+        if current_counter > expected_counter:
+            missing_count = current_counter - expected_counter
+
+            logger.warning(
+                "gap_detected",
+                execution_id=self.execution_id,
+                last_counter=last_counter,
+                current_counter=current_counter,
+                missing_count=missing_count
+            )
+
+            return {
+                "gap_detected": True,
+                "last_counter": last_counter,
+                "current_counter": current_counter,
+                "missing_count": missing_count,
+                "reason": f"Missing {missing_count} events between {last_counter} and {current_counter}"
+            }
+
+        return None
+
+    def check_buffer_miss(self, last_event_id: str) -> Optional[Dict[str, Any]]:
+        """
+        Check if last_event_id is no longer in buffer (too old).
+
+        Args:
+            last_event_id: Last event ID received by client
+
+        Returns:
+            Buffer miss information if events are too old, None otherwise:
+            {
+                "buffer_miss": True,
+                "last_known_id": str,
+                "buffer_oldest": str,
+                "reason": str
+            }
+        """
+        if not self.buffer:
+            return {
+                "buffer_miss": True,
+                "last_known_id": last_event_id,
+                "buffer_oldest": None,
+                "reason": "Event buffer is empty"
+            }
+
+        last_counter = self._parse_event_id(last_event_id)
+        if last_counter is None:
+            return None
+
+        # Check if any buffered event is after the last_event_id
+        has_newer_events = False
+        for buf_event_id, _, _, _ in self.buffer:
+            buf_counter = self._parse_event_id(buf_event_id)
+            if buf_counter is not None and buf_counter > last_counter:
+                has_newer_events = True
+                break
+
+        # If we have newer events but couldn't replay (empty replay_from_id),
+        # it means the last_event_id is older than our oldest buffered event
+        if not has_newer_events and last_counter > 0:
+            oldest_event_id = self.buffer[0][0] if self.buffer else None
+
+            logger.warning(
+                "gap_detected_buffer_miss",
+                execution_id=self.execution_id,
+                last_counter=last_counter,
+                buffer_size=len(self.buffer),
+                buffer_oldest=oldest_event_id
+            )
+
+            return {
+                "buffer_miss": True,
+                "last_known_id": last_event_id,
+                "buffer_oldest": oldest_event_id,
+                "reason": "Event buffer miss - events too old"
+            }
+
+        return None
+
+    def _parse_event_id(self, event_id: str) -> Optional[int]:
+        """
+        Extract sequence counter from event ID format: {execution_id}_{counter}_{timestamp}.
+
+        Args:
+            event_id: Event ID string to parse
+
+        Returns:
+            Sequence counter as integer, or None if parsing fails
+        """
+        try:
+            parts = event_id.split("_")
+
+            # Format: {execution_id}_{counter}_{timestamp_micros}
+            # Validate execution_id matches (first part)
+            if len(parts) >= 2 and parts[0] == self.execution_id:
+                return int(parts[1])
+
+            # If execution_id doesn't match, still try to parse counter
+            # for compatibility with different ID formats
+            if len(parts) >= 2:
+                return int(parts[1])
+
+        except (ValueError, IndexError) as e:
+            logger.warning(
+                "invalid_event_id_format",
+                execution_id=self.execution_id,
+                event_id=event_id,
+                error=str(e)
+            )
+
+        return None
+
+    def _estimate_size(self, data: Dict[str, Any]) -> int:
+        """
+        Estimate size of event data in bytes.
+
+        Args:
+            data: Dictionary of event data
+
+        Returns:
+            Estimated size in bytes
+        """
+        try:
+            return len(json.dumps(data))
+        except Exception:
+            # Fallback: rough estimate based on string representation
+            return len(str(data))
+
+    def get_buffer_size(self) -> int:
+        """
+        Get current buffer size in bytes.
+
+        Returns:
+            Total size of buffered event data in bytes
+        """
+        return self._current_size
+
+    def get_buffer_info(self) -> Dict[str, Any]:
+        """
+        Get buffer statistics for debugging.
+
+        Returns:
+            Dictionary with buffer statistics:
+            {
+                "execution_id": str,
+                "event_count": int,
+                "total_size_bytes": int,
+                "max_events": int,
+                "max_size_bytes": int,
+                "oldest_event_id": str,
+                "newest_event_id": str,
+                "utilization_percent": float
+            }
+        """
+        oldest_event_id = None
+        newest_event_id = None
+
+        if self.buffer:
+            oldest_event_id = self.buffer[0][0]
+            newest_event_id = self.buffer[-1][0]
+
+        event_count = len(self.buffer)
+        utilization = (event_count / self.MAX_BUFFER_EVENTS) * 100
+
+        return {
+            "execution_id": self.execution_id,
+            "event_count": event_count,
+            "total_size_bytes": self._current_size,
+            "max_events": self.MAX_BUFFER_EVENTS,
+            "max_size_bytes": self.MAX_BUFFER_SIZE,
+            "oldest_event_id": oldest_event_id,
+            "newest_event_id": newest_event_id,
+            "utilization_percent": round(utilization, 2)
+        }
+
+    def clear(self) -> None:
+        """Clear all buffered events."""
+        self.buffer.clear()
+        self._current_size = 0
+
+        logger.debug(
+            "event_buffer_cleared",
+            execution_id=self.execution_id
+        )