datarobot_genai.drmcpbase.fastmcp_transforms.utils

utils

get_header_value

get_header_value(headers: Mapping[str, str], name: str) -> str | None

Read a header by lowercase name (FastMCP normalizes keys; scan only as fallback).

Source code in datarobot_genai/drmcpbase/fastmcp_transforms/utils.py

def get_header_value(headers: Mapping[str, str], name: str) -> str | None:
    """Read a header by lowercase name (FastMCP normalizes keys; scan only as fallback)."""
    value = headers.get(name)
    if value is not None:
        return value
    target = name.casefold()
    for key, header_value in headers.items():
        if key.casefold() == target:
            return header_value
    return None

parse_bool_header

parse_bool_header(raw: str | None, *, default: bool = True) -> bool

Parse an optional boolean header; unrecognized values fall back to default.

Source code in datarobot_genai/drmcpbase/fastmcp_transforms/utils.py

def parse_bool_header(raw: str | None, *, default: bool = True) -> bool:
    """Parse an optional boolean header; unrecognized values fall back to *default*."""
    if raw is None:
        return default
    token = raw.strip().casefold()
    if token in _TRUE_HEADER_VALUES:
        return True
    if token in _FALSE_HEADER_VALUES:
        return False
    return default

parse_disabled_categories

parse_disabled_categories(headers: Mapping[str, str]) -> frozenset[str]

Resolve the category-gate headers to the set of disabled category names.

Every gate defaults to enabled; only an explicit false disables its category for this request.

Source code in datarobot_genai/drmcpbase/fastmcp_transforms/utils.py

def parse_disabled_categories(headers: Mapping[str, str]) -> frozenset[str]:
    """Resolve the category-gate headers to the set of disabled category names.

    Every gate defaults to enabled; only an explicit ``false`` disables its
    category for this request.
    """
    disabled: set[str] = set()
    for header_name, category_names in CATEGORY_GATE_HEADERS.items():
        if not parse_bool_header(get_header_value(headers, header_name)):
            disabled.update(category_names)
    return frozenset(disabled)

register_toolsets_allowlist_expander

register_toolsets_allowlist_expander(expander: ToolsetsAllowlistExpander | None) -> None

Install async expansion of toolset bundle names → tool function names.

Global-mcp registers a mongo-backed expander once at startup. User-mcp leaves the default, where a toolsets header resolves to no tools at all — see :func:effective_tool_allowlist for what that means for the request. Pass None to reset (tests).

Source code in datarobot_genai/drmcpbase/fastmcp_transforms/utils.py

def register_toolsets_allowlist_expander(expander: ToolsetsAllowlistExpander | None) -> None:
    """Install async expansion of toolset bundle names → tool function names.

    Global-mcp registers a mongo-backed expander once at startup. User-mcp leaves the
    default, where a toolsets header resolves to no tools at all — see
    :func:`effective_tool_allowlist` for what that means for the request. Pass ``None`` to
    reset (tests).
    """
    _toolsets_allowlist_expander[0] = expander
    # Swapping the expander invalidates anything the previous one produced. In production
    # this runs once at startup, before any request; in tests it is what keeps one case's
    # expansion from being served to the next.
    _toolset_expansion_cache.set(None)

expand_toolset_names_to_tools async

expand_toolset_names_to_tools(names: frozenset[str] | None) -> frozenset[str]

Expand parsed toolset bundle names to tool function names for allowlisting.

Every failure resolves to the empty set, which :func:effective_tool_allowlist keeps as an empty allowlist rather than "no allowlist" — degrading toward fewer tools, not more.

Source code in datarobot_genai/drmcpbase/fastmcp_transforms/utils.py

async def expand_toolset_names_to_tools(names: frozenset[str] | None) -> frozenset[str]:
    """Expand parsed toolset bundle *names* to tool function names for allowlisting.

    Every failure resolves to the empty set, which :func:`effective_tool_allowlist` keeps
    as an *empty allowlist* rather than "no allowlist" — degrading toward fewer tools, not
    more.
    """
    expander = _toolsets_allowlist_expander[0]
    if not names:
        return frozenset()
    if expander is None:
        logger.warning(
            "Received %s but this server has no Tool Sets expander registered "
            "(Tool Sets are global-mcp only); the request is allowed no tools.",
            MCP_TOOLSETS_HEADER,
        )
        return frozenset()
    try:
        return await expander(names)
    except Exception:
        logger.warning(
            "Tool Sets expander failed; the request is allowed no tools from %s",
            MCP_TOOLSETS_HEADER,
            exc_info=True,
        )
        return frozenset()

