feat: implement session runner for managing message processing and debounce logic

Signed-off-by: Frost Ming <me@frostming.com>

Author: frostming

src/bub/channels/base.py

@@ -6,8 +6,6 @@
 from collections.abc import Awaitable, Callable
 from typing import TYPE_CHECKING, Any
 
-from loguru import logger
-
 from bub.app.runtime import AppRuntime
 
 if TYPE_CHECKING:
@@ -30,26 +28,25 @@ def __init__(self, runtime: AppRuntime) -> None:
     async def start(self, on_receive: Callable[[T], Awaitable[None]]) -> None:
         """Start the channel and set up the receive callback."""
 
+    @property
+    def output_channel(self) -> str:
+        """The name of the channel to send outputs to. Defaults to the same channel."""
+        return self.name
+
+    @abstractmethod
+    def is_mentioned(self, message: T) -> bool:
+        """Determine if the message is relevant to this channel."""
+
     @abstractmethod
-    async def get_session_prompt(self, message: T) -> tuple[str, str] | None:
-        """Get the session id and prompt text for the given message.
-        If None is returned, the message will be ignored.
-        """
+    async def get_session_prompt(self, message: T) -> tuple[str, str]:
+        """Get the session id and prompt text for the given message."""
         pass
 
+    async def run_prompt(self, session_id: str, prompt: str) -> LoopResult:
+        """Run the given prompt through the runtime and return the result."""
+        return await self.runtime.handle_input(session_id, prompt)
+
     @abstractmethod
     async def process_output(self, session_id: str, output: LoopResult) -> None:
         """Process the output returned by the LLM."""
         pass
-
-    async def run_prompt(self, message: T) -> None:
-        """Run a prompt based on the received message."""
-        try:
-            result = await self.get_session_prompt(message)
-            if result is None:
-                return
-            session_id, prompt = result
-            output = await self.runtime.handle_input(session_id, prompt)
-            await self.process_output(session_id, output)
-        except Exception:
-            logger.exception("{}.agent.error", self.name)

src/bub/channels/discord.py

@@ -93,7 +93,7 @@ async def on_message(message: discord.Message) -> None:
             self._bot = None
             logger.info("discord.stopped")
 
-    async def get_session_prompt(self, message: discord.Message) -> tuple[str, str] | None:
+    async def get_session_prompt(self, message: discord.Message) -> tuple[str, str]:
         channel_id = str(message.channel.id)
         session_id = f"{self.name}:{channel_id}"
         content, media = self._parse_message(message)
@@ -125,14 +125,10 @@ async def get_session_prompt(self, message: discord.Message) -> tuple[str, str]
             metadata["reply_to_message"] = reply_meta
 
         metadata_json = json.dumps(
-            {"channel": "discord", "channel_id": channel_id, **exclude_none(metadata)}, ensure_ascii=False
-        )
-        prompt = (
-            "IMPORTANT: Please reply to this $discord message unless otherwise instructed.\n\n"
-            f"{content}\n———————\n{metadata_json}"
+            {"message": content, "channel_id": channel_id, **exclude_none(metadata)}, ensure_ascii=False
         )
         self._latest_message_by_session[session_id] = message
-        return session_id, prompt
+        return session_id, metadata_json
 
     async def process_output(self, session_id: str, output: LoopResult) -> None:
         parts = [part for part in (output.immediate_output, output.assistant_output) if part]
@@ -163,8 +159,6 @@ async def process_output(self, session_id: str, output: LoopResult) -> None:
     async def _on_message(self, message: discord.Message) -> None:
         if message.author.bot:
             return
-        if not self._allow_message(message):
-            return
         if self._on_receive is None:
             logger.warning("discord.inbound no handler for received messages")
             return
@@ -194,7 +188,7 @@ async def _resolve_channel(self, session_id: str) -> discord.abc.Messageable | N
                 return fetched
         return None
 
