code/agents/__init__.py

"""Custom Agent system — AI-generated, user-saved agent personas.

Inspired by Claude Code's sub-agent / agent-customization patterns. A custom
agent is a markdown file (`AGENT.md`) with YAML frontmatter that defines:

  - name          : unique kebab-case identifier
  - description   : one-line summary
  - tools         : comma-separated subset of the agent tool registry
                    (default: all tools)
  - skills        : comma-separated skill names to auto-load
  - temperature   : float 0.0–1.0 (optional)
  - max_iterations: int (optional, overrides default)
  - tags          : comma-separated tags
  - author        : "user" | "AI-generated"
  - created       : ISO date

The body of the file is a system-prompt extension that is appended to the
base SoniCoder system prompt when the agent is active.

Storage layout:
    workspace/.sonicoder/agents/<agent-name>/AGENT.md

Built-in agents live in `code/agents/builtins/` (read-only).
"""

from __future__ import annotations

import logging
import os
import re
from datetime import date
from typing import Any

from code.skills import _parse_frontmatter  # reuse the same simple YAML parser

logger = logging.getLogger(__name__)

# ─── Discovery roots ────────────────────────────────────────────────────

_BUILTIN_AGENTS_DIR = os.path.join(os.path.dirname(__file__), "builtins")
_USER_AGENTS_DIRNAME = ".sonicoder/agents"  # relative to workspace root

# All tools available to the agent — kept in sync with code/agent.TOOL_REGISTRY
ALL_TOOLS: tuple[str, ...] = (
    "read_file",
    "write_file",
    "edit_file",
    "multi_edit",
    "list_dir",
    "glob",
    "grep",
    "bash",
    "todo_read",
    "todo_write",
    "todo_update",
)

# Session state: the currently active agent (None = default SoniCoder)
_active_agent: str | None = None


# ─── Helpers ────────────────────────────────────────────────────────────

def _agent_dirs() -> list[str]:
    """Return all directories to search for agents (builtins first)."""
    dirs = [_BUILTIN_AGENTS_DIR]
    try:
        from code.tools.fs import get_workspace_root
        user_dir = os.path.join(get_workspace_root(), _USER_AGENTS_DIRNAME)
        if os.path.isdir(user_dir):
            dirs.append(user_dir)
    except Exception:
        pass
    return dirs


def _user_agents_dir() -> str:
    """Return the user agents directory (creating it if missing)."""
    from code.tools.fs import get_workspace_root
    user_dir = os.path.join(get_workspace_root(), _USER_AGENTS_DIRNAME)
    os.makedirs(user_dir, exist_ok=True)
    return user_dir


def _safe_agent_name(name: str) -> str:
    """Sanitize an agent name to kebab-case."""
    name = name.strip().lower()
    name = re.sub(r"[^a-z0-9-]+", "-", name)
    name = re.sub(r"-+", "-", name).strip("-")
    if not name:
        raise ValueError("Agent name must contain at least one alphanumeric character")
    if len(name) > 64:
        name = name[:64].rstrip("-")
    return name


def _load_agent(agent_dir: str) -> dict[str, Any] | None:
    """Load a single agent definition from a directory containing AGENT.md."""
    agent_md = os.path.join(agent_dir, "AGENT.md")
    if not os.path.isfile(agent_md):
        return None

    try:
        with open(agent_md, "r", encoding="utf-8") as f:
            content = f.read()
    except Exception as exc:
        logger.warning("Failed to read %s: %s", agent_md, exc)
        return None

    meta, body = _parse_frontmatter(content)

    # Parse comma-separated list fields
    def _split_csv(v: str) -> list[str]:
        return [s.strip() for s in (v or "").split(",") if s.strip()]

    tools = _split_csv(meta.get("tools", ""))
    if not tools:
        tools = list(ALL_TOOLS)
    # Filter to known tools (defensive)
    tools = [t for t in tools if t in ALL_TOOLS]

    skills = _split_csv(meta.get("skills", ""))
    tags = _split_csv(meta.get("tags", ""))

    try:
        temperature = float(meta["temperature"]) if meta.get("temperature") else None
    except (ValueError, TypeError):
        temperature = None

    try:
        max_iterations = int(meta["max_iterations"]) if meta.get("max_iterations") else None
    except (ValueError, TypeError):
        max_iterations = None

    return {
        "name": meta.get("name") or os.path.basename(agent_dir),
        "description": meta.get("description", ""),
        "tools": tools,
        "skills": skills,
        "temperature": temperature,
        "max_iterations": max_iterations,
        "tags": tags,
        "author": meta.get("author", "user"),
        "created": meta.get("created", ""),
        "body": body.strip(),
        "path": agent_dir,
    }


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

