agent.py

"""Wcode coding agent — main orchestrator.

The :class:`WodeApp` class is the central application entry point.
It composes all subsystems (session, memory, skills, hooks, cron, tools,
MCP, background, subagents) and provides a clean ``start()`` → ``run_turn()``
→ ``shutdown()`` lifecycle.

Domain logic lives in:
- config.py      — configuration, WodeConfig dataclass
- utils.py       — parse_frontmatter, tool_def, safe_path, get_final_answer
- skills.py      — SkillRegistry
- memory.py      — MemoryManager
- messaging.py   — SubagentManager
- tools.py       — ToolRegistry + tool functions
- background.py  — BackgroundExecutor
- hooks.py       — HookSystem
- context.py     — context compaction & summarisation
- session.py     — SessionManager
- cron_system.py — CronScheduler
- mcp_manager.py — MCPManager
- task_system.py — task CRUD
"""

import json
import threading
import time

from rich.console import Console
from rich.markdown import Markdown
from rich.text import Text

import litellm
from litellm.exceptions import (
    ContextWindowExceededError,
    RateLimitError,
    ServiceUnavailableError,
)

from background import BackgroundExecutor, should_run_background
from config import WodeConfig, MY_MODEL, MAX_TOKENS, WORKDIR, MCP_SERVERS, get_logger
from wcode_errors import WcodeError
from context import compact_context_by_llm, compact_context_fast
from cron_system import CronScheduler
from hooks import HookSystem, HookType
from memory import MemoryManager
from messaging import SubagentManager
from mcp_manager import MCPManager
from session import SessionManager
from skills import SkillRegistry
from tools import (
    ToolRegistry,
    run_bash,
    run_read_file,
    run_write_file,
    run_edit_file,
    run_glob,
)
from utils import get_final_answer

logger = get_logger(__name__)
console = Console()


# ═══════════════════════════════════════════════════════════════════════════
# WodeApp
# ═══════════════════════════════════════════════════════════════════════════


