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

A reply can also carry a voice (`voice_id`/`voice_name`), which is how the
server's `speak_as` marker reaches us: Bolt searched the ElevenLabs voice
library, picked one, and tagged the reply with it — the client is what
actually speaks in it. See `Reply` and controller._apply_voice.

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

from __future__ import annotations

import os
import signal
import subprocess
from pathlib import Path
from typing import Callable, NamedTuple, Optional

import requests

from . import config, sudo_askpass

_MAX_RELAY_HOPS = 16

# Command output handed back up the relay is capped: it becomes part of the
# server's prompt, and a runaway `find /` would blow the context window.
_MAX_COMMAND_OUTPUT = 6000


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


class Reply(NamedTuple):
    """One final reply from the desk API. *voice_id* is set only when the
    server tagged this reply with a `speak_as` voice; *voice_name* is the
    human-readable name that came with it (may be empty even when the id
    isn't). Both empty means "say it in the usual voice"."""

    text: str
    voice_id: str = ""
    voice_name: str = ""


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:
        # start_new_session puts the shell in its own process group so a timeout
        # can kill the whole tree. subprocess.run() would only SIGKILL the `sh`
        # itself, leaving whatever it spawned (a build, a `tail -f`, an ffmpeg)
        # running forever with no parent watching — one relayed command that
        # hangs shouldn't leak a process for the rest of the session.
        process = subprocess.Popen(
            command, shell=True, cwd=str(Path.home()), env=env, text=True,
            stdout=subprocess.PIPE, stderr=subprocess.PIPE,
            start_new_session=(os.name == "posix"),
        )
    except Exception as exc:
        return f"[command failed: {exc}]"

    try:
        stdout, stderr = process.communicate(timeout=timeout)
        return _command_output(f"[exit {process.returncode}]", stdout, stderr)
    except subprocess.TimeoutExpired:
        stdout, stderr = _terminate(process)
        # Whatever it managed to print before it hung is the useful part — a
        # bare "timed out" tells the model nothing it can act on, and the last
        # line of output usually says exactly what it was stuck waiting for.
        return _command_output(f"[command timed out after {timeout}s]", stdout, stderr)
    except Exception as exc:
        _terminate(process)
        return f"[command failed: {exc}]"


def _terminate(process: subprocess.Popen) -> tuple[str, str]:
    """Kill a timed-out command's whole process group and collect what it wrote.

    SIGTERM first so a shell script can clean up, SIGKILL a moment later for
    anything that ignores it. The final drain is itself time-boxed: a
    grandchild holding the pipe open must not turn a timeout into a hang."""
    try:
        if os.name == "posix":
            group = os.getpgid(process.pid)
            os.killpg(group, signal.SIGTERM)
            try:
                process.wait(timeout=2)
            except subprocess.TimeoutExpired:
                os.killpg(group, signal.SIGKILL)
        else:
            process.kill()
    except (ProcessLookupError, PermissionError, OSError):
        pass  # already gone, or never had its own group
    try:
        return process.communicate(timeout=2)
    except Exception:
        return "", ""


def _command_output(header: str, stdout: Optional[str], stderr: Optional[str]) -> str:
    body = (stdout or "") + (stderr or "")
    return f"{header}\n{body}"[:_MAX_COMMAND_OUTPUT]


def converse(
    text: str,
    on_command: Callable[[str], str] = run_local_command,
    timeout: float = 120.0,
) -> Reply:
    """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 Reply(
            text=str(payload.get("text") or ""),
            voice_id=str(payload.get("voice_id") or ""),
            voice_name=str(payload.get("voice_name") or ""),
        )
    if payload.get("type") == "command":
        # Fell out of the loop still being handed commands. Worth its own
        # message: "unknown server response" sent everyone looking at the
        # payload shape, when what actually happened is a model that kept
        # calling tools and never answered.
        raise ServerError(
            f"the server kept relaying commands past the {_MAX_RELAY_HOPS}-hop cap "
            "without producing a reply"
        )
    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