From 4026aa6575ab13b86e08e9d149d64e5aeb24e232 Mon Sep 17 00:00:00 2001
From: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
Date: Thu, 27 Aug 2026 20:58:14 -0700
Subject: [PATCH] fix(a2a): merge fresh agent vectors into the live cache and
 drop entries of another dimension

---
 litellm/proxy/agent_endpoints/agent_search.py | 12 +++++++---
 .../agent_endpoints/test_agent_search.py      | 22 +++++++++++++++++++
 2 files changed, 31 insertions(+), 3 deletions(-)

diff --git a/litellm/proxy/agent_endpoints/agent_search.py b/litellm/proxy/agent_endpoints/agent_search.py
index c249079ccac..46ab36d7b72 100644
--- a/litellm/proxy/agent_endpoints/agent_search.py
+++ b/litellm/proxy/agent_endpoints/agent_search.py
@@ -193,6 +193,14 @@ class AgentSearchIndex:
     def __init__(self) -> None:
         self._vectors: Mapping[str, Mapping[str, Vector]] = MappingProxyType({})
 
+    def _merged(self, embedding_model: str, embedded: _Embedded) -> Mapping[str, Vector]:
+        kept: Final = {
+            text: vector
+            for text, vector in self._vectors.get(embedding_model, _NO_VECTORS).items()
+            if len(vector) == len(embedded.query_vector)
+        }
+        return MappingProxyType({**kept, **embedded.vectors})
+
     async def search(
         self, query: str, agents: Sequence[AgentResponse], top_k: int, embed: Embedder, embedding_model: str
     ) -> AgentSearchHits | AgentSearchEmbeddingFailed:
@@ -207,9 +215,7 @@ class AgentSearchIndex:
             return AgentSearchEmbeddingFailed(
                 reason=f"embedding model {embedding_model} returned vectors of mixed dimensions"
             )
-        self._vectors = MappingProxyType(
-            {**self._vectors, embedding_model: MappingProxyType({**cached, **embedded.vectors})}
-        )
+        self._vectors = MappingProxyType({**self._vectors, embedding_model: self._merged(embedding_model, embedded)})
         ranked: Final = sorted(
             (
                 AgentSearchHit(agent=agent, score=cosine_similarity(embedded.query_vector, embedded.vectors[text]))
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 93f29c4a1d2..3fb09076e5f 100644
--- a/tests/test_litellm/proxy/agent_endpoints/test_agent_search.py
+++ b/tests/test_litellm/proxy/agent_endpoints/test_agent_search.py
@@ -1,3 +1,4 @@
+import asyncio
 from collections.abc import Sequence
 from types import MappingProxyType
 from typing import Final
@@ -84,6 +85,7 @@ class FixedDimensionEmbedder:
 
     async def __call__(self, texts: Sequence[str]) -> Sequence[Vector]:
         self.calls.append(tuple(texts))
+        await asyncio.sleep(0)
         return tuple((1.0,) * self.dimensions for _ in texts)
 
 
@@ -156,6 +158,26 @@ class TestAgentSearchIndex:
             ("language translation", *(agent_search_text(agent) for agent in AGENTS)),
         ]
 
+    @pytest.mark.asyncio
+    async def test_re_embedding_a_subset_drops_the_other_agents_old_vectors(self) -> None:
+        index = AgentSearchIndex()
+        await index.search("language translation", AGENTS, top_k=5, embed=FakeEmbedder(), embedding_model="m")
+        wide = FixedDimensionEmbedder(2)
+        await index.search("language translation", AGENTS[:1], top_k=5, embed=wide, embedding_model="m")
+        await index.search("language translation", AGENTS, top_k=5, embed=wide, embedding_model="m")
+        assert wide.calls[-1] == ("language translation", *(agent_search_text(agent) for agent in AGENTS[1:]))
+
+    @pytest.mark.asyncio
+    async def test_concurrent_searches_keep_each_others_vectors(self) -> None:
+        index = AgentSearchIndex()
+        embedder = FixedDimensionEmbedder(3)
+        await asyncio.gather(
+            index.search("q", AGENTS[:1], top_k=5, embed=embedder, embedding_model="m"),
+            index.search("q", AGENTS[1:], top_k=5, embed=embedder, embedding_model="m"),
+        )
+        await index.search("q", AGENTS, top_k=5, embed=embedder, embedding_model="m")
+        assert embedder.calls[-1] == ("q",)
+
     @pytest.mark.asyncio
     async def test_mixed_dimensions_in_one_batch_become_embedding_failed(self) -> None:
         async def mixed(texts: Sequence[str]) -> Sequence[Vector]: