insights/services/llm-gateway/tests/test_modal.py

from typing import Any
from unittest.mock import AsyncMock, MagicMock, patch

import pytest
from fastapi import HTTPException

from llm_gateway.config import Settings
from llm_gateway.modal import (
    MODAL_ALLOWED_MODELS,
    MODAL_MODELS,
    _inject_modal_params,
    ensure_modal_model_allowed,
    ensure_modal_model_configured,
    make_modal_responses_call,
    should_route_glm_to_modal,
)
from llm_gateway.modal_routing import send_modal_anthropic_messages
from llm_gateway.rate_limiting.cost_refresh import ALIAS_METRIC_LABELS, COST_ALIASES
from llm_gateway.rate_limiting.model_cost_overrides import MODEL_COST_OVERRIDES

GLM_MODEL = "@cf/zai-org/glm-5.2"
KIMI_MODEL = "moonshotai/kimi-k3"


def _modal_settings(**overrides: Any) -> Settings:
    base: dict[str, Any] = {
        "modal_api_base": "https://insights--glm.us-east.modal.direct/v1",
        "modal_kimi_api_base": "https://insights--kimi.us-west.modal.direct/v1",
        "modal_key": "wk-test",
        "modal_secret": "ws-test",
    }
    base.update(overrides)
    return Settings(**base)


def test_inject_modal_params_maps_model_and_pins_proxy_auth() -> None:
    # Caller-supplied header dicts (`headers` or `extra_headers`, any casing) must never reach the
    # authenticated upstream request — litellm merges both into the outbound headers, and a
    # forwarded Host header would exfiltrate the proxy-token pair.
    kwargs: dict = {
        "model": GLM_MODEL,
        "headers": {"Host": "attacker.example"},
        "extra_headers": {"MODAL-KEY": "attacker", "host": "attacker.example", "X-Other": "no"},
    }
    _inject_modal_params(kwargs, "https://modal.test/v1", "wk", "ws")

    assert kwargs["model"] == "openai/zai-org/GLM-5.2-FP8"
    assert kwargs["api_base"] == "https://modal.test/v1"
    # litellm's OpenAI client rejects a missing api_key even though Modal auth is header-based.
    assert kwargs["api_key"]
    assert "headers" not in kwargs
    assert kwargs["extra_headers"] == {"Modal-Key": "wk", "Modal-Secret": "ws"}
    # Load-bearing: if this fails, the litellm model key is unpriced/unlabeled and cost lookup
    # silently falls to default_fallback_cost_usd (and metrics lose the modal provider label).
    assert kwargs["model"] in COST_ALIASES
    assert kwargs["model"] in ALIAS_METRIC_LABELS


@pytest.mark.parametrize(("initial", "expected"), [({}, True), ({"drop_params": False}, False)])
def test_inject_modal_params_drop_params(initial: dict, expected: bool) -> None:
    # drop_params must default on (the OpenAI-compatible surface 400s on Anthropic-only params)
    # while an explicit caller opt-out keeps its value.
    kwargs: dict = {"model": GLM_MODEL, **initial}
    _inject_modal_params(kwargs, "https://modal.test/v1", "wk", "ws")
    assert kwargs["drop_params"] is expected


def test_modal_model_map_served_names_all_priced_and_labeled() -> None:
    # An unpriced Modal model would bill the flat fallback cost; the label must keep the public id
    # so one model id slices across backends in dashboards.
    for public_id, config in MODAL_MODELS.items():
        litellm_key = f"openai/{config.served_model}"
        assert litellm_key in COST_ALIASES or litellm_key in MODEL_COST_OVERRIDES
        provider, metric_model = ALIAS_METRIC_LABELS[litellm_key]
        assert provider == "modal"
        assert metric_model == public_id


@pytest.mark.parametrize("model", [GLM_MODEL, KIMI_MODEL])
def test_ensure_modal_model_allowed_accepts_mapped_model(model: str) -> None:
    ensure_modal_model_allowed(model)
    assert model in MODAL_ALLOWED_MODELS


def test_kimi_uses_its_endpoint_with_shared_credentials() -> None:
    assert ensure_modal_model_configured(KIMI_MODEL, _modal_settings()) == (
        "https://insights--kimi.us-west.modal.direct/v1",
        "wk-test",
        "ws-test",
    )


async def test_modal_anthropic_drops_clear_thinking_when_thinking_is_disabled() -> None:
    request = {
        "model": KIMI_MODEL,
        "messages": [{"role": "user", "content": "Hello"}],
        "thinking": {"type": "disabled"},
        "context_management": {"edits": [{"type": "clear_thinking_20251015", "keep": "all"}]},
    }

    with patch("llm_gateway.modal_routing.send_modal_request", new=AsyncMock(return_value={})) as send_request:
        await send_modal_anthropic_messages(request, MagicMock(), False, "insights_code")

    forwarded_request = send_request.call_args.args[0]
    assert "context_management" not in forwarded_request
    assert request["context_management"] == {"edits": [{"type": "clear_thinking_20251015", "keep": "all"}]}