-    def _allow_message(self, message: discord.Message) -> bool:
+    def is_mentioned(self, message: discord.Message) -> bool:
         channel_id = str(message.channel.id)
         if self._config.allow_channels and channel_id not in self._config.allow_channels:
             return False

src/bub/channels/manager.py

@@ -4,11 +4,13 @@
 
 import asyncio
 import contextlib
+import functools
 
 from loguru import logger
 
 from bub.app.runtime import AppRuntime
 from bub.channels.base import BaseChannel
+from bub.channels.runner import SessionRunner
 
 
 class ChannelManager:
@@ -18,6 +20,7 @@ def __init__(self, runtime: AppRuntime) -> None:
         self.runtime = runtime
         self._channels: dict[str, BaseChannel] = {}
         self._channel_tasks: list[asyncio.Task[None]] = []
+        self._session_runners: dict[str, SessionRunner] = {}
         for channel_cls in self.default_channels():
             self.register(channel_cls)
         runtime.install_hooks(self)
@@ -33,9 +36,7 @@ def channels(self) -> dict[str, BaseChannel]:
     async def run(self) -> None:
         logger.info("channel.manager.start channels={}", self.enabled_channels())
         for channel in self._channels.values():
-            # XXX: Currently we just call the same message handler with itself.
-            # But it will be likely decoupled later
-            task = asyncio.create_task(channel.start(channel.run_prompt))
+            task = asyncio.create_task(channel.start(functools.partial(self._process_input, channel)))
             self._channel_tasks.append(task)
         try:
             await asyncio.gather(*self._channel_tasks)
@@ -63,3 +64,14 @@ def default_channels(self) -> list[type[BaseChannel]]:
 
             result.append(DiscordChannel)
         return result
+
+    async def _process_input[T](self, channel: BaseChannel[T], message: T) -> None:
+        session_id, _ = await channel.get_session_prompt(message)
+        if session_id not in self._session_runners:
+            self._session_runners[session_id] = SessionRunner(
+                channel,
+                session_id,
+                self.runtime.settings.message_debounce_seconds,
+                self.runtime.settings.message_delay_seconds,
+            )
+        await self._session_runners[session_id].process_message(message)

src/bub/channels/runner.py

@@ -0,0 +1,60 @@
+import asyncio
+from typing import Any
+
+from loguru import logger
+
+from bub.channels.base import BaseChannel
+
+
+class SessionRunner:
+    def __init__(
+        self, channel: BaseChannel, session_id: str, debounce_seconds: int, message_delay_seconds: int
+    ) -> None:
+        self.session_id = session_id
+        self.channel = channel
+        self.debounce_seconds = debounce_seconds
+        self.message_delay_seconds = message_delay_seconds
+        self._prompts: list[str] = []
+        self._event = asyncio.Event()
+        self._timer: asyncio.TimerHandle | None = None
+        self._last_received_at: float | None = None
+        self._running_task: asyncio.Task[None] | None = None
+        self._loop = asyncio.get_running_loop()
+
+    async def _run(self) -> None:
+        await self._event.wait()
+        prompt = f"channel: ${self.channel.output_channel}\n" + "\n".join(self._prompts)
+        self._prompts.clear()
+        self._last_received_at = None
+        self._running_task = None
+        try:
+            result = await self.channel.run_prompt(self.session_id, prompt)
+            await self.channel.process_output(self.session_id, result)
+        except Exception:
+            logger.exception("session.run.error session_id={}", self.session_id)
+
+    def reset_timer(self, timeout: int) -> None:
+        self._event.clear()
+        if self._timer:
+            self._timer.cancel()
+        self._timer = self._loop.call_later(timeout, self._event.set)
+
+    async def process_message(self, message: Any) -> None:
+        is_mentioned = self.channel.is_mentioned(message)
+        _, prompt = await self.channel.get_session_prompt(message)
+        now = self._loop.time()
+        if not is_mentioned and self._last_received_at is None:
+            return
+        self._prompts.append(prompt)
+        if is_mentioned:
+            # wait at most 1 second to reply to mentioned messages.
+            self._last_received_at = now
+            self.reset_timer(self.debounce_seconds)
+            if self._running_task is None:
+                self._running_task = asyncio.create_task(self._run())
+            return await self._running_task
+        elif self._last_received_at is not None and self._running_task is None:
+            # Otherwise if bot is mentioned before, we will keep reading messages for at most 30 seconds.
+            self.reset_timer(self.message_delay_seconds)
+            self._running_task = asyncio.create_task(self._run())
+            return await self._running_task

src/bub/channels/telegram.py

@@ -5,7 +5,7 @@
 import asyncio
 import contextlib
 import json
-from collections.abc import Awaitable, Callable
+from collections.abc import AsyncGenerator, Awaitable, Callable
 from dataclasses import dataclass
 from typing import Any, ClassVar
 
@@ -124,9 +124,11 @@ def __init__(self, runtime: AppRuntime) -> None:
             proxy=settings.telegram_proxy,
         )
         self._app: Application | None = None
-        self._typing_tasks: dict[str, asyncio.Task[None]] = {}
         self._on_receive: Callable[[Message], Awaitable[None]] | None = None
 
+    def is_mentioned(self, message: Message) -> bool:
+        return bool(MESSAGE_FILTER.filter(message))
+
     async def start(self, on_receive: Callable[[Message], Awaitable[None]]) -> None:
         self._on_receive = on_receive
         proxy, _ = resolve_proxy(self._config.proxy)
@@ -153,11 +155,6 @@ async def start(self, on_receive: Callable[[Message], Awaitable[None]]) -> None:
         try:
             await asyncio.Event().wait()  # Keep running until stopped
         finally:
-            for task in self._typing_tasks.values():
-                task.cancel()
-            with contextlib.suppress(asyncio.CancelledError):
-                await asyncio.gather(*self._typing_tasks.values())
-            self._typing_tasks.clear()
             updater = self._app.updater
             with contextlib.suppress(Exception):
                 if updater is not None and updater.running:
@@ -167,9 +164,7 @@ async def start(self, on_receive: Callable[[Message], Awaitable[None]]) -> None:
             self._app = None
             logger.info("telegram.stopped")
 
-    async def get_session_prompt(self, message: Message) -> tuple[str, str] | None:
-        if MESSAGE_FILTER.filter(message) is False:
-            return None
+    async def get_session_prompt(self, message: Message) -> tuple[str, str]:
         chat_id = str(message.chat_id)
         session_id = f"{self.name}:{chat_id}"
         content, media = self._parse_message(message)
@@ -208,9 +203,8 @@ async def get_session_prompt(self, message: Message) -> tuple[str, str] | None:
         if reply_meta:
             metadata["reply_to_message"] = reply_meta
 
-        metadata_json = json.dumps({"channel": f"${self.name}", "chat_id": chat_id, **metadata}, ensure_ascii=False)
-        prompt = f"{content}\n———————\n{metadata_json}"
-        return session_id, prompt
+        metadata_json = json.dumps({"message": content, "chat_id": chat_id, **metadata}, ensure_ascii=False)
+        return session_id, metadata_json
 
     async def process_output(self, session_id: str, output: LoopResult) -> None:
         parts = [part for part in (output.immediate_output, output.assistant_output) if part]
@@ -259,30 +253,24 @@ async def _on_text(self, update: Update, _context: ContextTypes.DEFAULT_TYPE) ->
         if self._on_receive is None:
             logger.warning("telegram.inbound no handler for received messages")
             return
-        await self._start_typing(chat_id)
-        try:
+        async with self._start_typing(chat_id):
             await self._on_receive(update.message)
-        finally:
-            await self._stop_typing(chat_id)
 
