stackpilot/backend/services/agent_service.py

"""Talk to remote stackpilot-agent hosts over HTTP.

The central app stores an ``Agent`` row per remote host and proxies stack /
system calls to it using the agent's shared token. Connectivity state
(``status``, ``hostname``, ``last_seen``) is refreshed on every successful or
failed call so the UI can show a live dot per host.
"""
from __future__ import annotations

import logging
from datetime import datetime, timezone
from typing import Any, Optional

import httpx
from sqlmodel import Session

from models.agent import Agent

logger = logging.getLogger("stackpilot.agent_proxy")

_TIMEOUT = 30.0


class AgentError(Exception):
    def __init__(self, status: int, error: str, detail: str = ""):
        self.status = status
        self.error = error
        self.detail = detail
        super().__init__(f"{error}: {detail}" if detail else error)


def _now() -> datetime:
    return datetime.now(timezone.utc)


def _mark(session: Session, agent: Agent, status: str, hostname: Optional[str] = None) -> None:
    agent.status = status
    if status == "online":
        agent.last_seen = _now()
        if hostname:
            agent.hostname = hostname
    session.add(agent)
    session.commit()
    session.refresh(agent)


async def _request(
    agent: Agent,
    method: str,
    path: str,
    *,
    params: Optional[dict] = None,
    json: Any = None,
) -> httpx.Response:
    url = agent.url.rstrip("/") + path
    headers = {"Authorization": f"Bearer {agent.token}"}
    async with httpx.AsyncClient(follow_redirects=True) as client:
        return await client.request(
            method, url, headers=headers, params=params, json=json, timeout=_TIMEOUT
        )


async def call(
    session: Session,
    agent: Agent,
    method: str,
    path: str,
    *,
    params: Optional[dict] = None,
    json: Any = None,
) -> Any:
    """Proxy a request to the agent, updating its status, returning parsed JSON."""
    try:
        resp = await _request(agent, method, path, params=params, json=json)
    except httpx.HTTPError as exc:
        _mark(session, agent, "offline")
        raise AgentError(502, "agent_unreachable", str(exc)) from exc

    if resp.status_code in (401, 403):
        _mark(session, agent, "unauthorized")
        raise AgentError(resp.status_code, "agent_unauthorized", "Invalid agent token")

    _mark(session, agent, "online")

    if resp.status_code >= 400:
        detail = ""
        try:
            body = resp.json()
            detail = body.get("detail") if isinstance(body, dict) else str(body)
            if isinstance(detail, dict):
                detail = detail.get("detail") or detail.get("error") or str(detail)
        except ValueError:
            detail = resp.text[:500]
        raise AgentError(resp.status_code, "agent_error", str(detail))

    if resp.content:
        try:
            return resp.json()
        except ValueError:
            return resp.text
    return None


def _handle_status(session: Session, agent: Agent, status_code: int, body_text: str = "") -> None:
    """Update agent status from a response code; raise AgentError on failure."""
    if status_code in (401, 403):
        _mark(session, agent, "unauthorized")
        raise AgentError(status_code, "agent_unauthorized", "Invalid agent token")
    _mark(session, agent, "online")
    if status_code >= 400:
        raise AgentError(status_code, "agent_error", body_text[:500])


async def download_to_file(
    session: Session,
    agent: Agent,
    path: str,
    dest_path: str,
    *,
    params: Optional[dict] = None,
) -> None:
    """Stream a GET from the agent into ``dest_path``."""
    url = agent.url.rstrip("/") + path
    headers = {"Authorization": f"Bearer {agent.token}"}
    try:
        async with httpx.AsyncClient(follow_redirects=True) as client:
            async with client.stream("GET", url, headers=headers, params=params, timeout=None) as resp:
                if resp.status_code >= 400:
                    text = (await resp.aread()).decode("utf-8", "replace")
                    _handle_status(session, agent, resp.status_code, text)
                _handle_status(session, agent, resp.status_code)
                with open(dest_path, "wb") as fh:
                    async for chunk in resp.aiter_bytes(1024 * 256):
                        fh.write(chunk)
    except httpx.HTTPError as exc:
        _mark(session, agent, "offline")
        raise AgentError(502, "agent_unreachable", str(exc)) from exc


async def stream_download(
    session: Session,
    agent: Agent,
    path: str,
    *,
    params: Optional[dict] = None,
):
    """Stream a GET from the agent straight through, yielding chunks.

    Unlike :func:`download_to_file` this never buffers to disk, so a large
    response (e.g. a folder zip the agent builds on the fly) starts flowing to
    the browser immediately instead of being staged first.
    """
    url = agent.url.rstrip("/") + path
    headers = {"Authorization": f"Bearer {agent.token}"}
    try:
        async with httpx.AsyncClient(follow_redirects=True) as client:
            async with client.stream("GET", url, headers=headers, params=params, timeout=None) as resp:
                if resp.status_code >= 400:
                    text = (await resp.aread()).decode("utf-8", "replace")
                    _handle_status(session, agent, resp.status_code, text)
                _handle_status(session, agent, resp.status_code)
                async for chunk in resp.aiter_bytes(1024 * 256):
                    yield chunk
    except httpx.HTTPError as exc:
        _mark(session, agent, "offline")
        raise AgentError(502, "agent_unreachable", str(exc)) from exc


async def upload_file(
    session: Session,
    agent: Agent,
    path: str,
    file_path: str,
    filename: str,
    data: dict,
) -> Any:
    """Stream a multipart POST (file + form fields) to the agent, returning JSON."""
    url = agent.url.rstrip("/") + path
    headers = {"Authorization": f"Bearer {agent.token}"}
    try:
        async with httpx.AsyncClient(follow_redirects=True) as client:
            with open(file_path, "rb") as fh:
                files = {"file": (filename, fh, "application/gzip")}
                resp = await client.post(url, headers=headers, files=files, data=data, timeout=None)
    except httpx.HTTPError as exc:
        _mark(session, agent, "offline")
        raise AgentError(502, "agent_unreachable", str(exc)) from exc

    detail = ""
    if resp.status_code >= 400:
        try:
            body = resp.json()
            detail = body.get("detail") if isinstance(body, dict) else str(body)
        except ValueError:
            detail = resp.text[:500]
    _handle_status(session, agent, resp.status_code, str(detail))
    return resp.json() if resp.content else None


async def ping(session: Session, agent: Agent) -> dict:
    """Health-check an agent and refresh its status + hostname. Never raises."""
    try:
        data = await call(session, agent, "GET", "/agent/ping")
        if isinstance(data, dict) and data.get("hostname"):
            agent.hostname = data["hostname"]
            session.add(agent)
            session.commit()
            session.refresh(agent)
        return {"status": agent.status, "hostname": agent.hostname, "data": data}
    except AgentError:
        return {"status": agent.status, "hostname": agent.hostname, "data": None}