diff --git a/messaging/handler.py b/messaging/handler.py index 86a9e37..dc49a79 100644 --- a/messaging/handler.py +++ b/messaging/handler.py @@ -20,6 +20,34 @@ from .event_parser import parse_cli_event logger = logging.getLogger(__name__) +MDV2_SPECIAL_CHARS = set("\\_*[]()~`>#+-=|{}.!") + + +def escape_md_v2(text: str) -> str: + """Escape text for Telegram MarkdownV2.""" + return "".join(f"\\{ch}" if ch in MDV2_SPECIAL_CHARS else ch for ch in text) + + +def escape_md_v2_code(text: str) -> str: + """Escape text for Telegram MarkdownV2 code spans/blocks.""" + return text.replace("\\", "\\\\").replace("`", "\\`") + + +def mdv2_bold(text: str) -> str: + return f"*{escape_md_v2(text)}*" + + +def mdv2_code_inline(text: str) -> str: + return f"`{escape_md_v2_code(text)}`" + + +def format_status(emoji: str, label: str, suffix: Optional[str] = None) -> str: + base = f"{emoji} {mdv2_bold(label)}" + if suffix: + return f"{base} {escape_md_v2(suffix)}" + return base + + class ClaudeMessageHandler: """ Platform-agnostic handler for Claude interactions. @@ -137,8 +165,8 @@ class ClaudeMessageHandler: await self.platform.queue_edit_message( incoming.chat_id, status_msg_id, - f"šŸ“‹ **Queued** (position {queue_size}) - waiting...", - parse_mode="markdown", + format_status("šŸ“‹", "Queued", f"(position {queue_size}) - waiting..."), + parse_mode="MarkdownV2", ) async def _process_node( @@ -191,7 +219,7 @@ class ClaudeMessageHandler: if display and display != last_displayed_text: last_displayed_text = display await self.platform.queue_edit_message( - chat_id, status_msg_id, display, parse_mode="markdown" + chat_id, status_msg_id, display, parse_mode="MarkdownV2" ) try: @@ -210,7 +238,7 @@ class ClaudeMessageHandler: captured_session_id = session_or_temp_id except RuntimeError as e: components["errors"].append(str(e)) - await update_ui("ā³ **Session limit reached**", force=True) + await update_ui(format_status("ā³", "Session limit reached"), force=True) if tree: await tree.update_state( node_id, MessageState.ERROR, error_message=str(e) @@ -249,28 +277,28 @@ class ClaudeMessageHandler: for parsed in parsed_list: if parsed["type"] == "thinking": components["thinking"].append(parsed["text"]) - await update_ui("🧠 **Claude is thinking...**") + await update_ui(format_status("🧠", "Claude is thinking...")) elif parsed["type"] == "content": if parsed.get("text"): components["content"].append(parsed["text"]) - await update_ui("🧠 **Claude is working...**") + await update_ui(format_status("🧠", "Claude is working...")) elif parsed["type"] == "tool_start": names = [t.get("name") for t in parsed.get("tools", [])] components["tools"].extend(names) - await update_ui("ā³ **Executing tools...**") + await update_ui(format_status("ā³", "Executing tools...")) elif parsed["type"] == "subagent_start": tasks = parsed.get("tasks", []) components["subagents"].extend(tasks) - await update_ui("šŸ¤– **Subagent working...**") + await update_ui(format_status("šŸ¤–", "Subagent working...")) elif parsed["type"] == "complete": if not any(components.values()): components["content"].append("Done.") logger.info("HANDLER: Task complete, updating UI") - await update_ui("āœ… **Complete**", force=True) + await update_ui(format_status("āœ…", "Complete"), force=True) # Update node state and session if tree and captured_session_id: @@ -288,7 +316,7 @@ class ClaudeMessageHandler: ) components["errors"].append(error_msg) logger.info("HANDLER: Updating UI with error status") - await update_ui("āŒ **Error**", force=True) + await update_ui(format_status("āŒ", "Error"), force=True) if tree: await self._propagate_error_to_children( node_id, error_msg, "Parent task failed" @@ -297,7 +325,7 @@ class ClaudeMessageHandler: except asyncio.CancelledError: logger.warning(f"HANDLER: Task cancelled for node {node_id}") components["errors"].append("Task was cancelled") - await update_ui("āŒ **Cancelled**", force=True) + await update_ui(format_status("āŒ", "Cancelled"), force=True) if tree: await self._propagate_error_to_children( node_id, "Cancelled by user", "Parent task was stopped" @@ -308,7 +336,7 @@ class ClaudeMessageHandler: ) error_msg = str(e)[:200] components["errors"].append(error_msg) - await update_ui("šŸ’„ **Task Failed**", force=True) + await update_ui(format_status("šŸ’„", "Task Failed"), force=True) if tree: await self._propagate_error_to_children( node_id, error_msg, "Parent task failed" @@ -334,8 +362,8 @@ class ClaudeMessageHandler: self.platform.queue_edit_message( child.incoming.chat_id, child.status_message_id, - f"āŒ **Cancelled:** {child_status_text}", - parse_mode="markdown", + format_status("āŒ", "Cancelled:", child_status_text), + parse_mode="MarkdownV2", ) ) @@ -357,7 +385,9 @@ class ClaudeMessageHandler: if len(thinking_text) > 1000: thinking_text = "..." + thinking_text[-995:] - lines.append(f"šŸ’­ **Thinking:**\n```\n{thinking_text}\n```") + lines.append( + f"šŸ’­ {mdv2_bold('Thinking:')}\n```\n{escape_md_v2_code(thinking_text)}\n```" + ) # 2. Tools if components["tools"]: @@ -368,24 +398,26 @@ class ClaudeMessageHandler: unique_tools.append(str(t)) seen.add(t) if unique_tools: - lines.append(f"šŸ›  **Tools:** `{', '.join(unique_tools)}`") + lines.append( + f"šŸ›  {mdv2_bold('Tools:')} {mdv2_code_inline(', '.join(unique_tools))}" + ) # 3. Subagents if components["subagents"]: for task in components["subagents"]: - lines.append(f"šŸ¤– **Subagent:** `{task}`") + lines.append(f"šŸ¤– {mdv2_bold('Subagent:')} {mdv2_code_inline(task)}") # 4. Content if components["content"]: - lines.append("".join(components["content"])) + lines.append(escape_md_v2("".join(components["content"]))) # 5. Errors if components["errors"]: for err in components["errors"]: - lines.append(f"āš ļø **Error:** `{err}`") + lines.append(f"āš ļø {mdv2_bold('Error:')} {mdv2_code_inline(err)}") if not any(lines) and not status: - return "ā³ **Claude is working...**" + return format_status("ā³", "Claude is working...") # Telegram character limit is 4096. We leave buffer for status updates. LIMIT = 3900 @@ -400,7 +432,7 @@ class ClaudeMessageHandler: return ( main_text + status_text if main_text + status_text - else "ā³ **Claude is working...**" + else format_status("ā³", "Claude is working...") ) # If too long, truncate the start of the content (keep the end) @@ -408,7 +440,7 @@ class ClaudeMessageHandler: raw_truncated = main_text[-available_limit:].lstrip() # Check for unbalanced code blocks - prefix = "... (truncated)\n" + prefix = escape_md_v2("... (truncated)\n") if raw_truncated.count("```") % 2 != 0: prefix += "```\n" @@ -426,14 +458,20 @@ class ClaudeMessageHandler: # Reply to existing tree if self.tree_queue.is_node_tree_busy(parent_node_id): queue_size = self.tree_queue.get_queue_size(parent_node_id) + 1 - return f"šŸ“‹ **Queued** (position {queue_size}) - waiting..." - return "šŸ”„ **Continuing conversation...**" + return format_status( + "šŸ“‹", "Queued", f"(position {queue_size}) - waiting..." + ) + return format_status("šŸ”„", "Continuing conversation...") # New conversation stats = self.cli_manager.get_stats() if stats["active_sessions"] >= stats["max_sessions"]: - return f"ā³ **Waiting for slot...** ({stats['active_sessions']}/{stats['max_sessions']})" - return "ā³ **Launching new Claude CLI instance...**" + return format_status( + "ā³", + "Waiting for slot...", + f"({stats['active_sessions']}/{stats['max_sessions']})", + ) + return format_status("ā³", "Launching new Claude CLI instance...") async def stop_all_tasks(self) -> int: """ @@ -459,8 +497,8 @@ class ClaudeMessageHandler: self.platform.queue_edit_message( node.incoming.chat_id, node.status_message_id, - "ā¹ **Stopped.**", - parse_mode="markdown", + format_status("ā¹", "Stopped."), + parse_mode="MarkdownV2", ) ) @@ -476,7 +514,7 @@ class ClaudeMessageHandler: count = await self.stop_all_tasks() await self.platform.queue_send_message( incoming.chat_id, - f"ā¹ **Stopped.** Cancelled {count} pending or active requests.", + format_status("ā¹", "Stopped.", f"Cancelled {count} pending or active requests."), ) async def _handle_stats_command(self, incoming: IncomingMessage) -> None: @@ -485,5 +523,12 @@ class ClaudeMessageHandler: tree_count = self.tree_queue.get_tree_count() await self.platform.queue_send_message( incoming.chat_id, - f"šŸ“Š **Stats**\n• Active CLI: {stats['active_sessions']}\n• Max CLI: {stats['max_sessions']}\n• Message Trees: {tree_count}", + "šŸ“Š " + + mdv2_bold("Stats") + + "\n" + + escape_md_v2(f"• Active CLI: {stats['active_sessions']}") + + "\n" + + escape_md_v2(f"• Max CLI: {stats['max_sessions']}") + + "\n" + + escape_md_v2(f"• Message Trees: {tree_count}"), ) diff --git a/messaging/telegram.py b/messaging/telegram.py index 3d5365a..9b95cf2 100644 --- a/messaging/telegram.py +++ b/messaging/telegram.py @@ -19,6 +19,13 @@ from .models import IncomingMessage logger = logging.getLogger(__name__) +# Telegram MarkdownV2 escaping for inline strings. +MDV2_SPECIAL_CHARS = set("\\_*[]()~`>#+-=|{}.!") + + +def escape_md_v2(text: str) -> str: + return "".join(f"\\{ch}" if ch in MDV2_SPECIAL_CHARS else ch for ch in text) + # Optional import - python-telegram-bot may not be installed try: from telegram import Update, Bot @@ -132,8 +139,13 @@ class TelegramPlatform(MessagingPlatform): try: target = self.allowed_user_id if target: + startup_text = ( + f"šŸš€ *{escape_md_v2('Claude Code Proxy is online!')}* " + f"{escape_md_v2('(Bot API)')}" + ) await self.send_message( - target, "šŸš€ **Claude Code Proxy is online!** (Bot API)" + target, + startup_text, ) except Exception as e: logger.warning(f"Could not send startup message: {e}") @@ -201,52 +213,52 @@ class TelegramPlatform(MessagingPlatform): chat_id: str, text: str, reply_to: Optional[str] = None, - parse_mode: Optional[str] = "Markdown", + parse_mode: Optional[str] = "MarkdownV2", ) -> str: """Send a message to a chat.""" if not self._application or not self._application.bot: raise RuntimeError("Telegram application or bot not initialized") - async def _do_send(mode=parse_mode): + async def _do_send(parse_mode=parse_mode): bot = self._application.bot # type: ignore msg = await bot.send_message( chat_id=chat_id, text=text, reply_to_message_id=int(reply_to) if reply_to else None, - parse_mode=mode, + parse_mode=parse_mode, ) return str(msg.message_id) - return await self._with_retry(_do_send) + return await self._with_retry(_do_send, parse_mode=parse_mode) async def edit_message( self, chat_id: str, message_id: str, text: str, - parse_mode: Optional[str] = "Markdown", + parse_mode: Optional[str] = "MarkdownV2", ) -> None: """Edit an existing message.""" if not self._application or not self._application.bot: raise RuntimeError("Telegram application or bot not initialized") - async def _do_edit(mode=parse_mode): + async def _do_edit(parse_mode=parse_mode): bot = self._application.bot # type: ignore await bot.edit_message_text( chat_id=chat_id, message_id=int(message_id), text=text, - parse_mode=mode, + parse_mode=parse_mode, ) - await self._with_retry(_do_edit) + await self._with_retry(_do_edit, parse_mode=parse_mode) async def queue_send_message( self, chat_id: str, text: str, reply_to: Optional[str] = None, - parse_mode: Optional[str] = "Markdown", + parse_mode: Optional[str] = "MarkdownV2", fire_and_forget: bool = True, ) -> Optional[str]: """Enqueue a message to be sent (using limiter).""" @@ -268,7 +280,7 @@ class TelegramPlatform(MessagingPlatform): chat_id: str, message_id: str, text: str, - parse_mode: Optional[str] = "Markdown", + parse_mode: Optional[str] = "MarkdownV2", fire_and_forget: bool = True, ) -> None: """Enqueue a message edit.""" @@ -355,8 +367,9 @@ class TelegramPlatform(MessagingPlatform): try: await self.send_message( chat_id, - f"āŒ **Error:** {str(e)[:200]}", + f"āŒ *{escape_md_v2('Error:')}* {escape_md_v2(str(e)[:200])}", reply_to=incoming.message_id, + parse_mode="MarkdownV2", ) except Exception: pass diff --git a/tests/test_handler.py b/tests/test_handler.py index adb63d7..eb6fa26 100644 --- a/tests/test_handler.py +++ b/tests/test_handler.py @@ -22,7 +22,8 @@ async def test_handle_message_stop_command( handler.stop_all_tasks.assert_called_once() mock_platform.queue_send_message.assert_called_once_with( - incoming.chat_id, "ā¹ **Stopped.** Cancelled 5 pending or active requests." + incoming.chat_id, + "ā¹ *Stopped\\.* Cancelled 5 pending or active requests\\.", ) @@ -97,12 +98,12 @@ async def test_handle_message_queued(handler, mock_platform, incoming_message_fa await handler.handle_message(incoming) - mock_platform.queue_edit_message.assert_called_once_with( - incoming.chat_id, - "status_123", - "šŸ“‹ **Queued** (position 3) - waiting...", - parse_mode="markdown", - ) + mock_platform.queue_edit_message.assert_called_once_with( + incoming.chat_id, + "status_123", + "šŸ“‹ *Queued* \\(position 3\\) \\- waiting\\.\\.\\.", + parse_mode="MarkdownV2", + ) @pytest.mark.asyncio @@ -177,7 +178,7 @@ async def test_process_node_success_flow(handler, mock_cli_manager, mock_platfor # Note: update_ui is debounced, but COMPLETED/ERROR/CANCELLED are forced mock_platform.queue_edit_message.assert_called() last_call = mock_platform.queue_edit_message.call_args_list[-1] - assert "āœ… **Complete**" in last_call[0][2] + assert "āœ… *Complete*" in last_call[0][2] assert "Hello world" in last_call[0][2] @@ -216,5 +217,5 @@ async def test_process_node_error_flow(handler, mock_cli_manager, mock_platform) ) last_call = mock_platform.queue_edit_message.call_args_list[-1] - assert "āŒ **Error**" in last_call[0][2] + assert "āŒ *Error*" in last_call[0][2] assert "CLI crashed" in last_call[0][2] diff --git a/tests/test_handler_format.py b/tests/test_handler_format.py index d395321..a7094a2 100644 --- a/tests/test_handler_format.py +++ b/tests/test_handler_format.py @@ -1,6 +1,6 @@ import pytest from unittest.mock import MagicMock -from messaging.handler import ClaudeMessageHandler +from messaging.handler import ClaudeMessageHandler, escape_md_v2 @pytest.fixture @@ -20,7 +20,7 @@ def test_build_message_structure(handler): "content": ["Here is the file content."], "errors": ["Some error happened"], } - status = "āœ… **Complete**" + status = "āœ… *Complete*" msg = handler._build_message(components, status) @@ -31,23 +31,23 @@ def test_build_message_structure(handler): assert "list_files" in msg assert "read_file" in msg assert "Searching codebase..." in msg - assert "Here is the file content." in msg + assert escape_md_v2("Here is the file content.") in msg assert "Some error happened" in msg - assert "āœ… **Complete**" in msg + assert "āœ… *Complete*" in msg # Check headers - assert "šŸ’­ **Thinking:**" in msg - assert "šŸ›  **Tools:**" in msg - assert "šŸ¤– **Subagent:**" in msg - assert "āš ļø **Error:**" in msg + assert "šŸ’­ *Thinking:*" in msg + assert "šŸ›  *Tools:*" in msg + assert "šŸ¤– *Subagent:*" in msg + assert "āš ļø *Error:*" in msg # Check Order: Thinking -> Tools -> Subagents -> Content -> Errors -> Status p_thinking = msg.find("Thinking process...") - p_tools = msg.find("šŸ›  **Tools:**") - p_subagents = msg.find("šŸ¤– **Subagent:**") - p_content = msg.find("Here is the file content.") - p_errors = msg.find("āš ļø **Error:**") - p_status = msg.find("āœ… **Complete**") + p_tools = msg.find("šŸ›  *Tools:*") + p_subagents = msg.find("šŸ¤– *Subagent:*") + p_content = msg.find(escape_md_v2("Here is the file content.")) + p_errors = msg.find("āš ļø *Error:*") + p_status = msg.find("āœ… *Complete*") assert p_thinking < p_tools, "Thinking should come before Tools" assert p_tools < p_subagents, "Tools should come before Subagents" @@ -67,7 +67,7 @@ def test_build_message_simple(handler): } msg = handler._build_message(components, "Ready") - assert "Simple message." in msg + assert escape_md_v2("Simple message.") in msg assert "Ready" in msg assert "šŸ’­" not in msg assert "šŸ› " not in msg @@ -84,5 +84,5 @@ def test_subagents_formatting(handler): } msg = handler._build_message(components) - assert "šŸ¤– **Subagent:** `Task 1`" in msg - assert "šŸ¤– **Subagent:** `Task 2`" in msg + assert "šŸ¤– *Subagent:* `Task 1`" in msg + assert "šŸ¤– *Subagent:* `Task 2`" in msg diff --git a/tests/test_reliability.py b/tests/test_reliability.py index d4b8d81..5c351e3 100644 --- a/tests/test_reliability.py +++ b/tests/test_reliability.py @@ -2,7 +2,7 @@ import pytest from unittest.mock import AsyncMock, MagicMock, patch from messaging.telegram import TelegramPlatform from telegram.error import NetworkError, RetryAfter, TelegramError -from messaging.handler import ClaudeMessageHandler +from messaging.handler import ClaudeMessageHandler, format_status @pytest.fixture @@ -85,7 +85,7 @@ def test_handler_build_message_hardening(): "errors": [], } msg = handler._build_message(components) - assert msg == "ā³ **Claude is working...**" + assert msg == format_status("ā³", "Claude is working...") # Case 2: Truncation with code block closing long_thinking = "thought " * 200 # ~1400 chars diff --git a/tests/test_robust_formatting.py b/tests/test_robust_formatting.py index 4feb81b..43d1dcb 100644 --- a/tests/test_robust_formatting.py +++ b/tests/test_robust_formatting.py @@ -1,6 +1,6 @@ import pytest from unittest.mock import MagicMock -from messaging.handler import ClaudeMessageHandler +from messaging.handler import ClaudeMessageHandler, escape_md_v2 @pytest.fixture @@ -25,13 +25,13 @@ def test_truncation_closes_code_blocks(handler): "errors": [], } - msg = handler._build_message(components, "āœ… **Complete**") + msg = handler._build_message(components, "āœ… *Complete*") - assert "... (truncated)" in msg + assert escape_md_v2("... (truncated)") in msg # The limit is 3900. Our content + thinking is > 4000. # The backtick count must be even to be a valid block. assert msg.count("```") % 2 == 0 - assert msg.endswith("```") or "āœ… **Complete**" in msg.split("```")[-1] + assert msg.endswith("```") or "āœ… *Complete*" in msg.split("```")[-1] def test_truncation_preserves_status(handler): @@ -47,7 +47,7 @@ def test_truncation_preserves_status(handler): msg = handler._build_message(components, status) assert status in msg - assert "... (truncated)" in msg + assert escape_md_v2("... (truncated)") in msg def test_empty_components_with_status(handler): diff --git a/tests/test_telegram.py b/tests/test_telegram.py index 025e63b..df3b371 100644 --- a/tests/test_telegram.py +++ b/tests/test_telegram.py @@ -49,7 +49,10 @@ async def test_telegram_platform_send_message_success(telegram_platform): assert msg_id == "999" mock_bot.send_message.assert_called_once_with( - chat_id="chat_1", text="hello", reply_to_message_id=None, parse_mode="Markdown" + chat_id="chat_1", + text="hello", + reply_to_message_id=None, + parse_mode="MarkdownV2", ) @@ -62,7 +65,7 @@ async def test_telegram_platform_edit_message_success(telegram_platform): await telegram_platform.edit_message("chat_1", "999", "new text") mock_bot.edit_message_text.assert_called_once_with( - chat_id="chat_1", message_id=999, text="new text", parse_mode="Markdown" + chat_id="chat_1", message_id=999, text="new text", parse_mode="MarkdownV2" )