Newer
Older
cortex-hub / ai-hub / app / api / routes / agent_update.py
"""
Agent Update Distribution Endpoint.

Stable, frozen HTTP API — this contract NEVER changes shape.
  GET /api/v1/agent/version   →  {"version": "x.y.z"}
  GET /api/v1/agent/download  →  application/gzip  (agent-node tarball)

Auth: X-Agent-Token header must match the hub's secret key.
"""

import os
import io
import tarfile
import logging
from fastapi import APIRouter, HTTPException, Request
from fastapi.responses import StreamingResponse, JSONResponse

logger = logging.getLogger(__name__)

# Path to the agent-node source tree on the hub container.
# In Docker: mounted at /app/agent-node-source
# Overridable via env var for flexibility in other deployments.
_AGENT_NODE_DIR = os.environ.get(
    "AGENT_NODE_SRC_DIR",
    "/app/agent-node-source"
)
_SKILLS_DIR = os.environ.get(
    "SKILLS_SRC_DIR",
    "/app/skills"
)
_AGENT_NODE_DIR = os.path.abspath(_AGENT_NODE_DIR)
_SKILLS_DIR = os.path.abspath(_SKILLS_DIR)
_VERSION_FILE = os.path.join(_AGENT_NODE_DIR, "VERSION")

# Directories/files to exclude from the distributed tarball
_EXCLUDE_PATTERNS = {
    "__pycache__", ".git", "*.pyc", "*.pyo",
    "sync-node-1", "sync-node-2",  # Test workspace dirs
    "docker-compose.yml",           # Deployment-specific, not for generic clients
}


def _read_version() -> str:
    try:
        with open(_VERSION_FILE, "r") as f:
            return f.read().strip()
    except FileNotFoundError:
        return "0.0.0"


def _auth_ok(request: Request) -> bool:
    """Validates the agent auth token from the request header."""
    from app.config import settings
    token = request.headers.get("X-Agent-Token", "")
    return token == settings.SECRET_KEY


def _should_exclude(path: str) -> bool:
    """Returns True if the path should be excluded from the tarball."""
    parts = path.replace("\\", "/").split("/")
    for part in parts:
        if part in _EXCLUDE_PATTERNS:
            return True
        if part.endswith(".pyc") or part.endswith(".pyo"):
            return True
    return False


def _build_tarball() -> bytes:
    """
    Builds an in-memory gzipped tarball of the agent-node and skills.
    The tarball is flat (no agent-node-source/ prefix) so extraction is direct.
    """
    buf = io.BytesIO()
    with tarfile.open(fileobj=buf, mode="w:gz") as tar:
        # 1. Add agent-node files (Flat Root)
        if os.path.exists(_AGENT_NODE_DIR):
            for root, dirs, files in os.walk(_AGENT_NODE_DIR):
                dirs[:] = [d for d in dirs if not _should_exclude(os.path.join(root, d))]
                for filename in files:
                    abs_path = os.path.join(root, filename)
                    rel_path = os.path.relpath(abs_path, _AGENT_NODE_DIR)
                    if not _should_exclude(rel_path):
                        tar.add(abs_path, arcname=rel_path)
        
        # 2. Add skills directory (as 'skills/')
        if os.path.exists(_SKILLS_DIR):
            for root, dirs, files in os.walk(_SKILLS_DIR):
                dirs[:] = [d for d in dirs if not _should_exclude(os.path.join(root, d))]
                for filename in files:
                    abs_path = os.path.join(root, filename)
                    # Skills should be in a 'skills/' folder in the tarball
                    rel_path = os.path.join("skills", os.path.relpath(abs_path, _SKILLS_DIR))
                    if not _should_exclude(rel_path):
                        tar.add(abs_path, arcname=rel_path)

    return buf.getvalue()


def create_agent_update_router() -> APIRouter:
    router = APIRouter(prefix="/agent", tags=["Agent Update"])

    @router.get("/version", summary="Get current agent version")
    def get_agent_version(request: Request):
        """
        Returns the current agent-node version. Called by agent nodes at startup
        and periodically to detect if they need to self-update.
        """
        if not _auth_ok(request):
            raise HTTPException(status_code=401, detail="Unauthorized")

        version = _read_version()
        logger.info(f"[AgentUpdate] Version check → {version}")
        return JSONResponse({"version": version})

    @router.get("/download", summary="Download agent node tarball")
    def download_agent(request: Request):
        """
        Streams the current agent-node source as a gzipped tarball.
        Only called when an agent detects it is behind the hub's version.
        """
        if not _auth_ok(request):
            raise HTTPException(status_code=401, detail="Unauthorized")

        if not os.path.isdir(_AGENT_NODE_DIR):
            raise HTTPException(status_code=503, detail="Agent source not available on this hub.")

        version = _read_version()
        logger.info(f"[AgentUpdate] Serving agent tarball v{version}")

        try:
            tarball_bytes = _build_tarball()
        except Exception as e:
            logger.error(f"[AgentUpdate] Failed to build tarball: {e}")
            raise HTTPException(status_code=500, detail="Failed to package agent.")

        filename = f"cortex-agent-node-{version}.tar.gz"
        return StreamingResponse(
            iter([tarball_bytes]),
            media_type="application/gzip",
            headers={"Content-Disposition": f"attachment; filename={filename}"}
        )

    return router