litellm/tests/test_litellm_rust/cache/test_facade.py

import asyncio
import contextvars
import gc
import weakref
from types import SimpleNamespace
from typing import Final, cast

import pytest

import litellm
from litellm.caching.caching import Cache, disable_cache, enable_cache, update_cache
from litellm.caching.in_memory_cache import InMemoryCache
from litellm.rust_bridge import _native
from litellm.rust_bridge.catalog import CacheRule, Route, RouteRule, SecretManagerRule
from litellm.rust_bridge.configuration import Rollout
from litellm.rust_bridge.response_cache import ResponseCacheRuntime, resolve_response_cache
from litellm.types.caching import LiteLLMCacheType
from tests.test_litellm_rust.support.cache import CacheLookup, CacheTestHandle, CacheTestResolver, request
from tests.test_litellm_rust.support.isolation import rebound

pytestmark: Final = pytest.mark.requires_rust_extension


def test_existing_constructor_and_global_are_unchanged() -> None:
    facade: Final = Cache(type=LiteLLMCacheType.LOCAL)
    assert type(facade.cache) is InMemoryCache
    assert "_native_cache_handle" not in vars(facade)
    assert resolve_response_cache(facade) is None
    with rebound(litellm, "cache", facade):
        resolver: Final = CacheTestResolver(litellm)
        assert resolver.resolve().kind == "python_callback"
        resolver.resolve().store(None, {"answer": 7}, callback_kwargs={"cache_key": "key"})
        assert cast(CacheLookup, facade).get_cache(cache_key="key") == {"answer": 7}


async def test_catalog_constructs_native_runtime_from_public_cache_configuration() -> None:
    rules: Final = (
        RouteRule(Route.OCR, Rollout.PYTHON_ONLY),
        SecretManagerRule(Rollout.PYTHON_ONLY, systems=frozenset({"local"})),
        CacheRule(Rollout.RUST_REQUIRED, backends=frozenset({"local"})),
    )
    facade: Final = Cache(type=LiteLLMCacheType.LOCAL)
    runtime: Final = resolve_response_cache(facade, rules)
    assert isinstance(runtime, ResponseCacheRuntime)
    assert runtime.kind == "native"

    sync_request: Final = runtime.request(facade, {"cache_key": "sync"})
    assert sync_request is not None
    runtime.store(sync_request, {"answer": 1})
    assert runtime.lookup(sync_request) == {"answer": 1}
    assert facade.cache.get_cache("sync") is None

    async_request: Final = runtime.request(facade, {"cache_key": "async"})
    assert async_request is not None
    await runtime.async_store(async_request, {"answer": 2})
    assert await runtime.async_lookup(async_request) == {"answer": 2}
    assert await facade.cache.async_get_cache("async") is None

    requests: Final = (sync_request, async_request)
    expected: Final = {
        "values": [{"answer": 1}, {"answer": 2}],
        "missing_indices": [],
    }
    assert runtime.lookup_batch(requests) == expected
    assert await runtime.async_lookup_batch(requests) == expected

    await runtime.async_flush()
    assert runtime.lookup(sync_request) is None
    assert await runtime.async_lookup(async_request) is None


async def test_inference_resolver_uses_the_configured_native_cache_directly() -> None:
    rules: Final = (
        RouteRule(Route.OCR, Rollout.PYTHON_ONLY),
        SecretManagerRule(Rollout.PYTHON_ONLY, systems=frozenset({"local"})),
        CacheRule(Rollout.RUST_REQUIRED, backends=frozenset({"local"})),
    )
    facade: Final = Cache(type=LiteLLMCacheType.LOCAL)
    runtime: Final = resolve_response_cache(facade, rules)
    assert isinstance(runtime, ResponseCacheRuntime)
    facade._native_cache = runtime

    selected: Final = _native._CacheResolver(SimpleNamespace(cache=facade)).resolve()
    assert selected.kind == "native"
    request: Final = runtime.request(facade, {"cache_key": "inference-native"})
    assert request is not None
    await selected.async_store(request, {"answer": 42})
    assert await selected.async_lookup(request) == {"answer": 42}
    assert await runtime.async_lookup(request) == {"answer": 42}
    assert facade.cache.get_cache("inference-native") is None

    facade._native_cache = None
    fallback: Final = _native._CacheResolver(SimpleNamespace(cache=facade)).resolve()
    assert fallback.kind == "python_callback"
    await fallback.async_store(None, {"answer": 7}, callback_kwargs={"cache_key": "inference-python"})
    assert facade.get_cache(cache_key="inference-python") == {"answer": 7}
    assert facade.cache.get_cache("inference-python") is not None


async def test_inference_resolver_declines_a_native_runtime_whose_facade_changed() -> None:
    rules: Final = (
        RouteRule(Route.OCR, Rollout.PYTHON_ONLY),
        SecretManagerRule(Rollout.PYTHON_ONLY, systems=frozenset({"local"})),
        CacheRule(Rollout.RUST_REQUIRED, backends=frozenset({"local"})),
    )
    facade: Final = Cache(type=LiteLLMCacheType.LOCAL)
    runtime: Final = resolve_response_cache(facade, rules)
    assert isinstance(runtime, ResponseCacheRuntime)
    facade._native_cache = runtime
    stale_request: Final = runtime.request(facade, {"cache_key": "stale-only"})
    assert stale_request is not None
    await runtime.async_store(stale_request, {"answer": "stale"})

    replacement: Final = InMemoryCache()
    facade.cache = replacement
    with pytest.raises(_native.RustBridgeDeclined):
        _native._CacheResolver(SimpleNamespace(cache=facade)).resolve()
    assert await runtime.async_lookup(stale_request) == {"answer": "stale"}
    assert replacement.get_cache("stale-only") is None
    assert replacement.get_cache("swapped-backend") is None


def test_existing_global_lifecycle_remains_the_resolver_source_of_truth() -> None:
    resolver: Final = CacheTestResolver(litellm)

    enable_cache(type=LiteLLMCacheType.LOCAL, ttl=30)
    enabled: Final = litellm.cache
    assert isinstance(enabled, Cache)
    assert enabled.ttl == 30
    assert resolver.resolve().kind == "python_callback"

    enable_cache(type=LiteLLMCacheType.LOCAL, ttl=60)
    assert litellm.cache is enabled

    update_cache(type=LiteLLMCacheType.LOCAL, ttl=60)
    updated: Final = litellm.cache
    assert isinstance(updated, Cache)
    assert updated is not enabled
    assert updated.ttl == 60

    disable_cache()
    assert litellm.cache is None
    assert resolver.resolve().kind == "disabled"


async def test_native_bindings_survive_replacement_and_capture_writes_before_dispatch() -> None:
    namespace: Final = SimpleNamespace(cache=CacheTestHandle.memory())
    resolver: Final = CacheTestResolver(namespace)
    selected: Final = resolver.resolve()
    assert selected.kind == "native"
    selected.store(request(), {"answer": 1})
    assert await selected.async_lookup(request()) == {"answer": 1}
    with rebound(namespace, "cache", CacheTestHandle.memory()):
        replacement: Final = resolver.resolve()
        await selected.async_store(request(), {"answer": 2})
        assert replacement.lookup(request()) is None
        assert selected.lookup(request()) == {"answer": 2}
    with rebound(namespace, "cache", None):
        disabled: Final = resolver.resolve()
        assert disabled.kind == "disabled"
        assert disabled.lookup(None) is None
        await disabled.async_store(None, object())
        assert await disabled.async_lookup(None) is None
        assert selected.lookup(request()) == {"answer": 2}


async def test_python_callback_preserves_identity_caller_task_context_and_errors() -> None:
    context: Final = contextvars.ContextVar("cache_context", default="caller")
    caller: Final = asyncio.current_task()
    sentinel: Final = object()
    failure: Final = RuntimeError("callback failed")

    class CustomCache:
        async def async_get_cache(self, *, marker: object) -> object:
            assert marker is sentinel
            assert asyncio.current_task() is caller
            context.set("callback")
            return marker

        async def async_add_cache(self, response: object, *, marker: object) -> None:
            assert response is sentinel
            assert marker is sentinel
            raise failure

    namespace: Final = SimpleNamespace(cache=CustomCache())
    binding: Final = CacheTestResolver(namespace).resolve()
    assert binding.kind == "python_callback"
    assert await binding.async_lookup(None, callback_kwargs={"marker": sentinel}) is sentinel
    assert context.get() == "callback"
    with pytest.raises(RuntimeError) as caught:
        await binding.async_store(None, sentinel, callback_kwargs={"marker": sentinel})
    assert caught.value is failure


