deer-flow/backend/tests/test_provisioner_request_threading.py

"""Regression tests for provisioner request-path K8s IO threading."""

from __future__ import annotations

import asyncio
import inspect
import threading
import time
from contextlib import contextmanager
from types import SimpleNamespace

import httpx
import pytest
from blockbuster import BlockBuster
from kubernetes.client.rest import ApiException


class _RecordingCoreV1:
    def __init__(
        self,
        *,
        event_loop_thread_id: int,
        ready_after_service_reads: dict[str, int] | None = None,
        service_read_failures: dict[str, list[int]] | None = None,
    ) -> None:
        self.event_loop_thread_id = event_loop_thread_id
        self.thread_ids: list[int] = []
        self.service_sandboxes: set[str] = {"sandbox-existing"}
        self.ready_after_service_reads = ready_after_service_reads or {}
        self.service_read_failures = service_read_failures or {}
        self.service_read_counts: dict[str, int] = {}
        self.created_pods: list[str] = []
        self.created_pod_specs: dict[str, object] = {}
        self.created_services: list[str] = []

    def _record_k8s_call(self) -> None:
        thread_id = threading.get_ident()
        self.thread_ids.append(thread_id)
        time.sleep(0)
        if thread_id == self.event_loop_thread_id:
            raise AssertionError("Kubernetes client call ran on the ASGI event-loop thread")
        try:
            asyncio.get_running_loop()
        except RuntimeError:
            return
        raise AssertionError("Kubernetes client call ran inside an asyncio event loop")

    def read_namespaced_service(self, _name: str, _namespace: str):
        self._record_k8s_call()
        sandbox_id = _sandbox_id_from_service_name(_name)
        self.service_read_counts[sandbox_id] = self.service_read_counts.get(sandbox_id, 0) + 1
        failures = self.service_read_failures.get(sandbox_id) or []
        if failures:
            raise ApiException(status=failures.pop(0))
        ready_after_reads = self.ready_after_service_reads.get(sandbox_id, 1)
        if sandbox_id not in self.service_sandboxes or self.service_read_counts[sandbox_id] < ready_after_reads:
            raise ApiException(status=404)
        return _node_port_service(sandbox_id)

    def read_namespaced_pod(self, _name: str, _namespace: str):
        self._record_k8s_call()
        return SimpleNamespace(status=SimpleNamespace(phase="Running"))

    def create_namespaced_pod(self, _namespace: str, pod) -> None:
        self._record_k8s_call()
        sandbox_id = pod.metadata.labels["sandbox-id"]
        self.created_pods.append(sandbox_id)
        self.created_pod_specs[sandbox_id] = pod

    def create_namespaced_service(self, _namespace: str, service) -> None:
        self._record_k8s_call()
        sandbox_id = service.metadata.labels["sandbox-id"]
        self.created_services.append(sandbox_id)
        self.service_sandboxes.add(sandbox_id)

    def delete_namespaced_service(self, _name: str, _namespace: str) -> None:
        self._record_k8s_call()

    def delete_namespaced_pod(self, _name: str, _namespace: str) -> None:
        self._record_k8s_call()

    def list_namespaced_service(self, _namespace: str, *, label_selector: str):
        self._record_k8s_call()
        assert label_selector == "app=deer-flow-sandbox"
        return SimpleNamespace(items=[_node_port_service("sandbox-listed")])


def _node_port_service(sandbox_id: str):
    return SimpleNamespace(
        metadata=SimpleNamespace(labels={"sandbox-id": sandbox_id}),
        spec=SimpleNamespace(ports=[SimpleNamespace(name="http", port=8080, node_port=32123)]),
    )


def _sandbox_id_from_service_name(name: str) -> str:
    assert name.startswith("sandbox-")
    assert name.endswith("-svc")
    return name[len("sandbox-") : -len("-svc")]


@contextmanager
def _detect_provisioner_blocking_io(provisioner_module):
    detector = BlockBuster(scanned_modules=[provisioner_module.__name__])
    detector.activate()
    try:
        yield
    finally:
        detector.deactivate()


