vibe-remote 2.1.6__py3-none-any.whl

modules/agents/claude_agent.py

@@ -0,0 +1,344 @@
+import asyncio
+import logging
+import os
+from typing import Callable, Optional
+
+from claude_code_sdk import TextBlock, ToolUseBlock
+
+from modules.agents.base import AgentRequest, BaseAgent
+from modules.im import MessageContext
+
+logger = logging.getLogger(__name__)
+
+
+class ClaudeAgent(BaseAgent):
+    """Existing Claude Code integration extracted into an agent backend."""
+
+    name = "claude"
+
+    def __init__(self, controller):
+        super().__init__(controller)
+        self.session_handler = controller.session_handler
+        self.session_manager = controller.session_manager
+        self.receiver_tasks = controller.receiver_tasks
+        self.claude_sessions = controller.claude_sessions
+        self.claude_client = controller.claude_client
+        self._last_assistant_text: dict[str, str] = {}
+        self._pending_assistant_message: dict[str, str] = {}
+
+    async def handle_message(self, request: AgentRequest) -> None:
+        context = request.context
+
+        try:
+            client = await self.session_handler.get_or_create_claude_session(
+                context,
+                subagent_name=request.subagent_name,
+                subagent_model=request.subagent_model,
+                subagent_reasoning_effort=request.subagent_reasoning_effort,
+            )
+
+            await client.query(
+                request.message, session_id=request.composite_session_id
+            )
+            logger.info(
+                f"Sent message to Claude for session {request.composite_session_id}"
+            )
+
+            await self._delete_ack(context, request)
+
+            if (
+                request.composite_session_id not in self.receiver_tasks
+                or self.receiver_tasks[request.composite_session_id].done()
+            ):
+                self.receiver_tasks[request.composite_session_id] = asyncio.create_task(
+                    self._receive_messages(
+                        client, request.base_session_id, request.working_path, context
+                    )
+                )
+        except Exception as e:
+            logger.error(f"Error processing Claude message: {e}", exc_info=True)
+            await self.session_handler.handle_session_error(
+                request.composite_session_id, context, e
+            )
+        finally:
+            await self._delete_ack(context, request)
+
+    async def clear_sessions(self, settings_key: str) -> int:
+        """Clear Claude sessions scoped to the provided settings key."""
+        agent_map = self.settings_manager.sessions_store.get_agent_map(
+            settings_key, self.name
+        )
+        session_bases_to_clear = set(agent_map.keys())
+
+        self.settings_manager.clear_agent_sessions(settings_key, self.name)
+
+        sessions_to_clear = []
+        for session_key in list(self.claude_sessions.keys()):
+            base_part = session_key.split(":")[0] if ":" in session_key else session_key
+            if base_part in session_bases_to_clear:
+                sessions_to_clear.append(session_key)
+
+        for session_key in sessions_to_clear:
+            try:
+                client = self.claude_sessions[session_key]
+                if hasattr(client, "close"):
+                    await client.close()
+            except Exception as e:
+                logger.warning(f"Error closing Claude session {session_key}: {e}")
+            finally:
+                self.claude_sessions.pop(session_key, None)
+
+        # Legacy session manager cleanup (best-effort)
+        await self.session_manager.clear_session(settings_key)
+
+        return len(sessions_to_clear) or len(session_bases_to_clear)
+
+    async def handle_stop(self, request: AgentRequest) -> bool:
+        composite_key = request.composite_session_id
+        if composite_key not in self.claude_sessions:
+            return False
+
+        client = self.claude_sessions[composite_key]
+        await self.controller.emit_agent_message(
+            request.context, "notify", "🛑 Interrupting Claude session..."
+        )
+        try:
+            if hasattr(client, "interrupt"):
+                await client.interrupt()
+                return True
+            else:
+                await self.controller.emit_agent_message(
+                    request.context,
+                    "notify",
+                    "⚠️ This Claude session cannot be interrupted; consider /clear.",
+                )
+                return False
+        except Exception as err:
+            logger.error(f"Failed to interrupt Claude session {composite_key}: {err}")
+            await self.controller.emit_agent_message(
+                request.context,
+                "notify",
+                "⚠️ Failed to interrupt Claude session. Please try /clear.",
+            )
+            return False
+
+    async def _receive_messages(
+        self,
+        client,
+        base_session_id: str,
+        working_path: str,
+        context: MessageContext,
+    ):
+        """Receive messages from Claude SDK client."""
+        try:
+            settings_key = self.controller._get_settings_key(context)
+            composite_key = f"{base_session_id}:{working_path}"
+            async for message in client.receive_messages():
+                try:
+                    claude_session_id = self._maybe_capture_session_id(
+                        message, base_session_id, settings_key
+                    )
+                    if claude_session_id:
+                        logger.info(
+                            f"Captured Claude session id {claude_session_id} for {base_session_id}"
+                        )
+
+                    if self.claude_client._is_skip_message(message):
+                        continue
+
+                    message_type = self._detect_message_type(message)
+                    formatter = self.im_client.formatter
+
+                    if message_type == "assistant":
+                        toolcalls = []
+                        text_parts = []
+                        for block in getattr(message, "content", []) or []:
+                            if isinstance(block, ToolUseBlock):
+                                toolcalls.append(
+                                    formatter.format_toolcall(
+                                        block.name,
+                                        block.input,
+                                        get_relative_path=lambda path: self.get_relative_path(
+                                            path, context
+                                        ),
+                                    )
+                                )
+                            elif isinstance(block, TextBlock):
+                                text = block.text.strip() if block.text else ""
+                                if text:
+                                    text_parts.append(text)
+
+                        assistant_text = self._extract_text_blocks(message)
+                        if assistant_text:
+                            self._last_assistant_text[composite_key] = assistant_text
+
+                        pending = self._pending_assistant_message.pop(composite_key, None)
+                        if pending:
+                            await self.controller.emit_agent_message(
+                                context,
+                                "assistant",
+                                pending,
+                                parse_mode="markdown",
+                            )
+
+                        for toolcall in toolcalls:
+                            await self.controller.emit_agent_message(
+                                context,
+                                "toolcall",
+                                toolcall,
+                                parse_mode="markdown",
+                            )
+
+                        if text_parts:
+                            formatted_assistant = formatter.format_assistant_message(
+                                text_parts
+                            )
+                            self._pending_assistant_message[composite_key] = formatted_assistant
+                        continue
+
+                    if message_type == "system":
+                        formatted_message = self.claude_client.format_message(
+                            message,
+                            get_relative_path=lambda path: self.get_relative_path(
+                                path, context
+                            ),
+                        )
+                        if formatted_message and formatted_message.strip():
+                            await self.controller.emit_agent_message(
+                                context,
+                                "system",
+                                formatted_message,
+                                parse_mode="markdown",
+                            )
+                        continue
+
+                    if message_type == "result":
+                        pending = self._pending_assistant_message.pop(composite_key, None)
+                        result_text = getattr(message, "result", None)
+                        used_fallback = False
+                        if not result_text:
+                            fallback = self._last_assistant_text.get(composite_key)
+                            if fallback:
+                                result_text = fallback
+                                used_fallback = True
+
+                        if pending and not used_fallback:
+                            await self.controller.emit_agent_message(
+                                context,
+                                "assistant",
+                                pending,
+                                parse_mode="markdown",
+                            )
+
+                        await self.emit_result_message(
+                            context,
+                            result_text,
+                            subtype=getattr(message, "subtype", "") or "",
+                            duration_ms=getattr(message, "duration_ms", 0),
+                            parse_mode="markdown",
+                        )
+
+                        self._last_assistant_text.pop(composite_key, None)
+                        session = await self.session_manager.get_or_create_session(
+                            context.user_id, context.channel_id
+                        )
+                        if session:
+                            session.session_active[
+                                f"{base_session_id}:{working_path}"
+                            ] = False
+                        continue
+
+                    # Ignore UserMessage/tool results; toolcalls are emitted from ToolUseBlock.
+                    continue
+                except Exception as e:
+                    logger.error(
+                        f"Error processing message from Claude: {e}", exc_info=True
+                    )
+                    continue
+        except Exception as e:
+            composite_key = f"{base_session_id}:{working_path}"
+            logger.error(
+                f"Error in Claude receiver for session {composite_key}: {e}",
+                exc_info=True,
+            )
+            await self.session_handler.handle_session_error(composite_key, context, e)
+
+    async def _delete_ack(self, context: MessageContext, request: AgentRequest):
+        ack_id = request.ack_message_id
+        if ack_id and hasattr(self.im_client, "delete_message"):
+            try:
+                await self.im_client.delete_message(context.channel_id, ack_id)
+            except Exception as err:
+                logger.debug(f"Could not delete ack message: {err}")
+            finally:
+                request.ack_message_id = None
+
+    def get_relative_path(
+        self, abs_path: str, context: Optional[MessageContext] = None
+    ) -> str:
+        """Convert absolute path to relative path from working directory."""
+        try:
+            cwd = self.session_handler.get_working_path(context)
+            abs_path = os.path.abspath(os.path.expanduser(abs_path))
+            rel_path = os.path.relpath(abs_path, cwd)
+            if rel_path.startswith("../.."):
+                return abs_path
+            return rel_path
+        except Exception:
+            return abs_path
+
+    def _get_target_context(self, context: MessageContext) -> MessageContext:
+        """Return context for sending messages (respect Slack thread replies)."""
+        if self.im_client.should_use_thread_for_reply() and context.thread_id:
+            return MessageContext(
+                user_id=context.user_id,
+                channel_id=context.channel_id,
+                thread_id=context.thread_id,
+                message_id=context.message_id,
+                platform_specific=context.platform_specific,
+            )
+        return context
+
+    def _maybe_capture_session_id(
+        self,
+        message,
+        base_session_id: str,
+        settings_key: str,
+    ) -> Optional[str]:
+        """Capture session id from system init messages."""
+        if (
+            hasattr(message, "__class__")
+            and message.__class__.__name__ == "SystemMessage"
+            and getattr(message, "subtype", None) == "init"
+            and getattr(message, "data", None)
+        ):
+            session_id = message.data.get("session_id")
+            if session_id:
+                self.session_handler.capture_session_id(
+                    base_session_id, session_id, settings_key
+                )
+                return session_id
+        return None
+
+    def _extract_text_blocks(self, message) -> str:
+        """Extract text-only content blocks for result fallbacks."""
+        parts = []
+        for block in getattr(message, "content", []) or []:
+            if isinstance(block, TextBlock):
+                text = block.text.strip() if block.text else ""
+                if text:
+                    parts.append(self.claude_client.formatter.escape_special_chars(text))
+        return "\n\n".join(parts).strip()
+
+    def _detect_message_type(self, message) -> Optional[str]:
+        """Infer message type name from Claude SDK class."""
+        if not hasattr(message, "__class__"):
+            return None
+        class_name = message.__class__.__name__
+        mapping = {
+            "SystemMessage": "system",
+            "UserMessage": "user",
+            "AssistantMessage": "assistant",
+            "ResultMessage": "result",
+        }
+        return mapping.get(class_name)