Cortex-Inara/cortex/tools/aider.py

"""
Aider coding agent tool — invokes Aider AI pair programming as a subprocess.

Aider handles repo-map generation, file editing, git commits, and linting automatically.
It works with any OpenAI-compatible model — point it at DeepSeek, Ollama, OpenRouter, etc.
via AIDER_MODEL / AIDER_OPENAI_API_BASE env vars or the project's .aider.conf.yml.

Credentials are pulled automatically from the Cortex model registry:
  - Named cloud providers (OpenRouter, OpenAI, Groq, Anthropic, …) → --api-key slug=key
  - Generic OpenAI-compatible hosts (Open WebUI, Ollama, local) → --openai-api-base + key
  - Anthropic from providers.anthropic.credentials → --api-key anthropic=key

background=True runs the subprocess asynchronously and returns an agent_id immediately.
"""

import asyncio
import logging
import os
import shlex
from pathlib import Path

from google.genai import types

import agent_manager
from ._projects import PROJECT_ALIASES

logger = logging.getLogger(__name__)

_MAX_OUTPUT_CHARS = 12_000

# Maps URL fragments → Aider --api-key provider slug.
# Order matters: more specific patterns first.
_CLOUD_PROVIDER_URL_MAP: list[tuple[str, str]] = [
    ("openrouter.ai",    "openrouter"),
    ("api.openai.com",   "openai"),
    ("groq.com",         "groq"),
    ("api.together.xyz", "togetherai"),
    ("fireworks.ai",     "fireworks"),
    ("api.x.ai",         "xai"),
    ("api.deepseek.com", "deepseek"),
    ("api.mistral.ai",   "mistral"),
]


def _provider_slug(api_url: str) -> str | None:
    """Return the Aider --api-key provider slug for a known cloud URL, None for generic."""
    url_lower = api_url.lower()
    for fragment, slug in _CLOUD_PROVIDER_URL_MAP:
        if fragment in url_lower:
            return slug
    return None


def _host_flags(host: dict, model: str | None) -> tuple[list[str], str | None]:
    """Build Aider credential flags for a specific host entry.

    Returns (extra_args, adjusted_model). For generic (local) endpoints the model
    name may be prefixed with 'openai/' so Aider routes through the OpenAI client.
    """
    api_url  = (host.get("api_url") or "").rstrip("/")
    api_key  = host.get("api_key") or "none"
    host_type = host.get("host_type", "openai")
    slug     = _provider_slug(api_url)

    if slug:
        # Named cloud provider — Aider maps --api-key slug=key → SLUG_API_KEY env var
        flags = ["--api-key", f"{slug}={api_key}"] if api_key and api_key != "none" else []
        return flags, model

    # Generic OpenAI-compatible (local Open WebUI, Ollama, custom)
    base_url = api_url
    if host_type == "openwebui":
        # Open WebUI serves the chat endpoint at /api/chat/completions
        base_url = base_url + "/api"

    flags = ["--openai-api-base", base_url, "--openai-api-key", api_key]

    # Prefix model with 'openai/' for generic endpoints when no provider prefix is set
    adj_model = model
    if model and "/" not in model:
        adj_model = f"openai/{model}"

    return flags, adj_model


def _resolve_credentials(
    registry: dict,
    model: str | None,
    host_label: str | None,
) -> tuple[list[str], str | None]:
    """Determine Aider credential flags and (possibly adjusted) model name.

    Resolution order:
    1. Anthropic model hint (claude-* / anthropic/*) → Anthropic API key
    2. Explicit host_label → that host's credentials
    3. Model prefix hint (openrouter/*, groq/*, …) → matching host
    4. Default priority: OpenRouter → Anthropic → any keyed cloud host → local host

    Returns (extra_args, adjusted_model).
    """
    hosts = registry.get("hosts", [])

    # Extract Anthropic key from providers.anthropic.credentials (not a host entry)
    anthropic_key = None
    for cred in registry.get("providers", {}).get("anthropic", {}).get("credentials", []):
        if cred.get("api_key"):
            anthropic_key = cred["api_key"]
            break

    # ── 1. Anthropic model hint ────────────────────────────────────────────────
    if model and any(h in model.lower() for h in ("claude-", "anthropic/")):
        if anthropic_key:
            logger.debug("aider: Anthropic model detected — using Anthropic API key")
            return ["--api-key", f"anthropic={anthropic_key}"], model

    # ── 2. Explicit host_label override ───────────────────────────────────────
    if host_label:
        ll = host_label.lower()
        host = next((h for h in hosts if ll in h.get("label", "").lower()), None)
        if host:
            logger.debug("aider: using explicitly requested host '%s'", host.get("label"))
            return _host_flags(host, model)

    # ── 3. Model prefix hints ─────────────────────────────────────────────────
    if model:
        ml = model.lower()
        for fragment, slug in _CLOUD_PROVIDER_URL_MAP:
            if ml.startswith(slug + "/") or ml.startswith(fragment):
                host = next(
                    (h for h in hosts if fragment in h.get("api_url", "").lower()), None
                )
                if host:
                    logger.debug("aider: model prefix '%s' → host '%s'", slug, host.get("label"))
                    return _host_flags(host, model)

    # ── 4. Default priority ───────────────────────────────────────────────────
    # OpenRouter first (most model coverage)
    or_host = next((h for h in hosts if "openrouter.ai" in h.get("api_url", "")), None)
    if or_host and or_host.get("api_key"):
        logger.debug("aider: defaulting to OpenRouter")
        return _host_flags(or_host, model)

    # Anthropic API key (no model hint but it's configured)
    if anthropic_key:
        logger.debug("aider: defaulting to Anthropic API key")
        return ["--api-key", f"anthropic={anthropic_key}"], model

    # Any other keyed cloud host
    for host in hosts:
        slug = _provider_slug(host.get("api_url", ""))
        if slug and host.get("api_key"):
            logger.debug("aider: using keyed cloud host '%s'", host.get("label"))
            return _host_flags(host, model)

    # Generic / local host (no key or unknown provider)
    for host in hosts:
        flags, adj_model = _host_flags(host, model)
        if flags:
            logger.debug("aider: using local host '%s'", host.get("label"))
            return flags, adj_model

    logger.debug("aider: no credentials found in registry — relying on env vars / .aider.conf.yml")
    return [], model


async def aider_run(
    project: str,
    task: str,
    files: list[str] | None = None,
    model: str | None = None,
    host_label: str | None = None,
    auto_commit: bool = True,
    timeout: int = 300,
    background: bool = False,
    notify: bool = False,
) -> str:
    """Run Aider with a single task in a project directory, then exit.

    Credentials are resolved automatically from the Cortex model registry. Use
    host_label to pick a specific configured host (e.g. 'OpenRouter', 'Local').

    When background=True, fires the subprocess asynchronously and returns an agent_id
    immediately. Use agent_status(agent_id) to check progress; set notify=True to
    receive a push/Talk notification on completion.
    """
    proj_def = PROJECT_ALIASES.get(project)
    if proj_def is not None:
        cwd = Path(os.path.expanduser(proj_def.path))
        ssh_host = proj_def.ssh_host
    else:
        cwd = Path(os.path.expanduser(project))
        ssh_host = ""

    if not ssh_host and not cwd.is_dir():
        return f"Error: project directory '{cwd}' does not exist."

    timeout = min(max(int(timeout), 10), 600)

    # Resolve credentials before building the command (model name may be adjusted)
    user = "scott"
    extra_cred_flags: list[str] = []
    try:
        import model_registry
        from persona import get_user
        user = get_user() or "scott"
        registry = model_registry.get_registry(user)
        extra_cred_flags, model = _resolve_credentials(registry, model, host_label)
    except Exception as e:
        logger.debug("aider: credential resolution failed (%s) — relying on env", e)

    cmd: list[str] = [
        "aider",
        "--message", task,
        "--yes-always",
        "--no-pretty",
        "--no-stream",
        "--no-check-update",
        "--no-detect-urls",
        "--auto-commits" if auto_commit else "--no-auto-commits",
    ]

    cmd += extra_cred_flags

    if model:
        cmd += ["--model", model]

    for f in (files or []):
        cmd += ["--file", f]

    logger.info(
        "aider_run: project=%s ssh_host=%s model=%s host_label=%s auto_commit=%s background=%s task=%.120s",
        project, ssh_host or "local", model, host_label, auto_commit, background, task,
    )

    async def _run() -> str:
        if ssh_host:
            # Run aider natively on the remote host via a login shell so PATH
            # includes ~/.local/bin where aider is typically installed.
            inner_cmd = "cd " + shlex.quote(str(cwd)) + " && " + shlex.join(cmd)
            ssh_cmd = f"bash -l -c {shlex.quote(inner_cmd)}"
            proc = await asyncio.create_subprocess_exec(
                "ssh", ssh_host, ssh_cmd,
                stdout=asyncio.subprocess.PIPE,
                stderr=asyncio.subprocess.PIPE,
            )
        else:
            proc = await asyncio.create_subprocess_exec(
                *cmd,
                cwd=str(cwd),
                stdout=asyncio.subprocess.PIPE,
                stderr=asyncio.subprocess.PIPE,
            )
        stdout, stderr = await asyncio.wait_for(proc.communicate(), timeout=float(timeout))

        out = stdout.decode(errors="replace").strip()
        err = stderr.decode(errors="replace").strip()

        parts = []
        if out:
            parts.append(out)
        if err:
            parts.append(f"[stderr]\n{err}")
        combined = "\n".join(parts) if parts else "(no output)"

        if len(combined) > _MAX_OUTPUT_CHARS:
            half = _MAX_OUTPUT_CHARS // 2
            combined = (
                combined[:half]
                + f"\n\n[... {len(combined) - _MAX_OUTPUT_CHARS} chars trimmed ...]\n\n"
                + combined[-half:]
            )

        if proc.returncode not in (0, 1):
            return f"[exit {proc.returncode}]\n{combined}"
        return combined

    if background:
        rec = await agent_manager.register(
            user=user,
            role="aider",
            task=task,
            level=2,
            notify=notify,
        )

        async def _bg_task() -> None:
            try:
                result = await _run()
                await agent_manager.finish(rec.agent_id, result, "done")
                logger.info("aider_run [bg]: done %s", rec.agent_id[:8])
            except asyncio.CancelledError:
                await agent_manager.finish(rec.agent_id, "Cancelled.", "cancelled")
                raise
            except asyncio.TimeoutError:
                msg = f"Aider timed out after {timeout}s"
                logger.warning("aider_run [bg]: timeout %s", rec.agent_id[:8])
                await agent_manager.finish(rec.agent_id, msg, "timeout")
            except FileNotFoundError:
                msg = "Error: 'aider' not found in PATH — run: pip install aider-chat"
                await agent_manager.finish(rec.agent_id, msg, "failed")
            except Exception as e:
                logger.error("aider_run [bg]: failed %s: %s", rec.agent_id[:8], e)
                await agent_manager.finish(rec.agent_id, str(e), "failed")

        bg = asyncio.create_task(_bg_task())
        agent_manager.set_task_ref(rec.agent_id, bg)
        return (
            f"Aider task started in background. ID: {rec.agent_id}\n"
            f"Use agent_status('{rec.agent_id}') to monitor progress."
        )

    # Synchronous path
    try:
        return await _run()
    except asyncio.TimeoutError:
        return f"Error: aider timed out after {timeout}s"
    except FileNotFoundError:
        return "Error: 'aider' not found in PATH — run: pip install aider-chat"
    except Exception as e:
        logger.error("aider_run error: %s", e)
        return f"Error: {e}"


DECLARATIONS = [
    types.FunctionDeclaration(
        name="aider_run",
        description=(
            "Run the Aider AI coding agent on a project with a single task, then exit. "
            "Aider maps the repo, edits files, runs lint checks, and optionally commits. "
            "Credentials are resolved automatically from the Cortex model registry — "
            "OpenRouter, local Open WebUI/Ollama, Anthropic API, and other configured hosts "
            "are all supported. Use host_label to pick a specific host. "
            "aether_api, aether_frontend, and aether_container run aider natively on the "
            "workstation (scott-wks-main-i7) via SSH — aider must be installed there. "
            "Set background=True for long tasks — returns an agent_id immediately and sends "
            "a notification when done. ADMIN ONLY. Requires confirmation."
        ),
        parameters=types.Schema(
            type=types.Type.OBJECT,
            properties={
                "project": types.Schema(
                    type=types.Type.STRING,
                    description=(
                        "Project alias or absolute path. Known aliases: "
                        "'cortex' (this project), 'aether_api', 'aether_frontend', "
                        "'aether_container'. Or provide an absolute path."
                    ),
                ),
                "task": types.Schema(
                    type=types.Type.STRING,
                    description=(
                        "Full task description sent to Aider as --message. "
                        "Be specific — include file names, what to change, and why."
                    ),
                ),
                "files": types.Schema(
                    type=types.Type.ARRAY,
                    items=types.Schema(type=types.Type.STRING),
                    description=(
                        "Optional files to add explicitly to the editing context "
                        "(paths relative to project root). Aider builds a repo map "
                        "automatically — these get priority."
                    ),
                ),
                "model": types.Schema(
                    type=types.Type.STRING,
                    description=(
                        "Optional model override. Format depends on the provider: "
                        "'openrouter/anthropic/claude-3-5-haiku-20241022' (OpenRouter), "
                        "'claude-3-5-sonnet-20241022' (Anthropic direct), "
                        "'gemma-4-27b-it' or 'openai/gemma-4-27b-it' (local Open WebUI), "
                        "'deepseek/deepseek-chat' (DeepSeek via OpenRouter). "
                        "Defaults to the project's .aider.conf.yml model or AIDER_MODEL env var."
                    ),
                ),
                "host_label": types.Schema(
                    type=types.Type.STRING,
                    description=(
                        "Pick a specific configured host by label (partial match, case-insensitive). "
                        "Examples: 'OpenRouter', 'Local', 'scott-lt-i7-rtx'. "
                        "Overrides automatic credential resolution. "
                        "Omit to let credentials be chosen automatically."
                    ),
                ),
                "auto_commit": types.Schema(
                    type=types.Type.BOOLEAN,
                    description=(
                        "Auto-commit changes after edits (default: true). "
                        "Set to false to review diffs before committing manually."
                    ),
                ),
                "timeout": types.Schema(
                    type=types.Type.INTEGER,
                    description="Max seconds to wait for Aider to finish (default 300, max 600).",
                ),
                "background": types.Schema(
                    type=types.Type.BOOLEAN,
                    description=(
                        "Run asynchronously in the background (default: false). "
                        "Returns an agent_id immediately; use agent_status(agent_id) to monitor. "
                        "Recommended for tasks expected to take more than ~60 seconds."
                    ),
                ),
                "notify": types.Schema(
                    type=types.Type.BOOLEAN,
                    description=(
                        "Send a push/Talk notification when the background task completes "
                        "(default: false). Only applies when background=true."
                    ),
                ),
            },
            required=["project", "task"],
        ),
    )
]