async def test_callback_cancellation_stays_in_the_callers_task() -> None:
    entered: Final = asyncio.Event()
    finished: Final = asyncio.Event()

    class CustomCache:
        async def async_get_cache(self) -> None:
            entered.set()
            try:
                await asyncio.Future()
            finally:
                finished.set()

    binding: Final = CacheTestResolver(SimpleNamespace(cache=CustomCache())).resolve()

    async def lookup() -> object:
        return await binding.async_lookup(None, callback_kwargs={})

    task: Final = asyncio.create_task(lookup())
    await entered.wait()
    task.cancel()
    with pytest.raises(asyncio.CancelledError):
        await task
    assert finished.is_set()


def test_registered_facade_uses_native_and_instance_overrides_fall_back() -> None:
    facade: Final = Cache(type=LiteLLMCacheType.LOCAL)
    handle: Final = CacheTestHandle.memory()
    handle._bind_facade(facade)
    resolver: Final = CacheTestResolver(SimpleNamespace(cache=facade))
    native: Final = resolver.resolve()
    assert native.kind == "native"
    native.store(request(), {"source": "native"})
    assert native.lookup(request()) == {"source": "native"}
    assert cast(CacheLookup, facade).get_cache(cache_key="key") is None
    sentinel: Final = object()

    def outer_override(**_kwargs: object) -> object:
        return sentinel

    def backend_override(*_args: object, **_kwargs: object) -> dict[str, str]:
        return {"source": "override"}

    with rebound(facade, "get_cache", outer_override):
        fallback: Final = resolver.resolve()
        assert fallback.kind == "python_callback"
        assert fallback.lookup(None, callback_kwargs={"cache_key": "key"}) is sentinel
    assert resolver.resolve().kind == "python_callback"
    delattr(facade, "get_cache")
    assert resolver.resolve().kind == "native"
    with rebound(facade.cache, "get_cache", backend_override):
        backend_fallback: Final = resolver.resolve()
        assert backend_fallback.kind == "python_callback"
        assert backend_fallback.lookup(None, callback_kwargs={"cache_key": "key"}) == {"source": "override"}


def test_facade_subclasses_backend_replacement_and_configuration_changes_are_not_bypassed() -> None:
    class CustomCache(Cache):
        pass

    handle: Final = CacheTestHandle.memory()
    with pytest.raises(TypeError):
        handle._bind_facade(CustomCache(type=LiteLLMCacheType.LOCAL))
    facade: Final = Cache(type=LiteLLMCacheType.LOCAL)
    handle._bind_facade(facade)
    resolver: Final = CacheTestResolver(SimpleNamespace(cache=facade))
    with rebound(facade, "cache", InMemoryCache()):
        assert resolver.resolve().kind == "python_callback"
    with rebound(facade, "ttl", 12):
        assert resolver.resolve().kind == "python_callback"
    with rebound(facade, "semantic_cache_scope", "end_user"):
        assert resolver.resolve().kind == "python_callback"

    def custom_key(**_kwargs: object) -> str:
        return "custom"

    with rebound(facade, "get_cache_key", custom_key):
        assert resolver.resolve().kind == "python_callback"
    assert resolver.resolve().kind == "python_callback"
    delattr(facade, "get_cache_key")
    assert resolver.resolve().kind == "native"


def test_resolver_and_callback_cycles_can_be_collected() -> None:
    class CustomCache:
        pass

    def cyclic_reference() -> weakref.ReferenceType[CustomCache]:
        callback: Final = CustomCache()
        namespace: Final = SimpleNamespace(cache=callback)
        binding: Final = CacheTestResolver(namespace).resolve()
        setattr(callback, "binding", binding)
        return weakref.ref(callback)

    reference: Final = cyclic_reference()
    gc.collect()
    assert reference() is None


def test_invalid_duration_and_request_shape_fail_before_storage() -> None:
    binding: Final = CacheTestResolver(SimpleNamespace(cache=CacheTestHandle.memory())).resolve()
    for seconds in (-1.0, float("nan"), float("inf")):
        with pytest.raises(ValueError, match="cache durations must be finite and nonnegative"):
            binding.store({**request(), "ttl_seconds": seconds}, {"answer": 1})
    assert binding.lookup(request()) is None
    with pytest.raises(ValueError, match="cache durations must be finite and nonnegative"):
        CacheTestHandle.memory(ttl_seconds=-1)


