diff --git a/docs/concepts/mcp.md b/docs/concepts/mcp.md index d0c03489..f879f4a5 100644 --- a/docs/concepts/mcp.md +++ b/docs/concepts/mcp.md @@ -105,7 +105,7 @@ For configuration and the gateway boundary, see {doc}`/user_guide/how_to_guides/ ## Tool Surface -RedisVL MCP exposes up to three built-in tools, plus any configured [custom tool profiles](#custom-tool-profiles): +RedisVL MCP exposes up to three built-in tools, plus any configured [custom tool profiles](#custom-tool-profiles) and [code tools](#code-tools): - `list-indexes` enumerates the configured logical indexes for discovery - `search-records` searches a selected index using that index's server-owned search mode @@ -158,7 +158,7 @@ The discovery payload is deliberately minimal: The built-in tools expose the index generically, which leaves the model doing query engineering on every call: pick the index, understand the schema, build a filter, choose return fields. A **profile** moves those decisions into config. It is `search-records` with some arguments pre-filled and frozen and the rest still exposed, published under a name and description of your choosing. -Profiles are pure configuration. You add a `custom_tools` entry to the same YAML the server already loads and restart it; there is no Python to write. +Profiles are pure configuration. You add a `custom_tools` entry to the same YAML the server already loads and restart it; a profile needs no Python. For logic a profile cannot express, write a [code tool](#code-tools) instead. ```yaml custom_tools: @@ -243,14 +243,14 @@ Injection is checked at startup wherever the configuration alone can show it wou - authentication is not enabled, on any transport, including an unauthenticated loopback HTTP bind and any `--allow-unauthenticated` bind; - authentication is configured but the server runs over `stdio`, which is never authenticated (checked when the server starts through `rvl mcp` or `run_async`; an embedder that calls `startup()` directly is not, and every call is then refused at request time instead); -- the Redis index is also reachable without the tenant scope: through `search-records`, through `upsert-records` unless every binding over it is read-only, or through another custom tool on it that does not inject exactly the same entries, since a tool scoped by another field, or by the same field from another claim, reads across the tenants this one separates. Bindings that share a `redis_name` count as one index; +- the Redis index is also reachable without the tenant scope: through `search-records`, through `upsert-records` unless every binding over it is read-only, through another custom tool on it that does not inject exactly the same entries, since a tool scoped by another field, or by the same field from another claim, reads across the tenants this one separates, or through a [code tool](#code-tools) bound to it, since the server cannot see what a code tool queries. Bindings that share a `redis_name` count as one index; - the injected field is absent from the bound index, is not a tag field, or is declared `NOINDEX`; - an `inject` list is empty, names one field twice, or names a field that `lock.filter` also constrains; - `required` is anything but `true`, or `from` is anything but `claim`. An injected field must be a tag. Text equality is a phrase match over tokenised text, and text is tokenised on punctuation, so the phrase `acme-corp` would also match `acme-corp-eu`. -The tool set registers once per process. If a restart reloads a configuration that differs from the registered one, the server normally logs a warning and keeps the old tools. When injection is configured on either side of the change, startup fails instead, because keeping the old tenant scoping in force is not something a log line should report. +The tool set registers once per process. If a restart reloads a configuration that differs from the registered one, the server normally logs a warning and keeps the old tools. When injection is configured on either side of the change, startup fails instead, because keeping the old tenant scoping in force is not something a log line should report. Startup also fails when the change alters the code tools or the bindings they resolve to, since operator code would otherwise keep running against the old pairing. #### Threat Model @@ -267,6 +267,59 @@ Where tenants share one index, the injected filter is the only isolation boundar Listing the tenant claim under `auth.required_claims` is a cheap outer layer: the verifier then rejects a token that lacks the claim before any tool runs. That check confirms only that the claim is present. Its value is still validated by the profile on every call. +## Code Tools + +A **code tool** is an async Python function registered as an MCP tool. It is for logic a profile cannot express, such as re-ranking results, counting matches, or combining several queries into one answer. You register it on the server object from your own Python process: + +```python +from redisvl.mcp import MCPSettings, RedisVLMCPServer, ToolContext +from redisvl.query import CountQuery +from redisvl.query.filter import Tag + +server = RedisVLMCPServer(MCPSettings(config="/path/to/mcp_config.yaml")) + + +@server.code_tool(index="tickets") +async def count_open_tickets(ctx: ToolContext, product: str) -> dict: + """Count open support tickets for a product.""" + open_for_product = (Tag("status") == "open") & (Tag("product") == product) + return {"open": await ctx.index.query(CountQuery(open_for_product))} +``` + +The function's first parameter is the context and the rest are the tool's arguments, which the model supplies. The model sees the arguments and their types, and never the context. The docstring is the description the model chooses tools by, so a tool without one is refused. The tool's name is the function's name with hyphens for underscores (`count-open-tickets`), or `name=`; `namespace=` prefixes it. The same naming rules as profiles apply. + +`index` pins the binding the tool works on. It may be omitted only when exactly one index is configured, and the model never chooses it. + +### What a Code Tool Can Reach + +A code tool never receives the server, the real `AsyncSearchIndex`, or a Redis client. It receives a `ToolContext` with: + +| Attribute | What it is | +|---|---| +| `ctx.index` | A handle on the bound index offering `query`, `batch_query`, `search`, `batch_search`, `aggregate`, `fetch`, `exists` and `info`. A tool declared with `writes=True` also gets `load`, `drop_documents` and `expire_documents`. `document_id(key)` turns the full key a search result carries as `id` back into a document id. On JSON storage, a `FilterQuery` without `return_fields` returns whole documents, whose own `id` field replaces the key, so set `return_fields` when you need it. | +| `ctx.schema` | The bound index's schema. | +| `ctx.vectorizer` and `ctx.embed(text)` | The binding's vectorizer, if it has one, and a helper that embeds text with it. | +| `ctx.claims` | A read-only copy of the current request's verified token claims, or `None` when the server has no authentication or the request has no token. | + +Every call 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. So nothing on the handle takes a full Redis key, `load` returns document ids, and `query` refuses a `SQLQuery`, whose `FROM` clause names its own index. Either could otherwise reach another index's documents, including a tenant-scoped one. An index whose key prefix overlaps another's shares keys with it, which the handle cannot tell apart. A prefix that is empty, or made only of separator characters such as `:`, 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. + +Everything else on the index is out of reach: creating, dropping or clearing it, `info` for another index, `paginate`, and the filter-wide bulk writes `update_by_filter` and `drop_by_filter`, which a request timeout could cancel partway through. + +The server applies the same rails as for a built-in, without the author asking for them: + +- Every `ctx.index` call takes a slot under the server-wide concurrency cap and is bounded by the binding's `runtime.request_timeout_seconds`. The cap is taken per call, so a tool that queries several times never blocks itself. +- The whole tool call is bounded by the same timeout, so a tool cannot hold a request open indefinitely. A timeout returns a retryable `backend_unavailable` error. +- A caller needs the server's read scope, and the write scope as well for a tool declared with `writes=True`, whose handle can both read and write. A tool can require an extra scope with `scope=`, which a caller must hold as well, so a code tool can be stricter than the built-ins but never looser. Like every scope, it is enforced only when the server has authentication configured. +- An exception the tool raises becomes a stable MCP error: `ValueError` becomes `invalid_request`, a Redis error from a search, aggregation or fetch a retryable `backend_unavailable`, and most other exceptions `internal_error`. That includes any failure of `ctx.index.load`, which reports even a lost connection as a non-retryable `internal_error`. Raise `RedisVLMCPError` to choose the code and message yourself. + +The handle is a guardrail, not a sandbox. A code tool is your own code, running with the privileges of the process that started the server, and Python can still reach private attributes. What the handle does is make the bounded, policy-checked path the only obvious one. + +### What Fails for a Code Tool + +The function is checked when it is decorated: it must be async, take a `ToolContext` first, annotate every other parameter, and have a description. Registering a tool after the server has started is refused, because tools register once per process. + +The rest is checked at startup, against the loaded config: an `index` that is missing with several indexes configured or that names no binding; a name another custom tool already uses; a tool declared with `writes=True` on a read-only binding or server; and a code tool on a Redis index that a profile scopes with `lock.inject`. A code tool can run any query its handle allows, so nothing would show that it applies the tenant scope. Bind it to another index, or serve it from a separate server. + ## Why Use MCP Instead of Direct RedisVL Calls Use RedisVL MCP when you want a standard tool boundary for agent frameworks or assistants that already speak MCP. diff --git a/docs/user_guide/how_to_guides/mcp.md b/docs/user_guide/how_to_guides/mcp.md index 416a59db..7a918013 100644 --- a/docs/user_guide/how_to_guides/mcp.md +++ b/docs/user_guide/how_to_guides/mcp.md @@ -407,7 +407,7 @@ What the client sees for `search-customer-kb`: - A description whose field hints list `content` but not `org_id`. - Results from its own tenant only. A `filter` naming `org_id` ANDs with the injected value, so naming another tenant returns nothing. Results still carry `org_id`, always with the caller's own value, unless you lock `return_fields` to leave it out. -The server refuses to start while anything else can reach the same Redis index without the tenant scope: `search-records`, `upsert-records` unless every binding over the index is `read_only`, or another custom tool on that index whose `lock.inject` differs, including one with none. Bindings that share a `redis_name` are the same index, so a vector tool and a fulltext tool over one set of documents must inject the same scope. Each of those routes would hand every caller a way round the profile, and a write could retag another tenant's document as the writer's own. Ingest documents outside the server, stamping `org_id` exactly as the identity provider emits it. +The server refuses to start while anything else can reach the same Redis index without the tenant scope: `search-records`, `upsert-records` unless every binding over the index is `read_only`, another custom tool on that index whose `lock.inject` differs, including one with none, or any code tool bound to it. Bindings that share a `redis_name` are the same index, so a vector tool and a fulltext tool over one set of documents must inject the same scope. Each of those routes would hand every caller a way round the profile, and a write could retag another tenant's document as the writer's own. Ingest documents outside the server, stamping `org_id` exactly as the identity provider emits it. A tool meant to read across tenants, such as an internal support search, belongs on a separate server with its own authentication. Every tool on one server sits behind the same read scope, so any tenant could call it. @@ -415,6 +415,60 @@ Listing the tenant claim under `required_claims` makes the verifier reject a tok The server refuses to start an injecting profile without authentication, over `stdio`, or on a field the index does not hold as an indexed tag. The `stdio` check applies when the server starts through `rvl mcp` or `run_async`. For what the guarantee covers and what it rests on, read the threat model in {doc}`/concepts/mcp`. +### Code Tools Registered From Python + +When a tool needs logic a profile cannot express, write it as an async function and register it on the server before you run it. This suits a server you start from your own Python process: + +```python +from typing import cast + +from redisvl.mcp import MCPSettings, RedisVLMCPServer, ToolContext, WritableIndexHandle +from redisvl.query import FilterQuery +from redisvl.query.filter import Tag + +server = RedisVLMCPServer(MCPSettings(config="/path/to/mcp_config.yaml")) + + +@server.code_tool(index="tickets") +async def open_tickets_for_product(ctx: ToolContext, product: str, limit: int = 5) -> dict: + """List open support tickets for a product.""" + query = FilterQuery( + filter_expression=(Tag("status") == "open") & (Tag("product") == product), + return_fields=["subject", "created_at"], + num_results=limit, + ) + return {"tickets": await ctx.index.query(query)} + + +@server.code_tool(index="tickets", writes=True, scope="tickets.triage") +async def delete_spam_ticket(ctx: ToolContext, ticket_id: str) -> dict: + """Delete a support ticket that triage has marked as spam.""" + index = cast(WritableIndexHandle, ctx.index) + ticket = await index.fetch(ticket_id) + if ticket is None or ticket.get("status") != "spam": + raise ValueError(f"ticket {ticket_id} is not marked as spam") + return {"deleted": await index.drop_documents(ticket_id)} + + +server.run(transport="streamable-http", host="127.0.0.1", port=8000) +``` + +What the client sees: + +- `open-tickets-for-product`, taking `product` (required) and `limit`, described by its docstring. +- `delete-spam-ticket`, taking `ticket_id`. With authentication configured, a caller needs the server's read and write scopes and `tickets.triage`; without it, no scope is checked. Either way, the tool deletes nothing that is not already marked as spam. +- No argument for the context or the index, in either tool. + +Rules worth knowing: + +- **A code tool reaches Redis only through `ctx.index`.** It has the read methods, plus a handful of bounded writes for a tool registered with `writes=True`. Records are addressed by document id, never by full Redis key, so a tool stays under its index's key prefix. That does not separate indexes whose prefixes overlap, and a prefix that is empty or made only of separator characters covers every key. On an index with several prefixes, ids address the first prefix only. `load` returns document ids, and `ctx.index.document_id(hit["id"])` turns a search result's key into one. On JSON storage, give a `FilterQuery` `return_fields`, or a document's own `id` field replaces the key in its results. There is no Redis client and no way to create, drop or clear the index. +- **Writing tools need a cast for type checkers.** `ctx.index` is typed as `IndexHandle`. In a tool registered with `writes=True`, `cast(WritableIndexHandle, ctx.index)` gives a type checker the write methods, as `delete-spam-ticket` shows. Both types are exported from `redisvl.mcp`. +- **Every call is bounded.** Each `ctx.index` call takes a slot under the concurrency cap and the binding's `request_timeout_seconds`, and the whole tool call is bounded by the same timeout. +- **Raise to report a problem.** A `ValueError` reaches the client as `invalid_request`; raise `RedisVLMCPError` to choose the error code yourself. +- **Register before starting.** Tools register once, at the first startup, so a tool registered afterwards is refused. + +A code tool bound to an index that a profile scopes with `lock.inject` refuses to start, because nothing shows that it applies the tenant scope. For the full contract, read the code tools section in {doc}`/concepts/mcp`. + ## Tool Contracts RedisVL MCP exposes a small, implementation-owned contract. @@ -803,8 +857,12 @@ A profile with `lock.inject` refuses to start when authentication is not enabled ### Claim Injection Refuses an Unscoped Route -A profile with `lock.inject` refuses to start while its Redis index is also reachable through `search-records`, through `upsert-records` on any writable binding over it, or through another custom tool on it whose `lock.inject` differs, including one with none. Bindings that share a `redis_name` are the same index. The error names the Redis index and the bindings over it, each route it found, and each tool's injected scope. Disable the built-ins under `server.builtin_tools`, mark every binding over the index `read_only`, and give every custom tool on the index the same `lock.inject`. Move a tool meant to read across tenants to a separate server with its own authentication. +A profile with `lock.inject` refuses to start while its Redis index is also reachable through `search-records`, through `upsert-records` on any writable binding over it, through another custom tool on it whose `lock.inject` differs, including one with none, or through a code tool bound to it. Bindings that share a `redis_name` are the same index. The error names the Redis index and the bindings over it, each route it found, and each tool's injected scope. Disable the built-ins under `server.builtin_tools`, mark every binding over the index `read_only`, and give every custom tool on the index the same `lock.inject`. Bind a code tool to another index. Move a tool meant to read across tenants to a separate server with its own authentication. ### Claim Injection Fails Every Request With `forbidden` The token is verified but its claim is unusable: missing, empty, padded with whitespace, not a single string, or containing a control character or backtick. The error names the claim and the tool. A misspelled `claim` name in the config is the usual cause, since a JWT claim name such as `https://acme.example/org` must match exactly. + +### Code Tool Fails at Startup + +The error names the tool. Either its `index` is missing while several indexes are configured or names no configured binding, or another custom tool already uses its name, or it is registered with `writes=True` on a binding or server that is read-only. A code tool registered after the server has started is refused with a `RuntimeError`; register every tool before calling `run`, `run_async` or `startup`. A restart in the same process also fails with a `RuntimeError` if the code tools, or the indexes they bind to, have changed; restart the process instead. diff --git a/redisvl/mcp/__init__.py b/redisvl/mcp/__init__.py index 619c040c..a94bae1b 100644 --- a/redisvl/mcp/__init__.py +++ b/redisvl/mcp/__init__.py @@ -1,15 +1,19 @@ from redisvl.mcp.config import MCPConfig, MCPTransportSecurityConfig, load_mcp_config +from redisvl.mcp.context import IndexHandle, ToolContext, WritableIndexHandle from redisvl.mcp.errors import MCPErrorCode, RedisVLMCPError, map_exception from redisvl.mcp.server import RedisVLMCPServer from redisvl.mcp.settings import MCPSettings __all__ = [ + "IndexHandle", "MCPConfig", "MCPErrorCode", "MCPSettings", "MCPTransportSecurityConfig", "RedisVLMCPError", "RedisVLMCPServer", + "ToolContext", + "WritableIndexHandle", "load_mcp_config", "map_exception", ] diff --git a/redisvl/mcp/auth.py b/redisvl/mcp/auth.py index d2c8e462..666e8913 100644 --- a/redisvl/mcp/auth.py +++ b/redisvl/mcp/auth.py @@ -247,7 +247,6 @@ class InjectSpec(Protocol): A protocol rather than an import of the configuration model, so the security core has no dependency on the shape of the YAML that reaches it. - A v1.1 code tool assembles its own specs and gets the same guarantees. """ @property diff --git a/redisvl/mcp/server.py b/redisvl/mcp/server.py index ab550318..8a575071 100644 --- a/redisvl/mcp/server.py +++ b/redisvl/mcp/server.py @@ -1,5 +1,6 @@ import asyncio import logging +from collections.abc import Callable from contextlib import asynccontextmanager from enum import Enum, auto from importlib import import_module @@ -15,6 +16,12 @@ from redisvl.mcp.errors import MCPErrorCode, RedisVLMCPError from redisvl.mcp.runtime import BindingRuntime from redisvl.mcp.settings import MCPSettings +from redisvl.mcp.tools.code import ( + CodeToolSpec, + _ToolFn, + build_code_tool_spec, + register_code_tool, +) from redisvl.mcp.tools.list_indexes import register_list_indexes_tool from redisvl.mcp.tools.profiles import ( register_profile_tools, @@ -89,6 +96,12 @@ class RedisVLMCPServer(FastMCP): _LifecycleState = _LifecycleState + # Each code tool paired with the binding it resolved to, settled at every + # startup once the config is loaded, and the pairing the tools were + # registered with at the first. Empty until then. + _resolved_code_tools: tuple[tuple[CodeToolSpec, str], ...] = () + _registered_code_tools: tuple[tuple[CodeToolSpec, str], ...] = () + def __init__(self, settings: MCPSettings): """Create a server shell with lazy config and per-binding runtime state.""" self.mcp_settings = settings @@ -98,6 +111,7 @@ def __init__(self, settings: MCPSettings): self._tools_registered = False self._registered_tool_fingerprint = "" self._registered_tools_inject = False + self._code_tools: list[CodeToolSpec] = [] # Set by run_async, the only point in the process that knows it. None # means an embedder started the server without naming a transport. self._transport: str | None = None @@ -213,6 +227,57 @@ async def shutdown(self) -> None: finally: await self._mark_stopped() + def code_tool( + self, + *, + name: str | None = None, + namespace: str | None = None, + index: str | None = None, + writes: bool = False, + scope: str | None = None, + description: str | None = None, + ) -> Callable[[_ToolFn], _ToolFn]: + """Register an async function as a code tool on this server. + + The function takes a :class:`~redisvl.mcp.context.ToolContext` first and + the model's arguments after it; its docstring is the tool description + unless ``description`` is given. ``index`` pins the binding (required + when several are configured), ``writes`` grants the write methods and + requires the write scope on top of the read scope, and ``scope`` is an + extra scope a caller must hold on top of the server-wide ones; like + every scope, it is enforced only when authentication is configured. + The function is returned unchanged. + + The function is checked here; whether its binding exists and permits it + is checked at startup, where the config is known. + """ + + def decorator(fn: _ToolFn) -> _ToolFn: + # While starting, the code tools have already been read, so one + # added now would be silently left out. + if ( + self._tools_registered + or self._lifecycle_state is _LifecycleState.STARTING + ): + raise RuntimeError( + "MCP tools are registered once, at the first startup; " + "register code tools before starting the server" + ) + self._code_tools.append( + build_code_tool_spec( + fn, + name=name, + namespace=namespace, + index=index, + writes=writes, + scope=scope, + description=description, + ) + ) + return fn + + return decorator + def resolve_binding(self, index_id: str | None) -> BindingRuntime: """Resolve the runtime for a logical index id, honoring single-index defaults. @@ -389,7 +454,57 @@ def _verify_injection_has_a_token(self, profiles: Any) -> None: "streamable-http) so the configured auth applies." ) - def _verify_no_unscoped_route_to_injected_indexes(self, config: Any) -> None: + def _resolve_code_tools(self, config: Any) -> tuple[tuple[CodeToolSpec, str], ...]: + """Pair each code tool with its binding, refusing any that cannot hold. + + A code tool is declared in Python, so the rules the config model applies + to profiles are applied here instead, against the loaded config. + """ + taken = {profile.name for profile in config.custom_tools} + resolved: list[tuple[CodeToolSpec, str]] = [] + for spec in self._code_tools: + if spec.name in taken: + raise ValueError( + f"code tool name '{spec.name}' is already used by another " + "custom tool; tool names must be unique" + ) + taken.add(spec.name) + + if spec.index is None: + if len(config.indexes) > 1: + raise ValueError( + f"code tool '{spec.name}' must name its index because " + "several are configured: " + f"{', '.join(sorted(config.indexes))}" + ) + binding_id = next(iter(config.indexes)) + elif spec.index not in config.indexes: + raise ValueError( + f"code tool '{spec.name}' names unknown index '{spec.index}'; " + f"configured: {', '.join(sorted(config.indexes))}" + ) + else: + binding_id = spec.index + + # Refused here rather than per call, because the tool's writes were + # declared up front: it would fail on every call it exists for. + if spec.writes and self.mcp_settings.read_only: + raise ValueError( + f"code tool '{spec.name}' writes, but the server runs read-only" + ) + if spec.writes and config.indexes[binding_id].read_only: + raise ValueError( + f"code tool '{spec.name}' writes, but index '{binding_id}' is " + "read-only" + ) + resolved.append((spec, binding_id)) + return tuple(resolved) + + def _verify_no_unscoped_route_to_injected_indexes( + self, + config: Any, + code_tools: tuple[tuple[CodeToolSpec, str], ...] = (), + ) -> None: """Refuse a tool surface that reaches a tenant-scoped index unscoped. Injection isolates a tool, but tenant data lives in an index, and every @@ -406,6 +521,10 @@ def _verify_no_unscoped_route_to_injected_indexes(self, config: Any) -> None: tool scoped by another field, or by the same field from another claim, reads across the tenants the first one separates. + A code tool is a route too, and one the server cannot inspect: it can + run any query its handle allows, so nothing proves it applies the scope. + It is refused on a scoped index outright. + The index here is the Redis index, not the binding: two bindings with the same ``redis_name`` are two routes to one set of documents. The check sees only what this config names, so an alias, or a second index @@ -454,6 +573,19 @@ def _verify_no_unscoped_route_to_injected_indexes(self, config: Any) -> None: f"{', '.join(routes)}" ) + unscoped_code_tools = sorted( + spec.name + for spec, binding_id in code_tools + if config.indexes[binding_id].redis_name in injected + ) + if unscoped_code_tools: + problems.append( + "code tools " + f"{', '.join(repr(name) for name in unscoped_code_tools)} reach " + "it, and a code tool cannot carry an injected scope (bind it to " + "another index)" + ) + for redis_name in injected: by_tool = scopes[redis_name] if len(set(by_tool.values())) > 1: @@ -507,6 +639,17 @@ def _register_tools(self) -> None: # The dangerous direction is an operator disabling a tool or tightening # a lock and believing the restart applied it. if self._tools_registered: + # A code tool runs operator code against whatever binding it + # was registered on, so a restart that would pair it with a + # different one, or drop it, is never just a log line. + if self._resolved_code_tools != self._registered_code_tools: + raise RuntimeError( + "MCP code tools, or the indexes they are bound to, " + "changed since tools were registered. Tools register " + "once per process, so the previously registered code " + "tools would stay in effect. Restart the process to " + "apply the new configuration." + ) config = getattr(self, "config", None) current = self._tool_surface_fingerprint(config) if current != self._registered_tool_fingerprint: @@ -581,9 +724,13 @@ def _register_tools(self) -> None: register_upsert_tool(self, index_ids=unlisted_index_ids) registered.append("upsert-records") registered.extend(register_profile_tools(self)) + for spec, binding_id in self._resolved_code_tools: + register_code_tool(self, spec, binding_id) + registered.append(spec.name) self._warn_on_unusable_tool_surface(registered) self._registered_tool_fingerprint = self._tool_surface_fingerprint(config) + self._registered_code_tools = self._resolved_code_tools self._registered_tools_inject = _config_injects(config) self._tools_registered = True @@ -757,7 +904,10 @@ async def _initialize_runtime_resources(self) -> None: # Before any binding connects: both depend only on the loaded config, # so an unworkable one should not first need a reachable Redis. self._verify_injection_has_a_token(self.config.custom_tools) - self._verify_no_unscoped_route_to_injected_indexes(self.config) + self._resolved_code_tools = self._resolve_code_tools(self.config) + self._verify_no_unscoped_route_to_injected_indexes( + self.config, self._resolved_code_tools + ) # The semaphore is a single process-wide concurrency ceiling shared by # all bindings; take the max across bindings. This means the most # permissive binding sets the cap — e.g. five bindings each configured diff --git a/redisvl/mcp/tools/code.py b/redisvl/mcp/tools/code.py index 2c3bba83..38ea48d8 100644 --- a/redisvl/mcp/tools/code.py +++ b/redisvl/mcp/tools/code.py @@ -18,7 +18,7 @@ import typing from collections.abc import Callable from dataclasses import dataclass -from typing import Any +from typing import Any, TypeVar from pydantic import TypeAdapter @@ -27,6 +27,8 @@ from redisvl.mcp.context import ToolContext from redisvl.mcp.errors import MCPErrorCode, RedisVLMCPError, map_exception +_ToolFn = TypeVar("_ToolFn", bound=Callable[..., Any]) + _ACCEPTED_KINDS = ( inspect.Parameter.POSITIONAL_OR_KEYWORD, inspect.Parameter.KEYWORD_ONLY, diff --git a/tests/integration/test_mcp/test_code_tools.py b/tests/integration/test_mcp/test_code_tools.py new file mode 100644 index 00000000..fd61be67 --- /dev/null +++ b/tests/integration/test_mcp/test_code_tools.py @@ -0,0 +1,217 @@ +"""Code tools registered with ``server.code_tool``, against a real Redis. + +The unit suites own the validation and wrapping matrices. What lives here is +what depends on a started server: that a tool's handle reaches the real index, +that a writing tool really writes, that startup refuses what the config cannot +hold, and that the concurrency cap governs each Redis call without deadlocking +a tool that makes several. +""" + +import asyncio +from pathlib import Path + +import pytest +import yaml + +from redisvl.index import AsyncSearchIndex +from redisvl.mcp import ToolContext +from redisvl.mcp.server import RedisVLMCPServer +from redisvl.mcp.settings import MCPSettings +from redisvl.query import FilterQuery +from redisvl.query.filter import Tag +from redisvl.schema import IndexSchema + +fastmcp = pytest.importorskip( + "fastmcp", reason="fastmcp not installed (install redisvl[mcp])" +) +from fastmcp.server.auth.providers.jwt import RSAKeyPair + + +@pytest.fixture +async def tickets_index(async_client, worker_id): + schema = IndexSchema.from_dict( + { + "index": { + "name": f"mcp-code-tools-{worker_id}", + "prefix": f"mcp-code-tools:{worker_id}", + "storage_type": "hash", + }, + "fields": [ + {"name": "content", "type": "text"}, + {"name": "status", "type": "tag"}, + ], + } + ) + index = AsyncSearchIndex(schema=schema, redis_client=async_client) + await index.create(overwrite=True, drop=True) + await index.load( + [ + {"id": "t1", "content": "refund delayed", "status": "open"}, + {"id": "t2", "content": "refund issued", "status": "resolved"}, + {"id": "t3", "content": "login failure", "status": "open"}, + ], + id_field="id", + ) + yield index + await index.delete(drop=True) + + +@pytest.fixture +def config_path(tmp_path: Path, redis_url: str): + def factory( + redis_name: str, + *, + read_only: bool = False, + max_concurrency: int = 16, + custom_tools: list[dict] | None = None, + auth: dict | None = None, + ) -> str: + server: dict = {"redis_url": redis_url} + if auth is not None: + server["auth"] = auth + server["builtin_tools"] = { + "search-records": "disabled", + "upsert-records": "disabled", + } + config = { + "server": server, + "indexes": { + "tickets": { + "redis_name": redis_name, + "read_only": read_only, + "search": {"type": "fulltext"}, + "runtime": { + "text_field_name": "content", + "max_concurrency": max_concurrency, + }, + } + }, + "custom_tools": custom_tools or [], + } + path = tmp_path / f"{redis_name}-code-tools.yaml" + path.write_text(yaml.safe_dump(config), encoding="utf-8") + return str(path) + + return factory + + +@pytest.fixture +async def server_for(tickets_index, config_path): + servers: list[RedisVLMCPServer] = [] + + def build(**kwargs) -> RedisVLMCPServer: + server = RedisVLMCPServer( + MCPSettings(config=config_path(tickets_index.schema.index.name, **kwargs)) + ) + servers.append(server) + return server + + yield build + + for server in servers: + await server.shutdown() + + +async def _call(server: RedisVLMCPServer, name: str, **kwargs): + tool = await server.get_tool(name) + assert tool is not None, f"tool '{name}' is not registered" + return await tool.fn(**kwargs) + + +async def test_a_writing_tool_writes_through_its_handle(server_for): + server = server_for() + + @server.code_tool(writes=True) + async def close_ticket(ctx: ToolContext, ticket_id: str) -> dict: + """Mark a support ticket resolved.""" + record = await ctx.index.fetch(ticket_id) + await ctx.index.load([{**record, "status": "resolved"}], id_field="id") + return {"status": (await ctx.index.fetch(ticket_id))["status"]} + + await server.startup() + + assert await _call(server, "close-ticket", ticket_id="t1") == {"status": "resolved"} + + +async def test_ids_from_a_load_and_a_search_address_the_same_record(server_for): + # Both return full keys underneath, and the writes re-apply the prefix, so + # only document ids reach the record. + server = server_for() + + @server.code_tool(writes=True) + async def file_and_withdraw(ctx: ToolContext) -> dict: + """File a ticket under a generated id, find it again, then withdraw it.""" + [filed] = await ctx.index.load( + [{"content": "charged twice", "status": "draft"}] + ) + [hit] = await ctx.index.query( + FilterQuery(filter_expression=Tag("status") == "draft") + ) + found = ctx.index.document_id(hit["id"]) + return { + "same": found == filed, + "dropped": await ctx.index.drop_documents(found), + } + + await server.startup() + + assert await _call(server, "file-and-withdraw") == {"same": True, "dropped": 1} + + +async def test_a_code_tool_on_a_tenant_scoped_index_fails_startup(server_for): + key = RSAKeyPair.generate() + server = server_for( + auth={ + "type": "jwt", + "public_key": key.public_key, + "issuer": "https://auth.acme.example/", + "audience": "api://redisvl-mcp", + }, + custom_tools=[ + { + "name": "tenant-tickets", + "description": "Search this tenant's tickets.", + "lock": { + "inject": [{"field": "status", "from": "claim", "claim": "org"}] + }, + } + ], + ) + # Startup only checks for a token source on HTTP; naming one keeps the + # refusal under test the code-tool one. + server._transport = "streamable-http" + + @server.code_tool() + async def all_tickets(ctx: ToolContext) -> dict: + """List every ticket.""" + return {} + + with pytest.raises(ValueError, match="code tools 'all-tickets' reach it"): + await server.startup() + + +async def test_each_redis_call_takes_the_cap_without_deadlocking_the_tool( + server_for, +): + # With a cap of one, a tool holding the cap across its whole body would + # block its own second call. The cap is per Redis call, so two concurrent + # calls of a tool that queries twice each still finish. + server = server_for(max_concurrency=1) + + @server.code_tool() + async def ticket_counts(ctx: ToolContext) -> dict: + """Count open and resolved tickets.""" + counts = {} + for status in ("open", "resolved"): + query = FilterQuery(filter_expression=Tag("status") == status) + counts[status] = len(await ctx.index.query(query)) + return counts + + await server.startup() + + results = await asyncio.wait_for( + asyncio.gather(_call(server, "ticket-counts"), _call(server, "ticket-counts")), + timeout=10, + ) + + assert results == [{"open": 2, "resolved": 1}] * 2 diff --git a/tests/unit/test_mcp/test_server_unit.py b/tests/unit/test_mcp/test_server_unit.py index 70c72ec0..23999d50 100644 --- a/tests/unit/test_mcp/test_server_unit.py +++ b/tests/unit/test_mcp/test_server_unit.py @@ -4,9 +4,11 @@ import pytest from redisvl.mcp.config import MCPConfig, builtin_tool_names +from redisvl.mcp.context import ToolContext from redisvl.mcp.errors import MCPErrorCode, RedisVLMCPError from redisvl.mcp.runtime import BindingRuntime from redisvl.mcp.server import RedisVLMCPServer +from redisvl.mcp.tools.code import build_code_tool_spec class FakeClient: @@ -146,9 +148,26 @@ async def fake_close_resources(self, *, index, vectorizer): assert server._tools_registered is True -def _register_tools_with(monkeypatch, bindings: dict, *, config=None) -> list[str]: +def _register_tools_with( + monkeypatch, bindings: dict, *, config=None, code_tools=() +) -> list[str]: """Run _register_tools against the given bindings, returning registered names.""" + server, registered = _registration_server( + monkeypatch, bindings, config=config, code_tools=code_tools + ) + server._register_tools() + return registered + + +def _registration_server( + monkeypatch, bindings: dict, *, config=None, code_tools=() +) -> tuple[RedisVLMCPServer, list[str]]: + """A server ready for _register_tools, recording every name it registers.""" registered: list[str] = [] + monkeypatch.setattr( + "redisvl.mcp.server.register_code_tool", + lambda server, spec, binding_id: registered.append(spec.name), + ) monkeypatch.setattr( "redisvl.mcp.server.register_list_indexes_tool", lambda server: registered.append("list-indexes"), @@ -185,9 +204,8 @@ def fake_register_profile_tools(server): server.tool = object() server.config = config server.mcp_settings = SimpleNamespace(read_only=False) - - server._register_tools() - return registered + server._resolved_code_tools = tuple(code_tools) + return server, registered def _config_with(*, builtin_tools=None, custom_tools=None) -> MCPConfig: @@ -1042,3 +1060,255 @@ def test_tools_injecting_the_same_scope_share_an_index(): }, ], ) + + +# -------------------------------------------------------------------------- +# Code tools registered with server.code_tool +# -------------------------------------------------------------------------- + + +async def recent_tickets(ctx: ToolContext, query: str) -> dict: + """Find recent tickets.""" + return {} + + +def _code_tool_server(*, read_only=False) -> RedisVLMCPServer: + server = RedisVLMCPServer.__new__(RedisVLMCPServer) + server._code_tools = [] + server._tools_registered = False + server._lifecycle_state = server._LifecycleState.INITIAL + server.mcp_settings = SimpleNamespace(read_only=read_only) + return server + + +def _two_index_config(*, custom_tools=None, **knowledge_overrides) -> MCPConfig: + return MCPConfig.model_validate( + { + "server": {"redis_url": "redis://localhost:6379"}, + "indexes": { + "knowledge": { + "redis_name": "docs-index", + "search": {"type": "fulltext"}, + "runtime": {"text_field_name": "content"}, + **knowledge_overrides, + }, + "tickets": { + "redis_name": "tickets-index", + "search": {"type": "fulltext"}, + "runtime": {"text_field_name": "content"}, + }, + }, + "custom_tools": [ + {**tool, "index": tool.get("index", "knowledge")} + for tool in custom_tools or [] + ], + } + ) + + +def test_custom_tool_records_the_tool_and_returns_the_function(): + server = _code_tool_server() + + decorated = server.code_tool(namespace="support")(recent_tickets) + + assert decorated is recent_tickets + assert [spec.name for spec in server._code_tools] == ["support-recent-tickets"] + + +def test_custom_tool_checks_the_function_where_it_is_declared(): + server = _code_tool_server() + + def sync_tool(ctx: ToolContext, query: str) -> dict: + """Search.""" + return {} + + with pytest.raises(TypeError, match="must be an async function"): + server.code_tool()(sync_tool) + assert server._code_tools == [] + + +@pytest.mark.parametrize("moment", ["registered", "starting"]) +def test_custom_tool_after_registration_is_refused(moment): + server = _code_tool_server() + if moment == "registered": + server._tools_registered = True + else: + # Startup has read the code tools but not yet registered them. + server._lifecycle_state = server._LifecycleState.STARTING + + with pytest.raises(RuntimeError, match="register code tools before starting"): + server.code_tool()(recent_tickets) + assert server._code_tools == [] + + +def test_a_code_tool_on_a_single_index_server_binds_to_it(): + server = _code_tool_server() + server.code_tool()(recent_tickets) + + ((spec, binding_id),) = server._resolve_code_tools(_config_with()) + + assert (spec.name, binding_id) == ("recent-tickets", "knowledge") + + +def test_a_code_tool_binds_to_the_index_it_names(): + server = _code_tool_server() + server.code_tool(index="tickets")(recent_tickets) + + ((_, binding_id),) = server._resolve_code_tools(_two_index_config()) + + assert binding_id == "tickets" + + +@pytest.mark.parametrize( + "options, config, expected", + [ + pytest.param( + {}, + _two_index_config(), + "must name its index because several are configured: knowledge, tickets", + id="index-required", + ), + pytest.param( + {"index": "archive"}, + _two_index_config(), + "unknown index 'archive'; configured: knowledge, tickets", + id="unknown-index", + ), + pytest.param( + {"index": "knowledge", "name": "resolved-search"}, + _two_index_config( + custom_tools=[{"name": "resolved-search", "description": "Search."}] + ), + "'resolved-search' is already used by another custom tool", + id="profile-name", + ), + pytest.param( + {"index": "knowledge", "writes": True}, + _two_index_config(read_only=True), + "'recent-tickets' writes, but index 'knowledge' is read-only", + id="binding-read-only", + ), + ], +) +def test_a_code_tool_its_config_cannot_hold_fails_startup(options, config, expected): + server = _code_tool_server() + server.code_tool(**options)(recent_tickets) + + with pytest.raises(ValueError, match=expected): + server._resolve_code_tools(config) + + +def test_two_code_tools_cannot_share_a_name(): + server = _code_tool_server() + server.code_tool()(recent_tickets) + server.code_tool()(recent_tickets) + + with pytest.raises(ValueError, match="'recent-tickets' is already used"): + server._resolve_code_tools(_config_with()) + + +def test_a_writing_code_tool_is_refused_on_a_read_only_server(): + server = _code_tool_server(read_only=True) + server.code_tool(writes=True)(recent_tickets) + + with pytest.raises(ValueError, match="writes, but the server runs read-only"): + server._resolve_code_tools(_config_with()) + + +def test_a_reading_code_tool_may_run_on_a_read_only_server(): + server = _code_tool_server(read_only=True) + server.code_tool()(recent_tickets) + + assert len(server._resolve_code_tools(_config_with())) == 1 + + +def _code_tool_route_check(*, index, indexes): + server = _code_tool_server() + server.code_tool(index=index)(recent_tickets) + config = MCPConfig.model_validate( + { + "server": { + "redis_url": "redis://localhost:6379", + "builtin_tools": _BUILTINS_OFF, + }, + "indexes": { + index_id: { + "search": {"type": "fulltext"}, + "runtime": {"text_field_name": "content"}, + **overrides, + } + for index_id, overrides in indexes.items() + }, + "custom_tools": [{**_INJECTING, "index": "kb_semantic"}], + } + ) + server._verify_no_unscoped_route_to_injected_indexes( + config, server._resolve_code_tools(config) + ) + + +@pytest.mark.parametrize( + "index", + [ + pytest.param("kb_semantic", id="same-binding"), + # A second binding over the same Redis index reaches the same documents. + pytest.param("kb_keyword", id="sibling-binding"), + ], +) +def test_a_code_tool_on_a_tenant_scoped_index_is_refused(index): + with pytest.raises(ValueError, match="code tools 'recent-tickets' reach it"): + _code_tool_route_check( + index=index, + indexes={ + "kb_semantic": {"redis_name": "customer-kb"}, + "kb_keyword": {"redis_name": "customer-kb"}, + }, + ) + + +def test_a_code_tool_on_another_index_is_no_route(): + _code_tool_route_check( + index="public", + indexes={ + "kb_semantic": {"redis_name": "customer-kb"}, + "public": {"redis_name": "public-kb"}, + }, + ) + + +def test_register_tools_registers_code_tools_after_profiles(monkeypatch): + server = _code_tool_server() + server.code_tool()(recent_tickets) + config = _config_with( + custom_tools=[{"name": "resolved-search", "description": "Search."}] + ) + + registered = _register_tools_with( + monkeypatch, + {"knowledge": _binding_runtime("knowledge")}, + config=config, + code_tools=server._resolve_code_tools(config), + ) + + assert registered[-2:] == ["resolved-search", "recent-tickets"] + + +def _restart_with(monkeypatch, *, registered_on, restarted_on): + """Register a code tool on one binding, then restart pairing it with another.""" + spec = build_code_tool_spec(recent_tickets) + # The first startup records the pairing itself, as a real server does. + server, _ = _registration_server( + monkeypatch, {}, code_tools=[(spec, registered_on)] + ) + server._register_tools() + server._resolved_code_tools = ((spec, restarted_on),) + server._register_tools() + + +def test_a_restart_that_moves_a_code_tool_is_fatal(monkeypatch): + with pytest.raises(RuntimeError, match="code tools, or the indexes they are"): + _restart_with(monkeypatch, registered_on="knowledge", restarted_on="tickets") + + +def test_a_restart_with_the_same_code_tools_proceeds(monkeypatch): + _restart_with(monkeypatch, registered_on="knowledge", restarted_on="knowledge")