# encoding:utf-8 """ Unit tests for on-demand MCP tool retrieval. Covers the invariants the maintainer asked for in the feature review: * below threshold / degrade paths behave exactly like today (full injection), * the injected MCP tool set only ever grows within a run (never shrinks), plus the pure selection helpers (query building, cosine, top-k, fallbacks). """ import os import sys import unittest from unittest.mock import patch sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..")) from agent.tools.mcp.tool_retrieval import ( build_retrieval_query, cosine_similarity, select_mcp_tools, ) from agent.tools.mcp.mcp_tool import McpTool def _mcp_tool(name, description=""): """Build a real McpTool without needing a live MCP client.""" return McpTool( client=None, tool_schema={"name": name, "description": description, "inputSchema": {}}, server_name="test-server", ) class _FakeBuiltinTool: """Minimal stand-in for a built-in BaseTool.""" def __init__(self, name): self.name = name self.description = f"builtin {name}" self.params = {"type": "object", "properties": {}} def get_json_schema(self): return {} class _FakeToolManager: """Controls the vectors/query embeddings seen by the executor.""" def __init__(self, tool_vectors, query_vectors): self._tool_vectors = tool_vectors self._query_vectors = list(query_vectors) self._call = 0 def get_mcp_tool_vectors(self): return dict(self._tool_vectors) def embed_query(self, text): if self._call < len(self._query_vectors): vec = self._query_vectors[self._call] else: vec = self._query_vectors[-1] if self._query_vectors else None self._call += 1 return vec # -------------------------------------------------------------------------- # Pure helper: build_retrieval_query # -------------------------------------------------------------------------- class TestBuildRetrievalQuery(unittest.TestCase): def test_empty_messages(self): self.assertEqual(build_retrieval_query([]), "") def test_extracts_text_blocks(self): messages = [ {"role": "user", "content": [{"type": "text", "text": "hello"}]}, {"role": "assistant", "content": [{"type": "text", "text": "world"}]}, ] self.assertEqual(build_retrieval_query(messages), "hello\nworld") def test_skips_tool_result_and_tool_use(self): messages = [ {"role": "user", "content": [{"type": "text", "text": "do it"}]}, {"role": "assistant", "content": [ {"type": "tool_use", "name": "read", "input": {}}, ]}, {"role": "user", "content": [ {"type": "tool_result", "content": "huge payload " * 100}, ]}, ] self.assertEqual(build_retrieval_query(messages), "do it") def test_string_content_supported(self): messages = [{"role": "user", "content": "plain string"}] self.assertEqual(build_retrieval_query(messages), "plain string") def test_respects_recent_window(self): messages = [ {"role": "user", "content": [{"type": "text", "text": f"m{i}"}]} for i in range(10) ] # Only the last 3 messages should be kept. self.assertEqual(build_retrieval_query(messages, max_messages=3), "m7\nm8\nm9") # -------------------------------------------------------------------------- # Pure helper: cosine_similarity # -------------------------------------------------------------------------- class TestCosineSimilarity(unittest.TestCase): def test_identical_vectors(self): self.assertAlmostEqual(cosine_similarity([1.0, 0.0], [1.0, 0.0]), 1.0) def test_orthogonal_vectors(self): self.assertAlmostEqual(cosine_similarity([1.0, 0.0], [0.0, 1.0]), 0.0) def test_degenerate_inputs(self): self.assertEqual(cosine_similarity([], [1.0]), 0.0) self.assertEqual(cosine_similarity([0.0, 0.0], [1.0, 1.0]), 0.0) self.assertEqual(cosine_similarity([1.0], [1.0, 0.0]), 0.0) # -------------------------------------------------------------------------- # Pure helper: select_mcp_tools # -------------------------------------------------------------------------- class TestSelectMcpTools(unittest.TestCase): def setUp(self): self.vectors = { "a": [1.0, 0.0, 0.0], "b": [0.0, 1.0, 0.0], "c": [0.0, 0.0, 1.0], "d": [0.9, 0.1, 0.0], } def test_returns_top_k(self): selected = select_mcp_tools([1.0, 0.0, 0.0], self.vectors, top_k=2, already_selected=set()) self.assertEqual(selected, {"a", "d"}) def test_union_only_grows_across_turns(self): """The core invariant: a later turn never drops earlier selections.""" first = select_mcp_tools([1.0, 0.0, 0.0], self.vectors, top_k=2, already_selected=set()) second = select_mcp_tools([0.0, 1.0, 0.0], self.vectors, top_k=1, already_selected=first) self.assertTrue(first.issubset(second)) self.assertIn("b", second) def test_none_query_vector_falls_back(self): self.assertIsNone(select_mcp_tools(None, self.vectors, top_k=2, already_selected=set())) def test_empty_index_falls_back(self): self.assertIsNone(select_mcp_tools([1.0, 0.0, 0.0], {}, top_k=2, already_selected=set())) def test_dimension_mismatch_falls_back(self): mismatched = {"a": [1.0, 0.0]} # dim 2 vs query dim 3 self.assertIsNone(select_mcp_tools([1.0, 0.0, 0.0], mismatched, top_k=2, already_selected=set())) # -------------------------------------------------------------------------- # Executor integration: AgentStream._select_tools_for_injection # -------------------------------------------------------------------------- class TestSelectToolsForInjection(unittest.TestCase): """Exercise the executor decision without spinning up a real agent.""" def _make_self(self, mcp_count, builtins=("read", "write", "bash")): from types import SimpleNamespace tools = {} for name in builtins: tools[name] = _FakeBuiltinTool(name) for i in range(mcp_count): name = f"mcp_{i}" tools[name] = _mcp_tool(name, f"tool number {i}") return SimpleNamespace( tools=tools, messages=[{"role": "user", "content": [{"type": "text", "text": "hi"}]}], _retrieved_mcp_names=set(), ) def _call(self, fake_self): from agent.protocol.agent_stream import AgentStreamExecutor return AgentStreamExecutor._select_tools_for_injection(fake_self) def _conf(self, **overrides): cfg = { "mcp_tool_retrieval_enabled": True, "mcp_tool_retrieval_threshold": 20, "mcp_tool_retrieval_top_k": 2, } cfg.update(overrides) return cfg def test_disabled_returns_all_tools(self): fake = self._make_self(mcp_count=50) with patch("config.conf", return_value=self._conf(mcp_tool_retrieval_enabled=False)): result = self._call(fake) self.assertEqual(len(result), len(fake.tools)) def test_below_threshold_returns_all_tools(self): """Maintainer scenario 1: below threshold → behavior unchanged.""" fake = self._make_self(mcp_count=5) # <= threshold 20 with patch("config.conf", return_value=self._conf()): result = self._call(fake) self.assertEqual(len(result), len(fake.tools)) def test_degrade_no_provider_returns_all_tools(self): """Maintainer scenario 2: no embedding provider → full injection.""" fake = self._make_self(mcp_count=25) # > threshold fake_tm = _FakeToolManager(tool_vectors={}, query_vectors=[None]) with patch("config.conf", return_value=self._conf()), \ patch("agent.tools.ToolManager", return_value=fake_tm): result = self._call(fake) self.assertEqual(len(result), len(fake.tools)) def test_builtins_always_injected_and_set_grows(self): """Maintainer scenario 3: multi-turn MCP set only grows; builtins stay.""" fake = self._make_self(mcp_count=25) # Deterministic vectors: mcp_0 wins turn 1, mcp_1 wins turn 2. tool_vectors = {f"mcp_{i}": [0.1, 0.1, 0.1] for i in range(25)} tool_vectors["mcp_0"] = [1.0, 0.0, 0.0] tool_vectors["mcp_1"] = [0.0, 1.0, 0.0] fake_tm = _FakeToolManager( tool_vectors=tool_vectors, query_vectors=[[1.0, 0.0, 0.0], [0.0, 1.0, 0.0]], ) with patch("config.conf", return_value=self._conf(mcp_tool_retrieval_top_k=1)), \ patch("agent.tools.ToolManager", return_value=fake_tm): result1 = self._call(fake) names1 = {t.name for t in result1} result2 = self._call(fake) names2 = {t.name for t in result2} # Built-in tools present in both turns. for b in ("read", "write", "bash"): self.assertIn(b, names1) self.assertIn(b, names2) # Turn 1 selected mcp_0; turn 2 must still contain it (only-grows). self.assertIn("mcp_0", names1) self.assertIn("mcp_0", names2) self.assertIn("mcp_1", names2) self.assertTrue(fake._retrieved_mcp_names >= {"mcp_0", "mcp_1"}) if __name__ == "__main__": unittest.main()