-    async def _start_typing(self, chat_id: str) -> None:
-        await self._stop_typing(chat_id)
-        self._typing_tasks[chat_id] = asyncio.create_task(self._typing_loop(chat_id))
-
-    async def _stop_typing(self, chat_id: str) -> None:
-        task = self._typing_tasks.pop(chat_id, None)
-        if task is not None:
-            task.cancel()
+    @contextlib.asynccontextmanager
+    async def _start_typing(self, chat_id: str) -> AsyncGenerator[None, None]:
+        typing_task = asyncio.create_task(self._typing_loop(chat_id))
+        try:
+            yield
+        finally:
+            typing_task.cancel()
             with contextlib.suppress(asyncio.CancelledError):
-                await task
+                await typing_task
 
     async def _typing_loop(self, chat_id: str) -> None:
         try:
             while self._app is not None:
                 await self._app.bot.send_chat_action(chat_id=int(chat_id), action="typing")
                 await asyncio.sleep(4)
-        except asyncio.CancelledError:
-            return
         except Exception:
             logger.exception("telegram.typing_loop.error chat_id={}", chat_id)
             return

src/bub/config/settings.py

@@ -38,6 +38,8 @@ class Settings(BaseSettings):
     max_steps: int = Field(default=20, ge=1)
 
     proactive_response: bool = False
+    message_delay_seconds: int = 10
+    message_debounce_seconds: int = 1
 
     telegram_enabled: bool = False
     telegram_token: str | None = None

tests/test_channels.py

@@ -40,6 +40,10 @@ async def get_session_prompt(self, message: object) -> tuple[str, str]:
         _ = message
         return "session", "prompt"
 
+    def is_mentioned(self, message: object) -> bool:
+        _ = message
+        return True
+
     async def process_output(self, session_id: str, output) -> None:
         _ = (session_id, output)
 

tests/test_discord_filter.py

@@ -43,31 +43,31 @@ def _build_channel() -> DiscordChannel:
 def test_allow_message_when_content_contains_bub() -> None:
     channel = _build_channel()
     message = DummyMessage(content="please ask Bub to check this", channel=SimpleNamespace(id=100, name="general"))
-    assert channel._allow_message(message) is True  # type: ignore[arg-type]
+    assert channel.is_mentioned(message) is True  # type: ignore[arg-type]
 
 
 def test_allow_message_when_thread_name_starts_with_bub() -> None:
     channel = _build_channel()
     thread = SimpleNamespace(id=101, name="bub-help", parent=SimpleNamespace(name="forum"))
     message = DummyMessage(content="hello", channel=thread)
-    assert channel._allow_message(message) is True  # type: ignore[arg-type]
+    assert channel.is_mentioned(message) is True  # type: ignore[arg-type]
 
 
 def test_reject_message_when_only_parent_name_starts_with_bub() -> None:
     channel = _build_channel()
     thread = SimpleNamespace(id=102, name="question-1", parent=SimpleNamespace(name="bub-forum"))
     message = DummyMessage(content="hello", channel=thread)
-    assert channel._allow_message(message) is False  # type: ignore[arg-type]
+    assert channel.is_mentioned(message) is False  # type: ignore[arg-type]
 
 
 def test_reject_unrelated_message_without_bot_context() -> None:
     channel = _build_channel()
     message = DummyMessage(content="hello world", channel=SimpleNamespace(id=103, name="general"))
-    assert channel._allow_message(message) is False  # type: ignore[arg-type]
+    assert channel.is_mentioned(message) is False  # type: ignore[arg-type]
 
 
 def test_reject_empty_content_even_in_bub_thread() -> None:
     channel = _build_channel()
     thread = SimpleNamespace(id=104, name="bub-help", parent=SimpleNamespace(name="forum"))
     message = DummyMessage(content="   ", channel=thread)
-    assert channel._allow_message(message) is False  # type: ignore[arg-type]
+    assert channel.is_mentioned(message) is False  # type: ignore[arg-type]

tests/test_graceful_shutdown.py

@@ -34,6 +34,10 @@ async def get_session_prompt(self, message: object) -> tuple[str, str]:
         _ = message
         return "s", "p"
 
+    def is_mentioned(self, message: object) -> bool:
+        _ = message
+        return True
+
     async def process_output(self, session_id: str, output):
         _ = (session_id, output)