MCP Server Integration

import express from "express";
import { createRemoteJWKSet, jwtVerify } from "jose";
import { z } from "zod";
import { McpServer } from "@modelcontextprotocol/sdk/server/mcp.js";
import { StreamableHTTPServerTransport } from "@modelcontextprotocol/sdk/server/streamableHttp.js";
import { requireBearerAuth } from "@modelcontextprotocol/sdk/server/auth/middleware/bearerAuth.js";
import { InvalidTokenError } from "@modelcontextprotocol/sdk/server/auth/errors.js";
import type { OAuthTokenVerifier } from "@modelcontextprotocol/sdk/server/auth/provider.js";

const ISSUER = "https://auth.example.com";
const RESOURCE = "https://mcp.example.com"; // the application's url
const JWKS = createRemoteJWKSet(new URL(`${ISSUER}/api/auth/jwks`));

interface Claims {
  sub: string;
  roles: string[];
  features: Record<string, unknown>;
  org_id?: string;
}

const verifier: OAuthTokenVerifier = {
  async verifyAccessToken(token) {
    try {
      const { payload } = await jwtVerify(token, JWKS, {
        algorithms: ["EdDSA"],
        issuer: ISSUER,
        audience: RESOURCE, // matches when the `aud` array contains it
      });
      const claims: Claims = {
        sub: payload.sub ?? "",
        roles: Array.isArray(payload.roles) ? (payload.roles as string[]) : [],
        features:
          typeof payload.features === "object" && payload.features !== null
            ? (payload.features as Record<string, unknown>)
            : {},
        ...(typeof payload.org_id === "string" ? { org_id: payload.org_id } : {}),
      };
      return {
        token,
        clientId: String(payload.azp),
        scopes: String(payload.scope ?? "").split(" ").filter(Boolean),
        expiresAt: payload.exp,
        resource: new URL(RESOURCE),
        extra: { ...claims },
      };
    } catch (error) {
      throw new InvalidTokenError(error instanceof Error ? error.message : "Invalid token");
    }
  },
};

function createServer(): McpServer {
  const server = new McpServer({ name: "my-mcp-server", version: "1.0.0" });

  server.registerTool(
    "documents_list",
    {
      description: "List documents for the authenticated user",
      inputSchema: { page: z.number().int().min(1).default(1) },
    },
    async ({ page }, extra) => {
      const claims = extra.authInfo?.extra as unknown as Claims;

      // Role check
      if (!claims.roles.includes("editor")) {
        return { isError: true, content: [{ type: "text", text: "Requires the 'editor' role" }] };
      }
      // Feature flag from the user's subscription plan
      if (claims.features.apiAccess !== true) {
        return { isError: true, content: [{ type: "text", text: "API access is not part of your plan" }] };
      }

      // Scope every query by claims.sub and, when present, claims.org_id.
      const documents = [{ id: 1, page, owner: claims.sub, orgId: claims.org_id }];
      return { content: [{ type: "text", text: JSON.stringify(documents) }] };
    },
  );

  return server;
}

const app = express();
app.use(express.json());

app.post("/mcp", requireBearerAuth({ verifier }), async (req, res) => {
  const server = createServer();
  const transport = new StreamableHTTPServerTransport({ sessionIdGenerator: undefined });
  res.on("close", () => {
    void transport.close();
    void server.close();
  });
  await server.connect(transport);
  await transport.handleRequest(req, res, req.body);
});

app.listen(3000);

import asyncio
import json

import jwt  # PyJWT: pip install "pyjwt[crypto]"
from mcp.server.auth.middleware.auth_context import get_access_token
from mcp.server.auth.provider import AccessToken
from mcp.server.auth.settings import AuthSettings
from mcp.server.fastmcp import FastMCP

ISSUER = "https://auth.example.com"
RESOURCE = "https://mcp.example.com"  # the application's url

# PyJWKClient caches the key set and refetches it when it meets an unknown kid
jwks_client = jwt.PyJWKClient(f"{ISSUER}/api/auth/jwks")


def decode(token: str) -> dict:
    key = jwks_client.get_signing_key_from_jwt(token)
    return jwt.decode(
        token,
        key.key,
        algorithms=["EdDSA"],
        issuer=ISSUER,
        audience=RESOURCE,  # matches when the `aud` array contains it
    )


class AuthServiceVerifier:
    async def verify_token(self, token: str) -> AccessToken | None:
        try:
            # PyJWKClient does blocking I/O: keep it off the event loop
            claims = await asyncio.to_thread(decode, token)
        except jwt.PyJWTError:
            return None
        return AccessToken(
            token=token,
            client_id=claims["azp"],
            scopes=claims.get("scope", "").split(),
            expires_at=claims["exp"],
            resource=RESOURCE,
            subject=claims["sub"],
            claims=claims,
        )


mcp = FastMCP(
    "my-mcp-server",
    token_verifier=AuthServiceVerifier(),
    auth=AuthSettings(
        issuer_url=ISSUER,
        resource_server_url=RESOURCE,
        validate_token_resource=True,
    ),
)


@mcp.tool()
async def list_documents(page: int = 1) -> str:
    """List documents for the authenticated user."""
    access = get_access_token()
    assert access is not None and access.claims is not None
    claims = access.claims

    if "editor" not in claims.get("roles", []):
        raise PermissionError("Requires the 'editor' role")
    if claims.get("features", {}).get("apiAccess") is not True:
        raise PermissionError("API access is not part of your plan")

    # Scope every query by claims["sub"] and, when present, claims.get("org_id")
    documents = [{"page": page, "owner": claims["sub"], "org_id": claims.get("org_id")}]
    return json.dumps(documents)


if __name__ == "__main__":
    mcp.run(transport="streamable-http")