Skip to content

Conversation

@vl-kp
Copy link
Contributor

@vl-kp vl-kp commented Aug 26, 2025

Description

Contributors Checklist

Review Checklist

  • I have self-reviewed my changes
  • My Pull Request is ready for review

Follow up with #1192

@vl-kp vl-kp force-pushed the chore/add-schema-dereference-middleware branch from c7188ef to b5ec2be Compare August 26, 2025 14:09
@vl-kp vl-kp marked this pull request as ready for review August 26, 2025 14:14
@vl-kp vl-kp force-pushed the chore/add-schema-dereference-middleware branch from ea7ea72 to 5858dbe Compare September 8, 2025 01:04
@vl-kp
Copy link
Contributor Author

vl-kp commented Oct 1, 2025

Hi @jlowin could you review this PR

@vl-kp vl-kp force-pushed the chore/add-schema-dereference-middleware branch from 5858dbe to 75978f6 Compare October 1, 2025 06:46
@richardkmichael
Copy link
Contributor

Since I have a similar use-case and have my own hack locally, I tested this PR. I have comments on the PR below.

However, since I wanted slightly different behaviour, I also tried re-implementing it with jsonref (mentioned on #1192), and it's so simple that I'd be hesitant to add a dedicated middleware in repo. Maybe in contrib, but as-is, I'm not sure it's flexible enough.

First, here it is with jsonref -- it removes tools which fail dereferencing (to avoid LLM misuse) but logging a warning (i.e., fix the server):

from jsonref import JsonRefError, replace_refs

class DereferenceSchemaMiddleware(Middleware):
    def _dereference_tool(self, tool, remove_defs=True):
        """Dereference a tool's schema, returning the modified tool or None on failure.

        Args:
            tool: Tool object with parameters schema to dereference
            remove_defs: If True, remove $defs from schema after dereferencing (default: True)
        """
        try:
            # proxies=False returns plain dicts to ensure compatibility with FastMCP
            # lazy_load=False resolves immediately
            dereferenced = replace_refs(tool.parameters, proxies=False, lazy_load=False)

            # If requested, remove $defs (since all references have been resolved)
            if remove_defs and isinstance(dereferenced, dict) and "$defs" in dereferenced:
                dereferenced = {k: v for k, v in dereferenced.items() if k != "$defs"}

            tool.parameters = dereferenced
            return tool

        except JsonRefError as e:
            logger.warning(
                f"Failed to dereference schema for tool '{tool.name}': {e}. "
                "Tool will be hidden from client."
            )
            return None
        except Exception as e:
            logger.error(
                f"Unexpected error dereferencing schema for tool '{tool.name}': {e}. "
                "Tool will be hidden from client."
            )
            return None

    async def on_list_tools(self, context: MiddlewareContext, call_next):
        tools = await call_next(context)

        return [
            dereferenced_tool
            for tool in tools
            if (dereferenced_tool := self._dereference_tool(tool)) is not None
        ]

Even more minimal if tolerating a lingering $defs in the schema, and returning the unmodified tool if dereferencing fails (similar to the PR):

from jsonref import JsonRefError, replace_refs

class DereferenceSchemaMiddleware(Middleware):
    def _dereference_tool(self, tool):
        """Dereference a tool's schema, returning the unmodified tool if dereferencing fails.

        Args:
            tool: Tool object with parameters schema to dereference
        """
        try:
            # proxies=False returns plain dicts to ensure compatibility with FastMCP
            # lazy_load=False resolves immediately
            tool.parameters = replace_refs(tool.parameters, proxies=False, lazy_load=False)

        except JsonRefError:
            pass

        return tool


    async def on_list_tools(self, context: MiddlewareContext, call_next):
        tools = await call_next(context)

        return [ self._dereference_tool(tool) for tool in tools ]

Comments on the PR

  • I want deferencing conditional upon the client (i.e., broken ones).

The middleware is not conditional, so I tried both subclass and separate class importing the deref function; neither seem great options.

Since the deref function is module-level, a subclass can't self._deref(...), and needs to conditionally super().on_list_tools(...) -- this means for the else-branch (not deref'ing), it needs to know if the parent-class does anything in addition to deref and implement that part only.

If not subclassing, the deref function can be imported, but this essentially side-stepping this middleware entirely, and just using its "internals" as a utility function.

I think these could be helped by defining the deref function in the class, improving encapsulation and permitting a subclass to call it whenever needed.

Both mean re-implementing the hooks, adding the condition.

Conditionality could be provided by adding a callable to the constructor for simple cases, and also a template function (e.g., def self._should_dereference(): return True) to override for complex conditions (I needed this).

  • The PR seems to implement "best effort" deref, I want all-or-nothing with a warning/error.

I'm curious about the thinking here. The PR leaves alone refs in the corner-cases. I think this means tools failing to deref all arguments will still be potentially misused by the LLM (if it uses those args) -- making the problem less frequent, but still present. Also, there is no logging when this occurs.

  • The deref is in specific hooks, not at the schema-generation level.

This is debatable, i.e., a general vs. pragmatic and sufficient solution.

How were the specific hooks needing deref determined?

I don't think the deref in on_list_resource_templates() is actually needed because no JSON Schema is returned to the client for this call. I reviewed the MCP Schema, and I think only list_tools can contain $ref (input_schema is a JSON Schema fragment), but it's not straight-forward to determine this from the schema.

So a simple solution, for now, only needs to deref in on_list_tools(). Easy enough to add more later.

OTOH, a general approach would alter FastMCPs schema generation, and all future responses would be deref'd. (Tempting, to avoid MCP Schema interpretation.)

So my thought is that these knobs and design-decisions seem unnecessary, given the simplicity of the jsonref implementation. 🤷

I have other comments about the code itself, which I can add later if going forward.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants