226 lines
7.9 KiB
Python
226 lines
7.9 KiB
Python
from __future__ import annotations
|
|
|
|
from dataclasses import replace
|
|
|
|
from mcp_fixtures import build_mcp_fixture, mcp_session
|
|
|
|
from ragcore.application.audit.service import AuditEventService
|
|
from ragcore.application.limits.models import RateLimitRule
|
|
from ragcore.application.limits.rate_limiter import RateLimiter, RateLimitRegistry
|
|
from ragcore.domain.control import Scope
|
|
from ragcore.mcp.server import build_mcp_server
|
|
|
|
EXPECTED_TOOLS = {
|
|
"rag_search",
|
|
"rag_get_context",
|
|
"rag_answer",
|
|
"rag_list_spaces",
|
|
"rag_get_document",
|
|
"rag_get_citation",
|
|
}
|
|
|
|
|
|
async def test_tool_surface_is_exactly_the_agreed_read_only_tools() -> None:
|
|
fixture = await build_mcp_fixture()
|
|
async with mcp_session(fixture.server, token=fixture.token) as session:
|
|
listed = await session.list_tools()
|
|
|
|
names = {tool.name for tool in listed.tools}
|
|
assert names == EXPECTED_TOOLS
|
|
for tool in listed.tools:
|
|
assert tool.annotations is not None
|
|
assert tool.annotations.read_only_hint is True
|
|
assert tool.annotations.destructive_hint is False
|
|
|
|
|
|
async def test_rag_search_returns_typed_results() -> None:
|
|
fixture = await build_mcp_fixture()
|
|
async with mcp_session(fixture.server, token=fixture.token) as session:
|
|
result = await session.call_tool(
|
|
"rag_search", {"query": "refund", "spaces": [str(fixture.space_id)]}
|
|
)
|
|
|
|
assert result.is_error is not True
|
|
body = result.structured_content
|
|
assert body["effective_space_ids"] == [str(fixture.space_id)]
|
|
assert body["results"][0]["citation"]["excerpt_sha256"]
|
|
|
|
|
|
async def test_rag_get_context_returns_typed_blocks() -> None:
|
|
fixture = await build_mcp_fixture()
|
|
async with mcp_session(fixture.server, token=fixture.token) as session:
|
|
result = await session.call_tool(
|
|
"rag_get_context", {"query": "refund", "spaces": [str(fixture.space_id)]}
|
|
)
|
|
|
|
assert result.is_error is not True
|
|
body = result.structured_content
|
|
assert body["citations"]
|
|
assert body["blocks"][0]["citation_ids"]
|
|
|
|
|
|
async def test_rag_answer_returns_grounded_answer() -> None:
|
|
fixture = await build_mcp_fixture()
|
|
async with mcp_session(fixture.server, token=fixture.token) as session:
|
|
result = await session.call_tool(
|
|
"rag_answer", {"query": "refund window", "spaces": [str(fixture.space_id)]}
|
|
)
|
|
|
|
assert result.is_error is not True
|
|
body = result.structured_content
|
|
assert body["answer"]
|
|
assert body["citations"]
|
|
assert all(claim["citation_ids"] for claim in body["claims"])
|
|
|
|
|
|
async def test_rag_list_spaces_lists_only_granted_spaces() -> None:
|
|
fixture = await build_mcp_fixture()
|
|
async with mcp_session(fixture.server, token=fixture.token) as session:
|
|
result = await session.call_tool("rag_list_spaces", {})
|
|
|
|
assert result.is_error is not True
|
|
spaces = result.structured_content["result"]
|
|
assert [s["id"] for s in spaces] == [str(fixture.space_id)]
|
|
|
|
|
|
async def test_rag_get_document_returns_the_published_version() -> None:
|
|
from datetime import UTC, datetime
|
|
from uuid import uuid4
|
|
|
|
from ragcore.application.documents.models import (
|
|
DocumentRecord,
|
|
DocumentVersionRecord,
|
|
DocumentVersionState,
|
|
)
|
|
|
|
fixture = await build_mcp_fixture()
|
|
document = DocumentRecord(
|
|
id=uuid4(),
|
|
workspace_id=fixture.workspace.id,
|
|
space_id=fixture.space_id,
|
|
source_key="policy.md",
|
|
external_id="policy",
|
|
title="Refund policy",
|
|
)
|
|
version = DocumentVersionRecord(
|
|
id=uuid4(),
|
|
document_id=document.id,
|
|
source_fingerprint="fp",
|
|
content_hash="hash",
|
|
parser_version="v1",
|
|
chunk_profile_version="v1",
|
|
state=DocumentVersionState.PUBLISHED,
|
|
created_at=datetime(2026, 8, 1, tzinfo=UTC),
|
|
updated_at=datetime(2026, 8, 1, tzinfo=UTC),
|
|
)
|
|
fixture.document_repository.add_document(document)
|
|
fixture.document_repository.add_version(version)
|
|
|
|
async with mcp_session(fixture.server, token=fixture.token) as session:
|
|
result = await session.call_tool("rag_get_document", {"document_id": str(document.id)})
|
|
|
|
assert result.is_error is not True
|
|
body = result.structured_content
|
|
assert body["id"] == str(document.id)
|
|
assert body["version"]["id"] == str(version.id)
|
|
|
|
|
|
async def test_rag_get_document_hides_a_document_outside_granted_spaces() -> None:
|
|
from uuid import uuid4
|
|
|
|
from ragcore.application.documents.models import DocumentRecord
|
|
|
|
fixture = await build_mcp_fixture()
|
|
other_space_document = DocumentRecord(
|
|
id=uuid4(),
|
|
workspace_id=fixture.workspace.id,
|
|
space_id=uuid4(),
|
|
source_key="secret.md",
|
|
external_id="secret",
|
|
title="Not granted",
|
|
)
|
|
fixture.document_repository.add_document(other_space_document)
|
|
|
|
async with mcp_session(fixture.server, token=fixture.token) as session:
|
|
result = await session.call_tool(
|
|
"rag_get_document", {"document_id": str(other_space_document.id)}
|
|
)
|
|
|
|
assert result.is_error is True
|
|
|
|
|
|
async def test_rag_get_citation_returns_a_granted_citation() -> None:
|
|
fixture = await build_mcp_fixture()
|
|
(citation,) = fixture.citation_repository._citations.values()
|
|
|
|
async with mcp_session(fixture.server, token=fixture.token) as session:
|
|
result = await session.call_tool("rag_get_citation", {"citation_id": str(citation.id)})
|
|
|
|
assert result.is_error is not True
|
|
assert result.structured_content["id"] == str(citation.id)
|
|
|
|
|
|
async def test_rag_get_citation_hides_an_out_of_scope_citation() -> None:
|
|
from uuid import uuid4
|
|
|
|
fixture = await build_mcp_fixture()
|
|
outside = fixture.citation_repository
|
|
from mcp_fixtures import _citation
|
|
|
|
stray = _citation(fixture.workspace.id, uuid4())
|
|
outside.add(stray)
|
|
|
|
async with mcp_session(fixture.server, token=fixture.token) as session:
|
|
result = await session.call_tool("rag_get_citation", {"citation_id": str(stray.id)})
|
|
|
|
assert result.is_error is True
|
|
|
|
|
|
async def test_no_token_is_rejected() -> None:
|
|
fixture = await build_mcp_fixture()
|
|
async with mcp_session(fixture.server, token=None) as session:
|
|
result = await session.call_tool("rag_search", {"query": "refund"})
|
|
|
|
assert result.is_error is True
|
|
|
|
|
|
async def test_wrong_scope_token_is_rejected() -> None:
|
|
fixture = await build_mcp_fixture(grant_scopes={Scope.SEARCH})
|
|
async with mcp_session(fixture.server, token=fixture.token) as session:
|
|
result = await session.call_tool("rag_answer", {"query": "refund"})
|
|
|
|
assert result.is_error is True
|
|
|
|
|
|
async def test_malformed_token_is_rejected() -> None:
|
|
fixture = await build_mcp_fixture()
|
|
async with mcp_session(fixture.server, token="not-a-real-token") as session:
|
|
result = await session.call_tool("rag_search", {"query": "refund"})
|
|
|
|
assert result.is_error is True
|
|
|
|
|
|
async def test_rate_limit_is_enforced_via_the_shared_rate_limiter() -> None:
|
|
fixture = await build_mcp_fixture()
|
|
# A default rule (keyed by operation only, no subject) applies to any
|
|
# caller -- the same `RateLimitRegistry` primitive M6-02 built for REST,
|
|
# reused unchanged here. `AuditEventService`'s repository is never
|
|
# touched: only `set_override` (not used by this test) calls `.record()`.
|
|
registry = RateLimitRegistry(
|
|
default_rules={"mcp:rag_search": RateLimitRule(max_requests=1, window_seconds=60.0)},
|
|
audit=AuditEventService(repository=None), # type: ignore[arg-type]
|
|
)
|
|
limiter = RateLimiter(registry=registry)
|
|
limited_server = build_mcp_server(replace(fixture.deps, rate_limiter=limiter))
|
|
|
|
async with mcp_session(limited_server, token=fixture.token) as session:
|
|
first = await session.call_tool(
|
|
"rag_search", {"query": "refund", "spaces": [str(fixture.space_id)]}
|
|
)
|
|
second = await session.call_tool(
|
|
"rag_search", {"query": "refund", "spaces": [str(fixture.space_id)]}
|
|
)
|
|
|
|
assert first.is_error is not True
|
|
assert second.is_error is True
|