stackpilot/backend/routers/agents.py

"""Remote host (agent) management + proxied stack/system operations."""
from __future__ import annotations

from fastapi import APIRouter, Depends, HTTPException, Query, Request
from sqlmodel import Session, select

from auth import get_current_user, require_admin
from database import get_session
from models.agent import Agent, AgentCreate, AgentRead, AgentUpdate
from models.stack import StackCreate, StackUpdate
from models.user import User
from services import agent_service, audit_service
from services.agent_service import AgentError

router = APIRouter(prefix="/api/agents", tags=["agents"])

_ACTIONS = {"start", "stop", "restart", "pull", "update", "down"}


def _ip(request: Request) -> str:
    return request.client.host if request.client else "unknown"


def _to_read(a: Agent) -> AgentRead:
    return AgentRead(
        id=a.id,
        name=a.name,
        url=a.url,
        status=a.status,
        hostname=a.hostname,
        last_seen=a.last_seen,
        created_at=a.created_at,
        token_set=bool(a.token),
    )


def _get_or_404(session: Session, agent_id: int) -> Agent:
    agent = session.get(Agent, agent_id)
    if not agent:
        raise HTTPException(status_code=404, detail=f"Agent {agent_id} not found")
    return agent


def _raise(exc: AgentError):
    raise HTTPException(
        status_code=exc.status if exc.status >= 400 else 502,
        detail={"error": exc.error, "detail": exc.detail},
    )


async def _proxy(session: Session, agent: Agent, method: str, path: str, **kw):
    try:
        return await agent_service.call(session, agent, method, path, **kw)
    except AgentError as exc:
        _raise(exc)


# --------------------------------------------------------------------------- #
# CRUD
# --------------------------------------------------------------------------- #


@router.get("", response_model=list[AgentRead])
async def list_agents(
    refresh: bool = Query(True),
    session: Session = Depends(get_session),
    _user: User = Depends(get_current_user),
) -> list[AgentRead]:
    agents = session.exec(select(Agent).order_by(Agent.id)).all()
    if refresh:
        for agent in agents:
            await agent_service.ping(session, agent)
    return [_to_read(a) for a in agents]


@router.post("", response_model=AgentRead, status_code=201)
async def create_agent(
    body: AgentCreate,
    request: Request,
    session: Session = Depends(get_session),
    user: User = Depends(require_admin),
) -> AgentRead:
    agent = Agent(name=body.name, url=body.url.rstrip("/"), token=body.token)
    session.add(agent)
    session.commit()
    session.refresh(agent)
    # Validate connectivity immediately (best-effort; agent is saved regardless).
    await agent_service.ping(session, agent)
    audit_service.record(
        session, user=user.username, action="agent.create", target=agent.name,
        detail=agent.url, ip=_ip(request),
    )
    return _to_read(agent)


@router.put("/{agent_id}", response_model=AgentRead)
async def update_agent(
    agent_id: int,
    body: AgentUpdate,
    request: Request,
    session: Session = Depends(get_session),
    user: User = Depends(require_admin),
) -> AgentRead:
    agent = _get_or_404(session, agent_id)
    if body.name is not None:
        agent.name = body.name
    if body.url is not None:
        agent.url = body.url.rstrip("/")
    if body.token:
        agent.token = body.token
    agent.status = "unknown"
    session.add(agent)
    session.commit()
    session.refresh(agent)
    await agent_service.ping(session, agent)
    audit_service.record(
        session, user=user.username, action="agent.update", target=agent.name,
        ip=_ip(request),
    )
    return _to_read(agent)


@router.delete("/{agent_id}")
def delete_agent(
    agent_id: int,
    request: Request,
    session: Session = Depends(get_session),
    user: User = Depends(require_admin),
) -> dict:
    agent = _get_or_404(session, agent_id)
    name = agent.name
    session.delete(agent)
    session.commit()
    audit_service.record(
        session, user=user.username, action="agent.delete", target=name, ip=_ip(request),
    )
    return {"ok": True}


@router.post("/{agent_id}/ping")
async def ping_agent(
    agent_id: int,
    session: Session = Depends(get_session),
    _user: User = Depends(get_current_user),
) -> dict:
    agent = _get_or_404(session, agent_id)
    return await agent_service.ping(session, agent)


# --------------------------------------------------------------------------- #
# Proxied stack + system operations
# --------------------------------------------------------------------------- #