def test_sandbox_business_route_handlers_are_sync(provisioner_module) -> None:
    """FastAPI runs sync handlers in its worker pool, away from the event loop."""
    for handler in (
        provisioner_module.create_sandbox,
        provisioner_module.destroy_sandbox,
        provisioner_module.get_sandbox,
        provisioner_module.list_sandboxes,
    ):
        assert not inspect.iscoroutinefunction(handler)


@pytest.mark.asyncio
@pytest.mark.parametrize(
    ("method", "path", "json_body", "expected_created_sandbox"),
    [
        ("POST", "/api/sandboxes", {"sandbox_id": "sandbox-existing", "thread_id": "thread-1", "user_id": "user-1"}, None),
        ("POST", "/api/sandboxes", {"sandbox_id": "sandbox-new", "thread_id": "thread-1", "user_id": "user-1"}, "sandbox-new"),
        ("DELETE", "/api/sandboxes/sandbox-existing", None, None),
        ("GET", "/api/sandboxes/sandbox-existing", None, None),
        ("GET", "/api/sandboxes", None, None),
    ],
    ids=["create-existing", "create-new", "destroy", "get", "list"],
)
async def test_sandbox_business_routes_run_k8s_client_off_event_loop_thread(
    method: str,
    path: str,
    json_body: dict[str, str] | None,
    expected_created_sandbox: str | None,
    monkeypatch: pytest.MonkeyPatch,
    provisioner_module,
) -> None:
    fake_core_v1 = _RecordingCoreV1(
        event_loop_thread_id=threading.get_ident(),
        ready_after_service_reads={"sandbox-new": 3},
    )
    monkeypatch.setattr(provisioner_module, "core_v1", fake_core_v1)
    monkeypatch.setattr(provisioner_module, "PROVISIONER_API_KEY", "test-secret")

    with _detect_provisioner_blocking_io(provisioner_module):
        transport = httpx.ASGITransport(app=provisioner_module.app)
        async with httpx.AsyncClient(transport=transport, base_url="http://testserver") as client:
            headers = {"X-API-Key": "test-secret"}
            if json_body is None:
                response = await client.request(method, path, headers=headers)
            else:
                response = await client.request(method, path, json=json_body, headers=headers)

    assert response.status_code == 200
    assert fake_core_v1.thread_ids
    if expected_created_sandbox is not None:
        assert fake_core_v1.created_pods == [expected_created_sandbox]
        assert fake_core_v1.created_services == [expected_created_sandbox]


@pytest.mark.parametrize(
    ("include_legacy_skills", "expected_mount_names"),
    [
        (
            False,
            ["skills-public", "skills-custom", "user-data"],
        ),
        (
            True,
            ["skills-public", "skills-custom", "skills-legacy", "user-data"],
        ),
    ],
    ids=["without-legacy", "with-legacy"],
)
def test_create_sandbox_route_builds_expected_skills_mount_layout(
    include_legacy_skills: bool,
    expected_mount_names: list[str],
    monkeypatch: pytest.MonkeyPatch,
    provisioner_module,
) -> None:
    fake_core_v1 = _RecordingCoreV1(
        event_loop_thread_id=-1,
        ready_after_service_reads={"sandbox-layout": 1},
    )
    monkeypatch.setattr(provisioner_module, "core_v1", fake_core_v1)

    response = provisioner_module.create_sandbox(
        provisioner_module.CreateSandboxRequest(
            sandbox_id="sandbox-layout",
            thread_id="thread-1",
            user_id="user-1",
            include_legacy_skills=include_legacy_skills,
        )
    )

    assert response.status == "Running"
    pod = fake_core_v1.created_pod_specs["sandbox-layout"]
    volume_names = [volume.name for volume in pod.spec.volumes]
    mount_names = [mount.name for mount in pod.spec.containers[0].volume_mounts]
    assert volume_names == expected_mount_names
    assert mount_names == expected_mount_names


