tests/test_http_pull_worker.py

"""The outbound HTTP worker carries attested, strict results through the queue."""

import asyncio
import io
import json
import os
import time
import unittest
from concurrent.futures import ThreadPoolExecutor
from unittest.mock import patch

import httpx
from fastapi.testclient import TestClient

from app import create_app
from contract import MODEL
from direct_gateway import ARTIFACT_RESPONSE_HEADERS
from http_pull_worker import HTTPRuntime, HTTPWorker, runtime_from_environment
from model_registry import MODEL_ORDER, PROFILES
from pull_worker import Gateway, GatewayError
from relay import Relay

TOKEN = "t" * 32
MANIFEST = "a" * 64
REVISION = "b" * 40
CONTENT = "c" * 64
CANONICAL = PROFILES[MODEL]["repo_id"]
# The deployed Decision 1.0 runtime attests its pre-move Hub ID.
ATTESTED = PROFILES[MODEL]["runtime_model"]
ORIGIN = "http://127.0.0.1:18401"
ARTIFACT = {
    "repo_id": CANONICAL,
    "manifest_sha256": MANIFEST,
    "revision": REVISION,
    "content_sha256": CONTENT,
}


class NoTetris:
    def public_config(self):
        return {"competitors": []}


class LocalGateway:
    def __init__(self, client):
        self.client = client

    def post(self, action, payload):
        response = self.client.post(
            "/internal/worker/" + action,
            json=payload,
            headers={"Authorization": "Bearer " + TOKEN},
        )
        if response.status_code >= 400:
            raise GatewayError(response.status_code)
        return response.json()


def payload(*, batch=False, model=MODEL):
    request = {
        "model": model,
        "questions": {"route": {
            "type": "choice",
            "instructions": "Choose a route.",
            "criteria": {"left": None, "right": None},
        }},
    }
    if batch:
        request["states"] = [
            {"id": "one", "state": "First request"},
            {"id": "two", "state": "Second request"},
        ]
    else:
        request["state"] = "One request"
    return request


def answer():
    return {
        "type": "choice",
        "choice": "left",
        "confidence": 0.4,
        "probabilities": {"left": 0.7, "right": 0.3},
    }


def strict_response(request):
    usage = {"input_tokens": 12, "output_tokens": 0}
    if "states" in request:
        return {
            "model": request["model"],
            "results": [
                {"id": row["id"], "answers": {"route": answer()}, "usage": usage}
                for row in request["states"]
            ],
            "usage": {"input_tokens": 12 * len(request["states"]), "output_tokens": 0},
        }
    return {"model": request["model"], "answers": {"route": answer()}, "usage": usage}