@router.get("/{agent_id}/system")
async def agent_system(
    agent_id: int,
    session: Session = Depends(get_session),
    _user: User = Depends(get_current_user),
):
    agent = _get_or_404(session, agent_id)
    return await _proxy(session, agent, "GET", "/agent/system")


@router.get("/{agent_id}/stacks")
async def agent_stacks(
    agent_id: int,
    session: Session = Depends(get_session),
    _user: User = Depends(get_current_user),
) -> list[dict]:
    agent = _get_or_404(session, agent_id)
    stacks = await _proxy(session, agent, "GET", "/agent/stacks") or []
    for s in stacks:
        s["agent_id"] = agent.id
        s["agent_name"] = agent.name
    return stacks


@router.get("/{agent_id}/stacks/{stack_id}")
async def agent_stack_detail(
    agent_id: int,
    stack_id: str,
    session: Session = Depends(get_session),
    _user: User = Depends(get_current_user),
) -> dict:
    agent = _get_or_404(session, agent_id)
    data = await _proxy(session, agent, "GET", f"/agent/stacks/{stack_id}")
    data["agent_id"] = agent.id
    data["agent_name"] = agent.name
    return data


@router.get("/{agent_id}/stacks/{stack_id}/logs")
async def agent_stack_logs(
    agent_id: int,
    stack_id: str,
    tail: int = Query(200, le=2000),
    session: Session = Depends(get_session),
    _user: User = Depends(get_current_user),
) -> dict:
    agent = _get_or_404(session, agent_id)
    return await _proxy(
        session, agent, "GET", f"/agent/stacks/{stack_id}/logs", params={"tail": tail}
    )


@router.post("/{agent_id}/stacks", status_code=201)
async def agent_create_stack(
    agent_id: int,
    body: StackCreate,
    request: Request,
    session: Session = Depends(get_session),
    user: User = Depends(require_admin),
) -> dict:
    agent = _get_or_404(session, agent_id)
    result = await _proxy(
        session, agent, "POST", "/agent/stacks",
        json={"name": body.name, "yaml": body.yaml, "env": body.env},
    )
    audit_service.record(
        session, user=user.username, action="agent.stack.create",
        target=f"{agent.name}/{result.get('id')}", ip=_ip(request),
    )
    return result


@router.put("/{agent_id}/stacks/{stack_id}")
async def agent_update_stack(
    agent_id: int,
    stack_id: str,
    body: StackUpdate,
    request: Request,
    session: Session = Depends(get_session),
    user: User = Depends(require_admin),
) -> dict:
    agent = _get_or_404(session, agent_id)
    result = await _proxy(
        session, agent, "PUT", f"/agent/stacks/{stack_id}",
        json={"yaml": body.yaml, "env": body.env},
    )
    audit_service.record(
        session, user=user.username, action="agent.stack.update",
        target=f"{agent.name}/{stack_id}", ip=_ip(request),
    )
    return result


@router.delete("/{agent_id}/stacks/{stack_id}")
async def agent_delete_stack(
    agent_id: int,
    stack_id: str,
    request: Request,
    delete_files: bool = Query(True),
    session: Session = Depends(get_session),
    user: User = Depends(require_admin),
) -> dict:
    agent = _get_or_404(session, agent_id)
    result = await _proxy(
        session, agent, "DELETE", f"/agent/stacks/{stack_id}",
        params={"delete_files": delete_files},
    )
    audit_service.record(
        session, user=user.username, action="agent.stack.delete",
        target=f"{agent.name}/{stack_id}", ip=_ip(request),
    )
    return result


@router.post("/{agent_id}/stacks/{stack_id}/{action}")
async def agent_lifecycle(
    agent_id: int,
    stack_id: str,
    action: str,
    request: Request,
    session: Session = Depends(get_session),
    user: User = Depends(require_admin),
) -> dict:
    if action not in _ACTIONS:
        raise HTTPException(status_code=400, detail=f"Unknown action '{action}'")
    agent = _get_or_404(session, agent_id)
    result = await _proxy(session, agent, "POST", f"/agent/stacks/{stack_id}/{action}")
    audit_service.record(
        session, user=user.username, action=f"agent.stack.{action}",
        target=f"{agent.name}/{stack_id}", ip=_ip(request),
    )
    return result