agentic-mobile-control/apps/device-host-agent/tests/test_cloud_planner_client.py

from __future__ import annotations

import json

import httpx
import pytest

from host_agent.cloud_planner_client import (
    _CLOUD_PROXY_HTTP_TIMEOUT_SECONDS,
    CloudProxyToolCallingClient,
)
from host_agent.config import HostAgentConfig
from host_agent.planner_context import PlannerExecutionContext, _context
from runtime.tool_calling_client import ToolCallDecision, ToolCallUnavailable
from runtime.tool_specs import ToolSpec

_CONFIG = HostAgentConfig(
    control_plane_url="https://control-plane.example",
    host_id="host-a",
    token="token-a",
)

_TOOLS = [ToolSpec(name="tap", description="tap an element", parameters={})]


def _client(handler) -> CloudProxyToolCallingClient:
    transport = httpx.MockTransport(handler)
    http_client = httpx.Client(base_url=_CONFIG.control_plane_url, transport=transport)
    return CloudProxyToolCallingClient(_CONFIG, http_client=http_client)


def test_decide_returns_tool_call_decision_on_success() -> None:
    seen_requests: list[httpx.Request] = []

    def handler(request: httpx.Request) -> httpx.Response:
        seen_requests.append(request)
        return httpx.Response(
            200,
            json={"tool_name": "tap", "arguments": {"x": 1, "y": 2}},
        )

    client = _client(handler)

    decision = client.decide(
        system_prompt="you are a planner",
        user_prompt="tap login",
        screenshot=None,
        tools=_TOOLS,
        timeout=30.0,
    )

    assert decision == ToolCallDecision(tool_name="tap", arguments={"x": 1, "y": 2})
    assert len(seen_requests) == 1
    request = seen_requests[0]
    assert request.url.path == "/internal/v1/hosts/host-a/planner/decide"
    assert request.headers["Authorization"] == "Bearer token-a"
    body = json.loads(request.content)
    assert body["host_id"] == "host-a"
    assert body["system_prompt"] == "you are a planner"
    assert body["screenshot_base64"] is None
    assert body["timeout_seconds"] == 30.0


def test_decide_base64_encodes_screenshot() -> None:
    seen_requests: list[httpx.Request] = []

    def handler(request: httpx.Request) -> httpx.Response:
        seen_requests.append(request)
        return httpx.Response(200, json={"tool_name": "tap", "arguments": {}})

    client = _client(handler)

    client.decide(
        system_prompt="sp",
        user_prompt="up",
        screenshot=b"hello",
        tools=_TOOLS,
        timeout=10.0,
    )

    body = json.loads(seen_requests[0].content)
    assert body["screenshot_base64"] == "aGVsbG8="


def test_decide_clamps_legacy_timeout_and_waits_for_cloud_profile_timeout() -> None:
    seen_requests: list[httpx.Request] = []

    def handler(request: httpx.Request) -> httpx.Response:
        seen_requests.append(request)
        return httpx.Response(200, json={"tool_name": "tap", "arguments": {}})

    client = _client(handler)

    client.decide(
        system_prompt="sp",
        user_prompt="up",
        screenshot=None,
        tools=_TOOLS,
        timeout=180.0,
    )

    request = seen_requests[0]
    body = json.loads(request.content)
    assert body["timeout_seconds"] == 120.0
    assert request.extensions["timeout"]["read"] == _CLOUD_PROXY_HTTP_TIMEOUT_SECONDS


def test_decide_includes_bound_assignment_context() -> None:
    seen_requests: list[httpx.Request] = []

    def handler(request: httpx.Request) -> httpx.Response:
        seen_requests.append(request)
        return httpx.Response(200, json={"tool_name": "tap", "arguments": {}})

    client = _client(handler)
    token = _context.set(
        PlannerExecutionContext(task_id="task-a", attempt=2, lease_id="lease-a")
    )
    try:
        client.decide(
            system_prompt="sp",
            user_prompt="up",
            screenshot=None,
            tools=_TOOLS,
            timeout=10.0,
        )
    finally:
        _context.reset(token)

    body = json.loads(seen_requests[0].content)
    assert body["task_id"] == "task-a"
    assert body["attempt"] == 2
    assert body["lease_id"] == "lease-a"


def test_decide_raises_tool_call_unavailable_on_network_error() -> None:
    def handler(request: httpx.Request) -> httpx.Response:
        raise httpx.ConnectError("connection refused", request=request)

    client = _client(handler)

    with pytest.raises(ToolCallUnavailable):
        client.decide(
            system_prompt="sp",
            user_prompt="up",
            screenshot=None,
            tools=_TOOLS,
            timeout=5.0,
        )


def test_decide_raises_tool_call_unavailable_on_structured_error_response() -> None:
    def handler(request: httpx.Request) -> httpx.Response:
        return httpx.Response(
            502,
            json={"code": "planner_unavailable", "detail": "provider timed out"},
        )

    client = _client(handler)

    with pytest.raises(ToolCallUnavailable, match="provider timed out"):
        client.decide(
            system_prompt="sp",
            user_prompt="up",
            screenshot=None,
            tools=_TOOLS,
            timeout=5.0,
        )


def test_decide_raises_tool_call_unavailable_on_unstructured_error_response() -> None:
    def handler(request: httpx.Request) -> httpx.Response:
        return httpx.Response(500, text="internal server error")

    client = _client(handler)

    with pytest.raises(ToolCallUnavailable):
        client.decide(
            system_prompt="sp",
            user_prompt="up",
            screenshot=None,
            tools=_TOOLS,
            timeout=5.0,
        )


def test_close_closes_an_internally_constructed_http_client() -> None:
    client = CloudProxyToolCallingClient(_CONFIG)
    client.close()

    assert client._client.is_closed


def test_close_leaves_an_externally_supplied_http_client_open() -> None:
    """Only a client built internally (no ``http_client`` override) is closed
    by ``close()``; a caller-supplied ``http_client`` is left open for the
    caller to manage."""
    transport = httpx.MockTransport(lambda request: httpx.Response(200, json={}))
    external_http_client = httpx.Client(
        base_url=_CONFIG.control_plane_url, transport=transport
    )

    client = CloudProxyToolCallingClient(_CONFIG, http_client=external_http_client)
    client.close()

    assert not external_http_client.is_closed
    external_http_client.close()