class WodeApp:
    """Central application class for the Wcode coding agent.

    Composes all subsystems and provides the main entry points:

    - :meth:`start` — initialize everything and enter the REPL
    - :meth:`run_turn` — run a single turn of the agent loop
    - :meth:`shutdown` — gracefully shut down all subsystems

    Usage::

        config = WodeConfig(mcp_servers={"fetch": {...}})
        app = WodeApp(config)
        try:
            app.start()
        finally:
            app.shutdown()
    """

    def __init__(self, config: WodeConfig) -> None:
        self.config = config

        # ── Independent managers (no cross-dependencies) ──
        self.hooks = HookSystem()
        self.skills = SkillRegistry(config.skills_dir)
        self.memory = MemoryManager(config.memory_dir)
        self.session = SessionManager(config.sessions_dir)
        self.mcp = MCPManager()
        self.cron = CronScheduler()

        # ── Tool registry (depends on skills + cron) ──
        self.tools = ToolRegistry(
            skill_registry=self.skills,
            cron_scheduler=self.cron,
            subagent_manager=None,  # wired below after SubagentManager is created
        )

        # ── Subagent manager (depends on tool functions) ──
        self.subagent = SubagentManager(
            tool_functions=[run_bash, run_read_file, run_write_file, run_edit_file, run_glob],
            tool_handlers={
                "run_bash": self.tools.handlers["run_bash"],
                "run_read_file": self.tools.handlers["run_read_file"],
                "run_write_file": self.tools.handlers["run_write_file"],
                "run_edit_file": self.tools.handlers["run_edit_file"],
                "run_glob": self.tools.handlers["run_glob"],
            },
        )
        # Wire subagent into tool registry
        self.tools.handlers["run_spawn_subagent"] = self.subagent.spawn

        # ── Background executor (depends on tool execution) ──
        self.background = BackgroundExecutor(self._execute_tool)

        # ── Core agent state ──
        self.session_context: list[dict] = [{"role": "system", "content": ""}]
        self.rounds_since_todo = 0

    # ── Public API ─────────────────────────────────────────────────────

    def start(self) -> None:
        """Initialize all subsystems and enter the interactive REPL."""
        self._start_mcp()

        self.cron.start()
        threading.Thread(target=self._queue_processor, daemon=True).start()
        logger.info("[cron] queue processor started")

        session_id = self.session.init()
        logger.info(f"[session] active session: {session_id}")

        logger.info("Type a question, press Enter. Type q to quit.\n")
        self._repl()

    def run_turn(self, user_query: str | None = None) -> str:
        """Run one turn of the agent loop.

        Thread-safe — can be called from the queue-processor thread
        (which holds ``cron.agent_lock``) or from the main REPL.

        Args:
            user_query: Optional user message to append before running.

        Returns:
            The final assistant answer text.
        """
        if user_query is not None:
            self.session_context.append({"role": "user", "content": user_query})

        self._agent_loop(self.session_context)

        self.session.save(self.session_context)

        ans = get_final_answer(self.session_context)
        logger.info("FINAL_ANS:")
        logger.info(ans)
        return ans

    def shutdown(self) -> None:
        """Gracefully shut down all subsystems."""
        self.session.close()
        logger.info("[session] session closed")
        try:
            self.mcp.shutdown()
        except Exception:
            logger.debug("[mcp] shutdown error (non-critical)")

    # ── System prompt ──────────────────────────────────────────────────

    def _build_system_prompt(self) -> str:
        """Build the system prompt from memory index and skill list."""
        memories = self.memory.read_index()
        all_skills = self.skills.list_skills()
        prompt = (
            f"You are a coding agent at {WORKDIR}. \n"
            "When the user says 'remember' or expresses a clear preference, extract it as a memory. \n"
        )
        if memories:
            prompt += f"You have access to the following memories: {memories} \n"
        if all_skills:
            prompt += (
                f"You have access to the following skills: {all_skills} . "
                "Use load_skill(name) when a skill is relevant. \n"
            )
        return prompt

    # ── Agent loop ─────────────────────────────────────────────────────

    def _agent_loop(self, messages: list) -> None:
        """Main turn loop: call LLM, execute tools, handle errors & compaction."""
        memories_content = self.memory.load_relevant(messages)
        memory_turn = len(messages) - 1 if messages else None

        if messages[0].get("role") == "system":
            messages[0]["content"] = self._build_system_prompt()
        else:
            logger.error("[system prompt missing]")
            return

        retry_count = 0
        max_retries = 3
        base_delay = 3  # seconds

        while True:
            # ── Inject cron messages ──
            fired = self.cron.consume_queue()
            for job in fired:
                messages.append(
                    {"role": "user", "content": f"[Scheduled] {job.prompt}"}
                )
                logger.info(f"[inject cron] {job.prompt[:50]}")

            compact_context_fast(messages)

            # ── Periodic todo reminder ──
            if self.rounds_since_todo % 5 == 0:
                messages.append(
                    {"role": "user", "content": "<reminder>Update your todos.</reminder>"}
                )
            self.rounds_since_todo += 1

            # ── Inject relevant memories ──
            request_messages = messages
            if memories_content and memory_turn is not None and memory_turn < len(messages):
                request_messages = messages.copy()
                request_messages[memory_turn] = {
                    **messages[memory_turn],
                    "content": memories_content
                    + "\n\n"
                    + messages[memory_turn]["content"],
                }

            # ── LLM call ──
            try:
                response = litellm.completion(
                    messages=request_messages,
                    model=MY_MODEL,
                    tools=self.tools.tools,
                    max_tokens=MAX_TOKENS,
                )
            except ContextWindowExceededError:
                messages[:] = compact_context_by_llm(messages)
                logger.info(f"[ContextWindowExceededError compact] {messages}")
                continue
            except RateLimitError as e:
                retry_count += 1
                if retry_count > max_retries:
                    logger.exception(
                        f"[429 RateLimitError] {e}. "
                        f"Max retries ({max_retries}) exceeded, stopping."
                    )
                    break
                delay = base_delay * (2 ** (retry_count - 1))
                logger.warning(
                    f"[429 RateLimitError] {e}. "
                    f"Waiting {delay}s before retry {retry_count}/{max_retries}..."
                )
                time.sleep(delay)
                continue
            except ServiceUnavailableError as e:
                retry_count += 1
                if retry_count > max_retries:
                    logger.exception(
                        f"[529/503 ServiceUnavailableError] {e}. "
                        f"Max retries ({max_retries}) exceeded, stopping."
                    )
                    break
                delay = base_delay * (2 ** (retry_count - 1))
                logger.warning(
                    f"[529/503 ServiceUnavailableError] {e}. "
                    f"Waiting {delay}s before retry {retry_count}/{max_retries}..."
                )
                time.sleep(delay)
                continue

            retry_count = 0

            # ── Handle token-limit finish ──
            if response.choices[0].finish_reason == "length":
                messages[:] = compact_context_by_llm(messages)
                logger.info(f"[token limit compact] {messages}")
                continue

            msg = response.choices[0].message
            msg_dict = msg.model_dump()
            if not msg_dict.get("content") and msg_dict.get("reasoning_content"):
                msg_dict["content"] = msg_dict["reasoning_content"]
            msg_dict.pop("reasoning_content", None)
            messages.append(msg_dict)

            # ── Stop if no tool calls ──
            if not msg.tool_calls:
                self.hooks.trigger(HookType.Stop, messages)
                # TODO skip for test
                # self.memory.extract(messages)
                # self.memory.consolidate()
                break

            # ── Execute tool calls ──
            for tc in msg.tool_calls:
                name = tc.function.name
                args = json.loads(tc.function.arguments)

                blocked = self.hooks.trigger(HookType.PreToolUse, tc)
                if blocked:
                    tool_msg = {
                        "role": "tool",
                        "tool_call_id": tc.id,
                        "name": name,
                        "content": json.dumps(blocked, ensure_ascii=False),
                    }
                    messages.append(tool_msg)
                    continue

                tool_msg = {}
                tool_result = ""
                if should_run_background(name, args):
                    bg_id = self.background.start_task(tc)
                    tool_result = (
                        f"[Background task {bg_id} started] "
                        f"Command: {args.get('command')}. "
                        "Result will be available when complete."
                        "Don't query task through tools."
                    )
                    tool_msg = {
                        "role": "tool",
                        "tool_call_id": tc.id,
                        "name": name,
                        "content": tool_result,
                    }
                else:
                    tool_result = self.tools.execute(tc)
                    tool_msg = {
                        "role": "tool",
                        "tool_call_id": tc.id,
                        "name": name,
                        "content": json.dumps(tool_result, ensure_ascii=False),
                    }

                self.hooks.trigger(HookType.PostToolUse, tc, tool_result)
                messages.append(tool_msg)

            # ── Inject background & subagent notifications ──
            bg_notifications = self.background.collect_results()
            if bg_notifications:
                messages.append(
                    {"role": "user", "content": "\n".join(bg_notifications)}
                )
                logger.info(
                    f"[inject] {len(bg_notifications)} background notification(s)"
                )

            sub_notifications = self.subagent.collect_results()
            if sub_notifications:
                messages.append(
                    {"role": "user", "content": "\n".join(sub_notifications)}
                )
                logger.info(
                    f"[inject] {len(sub_notifications)} subagent notification(s)"
                )

    # ── Tool execution bridge (for BackgroundExecutor) ──────────────────

    def _execute_tool(self, name: str, args: dict) -> str:
        """Execute a tool by name with the given arguments.

        Used by :class:`BackgroundExecutor` to dispatch tool calls from
        background threads.  Catches :class:`WcodeError` and converts to
        a user-facing error string.
        """
        handler = self.tools.handlers.get(name)
        if handler is None:
            return f"Error: Unknown tool '{name}'"
        try:
            return handler(**args)
        except WcodeError as e:
            logger.warning(f"[tool] {name} error: {e}")
            return f"Error: {e}"

    # ── MCP startup ────────────────────────────────────────────────────

    def _start_mcp(self) -> None:
        """Start configured MCP servers and register their tools."""
        mcp_servers = self.config.mcp_servers
        if not mcp_servers:
            return

        logger.info("[mcp] starting %d server(s)…", len(mcp_servers))
        try:
            started = self.mcp.start_all(mcp_servers)
            if started:
                self.tools.ensure_mcp_tools(self.mcp)
                mcp_count = sum(
                    1
                    for t in self.tools.tools
                    if t.get("function", {}).get("name", "").startswith("mcp__")
                )
                logger.info(
                    "[mcp] %d server(s) ready, %d MCP tool(s) registered",
                    len(started),
                    mcp_count,
                )
        except Exception:
            logger.exception("[mcp] init failed — continuing without MCP tools")

    # ── Cron queue processor ───────────────────────────────────────────

    def _queue_processor(self) -> None:
        """Background thread: wake up when cron jobs fire and deliver them."""
        while True:
            time.sleep(1)
            if not self.cron.has_queue():
                continue
            if not self.cron.agent_lock.acquire(blocking=False):
                continue
            if not self.cron.has_queue():
                self.cron.agent_lock.release()
                continue
            logger.info("[queue processor] delivering scheduled work")
            try:
                self.run_turn()
            finally:
                self.cron.agent_lock.release()

    # ── REPL ───────────────────────────────────────────────────────────

    def _repl(self) -> None:
        """Interactive read-eval-print loop."""
        while True:
            try:
                query = input("\033[36m >> \033[0m")
            except (EOFError, KeyboardInterrupt):
                break
            if query.strip().lower() in ("q", "exit", ""):
                break
            if query.strip().lower() == "/sessions":
                self._print_sessions()
                continue
            if query.strip().lower().startswith("/load "):
                self._load_session_cmd(query)
                continue
            self.run_turn(query)

    def _print_sessions(self) -> None:
        """Print a formatted table of all saved sessions."""
        sessions = self.session.list_sessions()
        if not sessions:
            logger.info("No saved sessions found.")
            return
        lines = [
            f"\n{'SESSION ID':<30s} {'STATUS':<8s} {'MSGS':>5s}  TITLE",
            "-" * 80,
        ]
        for s in sessions:
            status = "active" if s.get("is_active") else "done"
            sid = s.get("session_id", "?")
            count = s.get("message_count", 0)
            title = s.get("title", "")
            lines.append(f"{sid:<30s} {status:<8s} {count:>5d}  {title}")
        current = self.session.get_current_id()
        lines.append(f"\nCurrent session: {current}")
        logger.info("\n".join(lines))

    def _load_session_cmd(self, query: str) -> None:
        """Load a saved session by ID, replacing the current session_context."""
        parts = query.strip().split(maxsplit=1)
        if len(parts) < 2:
            logger.info("Usage: /load <session_id>")
            return
        sid = parts[1].strip()
        msgs = self.session.load(sid)
        if msgs is None:
            logger.warning(f"Session not found: {sid}")
            return
        self.session.close()
        self.session_context = msgs
        self.session.resume(sid)
        logger.info(f"Loaded session {sid} ({len(msgs)} messages)")


# ═══════════════════════════════════════════════════════════════════════════
# Entry point
# ═══════════════════════════════════════════════════════════════════════════


def main():
    config = WodeConfig(mcp_servers=MCP_SERVERS)
    app = WodeApp(config)
    try:
        app.start()
    finally:
        app.shutdown()


if __name__ == "__main__":
    main()