agentic-mobile-control/tests/test_cloud_planner_decision_endpoint.py

from __future__ import annotations

from fastapi import FastAPI
from fastapi.testclient import TestClient

from cloud.auth import BearerCredential, ConfiguredBearerAuthProvider
from cloud.config import CloudConfig
from cloud.internal_api.api import create_internal_router
from cloud.pool import DevicePool
from cloud.store import CloudStore
from runtime.tool_calling_client import ToolCallDecision, ToolCallUnavailable
from runtime.tool_specs import ToolSpec


class _FakeToolCallingClient:
    """Stand-in ``ToolCallingClient`` for endpoint tests -- never calls a real
    LLM provider."""

    def __init__(
        self, decision: ToolCallDecision | None = None, error: str | None = None
    ):
        self.decision = decision
        self.error = error
        self.calls: list[dict[str, object]] = []

    def decide(
        self,
        *,
        system_prompt: str,
        user_prompt: str,
        screenshot: bytes | None,
        tools: list[ToolSpec],
        timeout: float,
    ) -> ToolCallDecision:
        self.calls.append(
            {
                "system_prompt": system_prompt,
                "user_prompt": user_prompt,
                "screenshot": screenshot,
                "tools": tools,
                "timeout": timeout,
            }
        )
        if self.error is not None:
            raise ToolCallUnavailable(self.error)
        assert self.decision is not None
        return self.decision


def _build_client(
    tmp_path, *, fake_client: _FakeToolCallingClient
) -> tuple[TestClient, _FakeToolCallingClient]:
    pool = DevicePool(
        CloudStore(tmp_path / "internal.sqlite3"),
        CloudConfig(stale_after_seconds=60),
    )
    auth_provider = ConfiguredBearerAuthProvider(
        [
            BearerCredential(principal_id="agent-a", token="token-a", host_id="host-a"),
            BearerCredential(principal_id="agent-b", token="token-b", host_id="host-b"),
        ]
    )
    app = FastAPI()
    app.include_router(
        create_internal_router(
            pool=pool,
            auth_provider=auth_provider,
            planner_client_factory=lambda: fake_client,
        )
    )
    return TestClient(app), fake_client


def _decision_payload(**overrides: object) -> dict[str, object]:
    payload = {
        "host_id": "host-a",
        "system_prompt": "you are a planner",
        "user_prompt": "tap the login button",
        "screenshot_base64": None,
        "tools": [
            {
                "name": "tap",
                "description": "tap an element",
                "parameters": {"type": "object", "properties": {}},
            }
        ],
        "timeout_seconds": 30.0,
    }
    payload.update(overrides)
    return payload


def test_authenticated_host_resolves_planner_decision(tmp_path) -> None:
    fake_client = _FakeToolCallingClient(
        decision=ToolCallDecision(tool_name="tap", arguments={"x": 1, "y": 2})
    )
    client, fake_client = _build_client(tmp_path, fake_client=fake_client)

    response = client.post(
        "/internal/v1/hosts/host-a/planner/decide",
        headers={"Authorization": "Bearer token-a"},
        json=_decision_payload(),
    )

    assert response.status_code == 200
    assert response.json() == {"tool_name": "tap", "arguments": {"x": 1, "y": 2}}
    assert len(fake_client.calls) == 1
    assert fake_client.calls[0]["system_prompt"] == "you are a planner"
    assert fake_client.calls[0]["timeout"] == 30.0


def test_planner_decision_decodes_screenshot_base64(tmp_path) -> None:
    fake_client = _FakeToolCallingClient(
        decision=ToolCallDecision(tool_name="tap", arguments={})
    )
    client, fake_client = _build_client(tmp_path, fake_client=fake_client)

    response = client.post(
        "/internal/v1/hosts/host-a/planner/decide",
        headers={"Authorization": "Bearer token-a"},
        json=_decision_payload(screenshot_base64="aGVsbG8="),
    )

    assert response.status_code == 200
    assert fake_client.calls[0]["screenshot"] == b"hello"


def test_invalid_screenshot_base64_is_rejected(tmp_path) -> None:
    fake_client = _FakeToolCallingClient(
        decision=ToolCallDecision(tool_name="tap", arguments={})
    )
    client, fake_client = _build_client(tmp_path, fake_client=fake_client)

    response = client.post(
        "/internal/v1/hosts/host-a/planner/decide",
        headers={"Authorization": "Bearer token-a"},
        json=_decision_payload(screenshot_base64="not-valid-base64!!"),
    )

    assert response.status_code == 422
    assert fake_client.calls == []


def test_unauthenticated_request_is_rejected(tmp_path) -> None:
    fake_client = _FakeToolCallingClient(
        decision=ToolCallDecision(tool_name="tap", arguments={})
    )
    client, fake_client = _build_client(tmp_path, fake_client=fake_client)

    response = client.post(
        "/internal/v1/hosts/host-a/planner/decide",
        json=_decision_payload(),
    )

    assert response.status_code == 401
    assert fake_client.calls == []


def test_foreign_host_token_cannot_request_another_hosts_decision(tmp_path) -> None:
    fake_client = _FakeToolCallingClient(
        decision=ToolCallDecision(tool_name="tap", arguments={})
    )
    client, fake_client = _build_client(tmp_path, fake_client=fake_client)

    response = client.post(
        "/internal/v1/hosts/host-a/planner/decide",
        headers={"Authorization": "Bearer token-b"},
        json=_decision_payload(),
    )

    assert response.status_code == 403
    assert fake_client.calls == []


def test_path_and_payload_host_id_mismatch_is_rejected(tmp_path) -> None:
    fake_client = _FakeToolCallingClient(
        decision=ToolCallDecision(tool_name="tap", arguments={})
    )
    client, fake_client = _build_client(tmp_path, fake_client=fake_client)

    response = client.post(
        "/internal/v1/hosts/host-a/planner/decide",
        headers={"Authorization": "Bearer token-a"},
        json=_decision_payload(host_id="host-b"),
    )

    assert response.status_code == 422
    assert fake_client.calls == []


def test_provider_failure_returns_structured_error_without_crashing(tmp_path) -> None:
    fake_client = _FakeToolCallingClient(error="anthropic timed out")
    client, fake_client = _build_client(tmp_path, fake_client=fake_client)

    response = client.post(
        "/internal/v1/hosts/host-a/planner/decide",
        headers={"Authorization": "Bearer token-a"},
        json=_decision_payload(),
    )

    assert response.status_code == 502
    assert response.json() == {
        "code": "planner_unavailable",
        "detail": "anthropic timed out",
    }