RFC: Expose initialization requests to middleware, and provide session storage

[richardkmichael]

This is a proof-of-concept I hacked up with Claude. I'm not proposing it for merge as-is.

I haven't thought carefully about the consequences of these modifications. This might be touching on stateful/less aspects of MCP; not sure and perhaps not any more so than progress tokens or other remembered data on either side of the connection. If there is interest, I'll take a closer look.

I was intrigued by a middleware based solution for #1193 and #1159, but this requires modifying FastMCP.

Middleware documentation mentions all requests are exposed to middleware, but this is not true. The initialization part of the lifecycle is handled in the session establishment and not available to middleware.

This hack adds:

• A new session class which exposes itself to middleware for on_initialize
• Adds the session to the MiddlewareContext so there is cross-request storage for data only available in specific requests

It allows a middleware like this:

class DereferenceForClaudeMiddleware(Middleware):
    async def on_initialize(self, context: MiddlewareContext, call_next):
        """Detect Claude clients and store on session object."""
        client_info = getattr(context.message, "clientInfo", None)
        client_name = (
            getattr(client_info, "name", "unknown") if client_info else "unknown"
        )

        is_claude_client = self._is_claude_client(client_name)

        # HACK: Store client type directly on session object
        if hasattr(context, "session") and context.session:
            context.session.is_claude_client = is_claude_client

        if is_claude_client:
            logger.info(f"Claude ({client_name}) detected - will dereference schemas")
        else:
            logger.info(
                f"Non-Claude ({client_name}) detected - schemas will remain unchanged"
            )

        return await call_next(context)

    async def on_list_tools(self, context: MiddlewareContext, call_next):
        """Dereference tool schemas for Claude clients."""
        tools = await call_next(context)

        # HACK: Read client type from session object (stored during initialization)
        is_claude_client = False
        if context.fastmcp_context and context.fastmcp_context.session:
            is_claude_client = getattr(
                context.fastmcp_context.session, "is_claude_client", False
            )

        if is_claude_client:
            for tool in tools:
                original_schema = deepcopy(tool.parameters)
                dereferenced_schema = self.dereference_json_schema(tool.parameters)

                if original_schema != dereferenced_schema:
                    tool.parameters = dereferenced_schema

        return tools

I think exposing initialization to middleware would also be useful for:

• logging
• analytics (i.e. server usage by specific MCP clients)
• future proofing for other [misbehaved] clients

[jlowin]

@richardkmichael thanks! I went down a rabbit hole of supporting every JSONRPC call at first but it required way too much rewriting of the low-level internals to work. I really like that you cut through that by overwriting just the one method and ensuring it's in place in the overwritten .run() method. This is pretty compelling! I want to think on it a little just to make sure we aren't inadvertently taking on a broader maintenance burden than intended (though I strongly suspect your solution limits that as much as possible)

[richardkmichael]

Thank you for the quick feedback!

I've just noticed #1160 -- but have not read it closely. Perhaps its state dictionary can be used for storage. Using the session was obviously a hack, but I needed that because there was no context at during initialization, IIRC. I certainly wouldn't want data stored in two places, so it would be worth looking at streamlining that aspect.

[jlowin]

My suggestion is lets look at #1160 for state and keep this PR focused on initializae middleware; I think 1160 may have stalled so I'll look at getting it over the line today.

(update: merged 1160)

[hopeful0]

Rewriting the run method of low-level server and extending ServerSession in Fastmcp is very useful. This gives us the opportunity to create a service session-level context to store some state that needs to be shared throughout the session, such as the session ID in #1242 and some stuff about proxy.

[jlowin]

@hopeful0 I agree. I have hesitated on rewriting low-level SDK at all because of the increased maintenance burden it places on FastMCP but I think we're at a place where the surface area is small enough and the benefit is large enough that this is worthwhile.

[rishid]

I have been toying with this myself and this is a good solution. I want to be able to get client info and collect metrics on this information. Even though documentation says on_message - Called for ALL MCP messages (both requests and notifications) it took me a while to figure out that initialize is not one of those message.

[dacamposol]

I love this proposal!

As @rishid mentioned, this would allow us to collect metrics on some information across the entire MCP session; plus there are some cases where the MCP serves as a shim to other upstream services, where being able to communicate directly on server communication whether it's available to be used or not to the MCP hosts would be helpful.

[jlowin]

I've added the middleware/initialize aspects of this PR to #1546 and will close this as I think it's become stale. Appreciate the head start though @richardkmichael!

[richardkmichael]

Thank you! I was going ask what to do here, but I've been busy. Thanks for running with it! 👍