vibe-remote 2.1.6__py3-none-any.whl

modules/agents/codex_agent.py

@@ -0,0 +1,368 @@
+import asyncio
+import json
+import logging
+import os
+import signal
+from asyncio.subprocess import Process
+from typing import Dict, Optional, Tuple
+
+from markdown_to_mrkdwn import SlackMarkdownConverter
+
+from modules.agents.base import AgentRequest, BaseAgent
+
+logger = logging.getLogger(__name__)
+
+STREAM_BUFFER_LIMIT = 8 * 1024 * 1024  # 8MB cap for Codex stdout/stderr streams
+
+class CodexAgent(BaseAgent):
+    """Codex CLI integration via codex exec JSON streaming mode."""
+
+    name = "codex"
+
+    def __init__(self, controller, codex_config):
+        super().__init__(controller)
+        self.codex_config = codex_config
+        self.active_processes: Dict[str, Tuple[Process, str]] = {}
+        self.base_process_index: Dict[str, str] = {}
+        self.composite_to_base: Dict[str, str] = {}
+        self._initialized_sessions: set[str] = set()
+        self._pending_assistant_messages: Dict[str, Tuple[str, Optional[str]]] = {}
+        self._slack_markdown_converter = (
+            SlackMarkdownConverter()
+            if getattr(self.controller.config, "platform", None) == "slack"
+            else None
+        )
+
+    async def handle_message(self, request: AgentRequest) -> None:
+        existing = self.base_process_index.get(request.base_session_id)
+        if existing and existing in self.active_processes:
+            await self.controller.emit_agent_message(
+                request.context,
+                "notify",
+                "⚠️ Codex is already processing a task in this thread. "
+                "Cancelling the previous run...",
+            )
+            await self._terminate_process(existing)
+            await self.controller.emit_agent_message(
+                request.context,
+                "notify",
+                "⏹ Previous Codex task cancelled. Starting the new request...",
+            )
+        resume_id = self.settings_manager.get_agent_session_id(
+            request.settings_key,
+            request.base_session_id,
+            agent_name=self.name,
+        )
+
+        if not os.path.exists(request.working_path):
+            os.makedirs(request.working_path, exist_ok=True)
+
+        cmd = self._build_command(request, resume_id)
+        try:
+            process = await asyncio.create_subprocess_exec(
+                *cmd,
+                stdout=asyncio.subprocess.PIPE,
+                stderr=asyncio.subprocess.PIPE,
+                cwd=request.working_path,
+                limit=STREAM_BUFFER_LIMIT,
+                **({"preexec_fn": os.setsid} if hasattr(os, "setsid") else {}),
+            )
+        except FileNotFoundError:
+            await self.controller.emit_agent_message(
+                request.context,
+                "notify",
+                "❌ Codex CLI not found. Please install it or set CODEX_CLI_PATH.",
+            )
+            return
+        except Exception as e:
+            logger.error(f"Failed to launch Codex CLI: {e}", exc_info=True)
+            await self.controller.emit_agent_message(
+                request.context, "notify", f"❌ Failed to start Codex CLI: {e}"
+            )
+            return
+
+        await self._delete_ack(request)
+
+        self.active_processes[request.composite_session_id] = (
+            process,
+            request.settings_key,
+        )
+        self.base_process_index[request.base_session_id] = request.composite_session_id
+        self.composite_to_base[request.composite_session_id] = request.base_session_id
+        logger.info(
+            f"Codex session {request.composite_session_id} started (pid={process.pid})"
+        )
+
+        stdout_task = asyncio.create_task(
+            self._consume_stdout(process, request)
+        )
+        stderr_task = asyncio.create_task(
+            self._consume_stderr(process, request)
+        )
+
+        try:
+            await process.wait()
+            await asyncio.gather(stdout_task, stderr_task)
+        finally:
+            self._unregister_process(request.composite_session_id)
+
+        if process.returncode != 0:
+            await self.controller.emit_agent_message(
+                request.context,
+                "notify",
+                "⚠️ Codex exited with a non-zero status. Review stderr for details.",
+            )
+
+    async def clear_sessions(self, settings_key: str) -> int:
+        self.settings_manager.clear_agent_sessions(settings_key, self.name)
+        # Terminate any active processes scoped to this settings key
+        terminated = 0
+        for key, (_, stored_key) in list(self.active_processes.items()):
+            if stored_key == settings_key:
+                await self._terminate_process(key)
+                terminated += 1
+        return terminated
+
+    async def handle_stop(self, request: AgentRequest) -> bool:
+        key = request.composite_session_id
+        if not await self._terminate_process(key):
+            key = self.base_process_index.get(request.base_session_id)
+            if not key or not await self._terminate_process(key):
+                return False
+        await self.controller.emit_agent_message(
+            request.context, "notify", "🛑 Terminated Codex execution."
+        )
+        logger.info(f"Codex session {key} terminated via /stop")
+        return True
+
+    def _unregister_process(self, composite_key: str):
+        self.active_processes.pop(composite_key, None)
+        self._pending_assistant_messages.pop(composite_key, None)
+        base_id = self.composite_to_base.pop(composite_key, None)
+        if base_id and self.base_process_index.get(base_id) == composite_key:
+            self.base_process_index.pop(base_id, None)
+
+    async def _terminate_process(self, composite_key: str) -> bool:
+        entry = self.active_processes.get(composite_key)
+        if not entry:
+            return False
+
+        proc, _ = entry
+        try:
+            if hasattr(os, "getpgid"):
+                try:
+                    os.killpg(os.getpgid(proc.pid), signal.SIGKILL)
+                except ProcessLookupError:
+                    pass
+            else:
+                proc.kill()
+            await proc.wait()
+        except ProcessLookupError:
+            pass
+
+        self._unregister_process(composite_key)
+        return True
+
+    def _build_command(self, request: AgentRequest, resume_id: Optional[str]) -> list:
+        cmd = [self.codex_config.binary, "exec", "--json"]
+        cmd += ["--dangerously-bypass-approvals-and-sandbox"]
+        cmd += ["--skip-git-repo-check"]
+
+        if self.codex_config.default_model:
+            cmd += ["--model", self.codex_config.default_model]
+
+        cmd += ["--cd", request.working_path]
+        cmd += self.codex_config.extra_args
+
+        if resume_id:
+            cmd += ["resume", resume_id]
+
+        cmd.append(request.message)
+
+        logger.info(f"Executing Codex command: {' '.join(cmd[:-1])} <prompt>")
+        return cmd
+
+    async def _consume_stdout(self, process: Process, request: AgentRequest):
+        assert process.stdout is not None
+        try:
+            while True:
+                try:
+                    line = await process.stdout.readline()
+                except (asyncio.LimitOverrunError, ValueError) as err:
+                    await self._notify_stream_error(
+                        request, f"Codex output too long; stream decode failed: {err}"
+                    )
+                    logger.exception("Codex stdout exceeded buffer limit")
+                    break
+                if not line:
+                    break
+                line = line.decode().strip()
+                if not line:
+                    continue
+                try:
+                    event = json.loads(line)
+                except json.JSONDecodeError:
+                    logger.debug(f"Codex emitted non-JSON line: {line}")
+                    continue
+                await self._handle_event(event, request)
+        except Exception as err:
+            await self._notify_stream_error(
+                request, f"Codex stdout 读取异常：{err}"
+            )
+            logger.exception("Unexpected Codex stdout error")
+
+    async def _consume_stderr(self, process: Process, request: AgentRequest):
+        assert process.stderr is not None
+        buffer = []
+        while True:
+            line = await process.stderr.readline()
+            if not line:
+                break
+            decoded = line.decode(errors="ignore").rstrip()
+            buffer.append(decoded)
+            logger.debug(f"Codex stderr: {decoded}")
+
+        if buffer:
+            joined = "\n".join(buffer[-10:])
+            await self.controller.emit_agent_message(
+                request.context,
+                "notify",
+                f"❗️ Codex stderr:\n```stderr\n{joined}\n```",
+                parse_mode="markdown",
+            )
+
+    async def _handle_event(
+        self, event: Dict, request: AgentRequest
+    ):
+        event_type = event.get("type")
+
+        if event_type == "thread.started":
+            thread_id = event.get("thread_id")
+            if thread_id:
+                self.settings_manager.set_agent_session_mapping(
+                    request.settings_key,
+                    self.name,
+                    request.base_session_id,
+                    thread_id,
+                )
+            session_key = request.composite_session_id
+            if session_key not in self._initialized_sessions:
+                self._initialized_sessions.add(session_key)
+                system_text = self.im_client.formatter.format_system_message(
+                    request.working_path, "init", thread_id
+                )
+                await self.controller.emit_agent_message(
+                    request.context,
+                    "system",
+                    system_text,
+                    parse_mode="markdown",
+                )
+            return
+
+        if event_type == "item.completed":
+            details = event.get("item", {})
+            item_type = details.get("type")
+
+            if item_type == "agent_message":
+                text = details.get("text", "")
+                if text:
+                    session_key = request.composite_session_id
+                    pending = self._pending_assistant_messages.get(session_key)
+                    if pending:
+                        pending_text, pending_parse_mode = pending
+                        await self.controller.emit_agent_message(
+                            request.context,
+                            "assistant",
+                            pending_text,
+                            parse_mode=pending_parse_mode or "markdown",
+                        )
+
+                    self._pending_assistant_messages[session_key] = self._prepare_last_message_payload(text)
+            elif item_type == "command_execution":
+                command = details.get("command")
+                status = details.get("status")
+                if command:
+                    toolcall = self.im_client.formatter.format_toolcall(
+                        "bash",
+                        {"command": command, "status": status},
+                    )
+                    await self.controller.emit_agent_message(
+                        request.context,
+                        "toolcall",
+                        toolcall,
+                        parse_mode="markdown",
+                    )
+            elif item_type == "reasoning":
+                text = details.get("text", "")
+                if text:
+                    await self.controller.emit_agent_message(
+                        request.context,
+                        "assistant",
+                        f"_🧠 {text}_",
+                        parse_mode="markdown",
+                    )
+            return
+
+        if event_type == "error":
+            message = event.get("message", "Unknown error")
+            await self.controller.emit_agent_message(
+                request.context, "notify", f"❌ Codex error: {message}"
+            )
+            return
+
+        if event_type == "turn.failed":
+            error = event.get("error", {}).get("message", "Turn failed.")
+            await self.controller.emit_agent_message(
+                request.context, "notify", f"⚠️ Codex turn failed: {error}"
+            )
+            self._pending_assistant_messages.pop(request.composite_session_id, None)
+            return
+
+        if event_type == "turn.completed":
+            pending = self._pending_assistant_messages.pop(
+                request.composite_session_id, None
+            )
+            if pending:
+                pending_text, pending_parse_mode = pending
+                await self.emit_result_message(
+                    request.context,
+                    pending_text,
+                    subtype="success",
+                    started_at=request.started_at,
+                    parse_mode=pending_parse_mode or "markdown",
+                )
+            else:
+                await self.emit_result_message(
+                    request.context,
+                    None,
+                    subtype="success",
+                    started_at=request.started_at,
+                    parse_mode="markdown",
+                )
+            return
+
+    async def _delete_ack(self, 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(
+                    request.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 _prepare_last_message_payload(
+        self, text: str
+    ) -> Tuple[str, Optional[str]]:
+        """Prepare cached assistant text for reuse in result messages."""
+        return text, "markdown"
+
+    async def _notify_stream_error(self, request: AgentRequest, message: str) -> None:
+        """Emit a notify message when Codex stdout handling fails."""
+        await self.controller.emit_agent_message(
+            request.context,
+            "notify",
+            f"⚠️ {message}\n请查看 `~/.vibe_remote/logs/vibe_remote.log` 获取更多细节。",
+        )