ReMe/tests/unit/test_local_embedding_store.py

"""Regression tests for LocalEmbeddingStore dimension and vector space handling."""

# pylint: disable=protected-access,too-many-lines

import asyncio
import json
from types import SimpleNamespace

import httpx
import numpy as np
import openai

from reme.components.as_embedding import DashScopeAsEmbedding, OllamaAsEmbedding, OpenAIAsEmbedding
from reme.components.embedding_store.base_embedding_store import BaseEmbeddingStore
from reme.components.embedding_store.local_embedding_store import LocalEmbeddingStore
from reme.schema import EmbNode


class FakeAsEmbedding:
    """Fake AgentScope embedding component."""

    dimensions = 2
    vector_space_id = "fakespace000"

    def initialize_model(self):
        """Mirror the real component's idempotent initialization hook."""

    async def __call__(self, texts: list[str], **_kwargs):
        return [[1.0] if text == "bad" else [1.0, 0.0] for text in texts]


class BadHealthAsEmbedding:
    """Fake provider whose health probe returns the wrong dimension."""

    dimensions = 2
    vector_space_id = "fakespace000"

    def initialize_model(self):
        """Mirror the real component's idempotent initialization hook."""

    async def __call__(self, _texts: list[str], **_kwargs):
        return [[1.0]]


class FailingHealthAsEmbedding:
    """Fake provider that records a failed health-check attempt."""

    dimensions = 2
    vector_space_id = "fakespace000"

    def __init__(self):
        self.calls = 0

    def initialize_model(self):
        """Mirror the real component's idempotent initialization hook."""

    async def __call__(self, _texts: list[str], **_kwargs):
        self.calls += 1
        raise ConnectionError("not ready")


class FakeProviderModel:
    """Stand-in for a constructed AgentScope embedding model object."""

    def __init__(
        self,
        model: str,
        dimensions: int = 2,
        base_url: str = "https://api.openai.com/v1",
    ):
        self.model = model
        self.dimensions = dimensions
        self.credential = SimpleNamespace(base_url=base_url)


class InsufficientQuotaError(Exception):
    """OpenAI-compatible quota error used without importing the provider SDK."""

    status_code = 429
    body = {"error": {"code": "insufficient_quota"}}


class RateLimitError(Exception):
    """OpenAI-compatible 429 error used without importing the provider SDK."""

    status_code = 429


class RateLimitedThenSuccessAsEmbedding:
    """Fail once with a 429, then return a valid embedding."""

    dimensions = 2

    def __init__(self):
        self.calls = 0

    async def __call__(self, texts: list[str], **_kwargs):
        self.calls += 1
        if self.calls == 1:
            raise RateLimitError("Requests are too frequent")
        return [[1.0, 0.0] for _ in texts]


class AlwaysRateLimitedAsEmbedding:
    """Always fail with a 429."""

    dimensions = 2

    def __init__(self):
        self.calls = 0

    async def __call__(self, _texts: list[str], **_kwargs):
        self.calls += 1
        raise RateLimitError("Requests are too frequent")


class QuotaThenSuccessAsEmbedding:
    """Fail once for quota, then return a valid embedding."""

    dimensions = 2

    def __init__(self):
        self.calls = 0

    async def __call__(self, texts: list[str], **_kwargs):
        self.calls += 1
        if self.calls == 1:
            raise InsufficientQuotaError("quota exhausted")
        return [[1.0, 0.0] for _ in texts]


class ContextLengthError(Exception):
    """OpenAI-compatible 400 raised when one item exceeds the provider's per-item window."""

    status_code = 400
    body = {"error": {"message": "the input length exceeds the context length"}}


class PerItemLimitAsEmbedding:
    """Reject a whole request when any single text is over the provider's per-item limit.

    Mirrors an OpenAI-compatible gateway: one over-limit item makes the provider answer
    HTTP 400 for the entire batch, so a client that gives up on the batch loses every
    vector in it.
    """

    dimensions = 2
    vector_space_id = "fakespace000"
    max_chars = 32

    def __init__(self):
        self.batch_sizes: list[int] = []

    def initialize_model(self):
        """Mirror the real component's idempotent initialization hook."""

    async def __call__(self, texts: list[str], **_kwargs):
        self.batch_sizes.append(len(texts))
        if any(len(text) > self.max_chars for text in texts):
            raise ContextLengthError("the input length exceeds the context length")
        return [[1.0, 0.0] for _ in texts]


class AlwaysInsufficientQuotaAsEmbedding:
    """Always fail with quota exhaustion, recording every requested batch size."""

    dimensions = 2

    def __init__(self):
        self.batch_sizes: list[int] = []

    async def __call__(self, texts: list[str], **_kwargs):
        self.batch_sizes.append(len(texts))
        raise InsufficientQuotaError("quota exhausted")


class PerItemLimitThenTransientErrorAsEmbedding:
    """Reject an over-limit batch, then fail the first single-item request once.

    Mirrors the review case: the batch is rejected because one item exceeds the per-item
    window, and a sibling item then hits a transient provider failure while retried alone.
    """

    dimensions = 2
    vector_space_id = "fakespace000"
    max_chars = 32

    def __init__(self, transient: Exception):
        self.transient = transient
        self.batch_sizes: list[int] = []
        self.transient_failures = 0

    def initialize_model(self):
        """Mirror the real component's idempotent initialization hook."""

    async def __call__(self, texts: list[str], **_kwargs):
        self.batch_sizes.append(len(texts))
        if any(len(text) > self.max_chars for text in texts):
            raise ContextLengthError("the input length exceeds the context length")
        if len(texts) == 1 and self.transient_failures == 0:
            self.transient_failures += 1
            raise self.transient
        return [[1.0, 0.0] for _ in texts]


class PerItemLimitThenQuotaAsEmbedding:
    """Reject an over-limit batch, recover the first isolated item, then exhaust the quota.

    Mirrors the review case: isolation starts, the first single-item request comes back, and
    every later call then fails provider-wide. Asking for the remaining items would only
    multiply the rejected requests, because no single item explains a quota failure.
    """

    dimensions = 2
    vector_space_id = "fakespace000"
    max_chars = 32

    def __init__(self):
        self.batch_sizes: list[int] = []
        self.singles = 0

    def initialize_model(self):
        """Mirror the real component's idempotent initialization hook."""

    async def __call__(self, texts: list[str], **_kwargs):
        self.batch_sizes.append(len(texts))
        if any(len(text) > self.max_chars for text in texts):
            raise ContextLengthError("the input length exceeds the context length")
        self.singles += 1
        if self.singles > 1:
            raise InsufficientQuotaError("quota exhausted")
        return [[1.0, 0.0] for _ in texts]


class AuthenticationError(Exception):
    """OpenAI-compatible 401 error used without importing the provider SDK."""

    status_code = 401


class PermissionDeniedError(Exception):
    """OpenAI-compatible 403 error used without importing the provider SDK."""

    status_code = 403


class InternalServerError(Exception):
    """OpenAI-compatible 500 error used without importing the provider SDK."""

    status_code = 500


class NotFoundError(Exception):
    """OpenAI-compatible 404 error, e.g. a missing model or endpoint."""

    status_code = 404


class APIError(Exception):
    """Base of the OpenAI SDK error hierarchy, reproduced without importing the SDK."""


class APIConnectionError(APIError):
    """OpenAI SDK connection failure: it does not inherit from the built-in ConnectionError."""


class APITimeoutError(APIConnectionError):
    """OpenAI SDK timeout, a subclass of APIConnectionError."""


class _RetryingPolicyModel:
    """Stand-in for AgentScope's OpenAI embedding model: it retries transport errors."""

    max_retries = 3

    @classmethod
    def _get_retryable_exceptions(cls):
        """Mirror the OpenAI policy the real wrapper consults."""
        return (APIConnectionError,)


class AlwaysProviderWideFailureAsEmbedding:
    """Always fail with the same provider-wide error, recording every requested batch size."""

    dimensions = 2

    def __init__(self, error: Exception):
        self.error = error
        self.batch_sizes: list[int] = []
        # This double stands for an error raised by the SDK beneath an OpenAI-style model, so it
        # advertises the same lower-layer retry policy: the store must not add a third loop.
        self.model = _RetryingPolicyModel()

    async def __call__(self, texts: list[str], **_kwargs):
        self.batch_sizes.append(len(texts))
        raise self.error


class TransientThenSuccessAsEmbedding:
    """Fail the first `failures` calls with one error, then return valid embeddings."""

    dimensions = 2
    vector_space_id = "fakespace000"

    def __init__(self, error: Exception, failures: int = 1):
        self.error = error
        self.failures = failures
        self.batch_sizes: list[int] = []

    def initialize_model(self):
        """Mirror the real component's idempotent initialization hook."""

    async def __call__(self, texts: list[str], **_kwargs):
        self.batch_sizes.append(len(texts))
        if len(self.batch_sizes) <= self.failures:
            raise self.error
        return [[1.0, 0.0] for _ in texts]


_REAL_OPENAI_ASYNC_CLIENT = openai.AsyncClient
_OFFLINE_OPENAI_BASE_URL = "http://127.0.0.1:9/v1"


class CountingOpenAIWrapper(OpenAIAsEmbedding):
    """Real OpenAI wrapper that records how often the store asked it for embeddings."""

    def __init__(self, **kwargs) -> None:
        super().__init__(**kwargs)
        self.calls = 0

    async def __call__(self, inputs, **kwargs):
        self.calls += 1
        return await super().__call__(inputs, **kwargs)


class OfflineOpenAIEndpoint:
    """Real openai.AsyncClient wired to an offline httpx transport that counts provider attempts.

    Gives the wrapper-and-SDK layers their real code path, so a regression test can assert how
    many attempts the provider would see instead of trusting the store's own loop.
    """

    def __init__(self, handler, sdk_max_retries: int = 2) -> None:
        self.requests = 0
        self.handler = handler
        self.client = _REAL_OPENAI_ASYNC_CLIENT(
            api_key="test-key",
            base_url=_OFFLINE_OPENAI_BASE_URL,
            http_client=httpx.AsyncClient(transport=httpx.MockTransport(self._handle)),
            max_retries=sdk_max_retries,
        )

    def _handle(self, request):
        self.requests += 1
        return self.handler(request)

    def install(self, monkeypatch) -> "OfflineOpenAIEndpoint":
        """Route every openai.AsyncClient built during this test to the offline client."""
        monkeypatch.setattr(openai, "AsyncClient", lambda **kwargs: self.client)
        return self


def offline_openai_endpoint(handler) -> OfflineOpenAIEndpoint:
    """Build an offline endpoint whose transport follows ``handler``."""
    return OfflineOpenAIEndpoint(handler)


def openai_wrapper(name: str) -> CountingOpenAIWrapper:
    """Build a real OpenAI embedding wrapper pointed at the offline endpoint."""
    return CountingOpenAIWrapper(
        name=name,
        model="text-embedding-v4",
        backend="openai",
        dimensions=2,
        credential={"base_url": _OFFLINE_OPENAI_BASE_URL, "api_key": "test-key"},
        retry_delay=0,
        pass_dimensions=False,
    )


def embedding_payload(request) -> httpx.Response:
    """Answer an embeddings request the way the provider would."""
    body = json.loads(request.content.decode("utf-8"))
    inputs = body["input"]
    if isinstance(inputs, str):
        inputs = [inputs]
    return httpx.Response(
        200,
        json={
            "object": "list",
            "data": [
                {"object": "embedding", "index": position, "embedding": [0.0, 1.0]} for position in range(len(inputs))
            ],
            "model": body.get("model", "text-embedding-v4"),
            "usage": {"prompt_tokens": len(inputs), "total_tokens": len(inputs)},
        },
    )


def context_length_rejection() -> httpx.Response:
    """Answer a batch the way an OpenAI-compatible provider reports an over-long input."""
    return httpx.Response(
        400,
        json={
            "error": {
                "message": (
                    "This model's maximum context length is 8191 tokens. " "Please reduce the length of the input."
                ),
                "type": "invalid_request_error",
                "param": None,
                "code": "context_length_exceeded",
            },
        },
    )


def read_timeout(request) -> httpx.Response:
    """Fail one attempt the way a read timeout reaches the SDK."""
    raise httpx.ReadTimeout("simulated read timeout", request=request)


class BadNodeEmbeddingStore(BaseEmbeddingStore):
    """Embedding store that returns wrong-dimensional vectors."""

    dimensions = 2

    async def health_check(self, timeout: float = 2.0) -> bool:
        return True

    async def get_embeddings(self, input_text: list[str], **_kwargs):
        return [np.array([1.0], dtype=np.float16) for _ in input_text]


def run(coro):
    """Run an async test body."""
    return asyncio.run(coro)


def test_truncate_uses_cjk_aware_integer_budget():
    """Truncation should preserve ASCII behavior and budget non-ASCII text."""
    store = BadNodeEmbeddingStore(name="t_base_embedding_truncate", max_input_length=10)

    assert store._truncate("abcdefghijk") == "abcdefghij"
    assert store._truncate("中文中文中文中文") == "中文中文中文"
    assert store._truncate("éabcdefghij") == "éabcdefgh"

    store.max_input_length = -1
    assert store._truncate("text") == ""
    assert store._truncate("中文") == ""


def test_compute_batch_rejects_embeddings_with_wrong_dimension():
    """Provider results with wrong dimensions are not padded, truncated, or cached."""

    async def go():
        store = LocalEmbeddingStore(name="t_local_embedding_dim")
        store.as_embedding = FakeAsEmbedding()

        results = await store._compute_batch(
            [
                (0, "ok", "ok-cache-key"),
                (1, "bad", "bad-cache-key"),
            ],
        )

        assert len(results) == 1
        assert results[0][0] == 0
        assert results[0][2].tolist() == [1.0, 0.0]
        assert isinstance(results[0][2], np.ndarray)

    run(go())


def test_base_get_node_embeddings_rejects_wrong_dimension():
    """Base node assignment should not accept wrong-dimensional vectors."""

    async def go():
        store = BadNodeEmbeddingStore(name="t_base_embedding_dim")
        node = EmbNode(text="bad")

        await store.get_node_embeddings([node])

        assert node.embedding is None

    run(go())


def test_health_check_rejects_wrong_dimension():
    """Health check should fail when the provider returns the wrong vector length."""

    async def go():
        store = LocalEmbeddingStore(name="t_local_embedding_health_dim")
        store.as_embedding = BadHealthAsEmbedding()

        assert await store.health_check() is False
        assert store.is_healthy is False

    run(go())


def test_health_check_starts_timeout_after_provider_initialization(monkeypatch):
    """One-time client construction must not consume the request timeout."""

    async def go():
        events = []

        class InitializingAsEmbedding(FakeAsEmbedding):
            """Record initialization and provider-call ordering."""

            def initialize_model(self):
                events.append("initialized")

            async def __call__(self, texts: list[str], **_kwargs):
                events.append("remote request")
                return [[1.0, 0.0] for _ in texts]

        original_wait_for = asyncio.wait_for

        async def checked_wait_for(awaitable, timeout):
            assert events == ["initialized"]
            assert timeout == 5.0
            return await original_wait_for(awaitable, timeout)

        store = LocalEmbeddingStore(name="t_local_embedding_health_timeout_scope")
        store.as_embedding = InitializingAsEmbedding()
        monkeypatch.setattr(asyncio, "wait_for", checked_wait_for)

        assert await store.health_check(timeout=5.0) is True
        assert events == ["initialized", "remote request"]

    run(go())


def test_health_check_makes_one_attempt():
    """A failed startup probe does not add hidden retries."""

    async def go():
        provider = FailingHealthAsEmbedding()
        store = LocalEmbeddingStore(
            name="t_local_embedding_health_retry",
            health_check_timeout=3.0,
        )
        store.as_embedding = provider

        assert await store.health_check() is False
        assert provider.calls == 1

    run(go())


