Bolt-Pet/bolt_pet/server_client.py

"""HTTP client for Bolt's desk API (ai/desk_api.py on the server).

Protocol is identical to desk_client/bolt_desk.py in the main tmn-api repo —
this pet is just another desk client, so it gets the exact same brain,
memory, tools, and persona as Discord chat and the Linux voice client:

    text  -> POST /desk/converse
    [server may relay a shell command back to run on THIS machine]
    ...   -> POST /desk/tool_result (repeat until the server sends a reply)
    reply <- returned to caller

Kept dependency-free beyond `requests` so it's easy to unit test with mocks.
"""

from __future__ import annotations

import subprocess
from pathlib import Path
from typing import Callable, Optional

import requests

from . import config, sudo_askpass

_MAX_RELAY_HOPS = 16


class ServerError(Exception):
    """Raised when the server responds with an error payload or unreachable."""


def _headers() -> dict:
    return {"X-Desk-Api-Key": config.API_KEY}


def check_health(timeout: float = 10.0) -> dict:
    response = requests.get(f"{config.SERVER_URL}/desk/health", headers=_headers(), timeout=timeout)
    response.raise_for_status()
    return response.json()


def run_local_command(command: str, timeout: int = None) -> str:
    """Execute a command relayed by the server, exactly as bolt_desk.py does —
    "full desktop control" for things like "open firefox" or "how full is my
    disk". Runs as the current desktop user. See README security notes.

    `sudo` gets special handling: the pet has no terminal, so sudo would sit
    waiting on a tty that nobody is looking at. With SUDO_ASKPASS_PROMPT on,
    bare `sudo` becomes `sudo -A` and the password is collected in a desktop
    dialog you have to answer — which also gives those commands a longer
    timeout, since a human has to notice the window and type."""
    env = None
    if config.SUDO_ASKPASS_PROMPT and sudo_askpass.needs_password_prompt(command):
        helper = sudo_askpass.find_helper()
        if helper:
            env = sudo_askpass.environment(helper)
            command = sudo_askpass.add_askpass_flag(command)
            timeout = timeout or config.SUDO_COMMAND_TIMEOUT_SECONDS
    timeout = timeout or config.COMMAND_TIMEOUT_SECONDS
    try:
        completed = subprocess.run(
            command, shell=True, capture_output=True, text=True,
            timeout=timeout, cwd=str(Path.home()), env=env,
        )
        output = (completed.stdout or "") + (completed.stderr or "")
        return f"[exit {completed.returncode}]\n{output}"[:6000]
    except subprocess.TimeoutExpired:
        return f"[command timed out after {timeout}s]"
    except Exception as exc:
        return f"[command failed: {exc}]"


def converse(
    text: str,
    on_command: Callable[[str], str] = run_local_command,
    timeout: float = 120.0,
) -> str:
    """Send one turn of conversation to the desk API, relaying any commands
    the server sends back until it produces a final reply.

    *on_command* is injectable for tests; defaults to actually running the
    command locally (matching bolt_desk.py's behavior).
    """
    headers = _headers()
    try:
        response = requests.post(
            f"{config.SERVER_URL}/desk/converse",
            json={"session_id": config.SESSION_ID, "text": text},
            headers=headers, timeout=timeout,
        )
        payload = response.json()
    except Exception as exc:
        raise ServerError(f"couldn't reach the server: {exc}") from exc

    for _ in range(_MAX_RELAY_HOPS):
        if payload.get("type") != "command":
            break
        output = on_command(str(payload.get("command") or ""))
        try:
            response = requests.post(
                f"{config.SERVER_URL}/desk/tool_result",
                json={
                    "session_id": config.SESSION_ID,
                    "token": payload.get("token"),
                    "output": output,
                },
                headers=headers, timeout=180,
            )
            payload = response.json()
        except Exception as exc:
            raise ServerError(f"couldn't reach the server during tool relay: {exc}") from exc

    if payload.get("type") == "reply":
        return str(payload.get("text") or "")
    raise ServerError(str(payload.get("error") or "unknown server response"))


def list_outbox_files(timeout: float = 15.0) -> list:
    """Files the server has queued for this session via its deliver_files
    tool (e.g. "send me that report" during a conversation) — each entry has
    id/name/size. Downloading one (download_outbox_file) dequeues it
    server-side, so a file is only ever handed out once."""
    try:
        response = requests.get(
            f"{config.SERVER_URL}/desk/files",
            params={"session_id": config.SESSION_ID},
            headers=_headers(), timeout=timeout,
        )
        response.raise_for_status()
        return list(response.json().get("files") or [])
    except Exception as exc:
        raise ServerError(f"couldn't list delivered files: {exc}") from exc


def download_outbox_file(file_id: str, timeout: float = 60.0) -> bytes:
    """Fetches and dequeues one file listed by list_outbox_files()."""
    try:
        response = requests.get(
            f"{config.SERVER_URL}/desk/files/{file_id}",
            params={"session_id": config.SESSION_ID},
            headers=_headers(), timeout=timeout,
        )
        response.raise_for_status()
        return response.content
    except Exception as exc:
        raise ServerError(f"couldn't download delivered file {file_id!r}: {exc}") from exc


def report_status(timeout: float = 15.0) -> Optional[str]:
    """Heartbeat — lets the desk API attach a pending spoken announcement
    (proactive nudges, reminders fired since the last heartbeat) that the pet
    can speak unprompted, exactly like the phone/desk clients. A pet has no
    battery/GPS to report, so the device-status fields
    (battery/is_charging/latitude/longitude/address, all optional server-side)
    are simply omitted.

    Returns the announcement text to speak, or None if there's nothing pending.
    """
    try:
        response = requests.post(
            f"{config.SERVER_URL}/desk/report_status",
            json={"session_id": config.SESSION_ID}, headers=_headers(), timeout=timeout,
        )
        response.raise_for_status()
        data = response.json()
    except Exception as exc:
        raise ServerError(f"heartbeat failed: {exc}") from exc
    reply = data.get("reply")
    return str(reply) if reply else None