claudetools/projects/discord-bot/bot/handlers/message_handler.py

"""Discord message handler. Maps each thread to a persistent Claude Agent session."""
from __future__ import annotations

import asyncio
import logging
import shutil
from pathlib import Path

import discord

from bot.claude.client import ClaudeAgentManager
from bot.config import settings

logger = logging.getLogger(__name__)

DISCORD_MAX_LEN = 2000
EDIT_THROTTLE_SECONDS = 0.75
MAX_FILE_BYTES = 25 * 1024 * 1024
MAX_TURN_BYTES = 100 * 1024 * 1024
ATTACHMENT_ROOT = settings.claudetools_root / "projects" / "discord-bot" / ".attachments"


class MessageHandler:
    def __init__(self, bot: discord.Client, agents: ClaudeAgentManager) -> None:
        self.bot = bot
        self.agents = agents

    async def handle_mention(self, message: discord.Message) -> None:
        if message.author == self.bot.user:
            return

        content = message.content
        for mention in message.mentions:
            if mention == self.bot.user:
                content = content.replace(f"<@{mention.id}>", "").replace(
                    f"<@!{mention.id}>", ""
                ).strip()

        if not content and not message.attachments:
            await message.reply("Hey! How can I help?")
            return

        if isinstance(message.channel, discord.Thread):
            thread = message.channel
        else:
            name = self._thread_name(content) if content else "Attachment"
            thread = await message.create_thread(
                name=name,
                auto_archive_duration=1440,
            )

        attachment_paths = await self._download_attachments(message, thread.id)
        if attachment_paths:
            lines = "\n".join(f"- {p}" for p in attachment_paths)
            content = (content + f"\n\nUser attached files:\n{lines}").strip()

        if not content:
            content = "User uploaded file(s) without a message."

        await self._run_turn(thread, content)

    @staticmethod
    async def _download_attachments(
        message: discord.Message, thread_id: int
    ) -> list[Path]:
        if not message.attachments:
            return []

        target_dir = ATTACHMENT_ROOT / str(thread_id)
        target_dir.mkdir(parents=True, exist_ok=True)

        saved: list[Path] = []
        running_total = 0
        for att in message.attachments:
            if att.size > MAX_FILE_BYTES:
                logger.warning(
                    "[WARNING] skipping %s: %d bytes exceeds per-file cap %d",
                    att.filename, att.size, MAX_FILE_BYTES,
                )
                continue
            if running_total + att.size > MAX_TURN_BYTES:
                logger.warning(
                    "[WARNING] skipping %s: would exceed per-turn cap %d",
                    att.filename, MAX_TURN_BYTES,
                )
                continue
            safe_name = Path(att.filename).name or f"attachment_{att.id}"
            target = target_dir / safe_name
            try:
                await att.save(target)
            except discord.HTTPException as e:
                logger.warning("[WARNING] failed to download %s: %s", att.filename, e)
                continue
            saved.append(target)
            running_total += att.size
            logger.info("[INFO] saved attachment %s (%d bytes)", target, att.size)
        return saved

    @staticmethod
    def _cleanup_attachments(thread_id: int) -> None:
        target_dir = ATTACHMENT_ROOT / str(thread_id)
        if target_dir.exists():
            try:
                shutil.rmtree(target_dir)
            except OSError as e:
                logger.warning("[WARNING] attachment cleanup failed for %d: %s", thread_id, e)

    @staticmethod
    def _thread_name(message: str) -> str:
        name = message[:50].replace("\n", " ").strip()
        if len(message) > 50:
            name += "..."
        return name or "ClaudeTools Conversation"

    async def _run_turn(self, thread: discord.Thread, user_message: str) -> None:
        # The Claude Code CLI cold-start can take a few seconds; show feedback.
        status_msg = await thread.send("[INFO] Thinking...")

        agent = await self.agents.get_or_create(thread.id)

        buffer: list[str] = []
        last_edit = 0.0
        lock = asyncio.Lock()

        async def on_text(chunk: str) -> None:
            nonlocal last_edit
            buffer.append(chunk)
            now = asyncio.get_event_loop().time()
            if now - last_edit < EDIT_THROTTLE_SECONDS:
                return
            async with lock:
                last_edit = asyncio.get_event_loop().time()
                preview = "".join(buffer)
                if len(preview) > DISCORD_MAX_LEN:
                    preview = preview[-DISCORD_MAX_LEN:]
                try:
                    await status_msg.edit(content=preview or "[INFO] Thinking...")
                except discord.errors.NotFound:
                    pass

        async def on_tool_use(tool_name: str) -> None:
            # Server-side log only — Discord-side notices clutter the thread
            # and ordering issues push them below the streaming answer.
            logger.info("[INFO] thread=%d tool=%s", thread.id, tool_name)

        try:
            full_text = await agent.send(user_message, on_text, on_tool_use)
        except Exception as e:
            logger.exception("[ERROR] Agent turn failed in thread %d", thread.id)
            await status_msg.edit(content=f"[ERROR] {e}")
            self._cleanup_attachments(thread.id)
            return

        await self._post_final(status_msg, thread, full_text or "[INFO] (no response)")
        self._cleanup_attachments(thread.id)

    async def _post_final(
        self,
        status_msg: discord.Message,
        thread: discord.Thread,
        text: str,
    ) -> None:
        chunks = self._split(text)
        try:
            await status_msg.edit(content=chunks[0])
        except discord.errors.NotFound:
            await thread.send(chunks[0])
        for chunk in chunks[1:]:
            await thread.send(chunk)

    @staticmethod
    def _split(content: str, max_len: int = DISCORD_MAX_LEN) -> list[str]:
        if len(content) <= max_len:
            return [content]

        chunks: list[str] = []
        current = ""
        for line in content.split("\n"):
            # A single line longer than the limit must be hard-split mid-line.
            while len(line) > max_len:
                if current:
                    chunks.append(current.rstrip())
                    current = ""
                chunks.append(line[:max_len])
                line = line[max_len:]
            if len(current) + len(line) + 1 > max_len:
                chunks.append(current.rstrip())
                current = line + "\n"
            else:
                current += line + "\n"
        if current:
            chunks.append(current.rstrip())
        return chunks