def test_insufficient_quota_waits_sixty_seconds_before_retry(monkeypatch):
    """Quota exhaustion uses the dedicated delay before ReMe retries."""

    async def go():
        sleeps = []

        async def fake_sleep(delay):
            sleeps.append(delay)

        store = LocalEmbeddingStore(
            name="t_local_embedding_quota",
            max_retries=2,
            quota_retry_delay=60.0,
        )
        embedding = QuotaThenSuccessAsEmbedding()
        store.as_embedding = embedding
        monkeypatch.setattr(asyncio, "sleep", fake_sleep)

        result = await store._call_with_retry(["text"])

        assert result == [[1.0, 0.0]]
        assert embedding.calls == 2
        assert sleeps == [60.0]

    run(go())


def test_insufficient_quota_does_not_retry_without_opt_in(monkeypatch):
    """The default store behavior remains unchanged for embedded consumers."""

    async def go():
        sleeps = []

        async def fake_sleep(delay):
            sleeps.append(delay)

        store = LocalEmbeddingStore(name="t_local_embedding_default_quota", max_retries=2)
        embedding = QuotaThenSuccessAsEmbedding()
        store.as_embedding = embedding
        monkeypatch.setattr(asyncio, "sleep", fake_sleep)

        result = await store._call_with_retry(["text"])

        assert result is None
        assert embedding.calls == 1
        assert not sleeps

    run(go())


def test_rate_limit_retries_with_backoff_without_opt_in(monkeypatch):
    """A 429 must retry like a network error, with no quota_retry_delay required."""

    async def go():
        sleeps = []

        async def fake_sleep(delay):
            sleeps.append(delay)

        store = LocalEmbeddingStore(name="t_local_embedding_rate_limit", max_retries=2)
        embedding = RateLimitedThenSuccessAsEmbedding()
        store.as_embedding = embedding
        monkeypatch.setattr(asyncio, "sleep", fake_sleep)

        result = await store._call_with_retry(["text"])

        assert result == [[1.0, 0.0]]
        assert embedding.calls == 2
        assert sleeps == [1.0]

    run(go())


def test_is_rate_limited_recognizes_status_code_or_body_and_nothing_else():
    """A 429 is recognized either by status_code or a body-embedded code, never by guesswork."""

    class StatusCodeError(Exception):
        """A 429 identified by an HTTP status code attribute."""

        status_code = 429

    class BodyCodeError(Exception):
        """A 429 identified only by a body-embedded error code."""

        body = {"error": {"code": "rate_limit_exceeded"}}

    assert LocalEmbeddingStore._is_rate_limited(StatusCodeError())
    assert LocalEmbeddingStore._is_rate_limited(BodyCodeError())
    assert not LocalEmbeddingStore._is_rate_limited(ValueError("unrelated"))


def test_is_rate_limited_defers_to_insufficient_quota_on_status_code_429():
    """An insufficient_quota error carrying status_code=429 is not the generic rate-limit case."""

    assert not LocalEmbeddingStore._is_rate_limited(InsufficientQuotaError("quota exhausted"))
    assert LocalEmbeddingStore._is_insufficient_quota(InsufficientQuotaError("quota exhausted"))


def test_rate_limit_exhausts_retries_and_reports_unhealthy(monkeypatch):
    """Repeated 429s still give up after max_retries, unlike an unclassified error."""

    async def go():
        sleeps = []

        async def fake_sleep(delay):
            sleeps.append(delay)

        store = LocalEmbeddingStore(name="t_local_embedding_rate_limit_exhausted", max_retries=3)
        embedding = AlwaysRateLimitedAsEmbedding()
        store.as_embedding = embedding
        monkeypatch.setattr(asyncio, "sleep", fake_sleep)

        result = await store._call_with_retry(["text"])

        assert result is None
        assert store.is_healthy is False
        assert embedding.calls == 3
        assert sleeps == [1.0, 2.0]

    run(go())


def test_vector_space_id_separates_models_of_equal_dimension():
    """Two models of the same width must not claim the same vector space."""
    common = {"backend": "openai", "dimensions": 1024, "credential": {"base_url": "https://example.com/v1"}}
    v3 = OpenAIAsEmbedding(name="t_space_v3", model="text-embedding-v3", **common)
    v4 = OpenAIAsEmbedding(name="t_space_v4", model="text-embedding-v4", **common)

    assert v3.dimensions == v4.dimensions
    assert v3.vector_space_id != v4.vector_space_id


def test_vector_space_id_separates_endpoints_of_one_model_name():
    """The same model name served by two endpoints is two vector spaces."""
    common = {"backend": "openai", "model": "text-embedding-v4", "dimensions": 1024}
    official = OpenAIAsEmbedding(name="t_space_a", credential={"base_url": "https://example.com/v1"}, **common)
    self_hosted = OpenAIAsEmbedding(name="t_space_b", credential={"base_url": "http://127.0.0.1:8000/v1"}, **common)

    assert official.vector_space_id != self_hosted.vector_space_id


def test_vector_space_id_ignores_trailing_slash_and_api_key():
    """Cosmetic and secret credential changes must not invalidate stored vectors."""
    common = {"backend": "openai", "model": "text-embedding-v4", "dimensions": 1024}
    first = OpenAIAsEmbedding(
        name="t_space_c",
        credential={"base_url": "https://example.com/v1", "api_key": "key-one"},
        **common,
    )
    second = OpenAIAsEmbedding(
        name="t_space_d",
        credential={"base_url": "https://example.com/v1/", "api_key": "key-two"},
        **common,
    )

    assert first.vector_space_id == second.vector_space_id


def test_vector_space_id_follows_a_model_swapped_in_after_start():
    """A provider replaced at runtime must win over the original kwargs."""
    embedding = OpenAIAsEmbedding(name="t_space_swap", backend="openai", model="v3", dimensions=2)
    before = embedding.vector_space_id

    # Mirrors Application.update_component("as_embedding", "default", model=<new provider>).
    embedding.model = FakeProviderModel("v4")

    assert embedding.vector_space_id != before


def test_vector_space_id_is_stable_across_lazy_provider_construction():
    """Constructing the configured provider must not look like a model switch."""
    embedding = OpenAIAsEmbedding(
        name="t_space_lazy",
        backend="openai",
        model="v3",
        dimensions=2,
        credential={"base_url": "https://example.com/v1"},
    )
    before = embedding.vector_space_id

    embedding.model = FakeProviderModel("v3", base_url="https://example.com/v1")

    assert embedding.vector_space_id == before


def test_vector_space_id_resolves_default_endpoint_before_lazy_construction():
    """Credential defaults must not change the cache namespace on the first request."""
    embedding = DashScopeAsEmbedding(
        name="t_space_default_endpoint",
        model="text-embedding-v3",
        dimensions=1024,
        credential={"api_key": "test"},
    )
    before = embedding.vector_space_id

    embedding._ensure_model()

    assert embedding.vector_space_id == before


def test_openai_vector_space_id_uses_sdk_resolved_endpoint(monkeypatch):
    """OPENAI_BASE_URL must separate caches and stay stable after lazy construction."""
    clients = []
    ids = []
    try:
        for endpoint in ("https://provider-a.example/v1", "https://provider-b.example/v1"):
            monkeypatch.setenv("OPENAI_BASE_URL", endpoint)
            embedding = OpenAIAsEmbedding(
                name="t_space_openai_env",
                backend="openai",
                model="text-embedding-3-small",
                dimensions=1536,
                credential={"api_key": "test"},
            )
            before = embedding.vector_space_id

            embedding._ensure_model()
            clients.append(embedding.model.client)

            assert str(embedding.model.client.base_url).rstrip("/") == endpoint
            assert embedding.vector_space_id == before
            ids.append(before)

        assert ids[0] != ids[1]
    finally:
        for client in clients:
            run(client.close())


def test_openai_vector_space_id_uses_sdk_default_endpoint(monkeypatch):
    """The SDK default URL must not change the namespace on first construction."""
    monkeypatch.delenv("OPENAI_BASE_URL", raising=False)
    embedding = OpenAIAsEmbedding(
        name="t_space_openai_default",
        backend="openai",
        model="text-embedding-3-small",
        dimensions=1536,
        credential={"api_key": "test"},
    )
    before = embedding.vector_space_id

    embedding._ensure_model()
    try:
        assert str(embedding.model.client.base_url).rstrip("/") == "https://api.openai.com/v1"
        assert embedding.vector_space_id == before
    finally:
        run(embedding.model.client.close())


def test_ollama_vector_space_id_uses_sdk_resolved_endpoint(monkeypatch):
    """OLLAMA_HOST must separate caches and stay stable after lazy construction."""
    ids = []
    for endpoint in ("http://provider-a.example:11434", "http://provider-b.example:11434"):
        monkeypatch.setenv("OLLAMA_HOST", endpoint)
        embedding = OllamaAsEmbedding(
            name="t_space_ollama_env",
            backend="ollama",
            model="nomic-embed-text",
            dimensions=768,
            credential={},
        )
        before = embedding.vector_space_id

        embedding._ensure_model()

        assert embedding.model.host is None
        assert embedding.vector_space_id == before
        ids.append(before)

    assert ids[0] != ids[1]


def test_ollama_vector_space_id_uses_sdk_default_endpoint(monkeypatch):
    """The Ollama SDK default URL must remain stable after lazy construction."""
    monkeypatch.delenv("OLLAMA_HOST", raising=False)
    embedding = OllamaAsEmbedding(
        name="t_space_ollama_default",
        backend="ollama",
        model="nomic-embed-text",
        dimensions=768,
        credential={},
    )
    before = embedding.vector_space_id

    embedding._ensure_model()

    assert embedding.vector_space[-1] == "http://127.0.0.1:11434"
    assert embedding.vector_space_id == before


def test_cache_is_saved_and_restored_per_vector_space(monkeypatch, tmp_path):
    """Switching models persists the old cache and restores it when switched back."""

    async def go():
        monkeypatch.setattr(
            LocalEmbeddingStore,
            "component_metadata_path",
            property(lambda _self: tmp_path),
        )
        embedding = OpenAIAsEmbedding(name="t_space_store", backend="openai", model="v3", dimensions=2)
        store = LocalEmbeddingStore(name="t_local_space")
        store.as_embedding = embedding
        await store.load()

        key = store._cache_key("hello")
        store._cache_put(key, np.array([1.0, 0.0], dtype=np.float16))
        v3_path = store.cache_path

        embedding.model = FakeProviderModel("v4")
        await store._sync_cache_space()

        assert store._cache_key("hello") == key
        assert store.cache_path != v3_path
        assert store._cache_get(store._cache_key("hello")) is None
        assert v3_path.exists()

        embedding.model = FakeProviderModel("v3")
        await store._sync_cache_space()

        np.testing.assert_array_equal(store._cache_get(key), np.array([1.0, 0.0], dtype=np.float16))

    run(go())


def test_cache_space_is_rechecked_after_async_load(monkeypatch, tmp_path):
    """A provider switch during disk I/O must not publish the stale namespace."""

    async def go():
        monkeypatch.setattr(
            LocalEmbeddingStore,
            "component_metadata_path",
            property(lambda _self: tmp_path),
        )
        embedding = OpenAIAsEmbedding(name="t_space_load_race", backend="openai", model="v3", dimensions=2)
        store = LocalEmbeddingStore(name="t_local_load_race")
        store.as_embedding = embedding
        await store.load()

        embedding.model = FakeProviderModel("v4")
        v4_space = embedding.vector_space_id
        np.savez(
            store._cache_path(v4_space),
            keys=np.array([store._cache_key("hello")]),
            embeddings=np.array([[4.0, 0.0]], dtype=np.float16),
        )
        original_to_thread = asyncio.to_thread

        async def switch_during_load(func, *args):
            if getattr(func, "__name__", "") == "_load_sync":
                embedding.model = FakeProviderModel("v3")
            return await original_to_thread(func, *args)

        monkeypatch.setattr(asyncio, "to_thread", switch_during_load)
        await store._sync_cache_space()

        assert store._cache_space == embedding.vector_space_id
        assert not store._cache

    run(go())


def test_whole_request_retries_after_vector_space_changes_between_batches():
    """Completed batches must be discarded when a later batch changes vector space."""

    async def go():
        embedding = FakeAsEmbedding()
        embedding.vector_space_id = "v3"
        store = LocalEmbeddingStore(name="t_local_write_race", max_batch_size=1, enable_cache=False)
        store.as_embedding = embedding
        store._cache_space = embedding.vector_space_id
        calls = 0

        async def switch_during_second_batch(batch, **_kwargs):
            nonlocal calls
            calls += 1
            if calls == 2:
                embedding.vector_space_id = "v4"
            idx, _text, key = batch[0]
            version = 3.0 if calls < 3 else 4.0
            return [(idx, key, np.array([version, 0.0], dtype=np.float16))]

        store._compute_batch = switch_during_second_batch
        results = await store.get_embeddings(["first", "second"])

        assert calls == 4
        assert store._cache_space == embedding.vector_space_id
        for result in results:
            np.testing.assert_array_equal(result, np.array([4.0, 0.0], dtype=np.float16))

    run(go())


def test_whole_request_rereads_cache_after_vector_space_changes(monkeypatch, tmp_path):
    """A cache hit from the old space must not survive a later provider switch."""

    async def go():
        monkeypatch.setattr(
            LocalEmbeddingStore,
            "component_metadata_path",
            property(lambda _self: tmp_path),
        )
        embedding = FakeAsEmbedding()
        embedding.vector_space_id = "v3"
        store = LocalEmbeddingStore(name="t_local_cache_race", enable_cache=True)
        store.as_embedding = embedding
        store._cache_space = embedding.vector_space_id
        first_key = store._cache_key("first")
        store._cache[first_key] = np.array([3.0, 0.0], dtype=np.float16)
        calls = 0

        async def switch_on_miss(batch, **_kwargs):
            nonlocal calls
            calls += 1
            if calls == 1:
                embedding.vector_space_id = "v4"
            version = 3.0 if calls == 1 else 4.0
            return [(idx, key, np.array([version, 0.0], dtype=np.float16)) for idx, _text, key in batch]

        store._compute_batch = switch_on_miss
        results = await store.get_embeddings(["first", "second"])

        assert calls == 2
        for result in results:
            np.testing.assert_array_equal(result, np.array([4.0, 0.0], dtype=np.float16))

    run(go())


def test_whole_request_stops_retrying_when_vector_space_keeps_changing():
    """Continuous configuration churn must discard the whole request instead of blocking forever."""

    async def go():
        embedding = FakeAsEmbedding()
        embedding.vector_space_id = "v3"
        store = LocalEmbeddingStore(name="t_local_write_churn")
        store.as_embedding = embedding
        store._cache_space = embedding.vector_space_id
        calls = 0

        async def change_space_every_time(_batch, **_kwargs):
            nonlocal calls
            calls += 1
            embedding.vector_space_id = f"v{calls + 3}"
            return [(0, "key", np.array([float(calls), 0.0], dtype=np.float16))]

        store._compute_batch = change_space_every_time
        results = await store.get_embeddings(["text"])

        assert calls == 3
        assert results == [None]
        assert "key" not in store._cache

    run(go())


def test_start_ignores_cache_file_without_vector_space_tag(monkeypatch, tmp_path):
    """An unattributable legacy cache is ignored without deleting derived data."""

    async def go():
        monkeypatch.setattr(
            LocalEmbeddingStore,
            "component_metadata_path",
            property(lambda _self: tmp_path),
        )
        store = LocalEmbeddingStore(name="t_local_untagged")
        store.as_embedding = FakeAsEmbedding()
        untagged = tmp_path / f"{store.name}_{store.cache_version}.npz"
        untagged.write_bytes(b"vectors from an unknown model")

        await store._start()

        assert untagged.exists()
        assert not store._cache

    run(go())


def test_batch_failure_falls_back_to_per_item_embedding():
    """One over-limit item must cost one vector, not its whole batch."""

    async def go():
        provider = PerItemLimitAsEmbedding()
        store = LocalEmbeddingStore(name="t_local_embedding_poison_batch", enable_cache=False)
        store.as_embedding = provider

        results = await store._call_with_retry(["ok", "x" * 64, "also ok"])

        assert results == [[1.0, 0.0], None, [1.0, 0.0]]
        assert provider.batch_sizes == [3, 1, 1, 1]

    run(go())


def test_poison_chunk_keeps_its_batch_mates():
    """A rejected chunk must not clear the vectors of the chunks beside it."""

    async def go():
        store = LocalEmbeddingStore(name="t_local_embedding_poison_nodes", enable_cache=False)
        store.as_embedding = PerItemLimitAsEmbedding()
        nodes = [EmbNode(text="first"), EmbNode(text="y" * 64), EmbNode(text="third")]

        await store.get_node_embeddings(nodes)

        assert nodes[0].embedding is not None
        assert nodes[1].embedding is None
        assert nodes[2].embedding is not None

    run(go())


def test_single_item_failure_still_reports_a_failed_batch():
    """A one-item request has nothing to isolate, so the existing contract stands."""

    async def go():
        store = LocalEmbeddingStore(name="t_local_embedding_single_poison", enable_cache=False)
        store.as_embedding = PerItemLimitAsEmbedding()

        assert await store._call_with_retry(["x" * 64]) is None

    run(go())


def test_rate_limited_batch_is_not_split_into_per_item_requests(monkeypatch):
    """A 429 is a provider-level failure; splitting it would multiply the load."""

    async def go():
        sleeps = []

        async def fake_sleep(delay):
            sleeps.append(delay)

        store = LocalEmbeddingStore(name="t_local_embedding_rate_limit_batch", max_retries=2)
        embedding = AlwaysRateLimitedAsEmbedding()
        store.as_embedding = embedding
        monkeypatch.setattr(asyncio, "sleep", fake_sleep)

        assert await store._call_with_retry(["a", "b"]) is None
        assert embedding.calls == 2
        assert sleeps == [1.0]

    run(go())


def test_insufficient_quota_batch_is_not_split_into_per_item_requests(monkeypatch):
    """Quota exhaustion is a service-level failure; splitting it would multiply the load."""

    async def go():
        sleeps = []

        async def fake_sleep(delay):
            sleeps.append(delay)

        store = LocalEmbeddingStore(name="t_local_embedding_quota_batch", max_retries=2)
        embedding = AlwaysInsufficientQuotaAsEmbedding()
        store.as_embedding = embedding
        monkeypatch.setattr(asyncio, "sleep", fake_sleep)

        assert await store._call_with_retry(["a", "b", "c"]) is None
        assert embedding.batch_sizes == [3]
        assert store.is_healthy is False
        assert not sleeps

    run(go())


def test_exhausted_quota_retries_still_do_not_split_the_batch(monkeypatch):
    """Once the opt-in quota retries run out, the store gives up instead of splitting."""

    async def go():
        sleeps = []

        async def fake_sleep(delay):
            sleeps.append(delay)

        store = LocalEmbeddingStore(
            name="t_local_embedding_quota_batch_exhausted",
            max_retries=3,
            quota_retry_delay=5.0,
        )
        embedding = AlwaysInsufficientQuotaAsEmbedding()
        store.as_embedding = embedding
        monkeypatch.setattr(asyncio, "sleep", fake_sleep)

        assert await store._call_with_retry(["a", "b"]) is None
        assert embedding.batch_sizes == [2, 2, 2]
        assert sleeps == [5.0, 5.0]

    run(go())


def test_per_item_fallback_keeps_a_rate_limited_vector(monkeypatch):
    """A 429 on one isolated item still gets the store's backoff and keeps its vector.

    Rate limiting is the failure class the store still owns: no lower layer reports it as a
    provider-wide error worth its own budget, so the retry lives here. Connection errors and
    timeouts are left to the wrapper and the SDK instead; see the transport tests above.
    """

    async def go():
        sleeps = []

        async def fake_sleep(delay):
            sleeps.append(delay)

        provider = PerItemLimitThenTransientErrorAsEmbedding(
            RateLimitError("Requests are too frequent"),
        )
        store = LocalEmbeddingStore(
            name="t_local_embedding_fallback_rate_limit",
            max_retries=3,
            enable_cache=False,
        )
        store.as_embedding = provider
        monkeypatch.setattr(asyncio, "sleep", fake_sleep)

        results = await store._call_with_retry(["ok", "x" * 64, "also ok"])

        assert results == [[1.0, 0.0], None, [1.0, 0.0]]
        assert provider.batch_sizes == [3, 1, 1, 1, 1]
        assert provider.transient_failures == 1
        assert sleeps == [1.0]

    run(go())


def test_provider_wide_credentials_errors_are_not_split_into_per_item_requests():
    """401/403 belong to the credentials: splitting ten items only multiplies the failures."""

    async def go():
        for error in (
            AuthenticationError("invalid api key"),
            PermissionDeniedError("no access to model"),
        ):
            provider = AlwaysProviderWideFailureAsEmbedding(error)
            store = LocalEmbeddingStore(name="t_local_embedding_provider_wide", max_retries=3)
            store.as_embedding = provider

            assert await store._call_with_retry([f"item{i}" for i in range(10)]) is None
            assert provider.batch_sizes == [10]
            assert store.is_healthy is False

    run(go())


def test_provider_wide_server_errors_are_not_retried_without_splitting(monkeypatch):
    """A 500 is left to the layers that own retries, so this store asks exactly once."""

    async def go():
        sleeps = []

        async def fake_sleep(delay):
            sleeps.append(delay)

        provider = AlwaysProviderWideFailureAsEmbedding(InternalServerError("upstream unavailable"))
        store = LocalEmbeddingStore(name="t_local_embedding_provider_server", max_retries=3)
        store.as_embedding = provider
        monkeypatch.setattr(asyncio, "sleep", fake_sleep)

        assert await store._call_with_retry([f"item{i}" for i in range(10)]) is None
        # One request only: the wrapper and the SDK already retry a 5xx, and retrying here
        # as well multiplied their attempts (12 -> 36 in the reported reproduction).
        assert provider.batch_sizes == [10]
        assert not sleeps
        assert store.is_healthy is False

    run(go())


def test_sdk_transport_errors_are_provider_wide_and_not_retried_here(monkeypatch):
    """SDK connection/timeout errors are provider-wide: never split, and not retried here.

    The wrapper and the SDK already retry a transport failure, so the store answers once and
    leaves the retry budget to those layers instead of multiplying their attempts.
    """

    async def go():
        sleeps = []

        async def fake_sleep(delay):
            sleeps.append(delay)

        monkeypatch.setattr(asyncio, "sleep", fake_sleep)

        for error in (APIConnectionError("connection reset"), APITimeoutError("timed out")):
            sleeps.clear()
            provider = AlwaysProviderWideFailureAsEmbedding(error)
            store = LocalEmbeddingStore(name="t_local_embedding_sdk_transport", max_retries=3)
            store.as_embedding = provider

            assert await store._call_with_retry([f"item{i}" for i in range(10)]) is None
            # One request only: the wrapper and the SDK already retry a transport failure
            # (4 wrapper attempts, each retried by the SDK), so another loop here only
            # multiplies them -- a persistent timeout used to cost 36 attempts instead of 12.
            assert provider.batch_sizes == [10]
            assert not sleeps
            assert store.is_healthy is False

    run(go())


def test_a_provider_without_a_retry_policy_keeps_the_store_backoff(monkeypatch):
    """A provider with no lower-layer retry policy still gets the store's backoff.

    This fake exposes no ``model`` with a retry policy, which is the same position Ollama is
    in: its embedding model inherits an empty retryable tuple and its SDK turns a connection
    error into a plain ``ConnectionError`` without retrying. The store's retry is therefore the
    only layer that can keep that vector.
    """

    async def go():
        sleeps = []

        async def fake_sleep(delay):
            sleeps.append(delay)

        for transient in (
            APIConnectionError("connection reset"),
            ConnectionError("connection reset"),
            TimeoutError("request timed out"),
        ):
            sleeps.clear()

            provider = PerItemLimitThenTransientErrorAsEmbedding(transient)
            store = LocalEmbeddingStore(
                name="t_local_embedding_sdk_isolated",
                max_retries=3,
                enable_cache=False,
            )
            store.as_embedding = provider
            monkeypatch.setattr(asyncio, "sleep", fake_sleep)

            results = await store._call_with_retry(["ok", "x" * 64, "also ok"])

            assert results == [[1.0, 0.0], None, [1.0, 0.0]]
            assert provider.batch_sizes == [3, 1, 1, 1, 1]
            assert sleeps == [1.0]

    run(go())


def test_a_persistent_transport_error_is_not_multiplied_by_the_store(monkeypatch):
    """A timeout is asked once: the wrapper and the SDK own the retry budget.

    The provider sees the 4 wrapper attempts x 3 SDK attempts those layers make, while the
    store's max_retries=3 no longer multiplies them -- a persistent timeout used to cost 36
    provider attempts instead of 12.
    """

    async def go():
        endpoint = offline_openai_endpoint(read_timeout).install(monkeypatch)
        provider = openai_wrapper("t_openai_transport_budget")
        store = LocalEmbeddingStore(
            name="t_local_embedding_sdk_transport_budget",
            max_retries=3,
            enable_cache=False,
        )
        store.as_embedding = provider

        assert await store._call_with_retry([f"item{i}" for i in range(3)]) is None
        assert provider.calls == 1
        assert endpoint.requests == 12
        assert store.is_healthy is False

    run(go())


def test_a_transient_transport_error_still_recovers_through_the_sdk(monkeypatch):
    """One timeout is answered by the SDK's own retry: the batch survives and never splits."""

    async def go():
        state = {"timeouts": 1}

        def handler(request):
            if state["timeouts"] > 0:
                state["timeouts"] -= 1
                return read_timeout(request)
            return embedding_payload(request)

        endpoint = offline_openai_endpoint(handler).install(monkeypatch)
        provider = openai_wrapper("t_openai_transport_recover")
        store = LocalEmbeddingStore(
            name="t_local_embedding_sdk_transport_recover",
            max_retries=3,
            enable_cache=False,
        )
        store.as_embedding = provider

        results = await store._call_with_retry([f"item{i}" for i in range(3)])

        assert results == [[0.0, 1.0]] * 3
        assert provider.calls == 1
        assert endpoint.requests == 2

    run(go())


def test_an_isolated_item_still_recovers_through_the_sdk(monkeypatch):
    """After a context-length rejection, a timeout on one remaining item is still recovered.

    The store asks the wrapper once per item; the retry that saves that item belongs to the SDK.
    """

    async def go():
        state = {"timeout_pending": True}

        def handler(request):
            body = json.loads(request.content.decode("utf-8"))
            inputs = body["input"]
            if isinstance(inputs, str):
                inputs = [inputs]
            if any(len(item) > 32 for item in inputs):
                return context_length_rejection()
            if len(inputs) == 1 and state["timeout_pending"]:
                state["timeout_pending"] = False
                return read_timeout(request)
            return embedding_payload(request)

        endpoint = offline_openai_endpoint(handler).install(monkeypatch)
        provider = openai_wrapper("t_openai_transport_isolated")
        store = LocalEmbeddingStore(
            name="t_local_embedding_sdk_transport_isolated",
            max_retries=3,
            enable_cache=False,
        )
        store.as_embedding = provider

        results = await store._call_with_retry(["ok", "x" * 64, "also ok"])

        assert results == [[0.0, 1.0], None, [0.0, 1.0]]
        assert provider.calls == 4
        assert endpoint.requests == 5

    run(go())


def test_provider_wide_statuses_exclude_input_dependent_bad_request():
    """404 is provider-wide, but an over-limit 400 must stay input-dependent."""

    async def go():
        store = LocalEmbeddingStore(name="t_local_embedding_classifier")
        for error in (
            AuthenticationError("bad key"),
            PermissionDeniedError("no access"),
            NotFoundError("model_not_found"),
            InternalServerError("boom"),
        ):
            assert store._is_provider_wide_failure(error) is True
        over_limit = ContextLengthError("the input length exceeds the context length")
        assert store._is_provider_wide_failure(over_limit) is False

    run(go())


def test_provider_wide_not_found_errors_are_not_split(monkeypatch):
    """A missing model or endpoint is provider-wide: one request, then failure, never a split."""

    async def go():
        sleeps = []

        async def fake_sleep(delay):
            sleeps.append(delay)

        provider = AlwaysProviderWideFailureAsEmbedding(NotFoundError("model_not_found"))
        store = LocalEmbeddingStore(name="t_local_embedding_not_found", max_retries=3)
        store.as_embedding = provider
        monkeypatch.setattr(asyncio, "sleep", fake_sleep)

        assert await store._call_with_retry([f"item{i}" for i in range(10)]) is None
        assert provider.batch_sizes == [10]
        assert not sleeps

    run(go())


class GenAiClientError(Exception):
    """Google GenAI client error: the HTTP status is an integer ``code``, not ``status_code``."""

    def __init__(self, code: int, message: str = "client error"):
        super().__init__(message)
        self.code = code


class GenAiServerError(Exception):
    """Google GenAI server error: the HTTP status is an integer ``code``, not ``status_code``."""

    def __init__(self, code: int, message: str = "server error"):
        super().__init__(message)
        self.code = code


class StringCodeError(Exception):
    """OpenAI-compatible error whose ``code`` names the error instead of carrying a status."""

    code = "invalid_api_key"


class BadRequestError(Exception):
    """OpenAI-compatible 400 that names no input, such as a malformed payload."""

    status_code = 400


class ResponseLike:
    """httpx-style wrapper that keeps the status on a ``response`` attribute."""

    def __init__(self, status_code: int):
        self.status_code = status_code


class WrappedTransportError(Exception):
    """Error that only exposes its status through ``response.status_code``."""

    def __init__(self, status_code: int):
        super().__init__("wrapped")
        self.response = ResponseLike(status_code)


class DashScopeRuntimeError(RuntimeError):
    """AgentScope turns every non-200 DashScope response into a plain ``RuntimeError``.

    The structured response is embedded in the message instead of being kept as an attribute,
    so a client that only reads exception attributes cannot classify it.
    """

    @staticmethod
    def for_status(status_code: int, code: str = "InvalidApiKey") -> "DashScopeRuntimeError":
        """Build the exception AgentScope raises for a non-200 DashScope response."""
        response = {
            "status_code": status_code,
            "request_id": "req-0",
            "code": code,
            "message": "request rejected",
            "output": None,
            "usage": None,
        }
        return DashScopeRuntimeError(f"Failed to get embedding from DashScope API: {json.dumps(response)}")


def test_http_status_reads_every_supported_exception_shape():
    """Read the status from status_code, an integer code, response.status_code or the message."""

    store = LocalEmbeddingStore(name="t_local_embedding_status_shapes")

    assert store._http_status(AuthenticationError("bad key")) == 401
    assert store._http_status(NotFoundError("model_not_found")) == 404
    assert store._http_status(GenAiClientError(401)) == 401
    assert store._http_status(GenAiServerError(503)) == 503
    assert store._http_status(WrappedTransportError(404)) == 404
    assert store._http_status(DashScopeRuntimeError.for_status(401)) == 401
    assert store._http_status(ValueError("no status here")) is None


def test_http_status_ignores_a_string_code():
    """A string ``code`` names the error; reading it as a status would misclassify quota errors."""

    store = LocalEmbeddingStore(name="t_local_embedding_status_string_code")

    assert store._http_status(StringCodeError("invalid key")) is None
    assert store._is_insufficient_quota(StringCodeError("invalid key")) is False
    assert store._http_status(InsufficientQuotaError("quota exhausted")) == 429
    assert store._is_insufficient_quota(InsufficientQuotaError("quota exhausted")) is True


