From bdc7fd5c1edce95ac4cb52f946738a20a2896a74 Mon Sep 17 00:00:00 2001 From: shuofengzhang Date: Mon, 9 Mar 2026 07:58:55 +0800 Subject: [PATCH 1/3] fix: tighten base64 data URL redaction in buffered logs --- src/bub/channels/handler.py | 22 ++++++++++++++++++---- tests/test_channels.py | 18 ++++++++++++++++++ 2 files changed, 36 insertions(+), 4 deletions(-) diff --git a/src/bub/channels/handler.py b/src/bub/channels/handler.py index 37b1e7a7..1a94fa43 100644 --- a/src/bub/channels/handler.py +++ b/src/bub/channels/handler.py @@ -40,11 +40,17 @@ async def _process(self) -> None: self._in_processing = None await self._handler(message) + @staticmethod + def prettify(content: str) -> str: + return MEDIA_DATA_URL_RE.sub("[media]", content) + async def __call__(self, message: ChannelMessage) -> None: now = self._loop.time() if message.content.startswith(","): logger.info( - "session.message received command session_id={}, content={}", message.session_id, message.content + "session.message received command session_id={}, content={}", + message.session_id, + self.prettify(message.content), ) await self._handler(message) return @@ -53,19 +59,27 @@ async def __call__(self, message: ChannelMessage) -> None: ): self._last_active_time = None logger.info( - "session.message received ignored session_id={}, content={}", message.session_id, message.content + "session.message received ignored session_id={}, content={}", + message.session_id, + self.prettify(message.content), ) return self._pending_messages.append(message) if message.is_active: self._last_active_time = now logger.info( - "session.message received active session_id={}, content={}", message.session_id, message.content + "session.message received active session_id={}, content={}", + message.session_id, + self.prettify(message.content), ) self._reset_timer(self.debounce_seconds) if self._in_processing is None: self._in_processing = asyncio.create_task(self._process()) elif self._last_active_time is not None and self._in_processing is None: - logger.info("session.receive followup session_id={} message={}", message.session_id, message.content) + logger.info( + "session.receive followup session_id={} message={}", + message.session_id, + self.prettify(message.content), + ) self._reset_timer(self.max_wait_seconds) self._in_processing = asyncio.create_task(self._process()) diff --git a/tests/test_channels.py b/tests/test_channels.py index 9debf1b7..c04a46d8 100644 --- a/tests/test_channels.py +++ b/tests/test_channels.py @@ -88,6 +88,24 @@ async def receive(message: ChannelMessage) -> None: assert handled == [",help"] +def test_buffered_handler_prettify_masks_base64_payload() -> None: + content = 'look data:image/png;base64,abcdef" end' + + assert BufferedMessageHandler.prettify(content) == 'look [media]" end' + + +def test_buffered_handler_prettify_masks_single_quoted_base64_payload() -> None: + content = "look 'data:image/png;base64,abcdef' end" + + assert BufferedMessageHandler.prettify(content) == "look '[media]' end" + + +def test_buffered_handler_prettify_preserves_following_text_for_unquoted_payload() -> None: + content = "look DATA:image/png;base64,abcdef tail" + + assert BufferedMessageHandler.prettify(content) == "look [media] tail" + + @pytest.mark.asyncio async def test_channel_manager_dispatch_uses_output_channel_and_preserves_metadata() -> None: cli_channel = FakeChannel("cli") From 76e8341bb49d463e3fd4f16a284aa2afabbd4ebc Mon Sep 17 00:00:00 2001 From: shuofengzhang Date: Mon, 16 Mar 2026 02:41:22 +0000 Subject: [PATCH 2/3] Fix direct tool registration to store wrapped handler --- src/bub/tools.py | 5 +++-- tests/test_tools.py | 14 ++++++++++++++ 2 files changed, 17 insertions(+), 2 deletions(-) diff --git a/src/bub/tools.py b/src/bub/tools.py index d723ec63..a57c9663 100644 --- a/src/bub/tools.py +++ b/src/bub/tools.py @@ -120,8 +120,9 @@ def tool( context=context, ) if isinstance(result, Tool): - REGISTRY[result.name] = result - return _add_logging(result) + tool_instance = _add_logging(result) + REGISTRY[tool_instance.name] = tool_instance + return tool_instance def decorator(func: Callable) -> Tool: tool_instance = _add_logging(result(func)) diff --git a/tests/test_tools.py b/tests/test_tools.py index 0c113f61..4d97e1e2 100644 --- a/tests/test_tools.py +++ b/tests/test_tools.py @@ -74,6 +74,20 @@ def failing_tool() -> str: assert errors[0].startswith("tool.call.error name=tests.failing_tool elapsed_time=") +@pytest.mark.asyncio +async def test_tool_direct_call_registers_wrapped_instance_in_registry() -> None: + tool_name = "tests.direct_call" + REGISTRY.pop(tool_name, None) + + def direct_call(value: str) -> str: + return value.upper() + + direct_tool = tool(direct_call, name=tool_name) + + assert REGISTRY[tool_name] is direct_tool + assert await REGISTRY[tool_name].run("hello") == "HELLO" + + def test_model_tools_rewrites_dotted_names_without_mutating_original() -> None: tool_name = "tests.rename_me" REGISTRY.pop(tool_name, None) From af9df0c52b4ddbdd09e7971ab6e7c5021b3d9486 Mon Sep 17 00:00:00 2001 From: Chojan Shang Date: Thu, 19 Mar 2026 01:37:58 +0800 Subject: [PATCH 3/3] fix: revert unrelated buffered handler changes --- src/bub/channels/handler.py | 25 ++++--------------------- tests/test_channels.py | 18 ------------------ 2 files changed, 4 insertions(+), 39 deletions(-) diff --git a/src/bub/channels/handler.py b/src/bub/channels/handler.py index 1a94fa43..24e664da 100644 --- a/src/bub/channels/handler.py +++ b/src/bub/channels/handler.py @@ -1,13 +1,10 @@ import asyncio -import re from loguru import logger from bub.channels.message import ChannelMessage from bub.types import MessageHandler -MEDIA_DATA_URL_RE = re.compile(r"data:[^;\s]+;base64,[^\"'\s]+", re.IGNORECASE) - class BufferedMessageHandler: """A message handler that buffers incoming messages and processes them in batch with debounce and active time window.""" @@ -40,17 +37,11 @@ async def _process(self) -> None: self._in_processing = None await self._handler(message) - @staticmethod - def prettify(content: str) -> str: - return MEDIA_DATA_URL_RE.sub("[media]", content) - async def __call__(self, message: ChannelMessage) -> None: now = self._loop.time() if message.content.startswith(","): logger.info( - "session.message received command session_id={}, content={}", - message.session_id, - self.prettify(message.content), + "session.message received command session_id={}, content={}", message.session_id, message.content ) await self._handler(message) return @@ -59,27 +50,19 @@ async def __call__(self, message: ChannelMessage) -> None: ): self._last_active_time = None logger.info( - "session.message received ignored session_id={}, content={}", - message.session_id, - self.prettify(message.content), + "session.message received ignored session_id={}, content={}", message.session_id, message.content ) return self._pending_messages.append(message) if message.is_active: self._last_active_time = now logger.info( - "session.message received active session_id={}, content={}", - message.session_id, - self.prettify(message.content), + "session.message received active session_id={}, content={}", message.session_id, message.content ) self._reset_timer(self.debounce_seconds) if self._in_processing is None: self._in_processing = asyncio.create_task(self._process()) elif self._last_active_time is not None and self._in_processing is None: - logger.info( - "session.receive followup session_id={} message={}", - message.session_id, - self.prettify(message.content), - ) + logger.info("session.receive followup session_id={} message={}", message.session_id, message.content) self._reset_timer(self.max_wait_seconds) self._in_processing = asyncio.create_task(self._process()) diff --git a/tests/test_channels.py b/tests/test_channels.py index c04a46d8..9debf1b7 100644 --- a/tests/test_channels.py +++ b/tests/test_channels.py @@ -88,24 +88,6 @@ async def receive(message: ChannelMessage) -> None: assert handled == [",help"] -def test_buffered_handler_prettify_masks_base64_payload() -> None: - content = 'look data:image/png;base64,abcdef" end' - - assert BufferedMessageHandler.prettify(content) == 'look [media]" end' - - -def test_buffered_handler_prettify_masks_single_quoted_base64_payload() -> None: - content = "look 'data:image/png;base64,abcdef' end" - - assert BufferedMessageHandler.prettify(content) == "look '[media]' end" - - -def test_buffered_handler_prettify_preserves_following_text_for_unquoted_payload() -> None: - content = "look DATA:image/png;base64,abcdef tail" - - assert BufferedMessageHandler.prettify(content) == "look [media] tail" - - @pytest.mark.asyncio async def test_channel_manager_dispatch_uses_output_channel_and_preserves_metadata() -> None: cli_channel = FakeChannel("cli")