Decorating class methods

[strawgate]

I was thinking a bit more about the class method decoration problem and ran through a little POC on doing delayed registration. The rough idea is that the decorator gains a delay_registration flag (or we add specific decorators for class instance methods, or we add a heuristic for detecting wrapping of a class method) and when this is called, the decorator does not register the tool/resource/etc but instead marks the class method as needing registration.

Finally, the MCP server provides a perform_delayed_registration method which takes a prefix and the object, identifies any wrapped methods, and performs registration.

This allows the continued use of Decorators, which makes the experience between class methods and other setups similar, while enabling the wrapping of class methods and enabling the instantiation of several objects, with each registered under its own prefix.

Here's a sample implementation in FastMCP for playing around with:
#175

The decorator:

        def decorator(fn: AnyFunction) -> AnyFunction:

             if delay_registration: # Mark for Delayed registration
                 fn.__perform_registration = lambda prefix, real_fn: self.add_tool(real_fn, name=f"{prefix}_{name or real_fn.__name__}", description=description, tags=tags)

             else: # Register the tool immediately
                 self.add_tool(fn, name=name, description=description, tags=tags)

             return fn

The method on the Server for registration:

    def perform_delayed_registration(self, prefix: str, object: object) -> None:
        """Perform delayed registration.

        This method is used to register tools/etc that were decorated with @tool(delay_registration=True).
        It allows them to be registered with a prefix that is determined at runtime.

        Args:
            prefix: The prefix to use for the registration
            object: The object containing methods marked with @tool(delay_registration=True)
        """
        
        methods_to_register = [
            getattr(object, method_name)
            for method_name in dir(object)
            if callable(getattr(object, method_name))
            if hasattr(getattr(object, method_name), "_FastMCP__perform_registration")
        ]
        for method_to_register in methods_to_register:
            method_to_register._FastMCP__perform_registration(prefix, method_to_register)

In he PR I provide a server that leverages this, which is what I took the screenshot of.

Here's a minimal viable implementation outside of FastMCP for playing around with. This sample is much more complex looking than the PR but is useful for confirming decorator behavior outside of FastMCP:

import inspect

# Global registry for tools
AVAILABLE_TOOLS = {}

print(f"Initial AVAILABLE_TOOLS: {AVAILABLE_TOOLS}")

def _register_item(key, item):
    """Adds item to AVAILABLE_TOOLS, handling warnings."""
    print(f" -> Registering item with key '{key}'.")
    if key in AVAILABLE_TOOLS:
        print(f"   Warning: Overwriting existing tool '{key}'")
    AVAILABLE_TOOLS[key] = item

def registertool(delay_registration=False):
    """
    Registers the function immediately or marks it for delayed registration.
    """
    def decorator(func):
        tool_name = func.__name__
        if delay_registration:
            # Mark for later registration
            print(f"Decorator @registertool marking '{tool_name}' for delayed registration.")
            func._needs_delayed_registration = True
            return func
        else:
            # Register immediately using the helper
            print(f"Decorator @registertool immediately registering '{tool_name}'.")
            _register_item(tool_name, func)
            return func
    return decorator

# Explicit Delayed Registration Function
def perform_delayed_registration(instance, prefix):
    """
    Finds marked methods on the instance and registers the bound methods
    in AVAILABLE_TOOLS using the provided prefix via the _register_item helper.
    """
    class_name = instance.__class__.__name__
    print(f"\n--- Running perform_delayed_registration for prefix '{prefix}' on {class_name} instance ---")
    found_tools = 0
    items_to_process = list(inspect.getmembers(instance))

    for name, member in items_to_process:
        # Check if it's a method originating from a marked function
        if inspect.ismethod(member) and getattr(member.__func__, '_needs_delayed_registration', False):
            tool_name = name
            registration_key = f"{prefix}_{tool_name}"
            found_tools += 1

            _register_item(registration_key, member) 

    if found_tools == 0:
        print(f" -> No methods marked for delayed registration found on this {class_name} instance.")
    print(f"--- Finished perform_delayed_registration for prefix '{prefix}' ---")


# --- Example Usage ---

@registertool()
def standalone_tool(x):
    """A simple standalone tool function."""
    print(f"Executing standalone_tool with {x}")
    return x * 10

print(f"\nAVAILABLE_TOOLS after defining standalone: {list(AVAILABLE_TOOLS.keys())}")

class ToolClass:
    def __init__(self, name):
        self.name = name
        print(f"ToolClass '{self.name}' instance created.")

    @registertool(delay_registration=True)
    def method_tool_one(self):
        """A tool method within a class."""
        print(f"Executing method_tool_one for instance '{self.name}'")
        return f"one_{self.name}"

    @registertool(delay_registration=True)
    def method_tool_two(self, value):
        """Another tool method."""
        print(f"Executing method_tool_two for instance '{self.name}' with value {value}")
        return f"two_{self.name}_{value}"

    def non_tool_method(self):
        print("This is just a regular method.")

print(f"\nAVAILABLE_TOOLS after class definition: {list(AVAILABLE_TOOLS.keys())}")

instance_one = ToolClass("Alpha")
instance_two = ToolClass("Beta")

print(f"\nAVAILABLE_TOOLS after creating instances: {list(AVAILABLE_TOOLS.keys())}")

# Explicitly Register Instance Tools
perform_delayed_registration(instance_one, prefix="alpha")
print(f"\nAVAILABLE_TOOLS after registering instance_one: {list(AVAILABLE_TOOLS.keys())}")

perform_delayed_registration(instance_two, prefix="beta")
print(f"\nAVAILABLE_TOOLS after registering instance_two: {list(AVAILABLE_TOOLS.keys())}")

print("\n--- Calling registered tools ---")

if 'standalone_tool' not in AVAILABLE_TOOLS:
    raise ValueError("Standalone tool not found in AVAILABLE_TOOLS")

print("Calling 'standalone_tool' via registry:")
AVAILABLE_TOOLS['standalone_tool'](5)

if 'alpha_method_tool_one' not in AVAILABLE_TOOLS:
    raise ValueError("Alpha's tool not found in AVAILABLE_TOOLS")

print("Calling 'alpha_method_tool_one' via registry:")
AVAILABLE_TOOLS['alpha_method_tool_one']()

if 'beta_method_tool_two' not in AVAILABLE_TOOLS:
    raise ValueError("Beta's tool not found in AVAILABLE_TOOLS")

print("Calling 'beta_method_tool_two' via registry:")
AVAILABLE_TOOLS['beta_method_tool_two'](42)

print(f"\nFinal AVAILABLE_TOOLS: {list(AVAILABLE_TOOLS.keys())}")

[strawgate]

second iteration #177