pallas/tests/test_mantle_shims.py

"""Tests for pallas.mantle_shims.

These tests exercise the module in isolation: they do not hit the network
and do not require fast-agent to be configured. They do, however, import
fast-agent so that the monkeypatch targets exist — so fast-agent-mcp must
be installed in the environment running pytest.
"""
from __future__ import annotations

import pytest

from pallas import mantle_shims


# ── is_mantle_base_url ───────────────────────────────────────────────────────

@pytest.mark.parametrize(
    "url,expected",
    [
        ("https://bedrock-mantle.us-east-1.api.aws/anthropic", True),
        ("https://bedrock-mantle.ca-central-1.api.aws/anthropic", True),
        ("https://api.anthropic.com", False),
        ("https://example.com/bedrock", False),
        ("", False),
        (None, False),
    ],
)
def test_is_mantle_base_url(url: str | None, expected: bool) -> None:
    assert mantle_shims.is_mantle_base_url(url) is expected


# ── install_wire_name_prefix ─────────────────────────────────────────────────

def test_install_wire_name_prefix_registers_prefixed_ids() -> None:
    from fast_agent.llm.model_database import ModelDatabase
    from fast_agent.llm.provider_types import Provider

    mantle_shims.install_wire_name_prefix()

    for fa_name, wire_name in mantle_shims.MANTLE_WIRE_NAMES.items():
        key = (Provider.ANTHROPIC, ModelDatabase.normalize_model_name(fa_name))
        assert ModelDatabase._PROVIDER_WIRE_MODEL_NAMES.get(key) == wire_name


def test_install_wire_name_prefix_is_idempotent() -> None:
    mantle_shims.install_wire_name_prefix()
    mantle_shims.install_wire_name_prefix()  # must not raise
    # Second call leaves the same mapping in place.
    from fast_agent.llm.model_database import ModelDatabase
    from fast_agent.llm.provider_types import Provider

    key = (Provider.ANTHROPIC, ModelDatabase.normalize_model_name("claude-opus-4-7"))
    assert ModelDatabase._PROVIDER_WIRE_MODEL_NAMES[key] == "anthropic.claude-opus-4-7"


# ── _strip_tool_use_caller ───────────────────────────────────────────────────

def test_strip_tool_use_caller_removes_caller_key() -> None:
    blocks = [
        {"type": "tool_use", "id": "t1", "name": "foo", "input": {}, "caller": None},
        {"type": "text", "text": "hello"},
        {"type": "tool_use", "id": "t2", "name": "bar", "input": {}},  # no caller
    ]
    result = mantle_shims._strip_tool_use_caller(blocks)

    assert "caller" not in result[0]
    assert result[1] == {"type": "text", "text": "hello"}
    assert "caller" not in result[2]


def test_strip_tool_use_caller_is_idempotent() -> None:
    blocks = [{"type": "tool_use", "id": "t1", "name": "foo", "input": {}, "caller": None}]
    mantle_shims._strip_tool_use_caller(blocks)
    mantle_shims._strip_tool_use_caller(blocks)  # second pass must be a no-op
    assert "caller" not in blocks[0]


def test_strip_tool_use_caller_ignores_non_dict_blocks() -> None:
    # Anthropic SDK model objects are sometimes passed instead of dicts;
    # the helper must leave those untouched.
    class Sentinel:
        type = "tool_use"

    s = Sentinel()
    blocks = [s]
    mantle_shims._strip_tool_use_caller(blocks)
    assert blocks[0] is s  # unchanged


# ── install_tool_use_caller_strip ────────────────────────────────────────────

def test_install_tool_use_caller_strip_patches_converter() -> None:
    from fast_agent.llm.provider.anthropic.multipart_converter_anthropic import (
        AnthropicConverter,
    )

    mantle_shims.install_tool_use_caller_strip()

    # The patched deserialize must strip caller from replayed tool_use dicts.
    channels: dict[str, list[dict]] = {"assistant_raw": []}
    # We can't easily stub the original's internal behaviour, so just
    # verify the patch is in place by round-tripping a destination list
    # through _append_server_tool_channel_blocks, which we *know* ends by
    # calling our strip helper on `destination`.
    destination: list[dict] = [
        {"type": "tool_use", "id": "t1", "name": "foo", "input": {}, "caller": None},
    ]
    AnthropicConverter._append_server_tool_channel_blocks(None, destination)
    assert "caller" not in destination[0]