def test_create_sandbox_retries_transient_service_read_errors(monkeypatch: pytest.MonkeyPatch, provisioner_module) -> None:
    fake_core_v1 = _RecordingCoreV1(
        event_loop_thread_id=-1,
        ready_after_service_reads={"sandbox-transient": 3},
        service_read_failures={"sandbox-transient": [503, 429]},
    )
    monkeypatch.setattr(provisioner_module, "core_v1", fake_core_v1)
    monkeypatch.setattr(provisioner_module.time, "sleep", lambda _seconds: None)

    response = provisioner_module.create_sandbox(
        provisioner_module.CreateSandboxRequest(
            sandbox_id="sandbox-transient",
            thread_id="thread-1",
            user_id="user-1",
        )
    )

    assert response.status == "Running"
    assert response.sandbox_url == provisioner_module._sandbox_url("sandbox-transient", node_port=32123)
    assert fake_core_v1.service_read_counts["sandbox-transient"] == 3


def test_sandbox_service_defaults_to_node_port_with_node_host_url(provisioner_module) -> None:
    provisioner_module.K8S_NAMESPACE = "mdv-sit"
    provisioner_module.SANDBOX_CONTAINER_PORT = 8080
    provisioner_module.SANDBOX_SERVICE_TYPE = "NodePort"
    provisioner_module.NODE_HOST = "node.example"

    service = provisioner_module._build_service("abc123")

    assert service.spec.type == "NodePort"
    assert service.spec.ports[0].port == 8080
    assert service.spec.ports[0].target_port == 8080
    assert provisioner_module._sandbox_url("abc123", node_port=32123) == "http://node.example:32123"


def test_sandbox_service_supports_cluster_ip_with_dns_url(provisioner_module) -> None:
    provisioner_module.K8S_NAMESPACE = "mdv-sit"
    provisioner_module.SANDBOX_CONTAINER_PORT = 8080
    provisioner_module.SANDBOX_SERVICE_TYPE = "ClusterIP"

    service = provisioner_module._build_service("abc123")

    assert service.spec.type == "ClusterIP"
    assert service.spec.ports[0].port == 8080
    assert service.spec.ports[0].target_port == 8080
    assert provisioner_module._sandbox_url("abc123") == ("http://sandbox-abc123-svc.mdv-sit.svc.cluster.local:8080")


@pytest.mark.asyncio
async def test_auth_middleware(monkeypatch: pytest.MonkeyPatch, provisioner_module) -> None:
    """Verify the X-API-Key middleware: /health is open; /api/* requires a correct key."""
    monkeypatch.setattr(provisioner_module, "PROVISIONER_API_KEY", "test-secret")
    fake_core_v1 = _RecordingCoreV1(event_loop_thread_id=-1)
    monkeypatch.setattr(provisioner_module, "core_v1", fake_core_v1)

    transport = httpx.ASGITransport(app=provisioner_module.app)
    async with httpx.AsyncClient(transport=transport, base_url="http://testserver") as client:
        # /health is always open — no key needed
        r = await client.get("/health")
        assert r.status_code == 200

        # /api/* with no header → 401
        r = await client.get("/api/sandboxes")
        assert r.status_code == 401

        # /api/* with wrong key → 401
        r = await client.get("/api/sandboxes", headers={"X-API-Key": "wrong-key"})
        assert r.status_code == 401

        # /api/* with correct key → not 401 (auth passed; handler runs with the K8s mock)
        r = await client.get("/api/sandboxes", headers={"X-API-Key": "test-secret"})
        assert r.status_code != 401


@pytest.mark.asyncio
async def test_auth_middleware_unset_key(monkeypatch: pytest.MonkeyPatch, provisioner_module) -> None:
    """When PROVISIONER_API_KEY is unset/empty, all /api/* routes return 401."""
    monkeypatch.setattr(provisioner_module, "PROVISIONER_API_KEY", "")
    fake_core_v1 = _RecordingCoreV1(event_loop_thread_id=-1)
    monkeypatch.setattr(provisioner_module, "core_v1", fake_core_v1)

    transport = httpx.ASGITransport(app=provisioner_module.app)
    async with httpx.AsyncClient(transport=transport, base_url="http://testserver") as client:
        # /health is always open even when key is unset
        r = await client.get("/health")
        assert r.status_code == 200

        # /api/* is always 401 when key is unset — even with a header
        r = await client.get("/api/sandboxes")
        assert r.status_code == 401

        r = await client.get("/api/sandboxes", headers={"X-API-Key": "anything"})
        assert r.status_code == 401