"""Composable message processing middleware. Each middleware is a standalone class that can be composed into a pipeline. They extract logic that was previously baked into the Channel base class, making it reusable across both legacy and plugin-based channels. Also contains the supporting data structures (DedupCache, GroupHistoryBuffer, TypingManager, PairingManager) that were previously in separate files. """ from __future__ import annotations import asyncio import dataclasses import logging import random import time from collections import OrderedDict, deque from collections.abc import Awaitable, Callable from dataclasses import dataclass from typing import Any from .base import RawIncoming from .bus.events import InboundMessage, OutboundMessage from .debug import emit_debug_event_if from .interaction import is_slash_command _logger = logging.getLogger(__name__) # ── Task cancellation helper ───────────────────────────────────────── async def _cancel_task(task: asyncio.Task) -> None: """Cancel an asyncio task and await its completion. Suppresses ``CancelledError`` from the cancelled *task* but re-raises if the **current** task was itself cancelled (to avoid swallowing an outer cancellation signal — required for correct behavior on Python 3.12+ where ``_must_cancel`` no longer auto-re-delivers). """ task.cancel() try: await task except asyncio.CancelledError: current = asyncio.current_task() if current is not None and current.cancelling() > 0: raise except Exception: pass # Already logged elsewhere; prevent unhandled propagation # ═══════════════════════════════════════════════════════════════════════ # Supporting data structures # ═══════════════════════════════════════════════════════════════════════ # ── Dedup cache ────────────────────────────────────────────────────── _DEDUP_MAX = 1000 _DEDUP_TRIM = 500 _DEDUP_TTL = 3600 # 1 hour class DedupCache: """Bounded ordered cache with TTL for detecting duplicate message IDs. Entries expire after *ttl_seconds* and are pruned lazily on each lookup. When the cache exceeds *max_size* entries it is trimmed down to *trim_to* by evicting the oldest entries. Accessed entries are moved to the end (LRU behavior). """ def __init__( self, max_size: int = _DEDUP_MAX, trim_to: int = _DEDUP_TRIM, ttl_seconds: float = _DEDUP_TTL, clock: Callable[[], float] | None = None, ) -> None: self._seen: OrderedDict[str, float] = OrderedDict() self._max = max_size self._trim = trim_to self._ttl = ttl_seconds self._clock = clock or time.monotonic # ── public API ────────────────────────────────────────────────── def is_duplicate(self, msg_id: str) -> bool: """Return ``True`` if *msg_id* has been seen before. First-time IDs are recorded and ``False`` is returned. Empty / falsy IDs are never considered duplicates. Expired entries are pruned before the check. """ if not msg_id: return False now = self._clock() self._prune(now) if msg_id in self._seen: # LRU: refresh position and timestamp self._seen.move_to_end(msg_id) self._seen[msg_id] = now return True self._seen[msg_id] = now if len(self._seen) > self._max: while len(self._seen) > self._trim: self._seen.popitem(last=False) return False def clear(self) -> None: """Remove all entries.""" self._seen.clear() @property def size(self) -> int: """Number of entries currently in the cache.""" return len(self._seen) # ── internal ──────────────────────────────────────────────────── def _prune(self, now: float | None = None) -> None: """Remove entries older than *ttl_seconds*.""" cutoff = (self._clock() if now is None else now) - self._ttl # OrderedDict is insertion-ordered; oldest entries are first. while self._seen: _key, ts = next(iter(self._seen.items())) if ts > cutoff: break self._seen.popitem(last=False) # ── Group history buffer ───────────────────────────────────────────── @dataclass class HistoryEntry: sender_id: str text: str timestamp: float message_id: str = "" class GroupHistoryBuffer: """Per-chat circular buffer of recent messages.""" def __init__(self, max_per_chat: int = 50, max_age_seconds: int = 3600): self._buffers: dict[str, deque[HistoryEntry]] = {} self._max = max_per_chat self._max_age = max_age_seconds def add(self, chat_id: str, entry: HistoryEntry) -> None: """Add a message to the chat's history buffer.""" if chat_id not in self._buffers: self._buffers[chat_id] = deque(maxlen=self._max) self._buffers[chat_id].append(entry) def get_recent(self, chat_id: str, limit: int = 20) -> list[HistoryEntry]: """Get recent messages for context injection, excluding expired ones.""" buf = self._buffers.get(chat_id) if not buf: return [] now = time.monotonic() recent = [e for e in buf if now - e.timestamp < self._max_age] return recent[-limit:] def format_context(self, chat_id: str, limit: int = 20) -> str: """Format recent messages as context block for the agent.""" entries = self.get_recent(chat_id, limit) if not entries: return "" lines = ["[Chat messages since your last reply - for context]"] for e in entries: lines.append(f"[from: {e.sender_id}] {e.text}") lines.append("[/Chat context]") return "\n".join(lines) def clear(self, chat_id: str) -> None: """Clear history for a chat (e.g., after the bot replies).""" self._buffers.pop(chat_id, None) # ── Typing indicator manager ───────────────────────────────────────── class TypingManager: """Manages background typing-indicator loops per chat_id. Args: send_action: Async callable that sends a single typing indicator for a given chat_id. interval: Seconds between typing indicator sends. """ def __init__( self, send_action: Callable[[str], Awaitable[None]], interval: float = 5.0, debug_trace: bool = False, channel_name: str = "unknown", ) -> None: self._send_action = send_action self._interval = interval self._tasks: dict[str, asyncio.Task] = {} self._debug_trace = debug_trace self._channel_name = channel_name async def start(self, chat_id: str) -> None: """Start a background typing-indicator loop for *chat_id*.""" await self.stop(chat_id) async def _loop() -> None: while True: try: await self._send_action(chat_id) except Exception as exc: _trace_named_event( "typing_error", enabled=self._debug_trace, channel_name=self._channel_name, chat_id=chat_id, error=str(exc), ) await asyncio.sleep(self._interval) self._tasks[chat_id] = asyncio.create_task(_loop()) async def stop(self, chat_id: str) -> None: """Cancel the typing-indicator loop for *chat_id*.""" task = self._tasks.pop(chat_id, None) if task: await _cancel_task(task) async def stop_all(self) -> None: """Cancel all active typing-indicator loops.""" for cid in list(self._tasks): await self.stop(cid) @property def active_chats(self) -> list[str]: """Return chat_ids with active typing loops.""" return list(self._tasks) # ── Pairing manager ───────────────────────────────────────────────── @dataclass class PairingRequest: sender_id: str channel: str code: str created_at: float approved: bool = False class PairingManager: """Manages DM pairing codes for channel access control.""" CODE_EXPIRY = 3600 # 1 hour MAX_PENDING = 50 # max pending requests def __init__(self): self._pending: dict[str, PairingRequest] = {} # code -> request self._approved: set[str] = set() # "channel:sender_id" keys def is_approved(self, channel: str, sender_id: str) -> bool: """Check if sender is already approved.""" return f"{channel}:{sender_id}" in self._approved def request_pairing(self, channel: str, sender_id: str) -> str: """Generate a pairing code for a new sender. Returns the code.""" # Check if already has pending request for code, req in list(self._pending.items()): if req.sender_id == sender_id and req.channel == channel: if time.monotonic() - req.created_at < self.CODE_EXPIRY: return code # return existing code else: del self._pending[code] break # Cleanup expired self._cleanup_expired() # Generate new code code = f"{random.randint(100000, 999999)}" while code in self._pending: code = f"{random.randint(100000, 999999)}" self._pending[code] = PairingRequest( sender_id=sender_id, channel=channel, code=code, created_at=time.monotonic(), ) _logger.info(f"Pairing code {code} generated for {channel}:{sender_id}") return code def approve(self, code: str) -> tuple[bool, str]: """Approve a pairing code. Returns (success, message).""" req = self._pending.get(code) if not req: return False, f"Unknown code: {code}" if time.monotonic() - req.created_at > self.CODE_EXPIRY: del self._pending[code] return False, f"Code {code} expired" key = f"{req.channel}:{req.sender_id}" self._approved.add(key) del self._pending[code] _logger.info(f"Approved pairing for {key}") return True, f"Approved {req.sender_id} on {req.channel}" def reject(self, code: str) -> tuple[bool, str]: """Reject a pairing code.""" if code in self._pending: del self._pending[code] return True, f"Rejected code {code}" return False, f"Unknown code: {code}" def list_pending(self) -> list[PairingRequest]: """List all pending (non-expired) requests.""" self._cleanup_expired() return list(self._pending.values()) def _cleanup_expired(self): now = time.monotonic() expired = [ c for c, r in self._pending.items() if now - r.created_at > self.CODE_EXPIRY ] for c in expired: del self._pending[c] # ═══════════════════════════════════════════════════════════════════════ # Middleware classes # ═══════════════════════════════════════════════════════════════════════ # ── Inbound middleware base ────────────────────────────────────────── class InboundMiddleware: """Base class for inbound message processing middleware.""" async def process_inbound( self, raw: RawIncoming, context: dict[str, Any], ) -> RawIncoming | None: """Process an inbound raw message. Return the (possibly modified) RawIncoming to continue the pipeline, or ``None`` to drop the message. """ return raw class OutboundMiddlewareBase: """Base class for outbound message processing middleware.""" async def process_outbound( self, message: OutboundMessage, context: dict[str, Any], ) -> OutboundMessage | None: """Process an outbound message. Return the (possibly modified) OutboundMessage to continue, or ``None`` to drop it. """ return message def _debug_trace_enabled(context: dict[str, Any]) -> bool: """Check whether channel-level debug tracing is enabled for this message.""" channel = context.get("channel") if channel is None: return False return channel.is_debug_trace_enabled() def _ctx_channel_name(context: dict[str, Any]) -> str: """Extract the channel name from middleware context.""" ch = context.get("channel") return getattr(ch, "name", "unknown") if ch else "unknown" def _trace_context_event( context: dict[str, Any], event: str, **fields: Any, ) -> None: """Emit a middleware trace event using the shared channel context.""" emit_debug_event_if( _logger, event, _debug_trace_enabled(context), channel=_ctx_channel_name(context), **fields, ) def _trace_named_event( event: str, *, enabled: bool, channel_name: str, **fields: Any, ) -> None: """Emit a trace event for helpers that already carry trace state.""" emit_debug_event_if( _logger, event, enabled, channel=channel_name, **fields, ) # ── Dedup ──────────────────────────────────────────────────────────── class DedupMiddleware(InboundMiddleware): """Message deduplication using a bounded TTL cache.""" def __init__( self, max_size: int = 1000, trim_to: int = 500, ttl_seconds: float = 3600.0, clock: Callable[[], float] | None = None, ) -> None: self._cache = DedupCache( max_size=max_size, trim_to=trim_to, ttl_seconds=ttl_seconds, clock=clock, ) async def process_inbound( self, raw: RawIncoming, context: dict[str, Any], ) -> RawIncoming | None: if raw.message_id and self._cache.is_duplicate(raw.message_id): _trace_context_event( context, "middleware_dedup_drop", message_id=raw.message_id, sender_id=raw.sender_id, ) return None return raw # ── Debounce ───────────────────────────────────────────────────────── class DebounceMiddleware: """Per-sender message batching with configurable timing. This middleware collects messages from the same sender and merges them after a debounce delay. It does not follow the simple process_inbound pattern because it needs to buffer across calls. Usage: call ``submit()`` for each message; merged results are delivered via the ``on_ready`` callback. """ def __init__( self, *, initial_debounce: float = 2.0, debounce_step: float = 0.5, max_debounce: float = 5.0, on_ready: Callable[[InboundMessage], Any] | None = None, ) -> None: self.initial_debounce = initial_debounce self.debounce_step = debounce_step self.max_debounce = max_debounce self.on_ready = on_ready self._buffers: dict[str, list[str]] = {} self._metadata: dict[str, dict] = {} self._media: dict[str, list[str]] = {} self._message_ids: dict[str, str] = {} self._tasks: dict[str, asyncio.Task] = {} self._channel_name: str = "" def set_channel_name(self, name: str) -> None: self._channel_name = name async def submit(self, msg: InboundMessage) -> None: """Buffer *msg* and schedule flush after debounce delay.""" sender = msg.sender_id if sender not in self._buffers: self._buffers[sender] = [] self._metadata[sender] = msg.metadata self._media[sender] = [] self._buffers[sender].append(msg.content) if msg.message_id: self._message_ids[sender] = msg.message_id if msg.media: self._media[sender].extend(msg.media) if sender in self._tasks: self._tasks[sender].cancel() count = len(self._buffers[sender]) wait = min( self.initial_debounce + (count - 1) * self.debounce_step, self.max_debounce, ) async def _flush(_s: str = sender, _w: float = wait) -> None: await asyncio.sleep(_w) await self._flush_sender(_s) self._tasks[sender] = asyncio.create_task(_flush()) async def _flush_sender(self, sender: str) -> None: messages = self._buffers.pop(sender, []) metadata = self._metadata.pop(sender, None) media = self._media.pop(sender, []) message_id = self._message_ids.pop(sender, "") self._tasks.pop(sender, None) if not messages: return merged = "\n".join(messages) chat_id = (metadata or {}).get("chat_id", sender) inbound = InboundMessage( channel=self._channel_name, sender_id=sender, chat_id=str(chat_id), content=merged, media=media, metadata=metadata or {}, message_id=message_id, ) if self.on_ready: await self.on_ready(inbound) async def cancel_all(self) -> None: """Cancel all pending debounce tasks and await their completion.""" tasks = list(self._tasks.values()) self._tasks.clear() for task in tasks: await _cancel_task(task) # ── Chunking ───────────────────────────────────────────────────────── class ChunkingMiddleware(OutboundMiddlewareBase): """Auto-split messages respecting format expansion. Wraps the existing ``chunking.chunk_text`` utility and the re-splitting logic from ``Channel._prepare_chunks``. """ def __init__(self, capabilities: Any) -> None: from .capabilities import ChannelCapabilities self._capabilities: ChannelCapabilities = capabilities def prepare_chunks( self, content: str, limit: int, format_fn: Callable[[str], str] | None = None, ) -> list[tuple[str, str]]: """Build ``(formatted, raw)`` pairs, re-splitting when needed. If *format_fn* is None, formatted == raw. """ from .base import chunk_text if format_fn is None: format_fn = lambda t: t # noqa: E731 raw_chunks = chunk_text(content, limit) pairs: list[tuple[str, str]] = [] for raw in raw_chunks: formatted = format_fn(raw) if len(formatted) <= limit: pairs.append((formatted, raw)) else: sub_limit = max(limit // 2, 500) for sub_raw in chunk_text(raw, sub_limit): sub_fmt = format_fn(sub_raw) if len(sub_fmt) <= limit: pairs.append((sub_fmt, sub_raw)) else: pairs.append((sub_raw, sub_raw)) return pairs # ── Formatting ─────────────────────────────────────────────────────── class FormattingMiddleware(OutboundMiddlewareBase): """Markdown -> channel format conversion. Uses ``UnifiedFormatter`` configured from capabilities. """ def __init__(self, capabilities: Any) -> None: from .capabilities import ChannelCapabilities from .formatter import UnifiedFormatter caps: ChannelCapabilities = capabilities self._formatter = UnifiedFormatter.for_channel(caps.format_type) def format(self, text: str) -> str: """Convert text to channel format.""" return self._formatter.format(text) async def process_outbound( self, message: OutboundMessage, context: dict[str, Any], ) -> OutboundMessage | None: formatted = self._formatter.format(message.content) return dataclasses.replace(message, content=formatted) # ── Retry ──────────────────────────────────────────────────────────── class RetryMiddleware: """Exponential backoff send retry. Wraps ``retry.retry_async`` with channel-appropriate configuration. """ def __init__(self, channel_name: str = "unknown") -> None: from .retry import DEFAULT_RETRY, RETRY_PRESETS self._config = RETRY_PRESETS.get(channel_name, DEFAULT_RETRY) self._channel_name = channel_name async def execute( self, coro_factory: Callable[[], Any], should_retry: Callable[[Exception, int], bool] | None = None, retry_after_s: Callable[[Exception], float | None] | None = None, ) -> Any: """Execute *coro_factory* with retry logic.""" from .retry import retry_async return await retry_async( coro_factory, config=self._config, should_retry=should_retry or (lambda exc, _: True), retry_after_s=retry_after_s, on_retry=lambda info: _logger.warning( f"{self._channel_name} retry {info.attempt}/{info.max_attempts} " f"in {info.delay_s:.2f}s: {info.error}" ), label=f"{self._channel_name}.send", ) # ── Typing ─────────────────────────────────────────────────────────── class TypingMiddleware: """Typing indicator management. Wraps ``TypingManager`` for use as a standalone middleware component. """ def __init__( self, send_typing_fn: Callable[[str], Any], interval: float = 5.0, debug_trace: bool = False, channel_name: str = "unknown", ) -> None: self._manager = TypingManager( send_typing_fn, interval=interval, debug_trace=debug_trace, channel_name=channel_name, ) async def start(self, chat_id: str) -> None: await self._manager.start(chat_id) async def stop(self, chat_id: str) -> None: await self._manager.stop(chat_id) async def stop_all(self) -> None: await self._manager.stop_all() # ── ACK Reaction ───────────────────────────────────────────────────── class AckReactionMiddleware: """ACK emoji reaction with configurable scope. Scope controls when reactions are sent: - ``"all"``: react to every message - ``"direct"``: react only in DMs - ``"group-all"``: react in group chats (all messages) - ``"group-mentions"``: react in groups only when mentioned - ``"off"``: disable reactions """ def __init__( self, *, scope: str = "all", emoji: str = "\U0001f440", remove_after_reply: bool = False, send_fn: Callable[[str, str, str], Any] | None = None, remove_fn: Callable[[str, str, str], Any] | None = None, debug_trace: bool = False, channel_name: str = "unknown", ) -> None: self.scope = scope self.emoji = emoji self.remove_after_reply = remove_after_reply self._send_fn = send_fn self._remove_fn = remove_fn self._pending: dict[str, str] = {} # chat_id -> message_id self._debug_trace = debug_trace self._channel_name = channel_name def should_react(self, *, is_group: bool, was_mentioned: bool) -> bool: if self.scope == "off": return False if self.scope == "all": return True if self.scope == "direct": return not is_group if self.scope == "group-all": return is_group if self.scope == "group-mentions": return is_group and was_mentioned return False async def send_ack(self, chat_id: str, message_id: str) -> None: if self._send_fn and message_id: try: await self._send_fn(chat_id, message_id, self.emoji) if self.remove_after_reply: self._pending[chat_id] = message_id except Exception as exc: _trace_named_event( "ack_send_error", enabled=self._debug_trace, channel_name=self._channel_name, chat_id=chat_id, message_id=message_id, error=str(exc), ) async def remove_ack(self, chat_id: str) -> None: message_id = self._pending.pop(chat_id, None) if message_id and self._remove_fn: try: await self._remove_fn(chat_id, message_id, self.emoji) except Exception as exc: _trace_named_event( "ack_remove_error", enabled=self._debug_trace, channel_name=self._channel_name, chat_id=chat_id, message_id=message_id, error=str(exc), ) # ── Mention Gating ─────────────────────────────────────────────────── class MentionGatingMiddleware(InboundMiddleware): """Filter messages based on mention policy. Policy values: - ``"always"``: require mention in all chats - ``"group"``: require mention only in groups (default) - ``"off"``: never require mention """ def __init__( self, require_mention: str = "group", strip_fn: Callable[[str], str] | None = None, ) -> None: self.require_mention = require_mention self._strip_fn = strip_fn async def process_inbound( self, raw: RawIncoming, context: dict[str, Any], ) -> RawIncoming | None: if not self._should_process(raw): _trace_context_event( context, "middleware_mention_drop", chat_id=raw.chat_id, policy=self.require_mention, ) return None # A slash command's platform target belongs only to its first token; # preserve mentions in its arguments. Ordinary group messages may # still carry a bot mention elsewhere and use the full-message strip. if self._strip_fn and is_slash_command(raw.text): text = raw.text token_start = len(text) - len(text.lstrip()) token_end = token_start while token_end < len(text) and not text[token_end].isspace(): token_end += 1 stripped_token = self._strip_fn(text[token_start:token_end]) raw = dataclasses.replace( raw, text=text[:token_start] + stripped_token + text[token_end:], ) elif self._strip_fn and raw.is_group: raw = dataclasses.replace(raw, text=self._strip_fn(raw.text)) return raw def _should_process(self, raw: RawIncoming) -> bool: if self.require_mention == "off": return True if self.require_mention == "always": return raw.was_mentioned # "group" — require mention only in groups if not raw.is_group: return True return raw.was_mentioned # ── AllowList ──────────────────────────────────────────────────────── class AllowListMiddleware(InboundMiddleware): """Sender and channel allow-list enforcement.""" def __init__( self, allowed_senders: set[str] | None = None, allowed_channels: set[str] | None = None, dm_policy: str = "allowlist", ) -> None: self.allowed_senders = allowed_senders self.allowed_channels = allowed_channels self.dm_policy = dm_policy async def process_inbound( self, raw: RawIncoming, context: dict[str, Any], ) -> RawIncoming | None: # Channel allow-list if self.allowed_channels and str(raw.chat_id) not in self.allowed_channels: _trace_context_event( context, "middleware_allowlist_drop", sender_id=raw.sender_id, chat_id=raw.chat_id, reason="chat_not_allowed", ) return None # Sender allow-list if not raw.is_group and self.dm_policy == "open": return raw # open DMs bypass sender checks if not self._is_sender_allowed(raw.sender_id): _trace_context_event( context, "middleware_allowlist_drop", sender_id=raw.sender_id, chat_id=raw.chat_id, reason="sender_not_allowed", ) return None return raw def _is_sender_allowed(self, sender: str) -> bool: if not self.allowed_senders: return True sender_str = str(sender) if sender_str in self.allowed_senders: return True if "|" in sender_str: for part in sender_str.split("|"): if part and part in self.allowed_senders: return True return False # ── Group History ──────────────────────────────────────────────────── class GroupHistoryMiddleware(InboundMiddleware): """Buffer non-mentioned group messages, inject as context when mentioned.""" def __init__( self, max_per_chat: int = 50, max_age_seconds: int = 3600, ) -> None: self._buffer = GroupHistoryBuffer( max_per_chat=max_per_chat, max_age_seconds=max_age_seconds, ) async def process_inbound( self, raw: RawIncoming, context: dict[str, Any], ) -> RawIncoming | None: if not raw.is_group: return raw # Use monotonic clock for consistent expiry calculation ts = time.monotonic() if not raw.was_mentioned: self._buffer.add( raw.chat_id, HistoryEntry( sender_id=raw.sender_id, text=raw.text, timestamp=ts, message_id=raw.message_id, ), ) # Don't drop here — let MentionGatingMiddleware handle that return raw # Slash commands must remain the leading content so channel command # dispatchers can recognize them. Keep buffered chatter for the next # normal mentioned message instead of injecting it ahead of a command. if is_slash_command(raw.text): return raw # Mentioned: inject history context history_context = self._buffer.format_context(raw.chat_id) if history_context: raw = dataclasses.replace( raw, text=history_context + "\n\n[Current message - respond to this]\n" + raw.text, ) self._buffer.clear(raw.chat_id) return raw # ── Pairing ────────────────────────────────────────────────────────── class PairingMiddleware(InboundMiddleware): """DM pairing flow management. When dm_policy is "pairing", unapproved DM senders receive a pairing code. Approved senders pass through normally. """ def __init__( self, channel_name: str, send_response_fn: Callable[[str, str], Any] | None = None, dm_policy: str = "allowlist", ) -> None: self._manager = PairingManager() self._channel_name = channel_name self._send_response_fn = send_response_fn self._dm_policy = dm_policy self._background_tasks: set[asyncio.Task] = set() async def process_inbound( self, raw: RawIncoming, context: dict[str, Any], ) -> RawIncoming | None: if raw.is_group: return raw # pairing only applies to DMs if self._dm_policy != "pairing": return raw if self._manager.is_approved(self._channel_name, raw.sender_id): return raw # Request pairing code = self._manager.request_pairing(self._channel_name, raw.sender_id) if self._send_response_fn: text = f"\U0001f510 Pairing required. Your code: {code}\nThis code expires in 1 hour." task = asyncio.create_task(self._send_response_fn(raw.chat_id, text)) # Track the task to prevent GC and handle exceptions self._background_tasks.add(task) task.add_done_callback(self._background_tasks.discard) _trace_context_event( context, "middleware_pairing_required", sender_id=raw.sender_id, ) _logger.info(f"Pairing required for {raw.sender_id}, code sent") return None