Telegram DM topic routing fix for text batches

Fixes on record

Merged to main. A stamp counts the issues a fix is linked to.

• #34342 fix(telegram): apply DM topic recovery before session keying

  merged to main

  files: gateway/platforms/base.py, gateway/platforms/telegram.py, gateway/run.py, tests/gateway/test_telegram_text_batching.py

  Show the diff

  • gateway/platforms/base.py

    @@ -472,6 +472,7 @@ def is_host_excluded_by_no_proxy(hostname: str, no_proxy_value: str | None = Non
         return False
     
     
    +import dataclasses
     from dataclasses import dataclass, field
     from datetime import datetime
     from pathlib import Path
    @@ -1561,6 +1562,10 @@ def __init__(self, config: PlatformConfig, platform: Platform):
             self.config = config
             self.platform = platform
             self._message_handler: Optional[MessageHandler] = None
    +        # Optional hook (e.g. Telegram DM topic recovery) that rewrites
    +        # ``event.source.thread_id`` before session keying. Returns the
    +        # corrected thread_id or None to leave the source untouched.
    +        self._topic_recovery_fn: Optional[Callable[[Any], Optional[str]]] = None
             self._running = False
             self._fatal_error_code: Optional[str] = None
             self._fatal_error_message: Optional[str] = None
    @@ -1816,6 +1821,40 @@ def set_message_handler(self, handler: MessageHandler) -> None:
             """
             self._message_handler = handler
     
    +    def set_topic_recovery_fn(
    +        self,
    +        fn: Optional[Callable[[Any], Optional[str]]],
    +    ) -> None:
    +        """Install a thread_id-recovery hook (Telegram DM topic mode).
    +
    +        The hook is called with ``event.source`` before session keying;
    +        a non-None return value replaces ``source.thread_id``. Pass
    +        ``None`` to clear the hook.
    +        """
    +        # Guard against subclasses that initialize via ``object.__new__`` in
    +        # tests and never run ``BasePlatformAdapter.__init__``.
    +        self._topic_recovery_fn = fn  # type: ignore[attr-defined]
    +
    +    def _apply_topic_recovery(self, event: MessageEvent) -> None:
    +        """Rewrite ``event.source.thread_id`` in place if the hook returns one."""
    +        recover = getattr(self, "_topic_recovery_fn", None)
    +        if recover is None:
    +            return
    +        source = getattr(event, "source", None)
    +        if source is None:
    +            return
    +        try:
    +            recovered = recover(source)
    +        except Exception:
    +            logger.debug("topic recovery hook failed", exc_info=True)
    +            return
    +        if recovered is None or str(recovered) == str(source.thread_id or ""):
    +            return
    +        try:
    +            event.source = dataclasses.replace(source, thread_id=str(recovered))
    +        except Exception:
    +            logger.debug("topic recovery rewrite failed", exc_info=True)
    +
         def set_busy_session_handler(self, handler: Optional[Callable[[MessageEvent, str], Awaitable[bool]]]) -> None:
             """Set an optional handler for messages arriving during active sessions."""
             self._busy_session_handler = handler
    @@ -3332,7 +3371,12 @@ async def handle_message(self, event: MessageEvent) -> None:
                 return
     
             coerce_plaintext_gateway_command(event)
    -        
    +
    +        # Rewrite ``event.source.thread_id`` via the installed recovery hook
    +        # (Telegram DM topic mode) so the session key, guard checks, and
    +        # downstream delivery all agree on the same lane.
    +        self._apply_topic_recovery(event)
    +
             session_key = build_session_key(
                 event.source,
                 group_sessions_per_user=self.config.extra.get("group_sessions_per_user", True),

  • gateway/platforms/telegram.py

    @@ -5027,8 +5027,14 @@ async def _handle_location_message(self, update: Update, context: ContextTypes.D
         # ------------------------------------------------------------------
     
         def _text_batch_key(self, event: MessageEvent) -> str:
    -        """Session-scoped key for text message batching."""
    +        """Session-scoped key for text message batching.
    +
    +        Applies the installed topic-recovery hook first so DM-topic batches
    +        coalesce on (and dispatch to) the recovered lane rather than the
    +        raw inbound ``message_thread_id`` Telegram may have attached.
    +        """
             from gateway.session import build_session_key
    +        self._apply_topic_recovery(event)
             return build_session_key(
                 event.source,
                 group_sessions_per_user=self.config.extra.get("group_sessions_per_user", True),

  • gateway/run.py

    @@ -4159,6 +4159,7 @@ async def start(self) -> bool:
                 adapter.set_fatal_error_handler(self._handle_adapter_fatal_error)
                 adapter.set_session_store(self.session_store)
                 adapter.set_busy_session_handler(self._handle_active_session_busy_message)
    +            adapter.set_topic_recovery_fn(self._recover_telegram_topic_thread_id)
                 adapter._busy_text_mode = self._busy_text_mode
                 
                 # Try to connect
    @@ -5864,6 +5865,7 @@ async def _platform_reconnect_watcher(self) -> None:
                         adapter.set_fatal_error_handler(self._handle_adapter_fatal_error)
                         adapter.set_session_store(self.session_store)
                         adapter.set_busy_session_handler(self._handle_active_session_busy_message)
    +                    adapter.set_topic_recovery_fn(self._recover_telegram_topic_thread_id)
                         adapter._busy_text_mode = self._busy_text_mode
     
                         success = await self._connect_adapter_with_timeout(adapter, platform)

  • tests/gateway/test_telegram_text_batching.py

    @@ -6,12 +6,14 @@
     """
     
     import asyncio
    +from types import SimpleNamespace
     from unittest.mock import AsyncMock, MagicMock, patch
     
     import pytest
     
     from gateway.config import Platform, PlatformConfig
     from gateway.platforms.base import MessageEvent, MessageType, SessionSource
    +from gateway.session import build_session_key
     
     
     def _make_adapter():
    @@ -119,3 +121,46 @@ async def test_batch_cleans_up_after_flush(self):
     
             assert len(adapter._pending_text_batches) == 0
             assert len(adapter._pending_text_batch_tasks) == 0
    +
    +    @pytest.mark.asyncio
    +    async def test_dm_topic_batching_recovers_thread_before_keying(self):
    +        """DM-topic text batches should use the recovered topic lane."""
    +        adapter = _make_adapter()
    +        adapter.set_topic_recovery_fn(
    +            lambda source: "222" if str(source.thread_id or "") == "1" else None
    +        )
    +        event = MessageEvent(
    +            text="hello from DM topic",
    +            message_type=MessageType.TEXT,
    +            source=SessionSource(
    +                platform=Platform.TELEGRAM,
    +                chat_id="12345",
    +                chat_type="dm",
    +                user_id="user-1",
    +                thread_id="1",
    +            ),
    +        )
    +
    +        adapter._enqueue_text_event(event)
    +
    +        def _key(thread_id: str) -> str:
    +            return build_session_key(
    +                SimpleNamespace(
    +                    platform=Platform.TELEGRAM,
    +                    chat_id="12345",
    +                    chat_type="dm",
    +                    thread_id=thread_id,
    +                ),
    +                group_sessions_per_user=True,
    +                thread_sessions_per_user=False,
    +            )
    +
    +        assert _key("222") in adapter._pending_text_batches
    +        assert _key("1") not in adapter._pending_text_batches
    +        assert event.source.thread_id == "222"
    +
    +        await asyncio.sleep(0.2)
    +
    +        adapter.handle_message.assert_called_once()
    +        dispatched = adapter.handle_message.call_args[0][0]
    +        assert dispatched.source.thread_id == "222"