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 json
import subprocess
from pathlib import Path
from typing import Callable, NamedTuple, 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."""


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 = ""
    # True when the sentences were already spoken as they streamed in. The
    # text is still carried — the follow-up rule needs to see whether the
    # answer ended on a question — it just must not be read out again.
    spoken: bool = False


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,
    on_say: Optional[Callable[[str], None]] = None,
) -> 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).

    *on_say* is called with a short holding line ("give me a sec") when the
    server sends one alongside a command. It is the difference between silence
    and an answer while a tool runs: the model's acknowledgement used to be
    discarded server-side, so the whole round trip was dead air and the model
    then repeated itself in the final reply. Optional, so an older pet against
    a newer server simply stays quiet as before.
    """
    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
        holding = str(payload.get("say") or "").strip()
        if holding and on_say is not None:
            # Spoken *before* the command runs — that is the whole point.
            try:
                on_say(holding)
            except Exception:
                pass  # a failed acknowledgement must not cost the tool call
        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 ""),
        )
    raise ServerError(str(payload.get("error") or "unknown server response"))


def converse_stream(
    text: str,
    on_say: Callable[[str], None],
    on_command: Callable[[str], str] = run_local_command,
    timeout: float = 180.0,
) -> Reply:
    """Same turn as converse(), but speaking each sentence as it arrives.

    Without this the pet waits out the *entire* model call before a single
    word is heard; with it the wait is time-to-first-sentence, which on a
    multi-sentence answer is most of the difference.

    Falls back by raising ServerError before anything has been spoken — the
    caller then retries the ordinary path and the user never finds out. Once a
    sentence *has* been spoken there is no going back, so late failures end the
    turn with whatever was said rather than repeating it.
    """
    spoke_anything = False
    try:
        response = requests.post(
            f"{config.SERVER_URL}/desk/converse_stream",
            json={"session_id": config.SESSION_ID, "text": text},
            headers=_headers(), timeout=timeout, stream=True,
        )
        response.raise_for_status()
        for raw in response.iter_lines(decode_unicode=True):
            if not raw:
                continue
            try:
                event = json.loads(raw)
            except (TypeError, ValueError):
                continue
            kind = str(event.get("type") or "")

            if kind == "say":
                line = str(event.get("text") or "").strip()
                if line:
                    spoke_anything = True
                    on_say(line)
            elif kind == "command":
                # The tool loop is request/response, so the rest of the turn
                # finishes through the ordinary relay rather than inside the
                # stream — one protocol for tools, not two.
                response.close()
                return _finish_relay(event, on_command, on_say)
            elif kind == "reply":
                return Reply(
                    text=str(event.get("text") or ""),
                    voice_id=str(event.get("voice_id") or ""),
                    voice_name=str(event.get("voice_name") or ""),
                    spoken=bool(event.get("already_spoken")),
                )
            elif kind == "error":
                raise ServerError(str(event.get("error") or "stream failed"))
    except ServerError:
        raise
    except Exception as exc:
        if spoke_anything:
            # Half a reply is out loud already; ending quietly beats saying it
            # all again through the fallback path.
            return Reply(text="", spoken=True)
        raise ServerError(f"streaming failed: {exc}") from exc
    raise ServerError("stream ended without a reply")


def _finish_relay(event: dict, on_command, on_say) -> Reply:
    """Run the tool the stream handed over, then continue the classic relay."""
    payload = dict(event)
    headers = _headers()
    for _ in range(_MAX_RELAY_HOPS):
        if payload.get("type") != "command":
            break
        holding = str(payload.get("say") or "").strip()
        if holding and on_say is not None:
            try:
                on_say(holding)
            except Exception:
                pass
        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 ""),
        )
    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