377 lines
13 KiB
Python
377 lines
13 KiB
Python
"""Fixtures for the MCP contract tests.
|
|
|
|
Mirrors the same fixture-construction pattern as `tests/contract/api/**`
|
|
and `tests/contract/sdk/python/**` (real `RetrievalAuthorizationService`,
|
|
a fixture executor/generator standing in only for the retrieval
|
|
backend/model provider) -- these tests exist to prove the MCP adapter's
|
|
own auth/mapping/rate-limit layer, not to re-prove authorization/
|
|
retrieval business logic already covered there.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import contextlib
|
|
from collections.abc import AsyncIterator
|
|
from dataclasses import dataclass, replace
|
|
from datetime import UTC, datetime
|
|
from typing import Any
|
|
from uuid import UUID, uuid4
|
|
|
|
import httpx2
|
|
from mcp.client.session import ClientSession
|
|
from mcp.client.streamable_http import streamable_http_client
|
|
from mcp.server.mcpserver import MCPServer
|
|
|
|
from ragcore.api.v1.answers.service import AnswerService
|
|
from ragcore.api.v1.context.service import (
|
|
ContextExecutionRequest,
|
|
ContextExecutionResult,
|
|
ContextService,
|
|
)
|
|
from ragcore.api.v1.search.service import (
|
|
AuthorizationInputs,
|
|
SearchHit,
|
|
SearchService,
|
|
)
|
|
from ragcore.application.identity.credentials import ServiceAccountCredentialService
|
|
from ragcore.application.retrieval.authorization import RetrievalAuthorizationService
|
|
from ragcore.domain.authorization import Principal
|
|
from ragcore.domain.citations import Citation
|
|
from ragcore.domain.control import Application, Role, Scope, ServiceAccount, Workspace
|
|
from ragcore.infrastructure.database.identity import InMemoryCredentialRepository
|
|
from ragcore.mcp.server import MCPDependencies, build_mcp_server
|
|
from ragcore.retrieval.context import ContextCandidate
|
|
from ragcore.retrieval.ranking import FusedCandidate
|
|
from ragcore.retrieval.search.models import Candidate, CandidateBackendKind, CandidateProvenance
|
|
|
|
EVIDENCE_EXCERPT = "The refund window is thirty days from the original purchase date."
|
|
CREDENTIAL_PEPPER = "test-credential-pepper-that-is-long-enough"
|
|
|
|
|
|
class InMemoryDocumentRepository:
|
|
def __init__(self) -> None:
|
|
self._documents: dict[UUID, Any] = {}
|
|
self._versions: dict[UUID, Any] = {}
|
|
|
|
def add_document(self, document: Any) -> None:
|
|
self._documents[document.id] = document
|
|
|
|
def add_version(self, version: Any) -> None:
|
|
self._versions[version.id] = version
|
|
|
|
async def get_document(self, document_id: UUID) -> Any:
|
|
return self._documents.get(document_id)
|
|
|
|
async def get_version(self, version_id: UUID) -> Any:
|
|
return self._versions.get(version_id)
|
|
|
|
async def list_searchable_versions(self, document_id: UUID) -> tuple[Any, ...]:
|
|
from ragcore.application.documents.models import DocumentVersionState
|
|
|
|
return tuple(
|
|
v
|
|
for v in self._versions.values()
|
|
if v.document_id == document_id and v.state is DocumentVersionState.PUBLISHED
|
|
)
|
|
|
|
|
|
class InMemoryCitationRepository:
|
|
def __init__(self) -> None:
|
|
self._citations: dict[UUID, Citation] = {}
|
|
|
|
def add(self, citation: Citation) -> None:
|
|
self._citations[citation.id] = citation
|
|
|
|
async def get_authorized(
|
|
self, citation_id: UUID, *, workspace_id: UUID, effective_space_ids: frozenset[UUID]
|
|
) -> Citation | None:
|
|
citation = self._citations.get(citation_id)
|
|
if citation is None or citation.workspace_id != workspace_id:
|
|
return None
|
|
if citation.space_id not in effective_space_ids:
|
|
return None
|
|
return citation
|
|
|
|
|
|
class _FixtureInputs:
|
|
"""Loads the one space plus a real workspace-level `Grant` tying it to
|
|
`application` -- `EffectiveSpaceResolver` requires an actual `Grant`
|
|
row (not just `Application.default_knowledge_space_ids`, which is only
|
|
consulted when the caller's own `requested_space_ids` is empty) to
|
|
intersect an *explicit* request against for a SERVICE_ACCOUNT
|
|
principal, exactly like a real workspace-level grant REST would use.
|
|
"""
|
|
|
|
def __init__(self, workspace: Workspace, space: Any, application: Application) -> None:
|
|
self.workspace = workspace
|
|
self.space = space
|
|
self.application = application
|
|
self.grant = workspace.grant(
|
|
id=uuid4(), application=application, knowledge_space=space, role=Role.EDITOR
|
|
)
|
|
|
|
async def load(self, principal: Principal) -> Any:
|
|
return type(
|
|
"Inputs",
|
|
(),
|
|
{
|
|
"spaces": (self.space,),
|
|
"grants": (self.grant,),
|
|
"user_grants": (),
|
|
"application": self.application,
|
|
},
|
|
)()
|
|
|
|
|
|
@dataclass
|
|
class McpFixture:
|
|
server: MCPServer
|
|
deps: MCPDependencies
|
|
workspace: Workspace
|
|
application: Application
|
|
space_id: UUID
|
|
token: str
|
|
document_repository: InMemoryDocumentRepository
|
|
citation_repository: InMemoryCitationRepository
|
|
|
|
|
|
def _citation(workspace_id: UUID, space_id: UUID) -> Citation:
|
|
return Citation.create(
|
|
workspace_id=workspace_id,
|
|
space_id=space_id,
|
|
document_id=uuid4(),
|
|
document_version_id=uuid4(),
|
|
chunk_id=uuid4(),
|
|
parent_chunk_id=None,
|
|
location={"page_start": 4, "section": "Refunds"},
|
|
excerpt=EVIDENCE_EXCERPT,
|
|
source_uri="https://example.test/policy.md",
|
|
acquired_at=datetime(2026, 8, 1, 12, tzinfo=UTC),
|
|
retrieval_run_id=uuid4(),
|
|
index_alias="rag_dense_hybrid_active",
|
|
profile_version="hybrid-v1",
|
|
)
|
|
|
|
|
|
async def build_mcp_fixture(*, grant_scopes: set[Scope] | None = None) -> McpFixture:
|
|
workspace = Workspace(uuid4(), "MCP workspace", "mcp-workspace")
|
|
space_id = uuid4()
|
|
space = workspace.create_knowledge_space(id=space_id, name="MCP space", slug="mcp-space")
|
|
application = workspace.create_application(
|
|
id=uuid4(),
|
|
name="MCP client",
|
|
slug="mcp-client",
|
|
scopes=frozenset(
|
|
grant_scopes
|
|
or {
|
|
Scope.SEARCH,
|
|
Scope.CONTEXT,
|
|
Scope.ANSWER,
|
|
Scope.DOCUMENTS_READ,
|
|
Scope.CITATIONS_READ,
|
|
}
|
|
),
|
|
default_knowledge_space_ids=frozenset({space_id}),
|
|
)
|
|
service_account = ServiceAccount(
|
|
id=uuid4(),
|
|
workspace_id=workspace.id,
|
|
application_id=application.id,
|
|
name="MCP service account",
|
|
scopes=application.scopes,
|
|
)
|
|
|
|
async def load_service_account(account_id: UUID) -> ServiceAccount | None:
|
|
return service_account if account_id == service_account.id else None
|
|
|
|
credential_service = ServiceAccountCredentialService(
|
|
InMemoryCredentialRepository(),
|
|
CREDENTIAL_PEPPER,
|
|
service_account_loader=load_service_account,
|
|
)
|
|
issued = await credential_service.issue(service_account, scopes=service_account.scopes)
|
|
|
|
citation = _citation(workspace.id, space_id)
|
|
candidate = Candidate(
|
|
point_id=uuid4(),
|
|
score=0.87,
|
|
rank=1,
|
|
backend=CandidateBackendKind.DENSE,
|
|
provenance=CandidateProvenance(
|
|
workspace_id=workspace.id,
|
|
space_id=space_id,
|
|
document_id=citation.document_id,
|
|
document_version_id=citation.document_version_id,
|
|
chunk_id=citation.chunk_id,
|
|
source_id=uuid4(),
|
|
parent_chunk_id=None,
|
|
status="published",
|
|
content_type="text/markdown",
|
|
language="en",
|
|
trust_level="trusted_internal",
|
|
),
|
|
)
|
|
hit = SearchHit(
|
|
fused=FusedCandidate(candidate, 1 / 61, 1, None),
|
|
text=EVIDENCE_EXCERPT,
|
|
title="Refund policy",
|
|
citation=citation,
|
|
dense_score=0.87,
|
|
)
|
|
|
|
class _SearchExecutor:
|
|
async def execute(self, request: Any) -> Any:
|
|
refreshed = replace(
|
|
hit,
|
|
citation=Citation.create(
|
|
workspace_id=hit.citation.workspace_id,
|
|
space_id=hit.citation.space_id,
|
|
document_id=hit.citation.document_id,
|
|
document_version_id=hit.citation.document_version_id,
|
|
chunk_id=hit.citation.chunk_id,
|
|
parent_chunk_id=hit.citation.parent_chunk_id,
|
|
location=hit.citation.location,
|
|
excerpt=hit.citation.excerpt,
|
|
source_uri=hit.citation.source_uri,
|
|
acquired_at=hit.citation.acquired_at,
|
|
retrieval_run_id=request.retrieval_run_id,
|
|
index_alias=hit.citation.index_alias,
|
|
profile_version=hit.citation.profile_version,
|
|
),
|
|
)
|
|
return type(
|
|
"Result",
|
|
(),
|
|
{
|
|
"hits": (refreshed,),
|
|
"degraded": False,
|
|
"degradation_codes": (),
|
|
"timings_ms": {},
|
|
},
|
|
)()
|
|
|
|
inputs = _FixtureInputs(workspace, space, application)
|
|
search_service = SearchService(
|
|
authorization=RetrievalAuthorizationService(),
|
|
inputs=inputs,
|
|
executor=_SearchExecutor(),
|
|
)
|
|
|
|
context_candidate = ContextCandidate(citation, EVIDENCE_EXCERPT, 1, 1.0)
|
|
|
|
class _ContextExecutor:
|
|
async def execute(self, request: ContextExecutionRequest) -> ContextExecutionResult:
|
|
refreshed = Citation.create(
|
|
workspace_id=citation.workspace_id,
|
|
space_id=citation.space_id,
|
|
document_id=citation.document_id,
|
|
document_version_id=citation.document_version_id,
|
|
chunk_id=citation.chunk_id,
|
|
parent_chunk_id=citation.parent_chunk_id,
|
|
location=citation.location,
|
|
excerpt=citation.excerpt,
|
|
source_uri=citation.source_uri,
|
|
acquired_at=citation.acquired_at,
|
|
retrieval_run_id=request.retrieval_run_id,
|
|
index_alias=citation.index_alias,
|
|
profile_version=citation.profile_version,
|
|
)
|
|
return ContextExecutionResult(
|
|
candidates=(replace(context_candidate, citation=refreshed),)
|
|
)
|
|
|
|
context_service = ContextService(
|
|
authorization=RetrievalAuthorizationService(), inputs=inputs, executor=_ContextExecutor()
|
|
)
|
|
|
|
import json
|
|
import re
|
|
|
|
citation_id_pattern = re.compile(r"citation_ids=([0-9a-fA-F-]{36})")
|
|
|
|
class _Generator:
|
|
async def generate(self, *, prompt: str, generation_profile: str, spaces: Any) -> str:
|
|
match = citation_id_pattern.search(prompt)
|
|
assert match is not None
|
|
return json.dumps(
|
|
{
|
|
"status": "answerable",
|
|
"answer": "The refund window is thirty days from purchase.",
|
|
"citations": [match.group(1)],
|
|
}
|
|
)
|
|
|
|
answer_service = AnswerService(
|
|
authorization=RetrievalAuthorizationService(),
|
|
inputs=inputs,
|
|
context_executor=_ContextExecutor(),
|
|
generator=_Generator(),
|
|
)
|
|
|
|
document_repository = InMemoryDocumentRepository()
|
|
citation_repository = InMemoryCitationRepository()
|
|
citation_repository.add(citation)
|
|
|
|
class _AuthorizationInputsForMcp:
|
|
async def load(self, principal: Principal) -> AuthorizationInputs:
|
|
loaded = await inputs.load(principal)
|
|
return AuthorizationInputs(
|
|
spaces=loaded.spaces,
|
|
grants=loaded.grants,
|
|
user_grants=loaded.user_grants,
|
|
application=loaded.application,
|
|
)
|
|
|
|
deps = MCPDependencies(
|
|
search_application=search_service,
|
|
context_application=context_service,
|
|
answer_application=answer_service,
|
|
authorization_inputs=_AuthorizationInputsForMcp(),
|
|
credential_service=credential_service,
|
|
citation_repository=citation_repository,
|
|
document_repository=document_repository,
|
|
)
|
|
server = build_mcp_server(deps)
|
|
|
|
return McpFixture(
|
|
server=server,
|
|
deps=deps,
|
|
workspace=workspace,
|
|
application=application,
|
|
space_id=space_id,
|
|
token=issued.token,
|
|
document_repository=document_repository,
|
|
citation_repository=citation_repository,
|
|
)
|
|
|
|
|
|
_MCP_TEST_URL = "http://127.0.0.1:8765"
|
|
|
|
|
|
@contextlib.asynccontextmanager
|
|
async def mcp_session(
|
|
server: MCPServer, *, token: str | None = None
|
|
) -> AsyncIterator[ClientSession]:
|
|
"""A real streamable-HTTP MCP session over `httpx2.ASGITransport` -- genuine
|
|
Starlette `Request`/header handling, no real TCP socket needed.
|
|
|
|
`ASGITransport` never runs the ASGI app's own lifespan protocol (that's
|
|
normally what starts the session manager's task group when a real
|
|
server like uvicorn owns the app), so `session_manager.run()` is
|
|
entered explicitly here instead, mirroring what a real deployment's
|
|
lifespan would do.
|
|
"""
|
|
|
|
app = server.streamable_http_app(stateless_http=True)
|
|
headers = {"authorization": f"Bearer {token}"} if token else {}
|
|
async with (
|
|
server.session_manager.run(),
|
|
httpx2.AsyncClient(
|
|
transport=httpx2.ASGITransport(app=app), base_url=_MCP_TEST_URL, headers=headers
|
|
) as http_client,
|
|
streamable_http_client(f"{_MCP_TEST_URL}/mcp", http_client=http_client) as (read, write),
|
|
ClientSession(read, write) as session,
|
|
):
|
|
await session.initialize()
|
|
yield session
|