Skip to content

Adding a tool, resource, template or prompt

Section 8 · From docs/MCP_INTERNALS.md on GitHub

The steps, then a complete module and its tests. The example was run from outside the repo (module and test in a scratch directory, -p tests.mcp.conftest to load the fixtures) and its four tests passed.

  1. Pick a category. A new category needs an entry in catalog/__init__.py:CATEGORIES and in OWN_CATEGORIES at the top of tests/mcp/test_catalog_registry.py; an existing one needs nothing. Prompts always use category prompts, whatever module defines them.
  2. Write the handler with ctx first and every argument annotated. Use the aliases from catalog/_common.py (DatasetArg, RefArg, Limit, Cursor) so arguments are described the same way across the catalog. Tools in the inventory, graph, findings, ontology and export categories take a dataset argument. Give the docstring the job description the model will read: what it returns, when to use it, an example call. The registry test wants more than 40 characters and a title.
  3. Declare honest metadata: sensitivity (what the output may contain), capabilities (what side effects it needs), and the hints when the derivation in 3.5 would be wrong. Policies act on these. The registry test requires read_only, idempotent and open_world to end up set, no write, cloud or exec capability on a read-only tool, and open_world true exactly when the tool has cloud_access.
  4. Return compact data. Use asset_brief / finding_brief and resource URIs instead of full objects, paginate lists with paginate and fingerprint, and attach ctx.link(...) for detail views.
  5. Fail with InvalidArgumentsError for bad input and ReferenceNotFoundError (or another NotFoundError) for misses, with suggestions in data=. The model gets both.
  6. Add an output schema only from a model whose fields are all optional with extra="allow" (like the *Out models), so transforms cannot break it, and add the tool with sample arguments to OUTPUT_CALLS in test_catalog_registry.py; that test calls it on the sample estate and validates the result against the schema.
  7. Template variables each need a completion function; required prompt arguments need one too, and every prompt argument needs a description.
  8. Register the module in catalog/__init__.py:_modules() and write tests with the call fixture.
  9. Run tests/mcp/test_catalog_registry.py; it checks the rules above for the whole catalog and validates every wire dict with the SDK's types when mcp is installed.

cloudg/mcp/catalog/owners.py:

"""Ownership tools: who owns what, by the ``owner`` tag."""

from __future__ import annotations

from typing import Annotated, Any

from pydantic import BaseModel, ConfigDict, Field

from cloudg.mcp.catalog._common import (
    DatasetArg,
    Limit,
    Cursor,
    asset_brief,
    fingerprint,
    paginate,
    ws_dataset,
)
from cloudg.mcp.core import (
    PromptArgument,
    PromptMessage,
    PromptResult,
    Registry,
    Sensitivity,
    TextContent,
)
from cloudg.mcp.schema import output_schema_for
from cloudg.mcp.state import ReferenceNotFoundError

CATEGORY = "ownership"


class OwnerAssetsOut(BaseModel):
    model_config = ConfigDict(extra="allow")
    owner: str | None = None
    total: int | None = None
    next_cursor: str | None = None
    items: list[dict[str, Any]] = Field(default_factory=list)


def _owners(ds: Any) -> dict[str, int]:
    counts: dict[str, int] = {}
    for a in ds.assets:
        owner = a.tags.get("owner")
        if owner:
            counts[owner] = counts.get(owner, 0) + 1
    return counts


def register(reg: Registry) -> None:
    @reg.tool(
        title="Assets by owner",
        category=CATEGORY,
        sensitivity=Sensitivity.CONFIDENTIAL,
        read_only=True,
        idempotent=True,
        open_world=False,
        output_schema=output_schema_for(OwnerAssetsOut),
    )
    def assets_by_owner(
        ctx: Any,
        owner: Annotated[str, Field(min_length=1, description="Value of the 'owner' tag.")],
        limit: Limit = 50,
        cursor: Cursor = "",
        dataset: DatasetArg = "",
    ) -> dict:
        """Assets whose 'owner' tag equals OWNER, as compact briefs."""
        ds = ws_dataset(ctx, dataset)
        if owner not in _owners(ds):
            raise ReferenceNotFoundError(
                f"No asset is tagged owner={owner!r}.",
                data={"owners": sorted(_owners(ds))},
            )
        matches = [a for a in ds.assets if a.tags.get("owner") == owner]
        page, env = paginate(matches, limit, cursor, fingerprint(ds, owner=owner))
        ctx.link(f"cloudg://owners/{owner}", f"owner {owner}")
        return {"owner": owner, **env, "items": [asset_brief(ds, a) for a in page]}

    def complete_owner(ctx: Any, partial: str) -> list[str]:
        try:
            ds = ws_dataset(ctx)
        except Exception:
            return []
        return sorted(o for o in _owners(ds) if o.startswith(partial))

    @reg.resource_template(
        "cloudg://owners/{owner}",
        title="Owner",
        category=CATEGORY,
        sensitivity=Sensitivity.INTERNAL,
        completions={"owner": complete_owner},
    )
    def owner_resource(ctx: Any, owner: str) -> dict:
        """Asset count for one owner."""
        ds = ws_dataset(ctx)
        return {"owner": owner, "assets": _owners(ds).get(owner, 0)}

    @reg.prompt(
        title="Owner review",
        category="prompts",
        arguments=[PromptArgument("owner", "Owner tag value", required=True,
                                  completion=complete_owner)],
    )
    def owner_review(ctx: Any, owner: str) -> PromptResult:
        """Review everything one team owns."""
        text = (f"Call assets_by_owner(owner={owner!r}), then findings_for_asset on each "
                "asset with open findings, and summarise the team's exposure.")
        return PromptResult([PromptMessage("user", TextContent(text))], "Owner review")