def list_agents() -> list[dict[str, Any]]:
    """List all available agents (builtins + user). Returns metadata only."""
    agents: list[dict[str, Any]] = []
    seen: set[str] = set()

    for agents_dir in _agent_dirs():
        if not os.path.isdir(agents_dir):
            continue
        for entry in sorted(os.listdir(agents_dir)):
            entry_path = os.path.join(agents_dir, entry)
            if not os.path.isdir(entry_path):
                continue
            agent = _load_agent(entry_path)
            if agent and agent["name"] not in seen:
                seen.add(agent["name"])
                agents.append({
                    "name": agent["name"],
                    "description": agent["description"],
                    "tools": agent["tools"],
                    "skills": agent["skills"],
                    "tags": agent["tags"],
                    "author": agent["author"],
                    "active": agent["name"] == _active_agent,
                })
    return agents


def get_agent(name: str) -> dict[str, Any] | None:
    """Get full agent definition by name (or None if not found)."""
    name = _safe_agent_name(name)
    for agents_dir in _agent_dirs():
        if not os.path.isdir(agents_dir):
            continue
        for entry in os.listdir(agents_dir):
            entry_path = os.path.join(agents_dir, entry)
            if not os.path.isdir(entry_path):
                continue
            agent = _load_agent(entry_path)
            if agent and agent["name"] == name:
                return agent
    return None


def agent_exists(name: str) -> bool:
    return get_agent(name) is not None


def save_agent(
    name: str,
    description: str,
    body: str,
    tools: list[str] | None = None,
    skills: list[str] | None = None,
    temperature: float | None = None,
    max_iterations: int | None = None,
    tags: list[str] | None = None,
    author: str = "user",
) -> dict[str, Any]:
    """Save (create or overwrite) an agent definition."""
    name = _safe_agent_name(name)

    # Validate tools
    if tools:
        invalid = [t for t in tools if t not in ALL_TOOLS]
        if invalid:
            return {"success": False, "error": f"Unknown tools: {invalid}", "valid_tools": list(ALL_TOOLS)}
        tools_str = ", ".join(tools)
    else:
        tools_str = ", ".join(ALL_TOOLS)

    skills_str = ", ".join(skills) if skills else ""
    tags_str = ", ".join(tags) if tags else ""
    temp_str = f"{temperature:.2f}" if temperature is not None else ""
    iter_str = str(max_iterations) if max_iterations is not None else ""

    # Build the markdown file
    frontmatter_lines = [
        "---",
        f"name: {name}",
        f"description: {description.strip()[:200]}",
        f"tools: {tools_str}",
        f"skills: {skills_str}",
    ]
    if temp_str:
        frontmatter_lines.append(f"temperature: {temp_str}")
    if iter_str:
        frontmatter_lines.append(f"max_iterations: {iter_str}")
    if tags_str:
        frontmatter_lines.append(f"tags: {tags_str}")
    frontmatter_lines.append(f"author: {author}")
    frontmatter_lines.append(f"created: {date.today().isoformat()}")
    frontmatter_lines.append("---")
    frontmatter_lines.append("")
    frontmatter_lines.append(body.strip())
    frontmatter_lines.append("")

    content = "\n".join(frontmatter_lines)

    # Write to workspace/.sonicoder/agents/<name>/AGENT.md
    agents_dir = _user_agents_dir()
    agent_dir = os.path.join(agents_dir, name)
    os.makedirs(agent_dir, exist_ok=True)
    agent_md = os.path.join(agent_dir, "AGENT.md")
    with open(agent_md, "w", encoding="utf-8") as f:
        f.write(content)

    logger.info("Saved agent '%s' to %s", name, agent_md)
    return {"success": True, "name": name, "path": agent_md}


def delete_agent(name: str) -> dict[str, Any]:
    """Delete a user-defined agent. Built-ins cannot be deleted."""
    name = _safe_agent_name(name)
    agent = get_agent(name)
    if not agent:
        return {"success": False, "error": f"Agent not found: {name}"}

    # Only allow deleting files under the user agents dir
    from code.tools.fs import get_workspace_root
    user_dir = os.path.join(get_workspace_root(), _USER_AGENTS_DIRNAME)
    if not agent["path"].startswith(user_dir + os.sep):
        return {"success": False, "error": f"Cannot delete built-in agent: {name}"}

    import shutil
    try:
        shutil.rmtree(os.path.dirname(agent["path"]), ignore_errors=True)
    except Exception as exc:
        return {"success": False, "error": str(exc)}

    if _active_agent == name:
        set_active_agent(None)

    logger.info("Deleted agent '%s'", name)
    return {"success": True, "name": name}


def set_active_agent(name: str | None) -> dict[str, Any]:
    """Set the active agent for subsequent prompts. None resets to default."""
    global _active_agent
    if name is None or name.strip() == "":
        _active_agent = None
        return {"success": True, "active_agent": None, "message": "Reset to default SoniCoder agent"}
    name = _safe_agent_name(name)
    if not agent_exists(name):
        return {"success": False, "error": f"Agent not found: {name}"}
    _active_agent = name
    return {"success": True, "active_agent": name}


def get_active_agent() -> str | None:
    """Return the name of the currently active agent, or None."""
    return _active_agent


def get_active_agent_config() -> dict[str, Any] | None:
    """Return the full config of the active agent, or None if default."""
    if _active_agent is None:
        return None
    return get_agent(_active_agent)


def build_agent_system_prompt_extension(agent_name: str) -> str:
    """Build the system-prompt extension for a custom agent.

    Returns empty string if the agent doesn't exist or has no body.
    """
    agent = get_agent(agent_name)
    if not agent:
        return ""
    parts: list[str] = []
    parts.append(f"# Active Custom Agent: {agent['name']}")
    if agent["description"]:
        parts.append(f"_{agent['description']}_")
    parts.append("")
    parts.append("You are operating under this custom agent persona. Follow its instructions and workflow.")
    parts.append("")
    parts.append("## Agent Persona & Instructions")
    parts.append("")
    parts.append(agent["body"])

    if agent["tools"] and len(agent["tools"]) < len(ALL_TOOLS):
        parts.append("")
        parts.append("## Tool Whitelist for this Agent")
        parts.append(
            f"You may ONLY use these tools: {', '.join(agent['tools'])}. "
            "Do not call any other tool."
        )

    if agent["skills"]:
        parts.append("")
        parts.append("## Auto-loaded Skills")
        parts.append(f"The following skills are pre-loaded for this agent: {', '.join(agent['skills'])}")

    return "\n".join(parts)


# ─── AI generation helper ───────────────────────────────────────────────

# The meta-prompt sent to the model when a user runs `/agent create <desc>`.
AGENT_GENERATION_PROMPT = """You are creating a custom agent definition for SoniCoder.

Based on the user's description, generate a complete AGENT.md file that defines a specialized agent persona.

## Available Tools (pick a subset, or all)
read_file, write_file, edit_file, multi_edit, list_dir, glob, grep, bash, todo_read, todo_write, todo_update

## Available Built-in Skills (pick zero or more)
frontend-design, feature-dev, code-review, debugging, fullstack-scaffold, commit-workflow

## Output Format
Use the @@FILE: multi-file format to write exactly ONE file:

@@FILE: .sonicoder/agents/<kebab-case-name>/AGENT.md
---
name: <kebab-case-name>
description: <one-line description, max 200 chars>
tools: <comma-separated subset of available tools>
skills: <comma-separated skill names, or empty>
temperature: <float 0.0-1.0 — lower for precise tasks, higher for creative>
max_iterations: <int 4-20>
tags: <comma-separated tags>
author: AI-generated
created: <today's date YYYY-MM-DD>
---

# <Agent Name in Title Case>

<Full system prompt extension. Include:>
<- Persona description (who the agent is)>
<- Core responsibilities>
<- Workflow / step-by-step approach>
<- Output format expectations>
<- Critical rules and constraints>

@@END@@

## Rules
- The agent name MUST be kebab-case (lowercase, hyphens only).
- Pick tools that match the agent's purpose — don't give a read-only reviewer `write_file`.
- The body should be 150-400 words, specific and actionable.
- Do NOT include any other files. Just the one AGENT.md.
- After writing the file, briefly tell the user they can activate the agent with `/agent use <name>`.

## User's Description
$ARGUMENTS
"""