Spaces:
Running
Running
| """Public FastMCP server: tool names, caps, no Cypher (no live Chroma).""" | |
| from __future__ import annotations | |
| import asyncio | |
| import unittest | |
| from unittest.mock import patch | |
| from langchain_core.documents import Document | |
| from starlette.middleware import Middleware | |
| from starlette.testclient import TestClient | |
| from src.deep_research_agent.tools import SearchResults | |
| from src.mcp_server import caps | |
| from src.mcp_server.rate_limit import IPRateLimitMiddleware | |
| from src.mcp_server.retrieval import documents_to_search_results, fetch, search | |
| from src.mcp_server.server import build_http_app, mcp | |
| def _doc(text: str, **meta) -> Document: | |
| return Document(page_content=text, metadata=meta) | |
| class TestCaps(unittest.TestCase): | |
| def test_clamp_k_caps_and_floors(self) -> None: | |
| self.assertEqual(caps.clamp_k(100, cap=8, default=5), 8) | |
| self.assertEqual(caps.clamp_k(0, cap=8, default=5), 1) | |
| self.assertEqual(caps.clamp_k(-3, cap=8, default=5), 1) | |
| self.assertEqual(caps.clamp_k(None, cap=8, default=5), 5) | |
| def test_snippet_limit_is_bounded(self) -> None: | |
| self.assertEqual(caps.clamp_snippet_limit(50_000), caps.MAX_SNIPPET_LIMIT) | |
| self.assertEqual(caps.clamp_snippet_limit(0), caps.MAX_SNIPPET_LIMIT) | |
| self.assertEqual(caps.clamp_snippet_limit(None), caps.DEFAULT_SNIPPET_LIMIT) | |
| self.assertGreaterEqual(caps.clamp_snippet_limit(1), 64) | |
| def test_query_and_id_are_bounded(self) -> None: | |
| self.assertEqual(len(caps.bound_query("x" * 5000)), caps.MAX_QUERY_CHARS) | |
| self.assertEqual(len(caps.bound_id("y" * 5000)), caps.MAX_ID_CHARS) | |
| def test_episode_id_variants(self) -> None: | |
| self.assertEqual(caps.episode_id_variants("S2E335"), ["S2E335"]) | |
| self.assertEqual(caps.episode_id_variants("S E18"), ["S E18", "SE18"]) | |
| self.assertIn("S E18", caps.episode_id_variants("SE18")) | |
| class TestPublicToolSurface(unittest.IsolatedAsyncioTestCase): | |
| async def test_only_search_and_fetch(self) -> None: | |
| from fastmcp import Client | |
| async with Client(mcp) as client: | |
| tools = await client.list_tools() | |
| names = sorted(t.name for t in tools) | |
| self.assertEqual(names, ["fetch", "search"]) | |
| blob = " ".join( | |
| f"{t.name} {t.description or ''} {getattr(t, 'input_schema', None) or getattr(t, 'inputSchema', '')}" | |
| for t in tools | |
| ).lower() | |
| self.assertNotIn("cypher", blob) | |
| self.assertNotIn("neo4j", blob) | |
| self.assertNotIn("query_knowledge_graph", names) | |
| async def test_search_returns_typed_results_and_clips_snippets(self) -> None: | |
| from fastmcp import Client | |
| long_text = "alpha " * 400 | |
| summaries = [_doc(long_text, type="key_takeaway", title="Ep", episode_id="S2E1", canonical_url="https://ex")] | |
| segments = [_doc(long_text, type="transcript_segment", title="Ep", episode_id="S2E1")] | |
| with patch("src.mcp_server.retrieval.retrieve_documents", return_value=(summaries, segments)): | |
| async with Client(mcp) as client: | |
| result = await client.call_tool( | |
| "search", | |
| {"query": "alpha", "max_segments": 99, "max_summaries": 99, "snippet_limit": 50_000}, | |
| ) | |
| payload = result.structured_content | |
| if isinstance(payload, SearchResults): | |
| data = payload.model_dump() | |
| else: | |
| data = payload | |
| self.assertIn("summaries", data) | |
| self.assertIn("segments", data) | |
| self.assertLessEqual(len(data["summaries"][0]["snippet"]), caps.MAX_SNIPPET_LIMIT + 20) | |
| self.assertEqual(data["summaries"][0]["episode_id"], "S2E1") | |
| self.assertEqual(data["summaries"][0]["canonical_url"], "https://ex") | |
| class TestRetrievalWrappers(unittest.TestCase): | |
| def test_search_calls_query_helpers(self) -> None: | |
| summaries = [_doc("sum", type="key_takeaway", title="T", canonical_url="https://a")] | |
| segments = [_doc("seg", type="transcript_segment", episode_id="S2E335")] | |
| with ( | |
| patch("src.mcp_server.retrieval._read_only_store", return_value="store"), | |
| patch("src.mcp_server.retrieval.query_summaries", return_value=summaries) as qs, | |
| patch("src.mcp_server.retrieval.query_segments", return_value=segments) as qg, | |
| ): | |
| result = search("parent agency", max_segments=3, max_summaries=2) | |
| self.assertIsInstance(result, SearchResults) | |
| qs.assert_called_once() | |
| qg.assert_called_once() | |
| self.assertEqual(result.summaries[0].canonical_url, "https://a") | |
| self.assertEqual(result.segments[0].episode_id, "S2E335") | |
| def test_fetch_does_not_accept_cypher(self) -> None: | |
| with ( | |
| patch("src.mcp_server.retrieval._read_only_store", return_value="store"), | |
| patch("src.mcp_server.retrieval.get_documents_where", return_value=[]), | |
| patch("src.mcp_server.retrieval.query_summaries", return_value=[]) as qs, | |
| patch("src.mcp_server.retrieval.query_segments", return_value=[]) as qg, | |
| ): | |
| result = fetch("MATCH (n) RETURN n") | |
| self.assertEqual(result.summaries, []) | |
| self.assertEqual(result.segments, []) | |
| for call in qs.call_args_list + qg.call_args_list: | |
| kwargs = call.kwargs | |
| filter_meta = kwargs.get("filter_metadata") or {} | |
| self.assertNotIn("cypher", str(filter_meta).lower()) | |
| def test_fetch_episode_returns_all_segments_in_time_order(self) -> None: | |
| docs = [ | |
| _doc("later", type="transcript_segment", episode_id="S2E335", start_time=120), | |
| _doc("first", type="transcript_segment", episode_id="S2E335", start_time=0), | |
| _doc("middle", type="transcript_segment", episode_id="S2E335", start_time=45), | |
| _doc("takeaway", type="key_takeaway", episode_id="S2E335"), | |
| ] | |
| with ( | |
| patch("src.mcp_server.retrieval._read_only_store", return_value="store"), | |
| patch("src.mcp_server.retrieval.get_documents_where", return_value=docs) as lookup, | |
| patch("src.mcp_server.retrieval.query_summaries") as qs, | |
| patch("src.mcp_server.retrieval.query_segments") as qg, | |
| ): | |
| result = fetch("S2E335", max_segments=2, snippet_limit=0) | |
| lookup.assert_called_once() | |
| self.assertEqual(lookup.call_args.args[0], {"episode_id": "S2E335"}) | |
| qs.assert_not_called() | |
| qg.assert_not_called() | |
| self.assertEqual([c.snippet for c in result.segments], ["first", "middle", "later"]) | |
| self.assertEqual(len(result.summaries), 1) | |
| def test_fetch_unknown_id_falls_back_to_similarity(self) -> None: | |
| summaries = [_doc("sum", type="article_summary_overview", doc_id="abc")] | |
| segments = [_doc("seg", type="article_text", doc_id="abc")] | |
| with ( | |
| patch("src.mcp_server.retrieval._read_only_store", return_value="store"), | |
| patch("src.mcp_server.retrieval.get_documents_where", return_value=[]), | |
| patch("src.mcp_server.retrieval.query_summaries", return_value=summaries), | |
| patch("src.mcp_server.retrieval.query_segments", return_value=segments), | |
| ): | |
| result = fetch("abc") | |
| self.assertEqual(len(result.segments), 1) | |
| self.assertEqual(result.segments[0].snippet, "seg") | |
| def test_snippet_truncation(self) -> None: | |
| docs = [_doc("n" * 5000, type="transcript_segment", title="x")] | |
| payload = documents_to_search_results([], docs, snippet_limit=200) | |
| chunk = payload.segments[0] | |
| self.assertTrue(chunk.truncated) | |
| self.assertEqual(chunk.stored_chars, 5000) | |
| self.assertEqual(chunk.omitted_chars, 4800) | |
| self.assertTrue(chunk.snippet.startswith("n" * 200)) | |
| self.assertIn("4800 of 5000 chars omitted", chunk.snippet) | |
| self.assertIn("snippet_limit=0", chunk.snippet) | |
| def test_snippet_limit_zero_returns_full_stored_chunk(self) -> None: | |
| text = "n" * 5000 | |
| docs = [_doc(text, type="article_text", title="x")] | |
| payload = documents_to_search_results([], docs, snippet_limit=0) | |
| chunk = payload.segments[0] | |
| self.assertEqual(chunk.snippet, text) | |
| self.assertFalse(chunk.truncated) | |
| self.assertEqual(chunk.omitted_chars, 0) | |
| self.assertEqual(chunk.stored_chars, 5000) | |
| class TestRateLimit(unittest.TestCase): | |
| def test_http_429_after_cap(self) -> None: | |
| async def inner_app(scope, receive, send): | |
| await send( | |
| { | |
| "type": "http.response.start", | |
| "status": 200, | |
| "headers": [(b"content-type", b"text/plain")], | |
| } | |
| ) | |
| await send({"type": "http.response.body", "body": b"ok"}) | |
| app = IPRateLimitMiddleware(inner_app, max_requests=2, window_seconds=60) | |
| async def request_once() -> int: | |
| status = {"code": 0} | |
| async def receive(): | |
| return {"type": "http.request"} | |
| async def send(message): | |
| if message["type"] == "http.response.start": | |
| status["code"] = message["status"] | |
| await app( | |
| { | |
| "type": "http", | |
| "asgi": {"version": "3.0"}, | |
| "http_version": "1.1", | |
| "method": "POST", | |
| "path": "/mcp", | |
| "headers": [(b"x-forwarded-for", b"203.0.113.9")], | |
| "client": ("127.0.0.1", 1234), | |
| "scheme": "http", | |
| "query_string": b"", | |
| }, | |
| receive, | |
| send, | |
| ) | |
| return status["code"] | |
| self.assertEqual(asyncio.run(request_once()), 200) | |
| self.assertEqual(asyncio.run(request_once()), 200) | |
| self.assertEqual(asyncio.run(request_once()), 429) | |
| def test_health_is_exempt(self) -> None: | |
| async def inner_app(scope, receive, send): | |
| await send( | |
| { | |
| "type": "http.response.start", | |
| "status": 200, | |
| "headers": [(b"content-type", b"text/plain")], | |
| } | |
| ) | |
| await send({"type": "http.response.body", "body": b"ok"}) | |
| app = IPRateLimitMiddleware(inner_app, max_requests=1, window_seconds=60) | |
| async def hit_health() -> int: | |
| status = {"code": 0} | |
| async def receive(): | |
| return {"type": "http.request"} | |
| async def send(message): | |
| if message["type"] == "http.response.start": | |
| status["code"] = message["status"] | |
| await app( | |
| { | |
| "type": "http", | |
| "asgi": {"version": "3.0"}, | |
| "http_version": "1.1", | |
| "method": "GET", | |
| "path": "/health", | |
| "headers": [], | |
| "client": ("10.0.0.1", 1), | |
| "scheme": "http", | |
| "query_string": b"", | |
| }, | |
| receive, | |
| send, | |
| ) | |
| return status["code"] | |
| self.assertEqual(asyncio.run(hit_health()), 200) | |
| self.assertEqual(asyncio.run(hit_health()), 200) | |
| def test_stateless_http_app_exposes_mcp_path(self) -> None: | |
| http_app = build_http_app() | |
| self.assertTrue(any("/mcp" in str(getattr(r, "path", "")) or "/mcp" in str(r) for r in http_app.routes)) | |
| with TestClient(http_app) as client: | |
| health = client.get("/health") | |
| self.assertEqual(health.status_code, 200) | |
| self.assertEqual(health.json()["status"], "ok") | |
| class TestMiddlewareWiring(unittest.TestCase): | |
| def test_rate_limit_middleware_is_configured(self) -> None: | |
| mw = [Middleware(IPRateLimitMiddleware)] | |
| self.assertEqual(mw[0].cls, IPRateLimitMiddleware) | |
| if __name__ == "__main__": | |
| unittest.main() | |