From 2ed2486bf40e442e9262067c8c2e3389dd75f637 Mon Sep 17 00:00:00 2001
From: Cursor Agent <cursoragent@cursor.com>
Date: Fri, 28 Aug 2026 02:33:19 +0000
Subject: [PATCH] fix(a2a): drop cached agent vectors whose dimension no longer
 matches the query

When the configured agent_search_embedding_model or a router fallback
switches to a model with a different embedding dimension, the query
vector shape stops matching cached agent vectors. cosine_similarity uses
zip(..., strict=True), so ranking raised ValueError outside the
try/except that treats embedding failures as AgentSearchEmbeddingFailed.
Because the cache never dropped the old vectors, every subsequent
GET /v1/agents?query= and agent_search MCP call returned 500 until the
worker restarted.

After embedding the query, filter the cache to only vectors of the same
dimension and re-embed any agent texts that got evicted.
---
 litellm/proxy/agent_endpoints/agent_search.py  | 18 ++++++++++++++++--
 .../proxy/agent_endpoints/test_agent_search.py | 13 +++++++++++++
 2 files changed, 29 insertions(+), 2 deletions(-)

diff --git a/litellm/proxy/agent_endpoints/agent_search.py b/litellm/proxy/agent_endpoints/agent_search.py
index 896a22fbf9d..fe4b255235f 100644
--- a/litellm/proxy/agent_endpoints/agent_search.py
+++ b/litellm/proxy/agent_endpoints/agent_search.py
@@ -164,10 +164,24 @@ class AgentSearchIndex:
             return AgentSearchEmbeddingFailed(
                 reason=f"embedding model returned {len(vectors)} vectors for {len(missing) + 1} inputs"
             )
-        self._vectors = MappingProxyType(dict(chain(self._vectors.items(), zip(missing, vectors[1:], strict=True))))
+        query_vector: Final = vectors[0]
+        fresh: Final = dict(zip(missing, vectors[1:], strict=True))
+        kept: Final = {text: vec for text, vec in self._vectors.items() if len(vec) == len(query_vector)}
+        stale: Final = tuple(dict.fromkeys(text for text in texts if text not in kept and text not in fresh))
+        try:
+            refreshed: Final = await embed(stale) if stale else ()
+        except (OpenAIError, ValueError, BudgetExceededError) as exc:
+            return AgentSearchEmbeddingFailed(reason=f"re-embedding stale agent texts failed: {exc}")
+        if len(refreshed) != len(stale):
+            return AgentSearchEmbeddingFailed(
+                reason=f"embedding model returned {len(refreshed)} vectors for {len(stale)} inputs"
+            )
+        self._vectors = MappingProxyType(
+            dict(chain(kept.items(), fresh.items(), zip(stale, refreshed, strict=True)))
+        )
         ranked: Final = sorted(
             (
-                AgentSearchHit(agent=agent, score=cosine_similarity(vectors[0], self._vectors[text]))
+                AgentSearchHit(agent=agent, score=cosine_similarity(query_vector, self._vectors[text]))
                 for agent, text in zip(agents, texts, strict=True)
             ),
             key=lambda hit: hit.score,
diff --git a/tests/test_litellm/proxy/agent_endpoints/test_agent_search.py b/tests/test_litellm/proxy/agent_endpoints/test_agent_search.py
index 4b674eb142b..c734d2fd9f3 100644
--- a/tests/test_litellm/proxy/agent_endpoints/test_agent_search.py
+++ b/tests/test_litellm/proxy/agent_endpoints/test_agent_search.py
@@ -147,6 +147,19 @@ class TestAgentSearchIndex:
         outcome = await AgentSearchIndex().search("q", AGENTS, top_k=5, embed=short)
         assert isinstance(outcome, AgentSearchEmbeddingFailed)
 
+    @pytest.mark.asyncio
+    async def test_dimension_change_invalidates_cached_agent_vectors(self) -> None:
+        index = AgentSearchIndex()
+        warm = FakeEmbedder()
+        await index.search("language translation", AGENTS, top_k=5, embed=warm)
+
+        async def wider(texts: Sequence[str]) -> Sequence[Vector]:
+            return tuple((1.0, 0.0, 0.0, 0.0) for _ in texts)
+
+        outcome = await index.search("language translation", AGENTS, top_k=5, embed=wider)
+        assert isinstance(outcome, AgentSearchHits)
+        assert len(outcome.hits) == len(AGENTS)
+
 
 class TestSearchAgents:
     @pytest.mark.asyncio