llm-connect/tests/test_messages_gate.py

"""Narrow Messages admission and wire behavior; deterministic, no provider SDK."""

import http.client
import json
from dataclasses import FrozenInstanceError, replace

import pytest

from llm_connect.messages_gate import (
    MessagesPolicy,
    MessagesServer,
    RequestRefused,
    _json,
    _Stream,
)


@pytest.fixture
def policy():
    return MessagesPolicy("fixture:no-live-tariff", "fixture-model", 1000, 100, 3, 15)


def body():
    return {
        "model": "fixture-model",
        "max_tokens": 100,
        "stream": True,
        "messages": [{"role": "user", "content": "fixture"}],
    }


def test_conservative_bound_ignores_unproven_input_estimates(policy):
    small = body()
    large = body()
    large["messages"][0]["content"] = "fixture " * 10000
    assert policy.validate(small, "") == policy.validate(large, "") == 4500
    cached = body()
    cached["cache_control"] = {"type": "ephemeral", "ttl": "1h"}
    assert policy.validate(cached, "") == 4500
    with pytest.raises(FrozenInstanceError):
        policy.context_tokens = 1
    assert replace(policy, context_tokens=1001).sha256 != policy.sha256


@pytest.mark.parametrize(
    "field,value",
    [
        ("context_tokens", True),
        ("context_tokens", 0),
        ("max_output_tokens", -1),
        ("input_microusd_per_token", 1.1),
        ("output_microusd_per_token", 0),
        ("timeout_seconds", 1000),
        ("max_body_bytes", 2000001),
        ("model", "x\nheader"),
        ("allowed_betas", ["unknown"]),
    ],
)
def test_invalid_policy_refused(policy, field, value):
    with pytest.raises(RequestRefused):
        replace(policy, **{field: value})


def test_beta_policy_never_implicitly_expands_context(policy):
    with pytest.raises(RequestRefused):
        policy.validate(body(), "context-1m-2025-08-07")
    admitted = replace(policy, allowed_betas=("fixture-beta",))
    assert admitted.validate(body(), "fixture-beta") == 4500
    with pytest.raises(RequestRefused):
        admitted.validate(body(), "fixture-beta,unknown")


@pytest.mark.parametrize(
    "raw", [b'{"max_tokens":1,"max_tokens":2}', b'{"a":NaN}', b'{"a":Infinity}']
)
def test_ambiguous_json_refused(raw):
    with pytest.raises(ValueError):
        _json(raw)


def event(stream, data):
    stream.event(b"data: " + json.dumps(data).encode())


def test_stream_usage_cumulative_cache_and_missing_terminal(policy):
    stream = _Stream(policy)
    event(
        stream,
        {
            "type": "message_start",
            "message": {
                "model": "fixture-model",
                "usage": {
                    "input_tokens": 10,
                    "output_tokens": 0,
                    "cache_creation_input_tokens": 20,
                    "cache_read_input_tokens": 30,
                },
            },
        },
    )
    event(stream, {"type": "message_delta", "delta": {}, "usage": {"output_tokens": 5}})
    event(
        stream,
        {
            "type": "message_delta",
            "delta": {"stop_reason": "end_turn"},
            "usage": {"output_tokens": 10},
        },
    )
    with pytest.raises(RequestRefused):
        stream.cost()
    event(stream, {"type": "message_stop"})
    assert stream.cost() == 330


@pytest.mark.parametrize(
    "extra",
    [
        {"type": "error"},
        {"type": "unknown"},
        {"type": "message_delta", "usage": {"input_tokens": 1}},
        {"type": "message_delta", "usage": {"server_tool_use": {"web_search_requests": 1}}},
        {"type": "content_block_start", "content_block": {"type": "fallback"}},
        {"type": "message_stop"},
    ],
)
def test_stream_unknown_fees_regression_and_premature_terminal_refuse(policy, extra):
    stream = _Stream(policy)
    event(
        stream,
        {
            "type": "message_start",
            "message": {
                "model": "fixture-model",
                "usage": {"input_tokens": 10, "output_tokens": 0},
            },
        },
    )
    with pytest.raises(RequestRefused):
        event(stream, extra)


class DeniedMeter:
    def __init__(self):
        self.calls = 0

    def reserve_request(self, *args):
        self.calls += 1
        raise RuntimeError("private credential-shaped exception must not escape")


def test_no_execute_route_no_header_forward_and_bounded_refusal(policy):
    meter = DeniedMeter()
    server = MessagesServer(policy, meter, provider_key="dummy-never-used")
    server.start()
    try:
        for path, extra, raw in [
            ("/execute", {}, json.dumps(body())),
            ("/v1/messages", {"Authorization": "Bearer dummy"}, json.dumps(body())),
            ("/v1/messages", {}, '{"model": "x", "model": "y"}'),
            ("/v1/messages", {}, json.dumps(body())),
        ]:
            conn = http.client.HTTPConnection("127.0.0.1", server.port, timeout=3)
            conn.request("POST", path, raw, {"Content-Type": "application/json", **extra})
            response = conn.getresponse()
            result = response.read()
            assert response.status in (400, 404)
            assert b"private" not in result and b"dummy" not in result
            conn.close()
        assert meter.calls == 1
    finally:
        server.stop()


@pytest.mark.parametrize(
    "url",
    [
        "http://example.com",
        "https://a:b@example.com",
        "https://example.com/execute",
        "https://example.com?token=x",
    ],
)
def test_fixed_origin_no_redirect_proxy_or_request_url(policy, url):
    with pytest.raises(RequestRefused):
        MessagesServer(
            policy, DeniedMeter(), provider_key="dummy", upstream_url=url, allow_test_http=True
        )