async def test_memory_size_policy_is_applied_by_the_native_host() -> None:
    handle: Final = CacheTestHandle.memory(capacity=2, max_entry_bytes=128)
    binding: Final = CacheTestResolver(SimpleNamespace(cache=handle)).resolve()
    small: Final = {"answer": "ok"}
    binding.store(request("small"), small)
    assert await binding.async_lookup(request("small")) == small
    await binding.async_store(request("large"), {"answer": "x" * 256})
    assert binding.lookup(request("large")) is None
    assert binding.lookup(request("small")) == small
    disabled: Final = CacheTestResolver(SimpleNamespace(cache=CacheTestHandle.memory(capacity=0))).resolve()
    await disabled.async_store(request(), small)
    assert await disabled.async_lookup(request()) is None


async def test_native_batch_lookup_and_store_report_partial_hits() -> None:
    binding: Final = CacheTestResolver(SimpleNamespace(cache=CacheTestHandle.memory())).resolve()
    requests: Final = [request("hit"), request("miss"), request("disabled")]
    requests[2]["controls"] = {
        "supported_call_type": True,
        "configured": True,
        "native_backend": True,
        "default_on": True,
        "caching": False,
        "no_cache": False,
        "no_store": False,
        "use_cache": False,
    }
    await binding.async_store_batch(requests, [{"value": 1}, {"value": 2}, {"value": 3}])

    partial: Final = await binding.async_lookup_batch(requests)

    assert partial == {
        "values": [{"value": 1}, {"value": 2}, None],
        "missing_indices": [2],
    }


async def test_python_batch_callbacks_use_the_builtin_cache_api() -> None:
    result: Final = object()
    marker: Final = object()

    class CustomCache(Cache):
        def get_cache(self, dynamic_cache_object: object = None, **kwargs: object) -> object:
            return ("sync", kwargs)

        async def async_get_cache(self, dynamic_cache_object: object = None, **kwargs: object) -> object:
            return ("async", kwargs)

        async def async_add_cache_pipeline(
            self, result: object, dynamic_cache_object: object = None, **kwargs: object
        ) -> object:
            return result, kwargs

    binding: Final = CacheTestResolver(SimpleNamespace(cache=CustomCache(type=LiteLLMCacheType.LOCAL))).resolve()
    assert binding.kind == "python_callback"
    requests: Final = [request("first"), request("second")]
    kwargs: Final = [{"cache_key": "first"}, {"cache_key": "second"}]

    assert binding.lookup_batch(requests, callback_kwargs=kwargs) == [("sync", kwargs[0]), ("sync", kwargs[1])]
    assert await binding.async_lookup_batch(requests, callback_kwargs=kwargs) == [
        ("async", kwargs[0]),
        ("async", kwargs[1]),
    ]
    with pytest.raises(ValueError, match="equal lengths"):
        binding.lookup_batch(requests, callback_kwargs=kwargs[:1])
    with pytest.raises(TypeError, match="callback_result"):
        await binding.async_store_batch(requests, [1, 2], callback_kwargs={"marker": marker})
    stored: Final = cast(
        tuple[object, dict[str, object]],
        await binding.async_store_batch(requests, [1, 2], callback_result=result, callback_kwargs={"marker": marker}),
    )
    assert stored[0] is result
    assert stored[1] == {"marker": marker}


async def test_unmodified_builtin_cache_callbacks_can_ping_and_flush() -> None:
    async def ping() -> str:
        return "pong"

    cache: Final = Cache(type=LiteLLMCacheType.LOCAL)
    cache.cache.set_cache("key", "value")
    binding: Final = CacheTestResolver(SimpleNamespace(cache=cache)).resolve()
    assert binding.kind == "python_callback"

    setattr(cache.cache, "ping", ping)
    assert await binding.ping() == "pong"
    await binding.async_flush()
    assert cache.cache.get_cache("key") is None


def test_facade_registration_rejects_mismatched_capacity() -> None:
    facade: Final = Cache(type=LiteLLMCacheType.LOCAL)
    with pytest.raises(TypeError, match="capacities must match"):
        CacheTestHandle.memory(capacity=7)._bind_facade(facade)