feat: support extending McpServer with authorization

[wangshijun]

Let's make some of the private members of McpServer accessible to its subclasses. This way, the community can easily build additional layers on top of the official version without needing to start from scratch.

Motivation and Context

We wanted to share some thoughts about the built-in McpServer. It's got some really handy methods like tool, prompt, and resource that make setting up an MCP server a breeze. However, it doesn't currently allow for any extensions, which makes it a bit tricky for us to implement server-side authorization for each tool call based on the current user session. We also feel that creating a whole new MCP server framework isn't the best route, especially since the official one works well in most scenarios.

With that in mind, we're suggesting a couple of changes:

• Allow some private members of McpServer to be accessible to its subclasses.
• Introduce a user property to the Transport and ensure it's accessible in the extra parameter during tool calls.

We did a bit of digging and found a potentially related issue here: #171

We'd love to hear your thoughts on this!

How Has This Been Tested?

Yes, we have tested this in: https://github.com/blocklet/mcp-server-demo

Breaking Changes

No, there are no any breaking changes to existing features.

Types of changes

• Bug fix (non-breaking change which fixes an issue)
• New feature (non-breaking change which adds functionality)
• Breaking change (fix or feature that would cause existing functionality to change)
• Documentation update

Checklist

• I have read the MCP Documentation
• My code follows the repository's style guidelines
• New and existing tests pass locally
• I have added appropriate error handling
• I have added or updated documentation as needed

Additional context

Example McpServer extension (with flexible access control policy):

import {
  CallToolRequestSchema,
  CallToolRequest,
  ListToolsRequestSchema,
  ListToolsResult,
  Tool,
} from '@modelcontextprotocol/sdk/types.js';
import { RequestHandlerExtra } from '@modelcontextprotocol/sdk/shared/protocol.js';
import { z, ZodRawShape } from 'zod';
import { McpServer, ToolCallback } from '@modelcontextprotocol/sdk/server/mcp.js';
import { zodToJsonSchema } from 'zod-to-json-schema';

export type AuthMethod = 'loginToken' | 'componentCall' | 'signedToken';

export interface SessionUser {
  did: string;
  role: string;
  provider: string;
  method?: AuthMethod;
  [key: string]: unknown;
}

export type AccessPolicy = {
  allow?: {
    dids?: string[];
    roles?: string[];
    providers?: string[];
    methods?: AuthMethod[];
  };
  deny?: {
    dids?: string[];
    roles?: string[];
    providers?: string[];
    methods?: AuthMethod[];
  };
};

interface RegisteredToolWithAuth {
  description?: string;
  inputSchema?: z.ZodObject<ZodRawShape>;
  callback: ToolCallback<ZodRawShape | undefined>;
  accessPolicy?: AccessPolicy;
}

export class McpServerWithAuth extends McpServer {
  protected override _registeredTools: { [name: string]: RegisteredToolWithAuth } = {};
  checkPermissions(user?: SessionUser, policy?: AccessPolicy): boolean {
    if (!policy) {
      return true;
    }

    if (!user) {
      return false;
    }

    // Check deny rules first
    if (policy.deny) {
      // Check denied DIDs
      if (policy.deny.dids?.includes(user.did)) {
        return false;
      }

      // Check denied roles
      if (policy.deny.roles?.includes(user.role)) {
        return false;
      }

      // Check denied providers
      if (policy.deny.providers?.includes(user.provider)) {
        return false;
      }

      // Check denied auth methods
      if (user.method && policy.deny.methods?.includes(user.method)) {
        return false;
      }
    }

    // Check allow rules
    if (policy.allow) {
      let isAllowed = false;

      // If no allow rules are specified, default to allowed
      if (!policy.allow.dids && !policy.allow.roles && !policy.allow.providers && !policy.allow.methods) {
        isAllowed = true;
      } else {
        // Check allowed DIDs
        if (policy.allow.dids?.includes(user.did)) {
          isAllowed = true;
        }

        // Check allowed roles
        if (policy.allow.roles?.includes(user.role)) {
          isAllowed = true;
        }

        // Check allowed providers
        if (policy.allow.providers?.includes(user.provider)) {
          isAllowed = true;
        }

        // Check allowed auth methods
        if (user.method && policy.allow.methods?.includes(user.method)) {
          isAllowed = true;
        }
      }

      return isAllowed;
    }

    // If no rules specified, default to allowed
    return true;
  }