def test_google_genai_service_failures_are_not_split_into_per_item_requests(monkeypatch):
    """A Gemini 401 / 429 / 503 describes the service, so the batch is never split per item.

    A 503 is answered once: the wrapper and the SDK own the retry for a 5xx, and retrying
    here as well multiplied their attempts. A 429 keeps its retry because the rate-limit
    branch above this classifier is the one place that reacts to it.
    """

    async def go():
        sleeps = []

        async def fake_sleep(delay):
            sleeps.append(delay)

        monkeypatch.setattr(asyncio, "sleep", fake_sleep)

        for error, expected_batches in (
            (GenAiClientError(401), [10]),
            (GenAiServerError(503), [10]),
            (GenAiClientError(429), [10, 10, 10]),
        ):
            sleeps.clear()
            provider = AlwaysProviderWideFailureAsEmbedding(error)
            store = LocalEmbeddingStore(
                name="t_local_embedding_genai",
                max_retries=3,
                enable_cache=False,
            )
            store.as_embedding = provider

            assert await store._call_with_retry([f"item{i}" for i in range(10)]) is None
            assert provider.batch_sizes == expected_batches

    run(go())


def test_dashscope_runtime_errors_are_not_split_into_per_item_requests(monkeypatch):
    """A DashScope failure arrives as a plain RuntimeError; the status lives in its message.

    As with the Gemini shapes, a 503 is answered once (the lower layers own that retry) and
    a 429 keeps the retry from the rate-limit branch.
    """

    async def go():
        sleeps = []

        async def fake_sleep(delay):
            sleeps.append(delay)

        monkeypatch.setattr(asyncio, "sleep", fake_sleep)

        for status, expected_batches in ((401, [10]), (503, [10]), (429, [10, 10, 10])):
            sleeps.clear()
            provider = AlwaysProviderWideFailureAsEmbedding(DashScopeRuntimeError.for_status(status))
            store = LocalEmbeddingStore(
                name="t_local_embedding_dashscope",
                max_retries=3,
                enable_cache=False,
            )
            store.as_embedding = provider

            assert await store._call_with_retry([f"item{i}" for i in range(10)]) is None
            assert provider.batch_sizes == expected_batches

    run(go())


def test_unrecognized_errors_are_not_split_into_per_item_requests():
    """Only an input-dependent error may split; anything unrecognized stays a batch failure."""

    async def go():
        provider = AlwaysProviderWideFailureAsEmbedding(RuntimeError("provider said no"))
        store = LocalEmbeddingStore(name="t_local_embedding_unrecognized", max_retries=3)
        store.as_embedding = provider

        assert await store._call_with_retry(["a", "b", "c"]) is None
        assert provider.batch_sizes == [3]

    run(go())


def test_input_dependent_classifier_pins_both_halves_of_the_boundary():
    """The isolation must keep working for the one error a single item can explain."""

    store = LocalEmbeddingStore(name="t_local_embedding_input_dependent")

    assert (
        store._is_input_dependent_failure(
            ContextLengthError("the input length exceeds the context length"),
        )
        is True
    )
    assert store._is_input_dependent_failure(ContextLengthError("too long")) is True
    assert store._is_input_dependent_failure(AuthenticationError("bad key")) is False
    assert store._is_input_dependent_failure(InternalServerError("boom")) is False
    assert store._is_input_dependent_failure(GenAiClientError(400)) is False
    assert store._is_input_dependent_failure(RuntimeError("provider said no")) is False


def test_bad_request_without_an_input_length_signal_is_not_split():
    """A bare 400 does not name an input, so the batch is not divided over it."""

    async def go():
        provider = AlwaysProviderWideFailureAsEmbedding(BadRequestError("malformed payload"))
        store = LocalEmbeddingStore(name="t_local_embedding_bare_400", max_retries=2)
        store.as_embedding = provider

        assert await store._call_with_retry(["a", "b", "c"]) is None
        assert provider.batch_sizes == [3]

    run(go())


def test_isolation_stops_when_a_provider_wide_failure_hits():
    """Once isolation meets a provider-wide failure, the remaining items are not requested.

    The vectors already recovered stay in the result: only the items that were never answered
    are dropped, and no further request is sent for them.
    """

    async def go():
        provider = PerItemLimitThenQuotaAsEmbedding()
        store = LocalEmbeddingStore(
            name="t_local_embedding_isolation_stops",
            max_retries=2,
            enable_cache=False,
        )
        store.as_embedding = provider

        results = await store._call_with_retry(["ok", "x" * 64, "also ok", "still ok"])

        assert results == [[1.0, 0.0], None, None, None]
        # One rejected batch, the recovered first item, the over-limit item, then the quota
        # failure. The fourth item is never requested once the quota failure is seen.
        assert provider.batch_sizes == [4, 1, 1, 1]

    run(go())


def test_partial_success_keeps_a_failed_provider_unhealthy():
    """Recovered vectors must not overwrite the health state of a provider-wide failure.

    Isolation can recover an early item and then meet a provider-wide failure such as
    exhausted quota. ``_call_with_retry`` marks the store unhealthy and stops, but the
    vectors recovered before it are still returned. Treating that partial success as a
    healthy provider would hide the outage, and it would break the unhealthy-to-healthy
    transition the file store uses to schedule missing-vector backfill. The existing
    isolation test drives ``_call_with_retry`` directly, so it never sees the overwrite.
    """

    async def go():
        provider = PerItemLimitThenQuotaAsEmbedding()
        store = LocalEmbeddingStore(
            name="t_local_embedding_partial_health",
            max_retries=2,
            enable_cache=False,
        )
        store.as_embedding = provider

        vectors = await store.get_embeddings(["ok", "x" * 64, "also ok", "still ok"])

        # The vector recovered before the outage is kept ...
        assert vectors[0] is not None
        assert vectors[1] is None and vectors[2] is None and vectors[3] is None
        # ... and the provider-wide failure still decides the health state.
        assert store.is_healthy is False
        # One rejected batch, the recovered first item, the over-limit item, then the quota
        # failure: the fourth item is never requested.
        assert provider.batch_sizes == [4, 1, 1, 1]

    run(go())


def test_isolation_keeps_going_while_only_inputs_fail():
    """An input-dependent failure is local to its item, so isolation must continue.

    This pins the other half of the boundary: stopping on every isolated failure would lose
    the vectors of the items that are perfectly embeddable on their own.
    """

    async def go():
        provider = PerItemLimitAsEmbedding()
        store = LocalEmbeddingStore(
            name="t_local_embedding_isolation_keeps_going",
            max_retries=2,
            enable_cache=False,
        )
        store.as_embedding = provider

        results = await store._call_with_retry(["ok", "x" * 64, "also ok", "still ok"])

        assert results == [[1.0, 0.0], None, [1.0, 0.0], [1.0, 0.0]]
        assert provider.batch_sizes == [4, 1, 1, 1, 1]

    run(go())


def test_dashscope_serialized_quota_errors_keep_their_quota_policy(monkeypatch):
    """A serialized DashScope quota error must classify as quota, not as a rate limit.

    AgentScope embeds the structured response in the message, so the business error code has to
    be read back out of it. Without that the opt-in quota retry policy is never applied and the
    error is retried with the generic backoff instead.
    """

    async def go():
        sleeps = []

        async def fake_sleep(delay):
            sleeps.append(delay)

        monkeypatch.setattr(asyncio, "sleep", fake_sleep)

        error = DashScopeRuntimeError.for_status(429, code="insufficient_quota")
        assert LocalEmbeddingStore._is_insufficient_quota(error) is True
        assert LocalEmbeddingStore._is_rate_limited(error) is False

        provider = AlwaysProviderWideFailureAsEmbedding(error)
        store = LocalEmbeddingStore(
            name="t_local_embedding_dashscope_quota",
            max_retries=3,
            quota_retry_delay=7.0,
            enable_cache=False,
        )
        store.as_embedding = provider

        assert await store._call_with_retry(["a", "b"]) is None
        assert provider.batch_sizes == [2, 2, 2]
        assert sleeps == [7.0, 7.0]

    run(go())


