- 
                Notifications
    
You must be signed in to change notification settings  - Fork 1.5k
 
feat: add schema dereference middleware #1641
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
base: main
Are you sure you want to change the base?
Conversation
c7188ef    to
    b5ec2be      
    Compare
  
    ea7ea72    to
    5858dbe      
    Compare
  
    | 
           Hi @jlowin could you review this PR  | 
    
5858dbe    to
    75978f6      
    Compare
  
    | 
           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  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  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
 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  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.,  
 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. 
 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  So a simple solution, for now, only needs to deref in  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.  | 
    
Description
Contributors Checklist
Review Checklist
Follow up with #1192