  override tool(name: string, cb: ToolCallback, accessPolicy?: AccessPolicy): void;
  override tool(name: string, description: string, cb: ToolCallback, accessPolicy?: AccessPolicy): void;
  override tool<Args extends ZodRawShape>(
    name: string,
    paramsSchema: Args,
    cb: ToolCallback<Args>,
    accessPolicy?: AccessPolicy
  ): void;
  override tool<Args extends ZodRawShape>(
    name: string,
    description: string,
    paramsSchema: Args,
    cb: ToolCallback<Args>,
    accessPolicy?: AccessPolicy
  ): void;
  override tool(name: string, ...rest: unknown[]): void {
    let description: string | undefined;
    let paramsSchema: ZodRawShape | undefined;
    let accessPolicy: AccessPolicy | undefined;
    let cb: ToolCallback<ZodRawShape | undefined>;

    // Parse arguments based on their types
    if (typeof rest[0] === 'function') {
      // Case: tool(name, cb, accessPolicy?)
      cb = rest[0] as ToolCallback<ZodRawShape | undefined>;
      accessPolicy = rest[1] as AccessPolicy | undefined;
    } else if (typeof rest[0] === 'string') {
      // Cases with description
      description = rest[0];
      if (typeof rest[1] === 'function') {
        // Case: tool(name, description, cb, accessPolicy?)
        cb = rest[1] as ToolCallback<ZodRawShape | undefined>;
        accessPolicy = rest[2] as AccessPolicy | undefined;
      } else {
        // Case: tool(name, description, paramsSchema, cb, accessPolicy?)
        paramsSchema = rest[1] as ZodRawShape;
        cb = rest[2] as ToolCallback<ZodRawShape>;
        accessPolicy = rest[3] as AccessPolicy | undefined;
      }
    } else {
      // Case: tool(name, paramsSchema, cb, accessPolicy?)
      paramsSchema = rest[0] as ZodRawShape;
      cb = rest[1] as ToolCallback<ZodRawShape>;
      accessPolicy = rest[2] as AccessPolicy | undefined;
    }

    // Register with base class
    const args: unknown[] = [name];
    if (description) args.push(description);
    if (paramsSchema) args.push(paramsSchema);
    args.push(cb);

    // Set up request handlers if not already initialized
    if (!this._toolHandlersInitialized) {
      this.server.assertCanSetRequestHandler(CallToolRequestSchema.shape.method.value);
      this.server.assertCanSetRequestHandler(ListToolsRequestSchema.shape.method.value);
      this.server.registerCapabilities({ tools: {} });

      // Add ListToolsRequestSchema handler
      this.server.setRequestHandler(ListToolsRequestSchema, (request, extra): ListToolsResult => {
        const user = extra.user as SessionUser | undefined;

        // Filter tools based on permissions
        const accessibleTools = Object.entries(this._registeredTools)
          .filter(([_, tool]) => this.checkPermissions(user, tool.accessPolicy))
          .map(
            ([name, tool]): Tool => ({
              name,
              description: tool.description,
              inputSchema: tool.inputSchema
                ? (zodToJsonSchema(tool.inputSchema, {
                    strictUnions: true,
                  }) as Tool['inputSchema'])
                : { type: 'object' },
            })
          );

        return { tools: accessibleTools };
      });

      this.server.setRequestHandler(
        CallToolRequestSchema,
        async (request: CallToolRequest, extra: RequestHandlerExtra) => {
          const tool = this._registeredTools[request.params.name];
          if (!tool) {
            throw new Error(`Tool ${request.params.name} not found`);
          }

          if (!this.checkPermissions(extra.user as SessionUser, tool.accessPolicy)) {
            throw new Error(`Access denied for tool: ${request.params.name}`);
          }

          if (tool.inputSchema) {
            const parseResult = await tool.inputSchema.safeParseAsync(request.params.arguments);
            if (!parseResult.success) {
              throw new Error(`Invalid arguments for tool ${request.params.name}: ${parseResult.error.message}`);
            }

            const args = parseResult.data;
            const cb = tool.callback as ToolCallback<ZodRawShape>;
            return await Promise.resolve(cb(args, extra));
          } else {
            const cb = tool.callback as ToolCallback<undefined>;
            return await Promise.resolve(cb(extra));
          }
        }
      );
      this._toolHandlersInitialized = true;
    }

    McpServer.prototype.tool.apply(this, args as Parameters<typeof McpServer.prototype.tool>);
    this._registeredTools[name].accessPolicy = accessPolicy;
  }
}

And usage example for above usage:

import { ResourceTemplate } from '@modelcontextprotocol/sdk/server/mcp.js';
import { z } from 'zod';

import { tellStory } from './agent';

const mcpServer = new McpServerWithAuth(
  {
    name: 'Example MCP Server on ArcBlock Platform',
    version: '1.0.0',
  },
  {
    capabilities: {
      tools: {},
      resources: {},
    },
  },
);

mcpServer.tool(
  'storytelling',
  'Tell a story about a given topic',
  {
    topic: z.string().describe('The topic to tell a story about'),
    language: z.string().describe('The language to tell the story in').default('zh-CN'),
  },
  async ({ topic, language }) => {
    const result = await tellStory(topic, language);
    return {
      content: [{ type: 'text', text: result }],
    };
  },
  {
    allow: {
      roles: ['admin', 'owner'],
    },
  },
);

[Copilot]

Copilot reviewed 4 out of 4 changed files in this pull request and generated no comments.

Comments suppressed due to low confidence (1)

src/server/mcp.ts:770

• [nitpick] If promptArgumentsFromSchema is intended only as an internal helper, consider not exporting it or renaming it (e.g. prefixing with an underscore) to reduce confusion regarding its public API status.

export function promptArgumentsFromSchema(

[cliffhall]

Hi @wangshijun. Looks good. Can you add some unit tests?