def test_install_tool_use_caller_strip_is_idempotent() -> None:
    mantle_shims.install_tool_use_caller_strip()
    mantle_shims.install_tool_use_caller_strip()  # must not raise or re-wrap


# ── install_fine_grained_tool_streaming_opt_out ──────────────────────────────

def test_fine_grained_beta_opt_out() -> None:
    from fast_agent.llm.provider.anthropic.llm_anthropic import AnthropicLLM

    mantle_shims.install_fine_grained_tool_streaming_opt_out()

    # The patched method never touches self, so a bare object suffices.
    stub = object()
    assert AnthropicLLM.supports_direct_anthropic_beta(stub, "fine_grained_tool_streaming") is False
    # Every other beta keeps the base-class answer (True).
    assert AnthropicLLM.supports_direct_anthropic_beta(stub, "interleaved_thinking") is True
    assert AnthropicLLM.supports_direct_anthropic_beta(stub, "long_context") is True


def test_fine_grained_beta_opt_out_is_idempotent() -> None:
    mantle_shims.install_fine_grained_tool_streaming_opt_out()
    mantle_shims.install_fine_grained_tool_streaming_opt_out()  # must not re-wrap

    from fast_agent.llm.provider.anthropic.llm_anthropic import AnthropicLLM

    assert (
        AnthropicLLM.supports_direct_anthropic_beta(object(), "fine_grained_tool_streaming")
        is False
    )


# ── install_max_tokens_clamp ─────────────────────────────────────────────────

@pytest.mark.parametrize(
    "initial,expected",
    [
        (128000, mantle_shims.MANTLE_MAX_OUTPUT_TOKENS),  # over the ceiling → clamped
        (None, mantle_shims.MANTLE_MAX_OUTPUT_TOKENS),    # unset → pinned to ceiling
        (4096, 4096),                                     # under the ceiling → untouched
    ],
)
def test_max_tokens_clamp(
    monkeypatch: pytest.MonkeyPatch, initial: int | None, expected: int
) -> None:
    from fast_agent.llm.provider.anthropic.llm_anthropic import AnthropicLLM
    from fast_agent.types import RequestParams

    # Stub the underlying initializer, then force a fresh wrap around it.
    monkeypatch.setattr(
        AnthropicLLM,
        "_initialize_default_params",
        lambda self, kwargs: RequestParams(maxTokens=initial),
    )
    monkeypatch.setattr(mantle_shims, "_max_tokens_clamp_installed", False)
    mantle_shims.install_max_tokens_clamp()

    params = AnthropicLLM._initialize_default_params(object(), {})
    assert params.maxTokens == expected


def test_max_tokens_clamp_is_idempotent() -> None:
    mantle_shims.install_max_tokens_clamp()
    mantle_shims.install_max_tokens_clamp()  # must not raise or re-wrap


# ── maybe_install ────────────────────────────────────────────────────────────

_INSTALLER_NAMES = [
    ("install_wire_name_prefix", "wire"),
    ("install_tool_use_caller_strip", "tool_use"),
    ("install_fine_grained_tool_streaming_opt_out", "beta_opt_out"),
    ("install_max_tokens_clamp", "max_tokens"),
]


def _patch_installers(monkeypatch: pytest.MonkeyPatch) -> list[str]:
    calls: list[str] = []
    for attr, label in _INSTALLER_NAMES:
        monkeypatch.setattr(
            mantle_shims, attr,
            lambda label=label: calls.append(label),
        )
    return calls


def test_maybe_install_installs_when_mantle(monkeypatch: pytest.MonkeyPatch) -> None:
    calls = _patch_installers(monkeypatch)

    installed = mantle_shims.maybe_install("https://bedrock-mantle.us-east-1.api.aws/anthropic")
    assert installed is True
    # "max_tokens" is intentionally absent: the 20 000 clamp is not installed
    # by default because the ceiling was never confirmed. See install_all().
    assert calls == ["wire", "tool_use", "beta_opt_out"]


def test_maybe_install_noop_for_non_mantle(monkeypatch: pytest.MonkeyPatch) -> None:
    calls = _patch_installers(monkeypatch)

    assert mantle_shims.maybe_install("https://api.anthropic.com") is False
    assert mantle_shims.maybe_install(None) is False
    assert mantle_shims.maybe_install("") is False
    assert calls == []