is_tool_allowed

is_tool_allowed(tool: Tool, allowlist: ToolAllowlist) -> bool

Whether tool survives allowlist, taking into account how it was named.

Three ways in, and which apply depends on what kind of tool this is:

• an explicit name admits anything — the client asked for this tool by name;
• a derived name (from expanding a static category) admits only built-ins, because the static taxonomy describes DataRobot's own tools. A user-authored tool that merely shares a name is not what dr_db meant;
• a bucket admits a tool whose marker puts it there.

Ordered so the common paths never read the tool's marker: an explicit hit returns at once, and a name in neither set with no bucket requested is rejected before the marker is touched. Only a name that could match — or a request that named a bucket — pays for the lookup.

Source code in datarobot_genai/drmcpbase/fastmcp_transforms/utils.py

def is_tool_allowed(tool: Tool, allowlist: ToolAllowlist) -> bool:
    """Whether *tool* survives *allowlist*, taking into account how it was named.

    Three ways in, and which apply depends on what kind of tool this is:

    - an **explicit** name admits anything — the client asked for this tool by name;
    - a **derived** name (from expanding a static category) admits only *built-ins*,
      because the static taxonomy describes DataRobot's own tools. A user-authored
      tool that merely shares a name is not what ``dr_db`` meant;
    - a **bucket** admits a tool whose marker puts it there.

    Ordered so the common paths never read the tool's marker: an explicit hit returns
    at once, and a name in neither set with no bucket requested is rejected before the
    marker is touched. Only a name that *could* match — or a request that named a
    bucket — pays for the lookup.
    """
    name = tool.name
    if name in allowlist.explicit:
        return True
    in_derived = name in allowlist.derived
    if not in_derived and not allowlist.buckets:
        return False
    kind = marked_kind(get_tool_category(tool))
    if kind is None:
        # A built-in: the static taxonomy speaks for it.
        return in_derived
    # Marker-classified (user / deployment / proxied). Its category comes from the
    # marker, so only a bucket can admit it — never a name a static category expanded
    # to. Proxied tools carry a marker but no category, so no bucket names them.
    bucket = kind["category"]
    return bucket is not None and bucket in allowlist.buckets

filter_tools_by_allowlist

filter_tools_by_allowlist(tools: Sequence[Tool], allowlist: ToolAllowlist) -> list[Tool]

Apply allowlist to a whole catalog.

The fast path is not a second copy of the rule — it is the one case where the rule provably collapses. With no derived names and no buckets, :func:is_tool_allowed can only return name in explicit, so the marker is never consulted and this is the same single set lookup per tool the allowlist has always cost.

Source code in datarobot_genai/drmcpbase/fastmcp_transforms/utils.py

def filter_tools_by_allowlist(
    tools: Sequence[Tool],
    allowlist: ToolAllowlist,
) -> list[Tool]:
    """Apply *allowlist* to a whole catalog.

    The fast path is not a second copy of the rule — it is the one case where the rule
    provably collapses. With no derived names and no buckets, :func:`is_tool_allowed`
    can only return ``name in explicit``, so the marker is never consulted and this is
    the same single set lookup per tool the allowlist has always cost.
    """
    if not allowlist.derived and not allowlist.buckets:
        explicit = allowlist.explicit
        return [tool for tool in tools if tool.name in explicit]
    return [tool for tool in tools if is_tool_allowed(tool, allowlist)]

get_tool_category

get_tool_category(tool: Tool) -> str | None

Category name stamped by the providers, or None for untagged (built-in) tools.

Source code in datarobot_genai/drmcpbase/fastmcp_transforms/utils.py

def get_tool_category(tool: Tool) -> str | None:
    """Category name stamped by the providers, or None for untagged (built-in) tools."""
    category = (tool.meta or {}).get("tool_category")
    return category if isinstance(category, str) else None

effective_tool_allowlist async

effective_tool_allowlist(ctx: MCPRequestContext) -> ToolAllowlist | None

Union x-datarobot-mcp-tools and expanded x-datarobot-mcp-toolsets names.

None means no allowlist — every tool is permitted — so it is only ever returned when neither header narrowed the request. An empty frozenset is the opposite answer: the caller asked to be narrowed and nothing survived, so nothing is permitted.

Collapsing those two is the failure this guards against. x-datarobot-mcp-toolsets expands to nothing whenever Mongo is unreachable, the Tool Sets feature is off, this server has no expander, or the named set was deleted — and if an empty expansion fell through to ctx.tool_allowlist (None when the toolsets header was the only filter sent), a session that asked for one bundle would silently receive the entire catalog. Precedence is unchanged: an allowlist is still a hard cap in every mode.

Source code in datarobot_genai/drmcpbase/fastmcp_transforms/utils.py

async def effective_tool_allowlist(ctx: MCPRequestContext) -> ToolAllowlist | None:
    """Union ``x-datarobot-mcp-tools`` and expanded ``x-datarobot-mcp-toolsets`` names.

    ``None`` means *no allowlist* — every tool is permitted — so it is only ever returned
    when neither header narrowed the request. An empty ``frozenset`` is the opposite
    answer: the caller asked to be narrowed and nothing survived, so nothing is permitted.

    Collapsing those two is the failure this guards against. ``x-datarobot-mcp-toolsets``
    expands to nothing whenever Mongo is unreachable, the Tool Sets feature is off, this
    server has no expander, or the named set was deleted — and if an empty expansion fell
    through to ``ctx.tool_allowlist`` (``None`` when the toolsets header was the *only*
    filter sent), a session that asked for one bundle would silently receive the entire
    catalog. Precedence is unchanged: an allowlist is still a hard cap in every mode.
    """
    if not ctx.toolset_names:
        return ctx.tool_allowlist
    expanded = await _expand_toolset_names_cached(ctx)
    if ctx.tool_allowlist is None:
        return ToolAllowlist(explicit=expanded)
    return ctx.tool_allowlist.with_explicit(expanded)

unfiltered_catalog_provider

unfiltered_catalog_provider(mcp: Any) -> Callable[[], Awaitable[Sequence[Tool]]]

Build a catalog provider that ignores the caller's x-datarobot-mcp-* headers.

For the describe-the-server REST routes (/static/tools/, /static/categories/, /metadata), which report what the server registers rather than what the current request may call. list_tools(run_middleware=False) is not enough on its own — FastMCP still applies registered catalog transforms, and :class:DataRobotMCPCatalogTransform is exactly where the session filter is enforced.

Pre-seeding the per-request context cache with a neutral context is what turns the transform into a pass-through, because every branch of it reads the context from there. The token is reset on the way out, so an MCP call sharing the same context (there is none today, but the routes and the protocol do share a process) still resolves its own headers afterwards.

Source code in datarobot_genai/drmcpbase/fastmcp_transforms/utils.py

def unfiltered_catalog_provider(mcp: Any) -> Callable[[], Awaitable[Sequence[Tool]]]:
    """Build a catalog provider that ignores the caller's ``x-datarobot-mcp-*`` headers.

    For the describe-the-server REST routes (``/static/tools/``,
    ``/static/categories/``,
    ``/metadata``), which report what the server registers rather than what the current
    request may call. ``list_tools(run_middleware=False)`` is not enough on its own —
    FastMCP still applies registered catalog transforms, and
    :class:`DataRobotMCPCatalogTransform` is exactly where the session filter is enforced.

    Pre-seeding the per-request context cache with a neutral context is what turns the
    transform into a pass-through, because every branch of it reads the context from there.
    The token is reset on the way out, so an MCP call sharing the same context (there is
    none today, but the routes and the protocol do share a process) still resolves its own
    headers afterwards.
    """

    async def provider() -> Sequence[Tool]:
        token = _request_context_cache.set(_UNFILTERED_REQUEST_CONTEXT)
        try:
            return await mcp.list_tools(run_middleware=False)
        finally:
            _request_context_cache.reset(token)

    return provider

is_category_disabled_for_request

is_category_disabled_for_request(category_name: str) -> bool

Return True when a category gate disables category_name for the current request.

Safe to call from providers, including outside an HTTP request (startup retrospection, in-process clients): any failure to read the request context means "no gates", preserving the default-enabled behavior.

Source code in datarobot_genai/drmcpbase/fastmcp_transforms/utils.py

def is_category_disabled_for_request(category_name: str) -> bool:
    """Return True when a category gate disables *category_name* for the current request.

    Safe to call from providers, including outside an HTTP request (startup
    retrospection, in-process clients): any failure to read the request context
    means "no gates", preserving the default-enabled behavior.
    """
    try:
        return category_name in get_request_context().disabled_categories
    except Exception:  # noqa: BLE001 — gates must never break the provider path
        return False