The sync handlers run in worker threads, which is fine for in-memory work. Write the handler async if it awaits I/O or must be cancellable on timeout.

tests/mcp/test_catalog_owners.py, using the fixtures from tests/mcp/conftest.py (workspace is the sample estate loaded as dataset sample; call is bound to the layer fixture, which this file overrides to add the new module):

"""Tests for the ownership catalog module."""

from __future__ import annotations

import json

import pytest

from cloudg.mcp.catalog import default_registry, owners
from cloudg.mcp.layer import CloudGMCPLayer


@pytest.fixture
def layer(workspace):
    reg = default_registry()
    owners.register(reg)        # not needed once owners is in _modules()
    return CloudGMCPLayer(policy="open", workspace=workspace, registry=reg)


async def test_assets_by_owner(call):
    out = await call("assets_by_owner", owner="platform", limit=2)
    assert out["owner"] == "platform"
    assert out["returned"] == len(out["items"]) <= 2
    assert all(i["uri"].startswith("cloudg://assets/") for i in out["items"])


async def test_unknown_owner_is_a_tool_error(call):
    text = await call.error("assets_by_owner", owner="nobody")
    assert "No asset is tagged owner='nobody'" in text
    raw = await call.raw("assets_by_owner", owner="nobody")
    assert raw.meta["cloudg/error_code"] == -32602
    assert "platform" in raw.meta["cloudg/error_data"]["owners"]


async def test_resource_completion_and_prompt(layer):
    contents = await layer.read_resource("cloudg://owners/platform")
    assert json.loads(contents[0].text)["assets"] >= 1
    done = await layer.complete({"type": "ref/resource", "uri": "cloudg://owners/{owner}"},
                                {"name": "owner", "value": "pl"})
    assert done["values"] == ["platform"]
    prompt = await layer.get_prompt("owner_review", {"owner": "platform"})
    assert "assets_by_owner" in prompt.messages[0].content.text


async def test_strict_policy_pseudonymises(workspace):
    reg = default_registry()
    owners.register(reg)
    strict = CloudGMCPLayer(policy="strict", workspace=workspace, registry=reg)
    res = await strict.call_tool("assets_by_owner", {"owner": "platform", "limit": 1})
    assert not res.is_error
    assert "111111111111" not in json.dumps(res.structured["items"][0])
    assert "outputSchema" not in [t for t in strict.tools_wire()
                                  if t["name"] == "assets_by_owner"][0]

The call fixture: await call(tool, **args) returns the structured result and asserts success; call.raw(...) returns the ToolResult; call.error(...) asserts an error and returns its text. The last test is worth copying for any tool that returns identifiers: it checks the tool under a pseudonymizing profile.

The run, from the repository root with the module and test in a scratch directory:

$ PYTHONPATH=scratch/ext .venv/bin/python -m pytest -p tests.mcp.conftest scratch/ext/test_owners.py \
    -q -p no:cacheprovider -o asyncio_mode=auto
4 passed, 8 warnings in 1.02s

To check the catalog invariants as if the module were registered, a scratch pytest plugin appended owners to _modules(), added ownership to CATEGORIES and OWN_CATEGORIES, and added "assets_by_owner": {"owner": "platform"} to OUTPUT_CALLS. Before the OUTPUT_CALLS entry, test_structured_results_match_output_schemas failed with AssertionError: {'assets_by_owner'}; with it, test_catalog_registry.py gave 221 passed, 4 warnings (re-run after the latest catalog changes).

(The warnings are pydantic deprecation notices from elsewhere in cloudg.) Inside the repo the -p and -o flags are unnecessary: tests/mcp/conftest.py and the pytest settings in pyproject.toml apply automatically.

For a primitive that stays in your own deployment, the registry test does not apply. Build the registry yourself (reg = default_registry(); my_module.register(reg)) and pass registry=reg to the layer, or point the CLI at a factory with --registry my_module:build_registry.