class ContextLength413Error(Exception):
    """An over-limit input reported as HTTP 413 instead of 400.

    Ollama v0.40.3 normalizes an embedding input-limit rejection to 413 and propagates that
    status through the embedding endpoint, so the same over-long item that arrives as a 400
    from an OpenAI-compatible gateway is a 413 for that provider.
    """

    status_code = 413
    body = {"error": {"message": "the input length exceeds the context length"}}


class RequestBodyTooLargeError(Exception):
    """A bare 413 that never names an input: the request body itself is too large."""

    status_code = 413


class PerItemLimitAs413Embedding:
    """Reject a whole request with HTTP 413 when any single text is over the per-item limit."""

    dimensions = 2
    vector_space_id = "fakespace000"
    max_chars = 32

    def __init__(self):
        self.batch_sizes: list[int] = []

    def initialize_model(self):
        """Mirror the real component's idempotent initialization hook."""

    async def __call__(self, texts: list[str], **_kwargs):
        self.batch_sizes.append(len(texts))
        if any(len(text) > self.max_chars for text in texts):
            raise ContextLength413Error("the input length exceeds the context length")
        return [[1.0, 0.0] for _ in texts]


def test_context_length_rejection_reported_as_http_413_is_isolated():
    """An over-limit input carried by a 413 must still reach the per-item isolation path.

    Ollama normalizes an embedding input-limit rejection to 413, so a classifier that only
    accepts 400 rejects the batch instead of isolating the item that caused it, and every
    healthy vector in the request is lost.
    """

    async def go():
        provider = PerItemLimitAs413Embedding()
        store = LocalEmbeddingStore(
            name="t_local_embedding_413_isolated",
            max_retries=2,
            enable_cache=False,
        )
        store.as_embedding = provider

        results = await store._call_with_retry(["ok", "x" * 64, "also ok"])

        assert results == [[1.0, 0.0], None, [1.0, 0.0]]
        assert provider.batch_sizes == [3, 1, 1, 1]

    run(go())


def test_bare_http_413_without_an_input_signal_is_not_split():
    """A 413 that names no input stays provider-wide, so the batch is not divided over it.

    The status alone is ambiguous: it can describe an unrelated request-body limit, which no
    single input explains, so only a 413 that blames the input itself may fan out per item.
    """

    async def go():
        provider = AlwaysProviderWideFailureAsEmbedding(
            RequestBodyTooLargeError("request entity too large"),
        )
        store = LocalEmbeddingStore(name="t_local_embedding_bare_413", max_retries=2)
        store.as_embedding = provider

        assert await store._call_with_retry(["a", "b", "c"]) is None
        assert provider.batch_sizes == [3]

    run(go())


def test_input_dependent_classifier_accepts_a_413_only_with_an_input_signal():
    """Pin both halves of the 413 boundary, and that the 400 behaviour is unchanged."""

    store = LocalEmbeddingStore(name="t_local_embedding_413_classifier")

    assert (
        store._is_input_dependent_failure(
            ContextLength413Error("the input length exceeds the context length"),
        )
        is True
    )
    assert store._is_input_dependent_failure(ContextLength413Error("too long")) is True
    assert store._is_input_dependent_failure(RequestBodyTooLargeError("request entity too large")) is False
    assert store._is_input_dependent_failure(ContextLengthError("too long")) is True
    assert store._is_input_dependent_failure(GenAiClientError(400)) is False


class _NoRetryPolicyModel:
    """Stand-in for AgentScope's Ollama embedding model: an empty policy plus a numeric budget.

    The real model keeps the inherited ``max_retries=3`` next to the empty tuple, so the double
    has to carry both fields: omitting the number is what let the previous helper mistake the
    budget for a policy.
    """

    max_retries = 3

    @classmethod
    def _get_retryable_exceptions(cls):
        """Mirror the inherited empty policy."""
        return ()


class _PartialRetryPolicyModel:
    """Stand-in for AgentScope's DashScope embedding model: ``(RuntimeError,)`` and nothing else."""

    max_retries = 3

    @classmethod
    def _get_retryable_exceptions(cls):
        """Mirror the DashScope policy, which does not cover a connection error."""
        return (RuntimeError,)


class OllamaLikeTransientProvider(TransientThenSuccessAsEmbedding):
    """An Ollama-shaped provider: the bound model declares no retry policy."""

    def __init__(self, error: Exception, failures: int = 1):
        super().__init__(error, failures)
        self.model = _NoRetryPolicyModel()


def test_a_provider_with_an_empty_retry_policy_recovers_a_connection_error(monkeypatch):
    """The store is the retry owner when the model below declares no policy.

    Ollama's embedding model inherits an empty retryable tuple and its SDK converts
    ``httpx.ConnectError`` into a plain ``ConnectionError`` without retrying, so leaving the
    retry to the lower layers would turn one dropped connection into a lost vector. The base
    implementation answered that case with two requests; so must this one.
    """

    async def go():
        sleeps = []

        async def fake_sleep(delay):
            sleeps.append(delay)

        monkeypatch.setattr(asyncio, "sleep", fake_sleep)

        provider = OllamaLikeTransientProvider(ConnectionError("connection reset"))
        store = LocalEmbeddingStore(
            name="t_local_embedding_empty_retry_policy",
            max_retries=3,
            enable_cache=False,
        )
        store.as_embedding = provider

        results = await store._call_with_retry(["a", "b", "c"])

        assert results == [[1.0, 0.0]] * 3
        assert provider.batch_sizes == [3, 3]
        assert sleeps == [1.0]

    run(go())


def test_a_policy_that_does_not_cover_the_error_keeps_the_store_backoff(monkeypatch):
    """A nonempty policy is not ownership unless it covers the error that actually occurred.

    AgentScope's DashScope model declares ``(RuntimeError,)``, which does not cover a connection
    error raised by ``requests`` underneath it. Trusting the whole tuple would leave that retry
    to a layer that re-raises the error immediately, so one dropped connection would still cost
    the vector. The store has to answer it as the base implementation did.
    """

    async def go():
        sleeps = []

        async def fake_sleep(delay):
            sleeps.append(delay)

        monkeypatch.setattr(asyncio, "sleep", fake_sleep)

        provider = OllamaLikeTransientProvider(ConnectionError("connection reset"))
        provider.model = _PartialRetryPolicyModel()
        store = LocalEmbeddingStore(
            name="t_local_embedding_partial_policy",
            max_retries=3,
            enable_cache=False,
        )
        store.as_embedding = provider

        results = await store._call_with_retry(["a", "b", "c"])

        assert results == [[1.0, 0.0]] * 3
        assert provider.batch_sizes == [3, 3]
        assert sleeps == [1.0]

    run(go())


def test_an_isolated_item_does_not_multiply_sdk_transport_retries(monkeypatch):
    """Isolating an item must not add a third retry loop over the SDK's budget.

    A rejected batch costs one request; the first isolated item then hits a persistent timeout,
    which the wrapper and the SDK already try 12 times. The store must stop there (13 requests
    in total) instead of asking that item three more times (37), and the failure is provider
    wide, so the remaining items are left unanswered rather than probed one by one.
    """

    async def go():
        def handler(request):
            body = json.loads(request.content.decode("utf-8"))
            inputs = body["input"]
            if isinstance(inputs, str):
                inputs = [inputs]
            if any(len(item) > 32 for item in inputs):
                return context_length_rejection()
            return read_timeout(request)

        endpoint = offline_openai_endpoint(handler).install(monkeypatch)
        provider = openai_wrapper("t_openai_transport_isolated_budget")
        store = LocalEmbeddingStore(
            name="t_local_embedding_sdk_isolated_budget",
            max_retries=3,
            enable_cache=False,
        )
        store.as_embedding = provider

        results = await store._call_with_retry(["ok", "x" * 64, "also ok"])

        assert results == [None, None, None]
        assert provider.calls == 2
        assert endpoint.requests == 13
        assert store.is_healthy is False

    run(go())