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.
This commit is contained in:
Cursor Agent 2026-08-28 02:33:19 +00:00
parent e9cc9c9bc3
commit 2ed2486bf4
No known key found for this signature in database
2 changed files with 29 additions and 2 deletions

View file

@ -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,

View file

@ -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