class HTTPPullWorkerTests(unittest.TestCase):
    def setUp(self):
        self.calls = []
        self.bad_header = False
        self.bad_response = False
        self.offline = False

        def handler(request):
            if request.method == "GET" and request.url.path == "/api/status":
                return httpx.Response(200, json={
                    "status": "offline" if self.offline else "ready",
                    "artifact": {
                        "model": ATTESTED,
                        "revision": REVISION,
                        "manifest_sha256": MANIFEST,
                        "content_sha256": CONTENT,
                    },
                })
            body = json.loads(request.content)
            self.calls.append((request.url.path, body))
            response = strict_response(body)
            if self.bad_response:
                response["model"] = "another/model"
            headers = {
                header: {
                    "model": ATTESTED,
                    "revision": REVISION,
                    "manifest_sha256": "d" * 64 if self.bad_header else MANIFEST,
                    "content_sha256": CONTENT,
                }[field]
                for field, header in ARTIFACT_RESPONSE_HEADERS.items()
            }
            return httpx.Response(200, json=response, headers=headers)

        self.runtime_client = httpx.AsyncClient(transport=httpx.MockTransport(handler))
        self.runtime = HTTPRuntime(MODEL, ORIGIN, ARTIFACT, client=self.runtime_client)
        self.app_context = TestClient(create_app(
            mode="pull_queue",
            relay=Relay(TOKEN, MANIFEST, model=MODEL),
            registry=[{
                "id": MODEL,
                "label": "Kai",
                "version": "1.0",
                "manifest_sha256": MANIFEST,
            }],
            tetris_manager=NoTetris(),
        ))
        self.client = self.app_context.__enter__()
        self.worker = HTTPWorker(LocalGateway(self.client), self.runtime)
        self.runtime._load()
        self.worker.heartbeat()

    def tearDown(self):
        self.runtime.close()
        asyncio.run(self.runtime_client.aclose())
        self.app_context.__exit__(None, None, None)

    def process_one(self):
        claimed = self.worker.post("claim", {
            "worker_id": self.worker.worker_id,
            "wait_seconds": 0,
        })["job"]
        self.assertIsNotNone(claimed)
        self.worker.execute(claimed)
        return claimed

    def test_single_and_batch_jobs_preserve_canonical_runtime_response(self):
        for batch, expected_path in (
            (False, "/v1/systemone"),
            (True, "/v1/systemone/batches"),
        ):
            with self.subTest(batch=batch):
                submitted = self.client.post("/api/jobs", json=payload(batch=batch))
                self.assertEqual(submitted.status_code, 202)
                self.process_one()
                completed = self.client.get("/api/jobs/" + submitted.json()["id"])
                self.assertEqual(completed.status_code, 200)
                self.assertEqual(completed.json()["status"], "succeeded")
                self.assertEqual(completed.json()["result"],
                                 strict_response(payload(batch=batch, model=CANONICAL)))
                self.assertEqual(self.calls[-1],
                                 (expected_path, payload(batch=batch, model=ATTESTED)))

    def test_public_synchronous_route_returns_strict_result(self):
        with ThreadPoolExecutor(max_workers=1) as executor:
            pending = executor.submit(
                self.client.post, "/v1/systemone", json=payload(model=CANONICAL)
            )
            for _ in range(200):
                if self.client.get("/api/status").json()["queued"]:
                    break
                time.sleep(0.01)
            else:
                self.fail("The public request was not queued")
            self.process_one()
            response = pending.result(timeout=5)
        self.assertEqual(response.status_code, 200)
        self.assertEqual(response.json(), strict_response(payload(model=CANONICAL)))

    def test_studio_batch_routes_admit_over_256_kib_but_single_stays_bounded(self):
        large_state = "x" * (300 * 1024)
        batch_request = payload(batch=True)
        batch_request["states"] = [{"id": "large", "state": large_state}]
        single_request = payload()
        single_request["state"] = large_state
        for route in ("/api/jobs", "/api/evaluate"):
            with self.subTest(route=route):
                self.assertEqual(
                    self.client.post(route, json=single_request).status_code, 413
                )
        submitted = self.client.post("/api/jobs", json=batch_request)
        self.assertEqual(submitted.status_code, 202)
        self.process_one()
        completed = self.client.get("/api/jobs/" + submitted.json()["id"]).json()
        self.assertEqual(completed["status"], "succeeded")

        with ThreadPoolExecutor(max_workers=1) as executor:
            pending = executor.submit(
                self.client.post, "/api/evaluate", json=batch_request
            )
            for _ in range(200):
                if self.client.get("/api/status").json()["queued"]:
                    break
                time.sleep(0.01)
            else:
                self.fail("The Studio batch was not queued")
            self.process_one()
            response = pending.result(timeout=5)
        self.assertEqual(response.status_code, 200)
        self.assertEqual(
            response.json(), strict_response(dict(batch_request, model=CANONICAL))
        )

    def test_bad_artifact_or_response_fails_without_prediction(self):
        for field in ("bad_header", "bad_response"):
            with self.subTest(field=field):
                setattr(self, field, True)
                submitted = self.client.post("/api/jobs", json=payload())
                self.assertEqual(submitted.status_code, 202)
                self.process_one()
                completed = self.client.get("/api/jobs/" + submitted.json()["id"]).json()
                self.assertEqual(completed["status"], "failed")
                self.assertNotIn("result", completed)
                setattr(self, field, False)

    def test_offline_runtime_does_not_send_a_fresh_heartbeat(self):
        before = self.client.app.state.relay.worker["seen"]
        self.offline = True
        with self.assertRaises(RuntimeError):
            self.worker.heartbeat()
        self.assertEqual(self.client.app.state.relay.worker["seen"], before)

    def test_space_rejects_mismatched_strict_result(self):
        submitted = self.client.post("/api/jobs", json=payload()).json()
        claimed = self.worker.post("claim", {
            "worker_id": self.worker.worker_id,
            "wait_seconds": 0,
        })["job"]
        wrong = strict_response(payload(model=ATTESTED))
        wrong["model"] = "another/model"
        with self.assertRaises(GatewayError) as rejected:
            self.worker.post("result", {
                "worker_id": self.worker.worker_id,
                "id": claimed["id"],
                "lease_token": claimed["lease_token"],
                "result": {"kind": "http_runtime_v1", "response": wrong},
            })
        self.assertEqual(rejected.exception.code, 422)
        self.worker.execute(claimed)
        completed = self.client.get("/api/jobs/" + submitted["id"]).json()
        self.assertEqual(completed["status"], "succeeded")

    def test_strict_results_accept_reordered_question_maps(self):
        questions = {
            "beta": {"type": "noul", "instructions": "Is beta true?"},
            "alpha": {"type": "noul", "instructions": "Is alpha true?"},
        }
        answers = {
            "alpha": {"type": "noul", "noul": 0.2},
            "beta": {"type": "noul", "noul": 0.8},
        }
        for batch in (False, True):
            with self.subTest(batch=batch):
                request = {"model": MODEL, "questions": questions}
                if batch:
                    request["states"] = [
                        {"id": "first", "state": "First"},
                        {"id": "second", "state": "Second"},
                    ]
                    strict = {
                        "model": ATTESTED,
                        "results": [
                            {"id": row["id"], "answers": answers,
                             "usage": {"input_tokens": 12, "output_tokens": 0}}
                            for row in request["states"]
                        ],
                        "usage": {"input_tokens": 24, "output_tokens": 0},
                    }
                else:
                    request["state"] = "One request"
                    strict = {
                        "model": ATTESTED,
                        "answers": answers,
                        "usage": {"input_tokens": 12, "output_tokens": 0},
                    }
                submitted = self.client.post("/api/jobs", json=request).json()
                claimed = self.worker.post("claim", {
                    "worker_id": self.worker.worker_id,
                    "wait_seconds": 0,
                })["job"]
                self.worker.post("result", {
                    "worker_id": self.worker.worker_id,
                    "id": claimed["id"],
                    "lease_token": claimed["lease_token"],
                    "result": {"kind": "http_runtime_v1", "response": strict},
                })
                completed = self.client.get("/api/jobs/" + submitted["id"]).json()
                self.assertEqual(completed["status"], "succeeded")
                self.assertEqual(completed["result"], dict(strict, model=CANONICAL))


class WorkerConfigurationTests(unittest.TestCase):
    def test_claim_transport_accepts_a_large_valid_batch_body(self):
        gateway = Gateway("https://example.test", TOKEN)
        reply = json.dumps({"job": {"body": {"state": "x" * (1200 * 1024)}}}).encode()

        class Opener:
            def open(self, request, timeout):
                self.request = request
                return io.BytesIO(reply)

        gateway.opener = Opener()
        response = gateway.post("claim", {"worker_id": "1" * 32})
        self.assertEqual(len(response["job"]["body"]["state"]), 1200 * 1024)
        self.assertTrue(gateway.opener.request.full_url.endswith("/internal/worker/claim"))

    def test_every_wire_id_binds_to_its_exact_canonical_model(self):
        registry = [{
            "id": model,
            "label": model,
            "version": "1.0",
            "manifest_sha256": format(index + 1, "x") * 64,
            "revision": format(index + 1, "x") * 40,
        } for index, model in enumerate(MODEL_ORDER)]
        for model in MODEL_ORDER:
            with self.subTest(model=model), patch.dict(os.environ, {
                "DECISION_MODEL_REGISTRY_V2": json.dumps(registry),
                "DECISION_WORKER_MODEL": model,
                "DECISION_RUNTIME_URL": ORIGIN,
            }):
                runtime = runtime_from_environment()
                self.assertEqual(runtime.model, model)
                self.assertEqual(runtime.canonical_model, PROFILES[model].get(
                    "runtime_model", PROFILES[model]["repo_id"]))
                self.assertEqual(runtime.manifest, registry[MODEL_ORDER.index(model)]["manifest_sha256"])
        with patch.dict(os.environ, {
            "DECISION_MODEL_REGISTRY_V2": json.dumps(registry),
            "DECISION_WORKER_MODEL": "decision-nano-preview",
            "DECISION_RUNTIME_URL": ORIGIN,
        }), self.assertRaises(ValueError):
            runtime_from_environment()


class PullQueueTetrisReadinessTests(unittest.TestCase):
    def test_offline_queue_is_unavailable_in_config_and_race_admission(self):
        lux_manifest = "d" * 64
        relays = {
            MODEL: Relay(TOKEN, MANIFEST, model=MODEL),
            "decision-lux": Relay(
                TOKEN, lux_manifest, model="decision-lux",
                complete_input_tokens=16384,
            ),
        }
        registry = [
            {"id": MODEL, "label": "Kai", "version": "1.0",
             "manifest_sha256": MANIFEST},
            {"id": "decision-lux", "label": "Lux", "version": "1.0",
             "manifest_sha256": lux_manifest},
        ]
        with patch.dict(os.environ, {
            "TETRIS_LOCAL_API_URL": "http://127.0.0.1:7860/v1/systemone",
        }), TestClient(create_app(
            mode="pull_queue", relays=relays, registry=registry,
        )) as client:
            headers = {"Authorization": "Bearer " + TOKEN}
            ready = client.post("/internal/worker/heartbeat", json={
                "model": MODEL,
                "manifest_sha256": MANIFEST,
                "worker_id": "1" * 32,
                "phase": "ready",
                "capabilities": ["context_batch_v1"],
            }, headers=headers)
            self.assertEqual(ready.status_code, 200)
            config = client.get("/api/tetris/config").json()
            by_id = {item["id"]: item["ready"] for item in config["competitors"]}
            self.assertTrue(by_id["kai"])
            self.assertFalse(by_id["lux"])
            self.assertNotIn("nox", by_id)
            race = {"left": "kai", "right": "lux", "mode": "steps", "max_steps": 1}
            self.assertEqual(client.post("/api/tetris/races", json=race).status_code, 503)

            ready = client.post("/internal/worker/heartbeat", json={
                "model": "decision-lux",
                "manifest_sha256": lux_manifest,
                "worker_id": "2" * 32,
                "phase": "ready",
                "capabilities": ["context_batch_v1"],
            }, headers=headers)
            self.assertEqual(ready.status_code, 200)
            config = client.get("/api/tetris/config").json()
            by_id = {item["id"]: item["ready"] for item in config["competitors"]}
            self.assertTrue(by_id["lux"])


if __name__ == "__main__":
    unittest.main()