@pytest.mark.parametrize("model", ["@cf/moonshotai/kimi-k2.6", "@cf/unknown/model", "zai-org/GLM-5.2-FP8"])
def test_ensure_modal_model_allowed_rejects_unmapped_model(model: str) -> None:
    with pytest.raises(HTTPException) as exc_info:
        ensure_modal_model_allowed(model)
    assert exc_info.value.status_code == 400
    assert exc_info.value.detail["error"]["type"] == "invalid_request_error"
    assert model in exc_info.value.detail["error"]["message"]


@pytest.mark.parametrize(
    ("fraction", "product_fractions", "configured", "expected"),
    [
        (0.0, {}, True, False),
        (1.0, {}, True, True),
        # Missing credentials beat any fraction.
        (1.0, {}, False, False),
        # Per-product override beats the global fraction in both directions.
        (0.0, {"insights_code": 1.0}, True, True),
        (1.0, {"insights_code": 0.0}, True, False),
    ],
)
def test_should_route_glm_to_modal(
    fraction: float, product_fractions: dict[str, float], configured: bool, expected: bool
) -> None:
    overrides: dict[str, Any] = {
        "glm_modal_traffic_fraction": fraction,
        "glm_modal_product_traffic_fractions": product_fractions,
    }
    if not configured:
        overrides["modal_key"] = None
    settings = _modal_settings(**overrides)
    assert (
        should_route_glm_to_modal(GLM_MODEL, product="insights_code", user_key="user-1", settings=settings) is expected
    )


@pytest.mark.parametrize("product", ["twig", "array"])
def test_product_aliases_resolve_in_fraction_lookup(product: str) -> None:
    # Legacy aliases must read insights_code's per-product fraction, or its ramp/rollback would
    # silently skip alias traffic.
    settings = _modal_settings(
        glm_modal_traffic_fraction=0.0, glm_modal_product_traffic_fractions={"insights_code": 1.0}
    )
    assert should_route_glm_to_modal(GLM_MODEL, product=product, user_key="user-1", settings=settings) is True


def test_should_route_only_modal_served_models() -> None:
    # kimi is CF-allowlisted but has no Modal-served equivalent; routing it to Modal would 404.
    settings = _modal_settings(glm_modal_traffic_fraction=1.0)
    assert (
        should_route_glm_to_modal("@cf/moonshotai/kimi-k2.6", product="insights_code", user_key="u", settings=settings)
        is False
    )


def test_partial_fraction_is_sticky_and_monotonic() -> None:
    # Repeat calls must be identical, and a user routed to Modal at fraction f must stay routed at
    # any higher fraction — otherwise each ramp-up step would flap migrated users back.
    fractions = (0.25, 0.5, 0.75)
    settings = [_modal_settings(glm_modal_traffic_fraction=f) for f in fractions]
    for user in ("user-1", "user-2", "user-3", "user-4", "user-5"):
        decisions = [
            should_route_glm_to_modal(GLM_MODEL, product="insights_code", user_key=user, settings=s) for s in settings
        ]
        assert decisions == sorted(decisions)
        repeat = [
            should_route_glm_to_modal(GLM_MODEL, product="insights_code", user_key=user, settings=s) for s in settings
        ]
        assert repeat == decisions


def test_partial_fraction_splits_pinned_users() -> None:
    # Pinned user keys with known buckets ("3" -> ~0.25, "2" -> ~0.96) must land on opposite sides
    # of fraction 0.5 — a degenerate bucket (constant 0 or 1) passes the monotonicity test above
    # but fails here.
    settings = _modal_settings(glm_modal_traffic_fraction=0.5)
    assert should_route_glm_to_modal(GLM_MODEL, product="insights_code", user_key="3", settings=settings) is True
    assert should_route_glm_to_modal(GLM_MODEL, product="insights_code", user_key="2", settings=settings) is False


async def test_make_modal_responses_call_forces_bridge_and_ignores_smuggled_flag() -> None:
    # vLLM has no /responses route, so the Responses->chat/completions bridge must be forced and a
    # caller-smuggled use_chat_completions_api=False must not escape it (same bug class as CF).
    llm_call = make_modal_responses_call("https://modal.test/v1", "wk", "ws")

    with patch("llm_gateway.modal.litellm.aresponses", new=AsyncMock(return_value="ok")) as mock_aresponses:
        await llm_call(model=GLM_MODEL, input="hi", use_chat_completions_api=False)

    kwargs = mock_aresponses.call_args.kwargs
    assert kwargs["use_chat_completions_api"] is True
    assert kwargs["model"] == "openai/zai-org/GLM-5.2-FP8"
    assert kwargs["api_base"] == "https://modal.test/v1"
    assert kwargs["extra_headers"]["Modal-Key"] == "wk"
    assert kwargs["extra_headers"]["Modal-Secret"] == "ws"
    assert kwargs["input"] == "hi"