diff --git a/redisvl/mcp/context.py b/redisvl/mcp/context.py new file mode 100644 index 00000000..3e96e642 --- /dev/null +++ b/redisvl/mcp/context.py @@ -0,0 +1,282 @@ +"""The handle a code tool gets on Redis. + +A code tool never receives the server, its ``BindingRuntime``, the real +``AsyncSearchIndex`` or a Redis client. It receives a :class:`ToolContext`, +whose ``index`` is an allow-listed wrapper: every method on it is spelled out +here, routes through ``server.run_guarded`` (the process-wide concurrency cap +and the binding's request timeout), and write methods exist only on the +handle given to a tool that declared it writes. A method added to +``AsyncSearchIndex`` later is therefore unreachable until it is added here. + +Every method stays inside the bound index: searches and aggregations run +against its name, and records are addressed by document id, which the index +turns into a key under its own prefix. Nothing accepts a full Redis key or a +query that names its own target, because either could reach another index's +documents, including a tenant-scoped one. An index whose prefix overlaps +another's shares keys with it, which no handle can tell apart. A prefix that +is empty, or made only of separator characters, turns document ids into full +Redis keys, so every key is in reach. On an index with several prefixes, +document ids address records under the first one only. + +This is a guardrail, not a sandbox: Python can still reach private attributes. +It makes the safe path the only obvious one. +""" + +from collections.abc import Callable, Iterable, Mapping +from types import MappingProxyType +from typing import TYPE_CHECKING, Any + +from redisvl.index.index import DEFAULT_BULK_BATCH_SIZE +from redisvl.mcp.tools.search import _embed_query +from redisvl.query.sql import SQLQuery + +if TYPE_CHECKING: + from redis.commands.search.aggregation import AggregateResult + from redis.commands.search.result import Result + + from redisvl.index import AsyncSearchIndex + from redisvl.mcp.runtime import BindingRuntime + from redisvl.schema import IndexSchema + + +def _refuse_sql(query: Any) -> None: + if isinstance(query, SQLQuery): + raise ValueError( + "a code tool cannot run a SQLQuery: it names its own index, so it " + "could read beyond the one the tool is bound to" + ) + + +class IndexHandle: + """Read access to one bound index, bounded like every built-in tool call.""" + + __slots__ = ("_server", "_index", "_timeout_seconds") + + def __init__( + self, server: Any, index: "AsyncSearchIndex", *, timeout_seconds: float + ) -> None: + self._server = server + self._index = index + self._timeout_seconds = timeout_seconds + + async def _guarded(self, operation_name: str, awaitable: Any) -> Any: + return await self._server.run_guarded( + operation_name, awaitable, timeout_seconds=self._timeout_seconds + ) + + async def query(self, query: Any) -> Any: + """Run a RedisVL query object and return its processed results. + + A ``CountQuery`` returns the match count rather than records. A + ``SQLQuery`` is refused: its ``FROM`` clause, not the bound index, + decides which index it reads. + """ + _refuse_sql(query) + return await self._guarded("query", self._index.query(query)) + + async def batch_query( + self, queries: list[Any], batch_size: int = 10 + ) -> list[list[dict[str, Any]]]: + """Run several RedisVL query objects and return each one's results.""" + for query in queries: + _refuse_sql(query) + return await self._guarded( + "batch_query", self._index.batch_query(queries, batch_size=batch_size) + ) + + async def search(self, *args: Any, **kwargs: Any) -> "Result": + """Run a raw ``FT.SEARCH`` and return the redis-py result.""" + return await self._guarded("search", self._index.search(*args, **kwargs)) + + async def batch_search( + self, queries: list[Any], batch_size: int = 10 + ) -> list["Result"]: + """Run several raw searches and return each redis-py result.""" + return await self._guarded( + "batch_search", self._index.batch_search(queries, batch_size=batch_size) + ) + + async def aggregate(self, *args: Any, **kwargs: Any) -> "AggregateResult": + """Run a raw ``FT.AGGREGATE`` and return the redis-py result.""" + return await self._guarded("aggregate", self._index.aggregate(*args, **kwargs)) + + async def fetch(self, id: str) -> dict[str, Any] | None: + """Return one record by its id, or ``None`` when it does not exist.""" + return await self._guarded("fetch", self._index.fetch(id)) + + async def exists(self) -> bool: + """Return whether the index exists in Redis.""" + return await self._guarded("exists", self._index.exists()) + + async def info(self) -> dict[str, Any]: + """Return ``FT.INFO`` for the bound index.""" + return await self._guarded("info", self._index.info()) + + def document_id(self, key: str) -> str: + """Return the document id of a Redis key under this index's prefix. + + Search results carry each record's full key as ``id``, while the + handle addresses records by document id, so this is the bridge from a + result back to ``fetch`` or a write. One exception: on JSON storage, a + ``FilterQuery`` without ``return_fields`` returns whole documents, and + a document's own ``id`` field replaces the key; set ``return_fields`` + to keep it. On an index with several prefixes, ids address records + under the first prefix, the same one ``fetch`` and the writes use. + Raises ``ValueError`` for a key outside it. + """ + prefix = self._index.key("") + if not prefix: + return key + if not key.startswith(prefix) or key == prefix: + raise ValueError( + f"key {key!r} is not under {prefix!r}, the key prefix this " + "index's document ids address" + ) + return key[len(prefix) :] + + +class WritableIndexHandle(IndexHandle): + """An :class:`IndexHandle` that may also write, for tools declared as writing. + + Only writes whose cost is bounded by the input they are given are exposed; + the filter-wide bulk operations are not, because the request timeout can + cancel them partway through. + """ + + __slots__ = () + + async def load( + self, + data: Iterable[Any], + id_field: str | None = None, + ttl: int | None = None, + preprocess: Callable | None = None, + batch_size: int | None = None, + ) -> list[str]: + """Write records to the index and return their document ids. + + Each record's id is its ``id_field`` value, or a generated one when + ``id_field`` is not given. The ids, not full Redis keys, are what + ``fetch``, ``drop_documents`` and ``expire_documents`` take. + """ + keys = await self._guarded( + "load", + self._index.load( + data, + id_field=id_field, + ttl=ttl, + preprocess=preprocess, + batch_size=batch_size, + ), + ) + return [self.document_id(key) for key in keys] + + async def drop_documents( + self, ids: str | list[str], batch_size: int = DEFAULT_BULK_BATCH_SIZE + ) -> int: + """Delete records by document id and return how many were removed.""" + return await self._guarded( + "drop_documents", self._index.drop_documents(ids, batch_size=batch_size) + ) + + async def expire_documents(self, ids: str | list[str], ttl: int) -> int | list[int]: + """Set a time-to-live, in seconds, on records by document id.""" + keys = ( + self._index.key(ids) + if isinstance(ids, str) + else [self._index.key(id) for id in ids] + ) + return await self._guarded( + "expire_documents", self._index.expire_keys(keys, ttl) + ) + + +def _freeze(value: Any) -> Any: + """Return a read-only copy: mappings become proxies, lists become tuples.""" + if isinstance(value, Mapping): + return MappingProxyType({key: _freeze(item) for key, item in value.items()}) + if isinstance(value, (list, tuple)): + return tuple(_freeze(item) for item in value) + return value + + +def _current_claims(server: Any) -> Mapping[str, Any] | None: + """The verified claims of the current request, or ``None`` without a token.""" + # Without auth no token is verified. Returning first also keeps code tools + # working on FastMCP releases before 2.10, which lack get_access_token. + if not getattr(server, "_auth_enabled", False): + return None + # Imported here because this module stays importable without the `mcp` extra. + from fastmcp.server.dependencies import get_access_token + + token = get_access_token() + if token is None: + return None + return _freeze(getattr(token, "claims", None) or {}) + + +class ToolContext: + """Everything a code tool may use to reach Redis, and nothing else. + + ``index`` is the bound index's handle. ``schema`` is a copy of the + binding's schema, and ``vectorizer`` is the binding's. ``claims`` is a + read-only view of the current request's verified token claims, or ``None`` + when the request carries no token. + """ + + __slots__ = ("_index", "_schema", "_vectorizer", "_claims") + + def __init__( + self, + *, + index: IndexHandle, + schema: "IndexSchema", + vectorizer: Any | None, + claims: Mapping[str, Any] | None, + ) -> None: + self._index = index + self._schema = schema + self._vectorizer = vectorizer + self._claims = claims + + @classmethod + def for_request( + cls, server: Any, runtime: "BindingRuntime", *, writes: bool + ) -> "ToolContext": + """Build the context for one call of a tool bound to ``runtime``.""" + handle_type = WritableIndexHandle if writes else IndexHandle + return cls( + index=handle_type( + server, + runtime.index, + timeout_seconds=runtime.binding.runtime.request_timeout_seconds, + ), + # A copy, because the index holds the same object: a tool that + # edited it would retarget every later call on the binding, + # built-in tools included. + schema=runtime.schema.model_copy(deep=True), + vectorizer=runtime.vectorizer, + claims=_current_claims(server), + ) + + @property + def index(self) -> IndexHandle: + return self._index + + @property + def schema(self) -> "IndexSchema": + return self._schema + + @property + def vectorizer(self) -> Any | None: + return self._vectorizer + + @property + def claims(self) -> Mapping[str, Any] | None: + return self._claims + + async def embed(self, text: str) -> Any: + """Embed ``text`` with the binding's vectorizer.""" + if self._vectorizer is None: + raise RuntimeError("MCP server vectorizer is not configured") + return await _embed_query(self._vectorizer, text) diff --git a/tests/unit/test_mcp/test_tool_context.py b/tests/unit/test_mcp/test_tool_context.py new file mode 100644 index 00000000..2e436b2b --- /dev/null +++ b/tests/unit/test_mcp/test_tool_context.py @@ -0,0 +1,390 @@ +"""The handle a code tool gets on Redis: an allow-list, every call guarded.""" + +import inspect +from types import SimpleNamespace + +import pytest +from conftest import _schema + +pytest.importorskip("fastmcp", reason="fastmcp not installed (install redisvl[mcp])") + +from redisvl.index import AsyncSearchIndex +from redisvl.mcp.config import MCPConfig +from redisvl.mcp.context import IndexHandle, ToolContext, WritableIndexHandle +from redisvl.mcp.runtime import BindingRuntime +from redisvl.query.sql import SQLQuery + +READS = { + "query", + "batch_query", + "search", + "batch_search", + "aggregate", + "fetch", + "exists", + "info", +} +WRITES = {"load", "drop_documents", "expire_documents"} +# Answered by the handle itself, with no Redis call. +LOCAL = {"document_id"} +TIMEOUT = 7.0 + + +class RecordingIndex: + """Answers any method with a record of how it was called.""" + + def __init__(self): + self.awaited: list[str] = [] + self.calls: list[tuple] = [] + + def key(self, id): + return f"doc:{id}" + + def __getattr__(self, name): + async def method(*args, **kwargs): + self.awaited.append(name) + self.calls.append((name, args, kwargs)) + if name == "load": + # The real load returns full Redis keys. + return [self.key(record["id"]) for record in args[0]] + return (name, args, kwargs) + + return method + + +class FakeServer: + def __init__(self, *, auth_enabled: bool = True): + self._auth_enabled = auth_enabled + self.guarded: list[tuple[str, float]] = [] + + async def run_guarded(self, operation_name, awaitable, *, timeout_seconds): + self.guarded.append((operation_name, timeout_seconds)) + return await awaitable + + +class FakeVectorizer: + async def aembed(self, text): + return [len(text), 0.0, 0.0] + + +def _runtime(vectorizer=None) -> BindingRuntime: + config = MCPConfig.model_validate( + { + "server": {"redis_url": "redis://localhost:6379"}, + "indexes": { + "knowledge": { + "redis_name": "docs-index", + "search": {"type": "fulltext"}, + "runtime": { + "text_field_name": "content", + "request_timeout_seconds": TIMEOUT, + }, + } + }, + } + ) + return BindingRuntime( + binding_id="knowledge", + binding=config.indexes["knowledge"], + index=RecordingIndex(), + schema=_schema(), + vectorizer=vectorizer, + supports_native_hybrid_search=False, + effective_read_only=False, + ) + + +def _as_caller(monkeypatch, token): + monkeypatch.setattr( + "fastmcp.server.dependencies.get_access_token", lambda: token, raising=False + ) + + +def _public(cls) -> set[str]: + return {name for name in dir(cls) if not name.startswith("_")} + + +def test_the_read_handle_exposes_exactly_the_read_methods(): + # Pinned exactly, so a method added to the handle is a deliberate, + # reviewed change rather than a quiet widening. + assert _public(IndexHandle) == READS | LOCAL + + +def test_the_writable_handle_adds_exactly_the_bounded_writes(): + assert _public(WritableIndexHandle) == READS | WRITES | LOCAL + + +QUERY = object() + +# (handle method, its arguments, the index method it reaches, what it forwards) +CALLS = [ + ("query", (QUERY,), {}, "query", ((QUERY,), {})), + ( + "batch_query", + ([QUERY],), + {"batch_size": 3}, + "batch_query", + (([QUERY],), {"batch_size": 3}), + ), + ( + "search", + ("*",), + {"query_params": {}}, + "search", + (("*",), {"query_params": {}}), + ), + ("batch_search", ([QUERY],), {}, "batch_search", (([QUERY],), {"batch_size": 10})), + ("aggregate", (QUERY,), {}, "aggregate", ((QUERY,), {})), + ("fetch", ("a1",), {}, "fetch", (("a1",), {})), + ("exists", (), {}, "exists", ((), {})), + ("info", (), {}, "info", ((), {})), + ( + "load", + ([{"id": "a1"}],), + {"id_field": "id"}, + "load", + ( + ([{"id": "a1"}],), + {"id_field": "id", "ttl": None, "preprocess": None, "batch_size": None}, + ), + ), + ( + "drop_documents", + (["a1"],), + {}, + "drop_documents", + ((["a1"],), {"batch_size": 500}), + ), + # Addressed by id, so the key is built under the index's own prefix. + ( + "expire_documents", + (["a1", "a2"], 60), + {}, + "expire_keys", + ((["doc:a1", "doc:a2"], 60), {}), + ), + ("expire_documents", ("a1", 60), {}, "expire_keys", (("doc:a1", 60), {})), +] +CALL_IDS = [f"{call[0]}-{n}" for n, call in enumerate(CALLS)] +# load hands back document ids, not the keys the index returns. +RETURNS = {"load": ["a1"]} + + +@pytest.mark.parametrize("method, args, kwargs, target, forwarded", CALLS, ids=CALL_IDS) +async def test_every_call_reaches_the_index_through_the_guard( + method, args, kwargs, target, forwarded +): + server = FakeServer() + index = RecordingIndex() + handle = WritableIndexHandle(server, index, timeout_seconds=TIMEOUT) + + result = await getattr(handle, method)(*args, **kwargs) + + assert index.calls == [(target, *forwarded)] + assert result == RETURNS.get(method, (target, *forwarded)) + assert server.guarded == [(method, TIMEOUT)] + + +@pytest.mark.parametrize("method, args, kwargs, target, forwarded", CALLS, ids=CALL_IDS) +def test_every_forwarded_call_fits_the_real_index_signature( + method, args, kwargs, target, forwarded +): + # The recording index accepts anything, so check what the handle forwards + # against the method it stands in for. + forwarded_args, forwarded_kwargs = forwarded + signature = inspect.signature(getattr(AsyncSearchIndex, target)) + + signature.bind(None, *forwarded_args, **forwarded_kwargs) + + +def test_the_covered_calls_are_the_whole_surface(): + assert {call[0] for call in CALLS} == READS | WRITES + + +async def test_a_guard_refusal_reaches_the_tool_without_touching_the_index(): + class StoppedServer: + async def run_guarded(self, operation_name, awaitable, *, timeout_seconds): + awaitable.close() + raise RuntimeError("MCP server is not running") + + index = RecordingIndex() + handle = IndexHandle(StoppedServer(), index, timeout_seconds=TIMEOUT) + + with pytest.raises(RuntimeError, match="not running"): + await handle.fetch("a1") + assert index.awaited == [] + + +@pytest.mark.parametrize("method", ["query", "batch_query"]) +async def test_a_sql_query_is_refused_because_it_names_its_own_index(method): + sql = SQLQuery("SELECT * FROM tenant_docs") + index = RecordingIndex() + handle = IndexHandle(FakeServer(), index, timeout_seconds=TIMEOUT) + + with pytest.raises(ValueError, match="cannot run a SQLQuery"): + await getattr(handle, method)(sql if method == "query" else [QUERY, sql]) + assert index.awaited == [] + + +@pytest.mark.parametrize( + "call", + [ + lambda handle: handle.info("other-index"), + lambda handle: handle.load([{"id": "1"}], keys=["other:1"]), + ], + ids=["info-names-an-index", "load-takes-full-keys"], +) +async def test_a_handle_never_takes_another_target(call): + handle = WritableIndexHandle( + FakeServer(), RecordingIndex(), timeout_seconds=TIMEOUT + ) + + with pytest.raises(TypeError): + await call(handle) + + +def _real_index(prefix: str | list[str]) -> AsyncSearchIndex: + return AsyncSearchIndex.from_dict( + { + "index": {"name": "tickets", "prefix": prefix}, + "fields": [{"name": "status", "type": "tag"}], + } + ) + + +@pytest.mark.parametrize("prefix", ["ticket", "ticket:", "", ["ticket", "archive"]]) +def test_a_key_from_a_result_or_a_load_maps_back_to_its_document_id(prefix): + # Results and load return full keys; fetch and the writes re-apply the + # prefix, so passing a key straight back would miss the record. + index = _real_index(prefix) + handle = IndexHandle(FakeServer(), index, timeout_seconds=TIMEOUT) + + assert handle.document_id(index.key("t1")) == "t1" + + +# A key under a second prefix is in the index, but outside the prefix that +# document ids, fetch and the writes address. +@pytest.mark.parametrize("key", ["other:t1", "ticket:", "archive:t1"]) +def test_a_key_outside_the_index_has_no_document_id(key): + index = _real_index(["ticket", "archive"]) + handle = IndexHandle(FakeServer(), index, timeout_seconds=TIMEOUT) + + with pytest.raises(ValueError, match="is not under 'ticket:'"): + handle.document_id(key) + + +def test_a_handle_offers_no_client_and_takes_no_new_attributes(): + handle = WritableIndexHandle(FakeServer(), RecordingIndex(), timeout_seconds=1) + + with pytest.raises(AttributeError): + handle.client + with pytest.raises(AttributeError): + handle.client = object() + + +@pytest.mark.parametrize( + "writes, handle_type", [(False, IndexHandle), (True, WritableIndexHandle)] +) +def test_only_a_writing_tool_gets_the_writable_handle(monkeypatch, writes, handle_type): + _as_caller(monkeypatch, None) + + ctx = ToolContext.for_request(FakeServer(), _runtime(), writes=writes) + + assert type(ctx.index) is handle_type + + +async def test_the_context_bounds_calls_by_the_bindings_timeout(monkeypatch): + _as_caller(monkeypatch, None) + server = FakeServer() + runtime = _runtime() + + ctx = ToolContext.for_request(server, runtime, writes=False) + await ctx.index.exists() + + assert server.guarded == [("exists", TIMEOUT)] + + +def test_the_context_schema_is_a_copy_a_tool_cannot_retarget_the_index_with( + monkeypatch, +): + _as_caller(monkeypatch, None) + runtime = _runtime() + + ctx = ToolContext.for_request(FakeServer(), runtime, writes=False) + ctx.schema.index.name = "tenant-index" + + assert runtime.schema.index.name == "docs-index" + + +def test_the_context_is_read_only(monkeypatch): + _as_caller(monkeypatch, None) + ctx = ToolContext.for_request(FakeServer(), _runtime(), writes=True) + + with pytest.raises(AttributeError): + ctx.index = WritableIndexHandle( + FakeServer(), RecordingIndex(), timeout_seconds=1 + ) + with pytest.raises(AttributeError): + ctx.anything = object() + + +def test_a_server_without_auth_never_looks_for_a_token(monkeypatch): + # FastMCP before 2.10 has no get_access_token, so reaching for it without + # auth would fail every code-tool call there. + def no_such_helper(): + raise AssertionError("get_access_token was called") + + monkeypatch.setattr( + "fastmcp.server.dependencies.get_access_token", no_such_helper, raising=False + ) + server = FakeServer(auth_enabled=False) + + assert ToolContext.for_request(server, _runtime(), writes=False).claims is None + + +def test_a_tokenless_request_has_no_claims(monkeypatch): + _as_caller(monkeypatch, None) + + assert ( + ToolContext.for_request(FakeServer(), _runtime(), writes=False).claims is None + ) + + +def test_a_token_without_claims_still_reads_as_a_token(monkeypatch): + _as_caller(monkeypatch, SimpleNamespace(claims=None)) + + claims = ToolContext.for_request(FakeServer(), _runtime(), writes=False).claims + + assert claims is not None and dict(claims) == {} + + +def test_claims_are_a_read_only_copy_all_the_way_down(monkeypatch): + claims = {"sub": "user-42", "org": {"id": "acme", "regions": ["eu", "us"]}} + _as_caller(monkeypatch, SimpleNamespace(claims=claims)) + + seen = ToolContext.for_request(FakeServer(), _runtime(), writes=False).claims + claims["org"]["id"] = "victim" + + assert seen["org"]["id"] == "acme" + assert seen["org"]["regions"] == ("eu", "us") + with pytest.raises(TypeError): + seen["sub"] = "someone-else" + with pytest.raises(TypeError): + seen["org"]["id"] = "victim" + + +async def test_embed_uses_the_bindings_vectorizer(monkeypatch): + _as_caller(monkeypatch, None) + ctx = ToolContext.for_request( + FakeServer(), _runtime(FakeVectorizer()), writes=False + ) + + assert await ctx.embed("refund") == [6, 0.0, 0.0] + + +async def test_embed_without_a_vectorizer_fails_clearly(monkeypatch): + _as_caller(monkeypatch, None) + ctx = ToolContext.for_request(FakeServer(), _runtime(), writes=False) + + with pytest.raises(RuntimeError, match="vectorizer is not configured"): + await ctx.embed("refund")