index.ts•3.02 kB
import OAuthProvider from "@cloudflare/workers-oauth-provider";
import { McpServer } from "@modelcontextprotocol/sdk/server/mcp.js";
import { McpAgent } from "agents/mcp";
import { Octokit } from "octokit";
import { z } from "zod";
import { GitHubHandler } from "./github-handler";
// Context from the auth process, encrypted & stored in the auth token
// and provided to the DurableMCP as this.props
type Props = {
login: string;
name: string;
email: string;
accessToken: string;
};
const PHONE_NUMBER = 919998881729;
const ALLOWED_USERNAMES = new Set<string>([
// Add GitHub usernames of users who should have access to the image generation tool
// For example: 'yourusername', 'coworkerusername'
]);
export class MyMCP extends McpAgent<Env, Record<string, never>, Props> {
server = new McpServer({
name: "Github OAuth Demo",
version: "1.0.0",
});
async init() {
this.server.tool("validate", "Validated this mcp server to be used by PuchAI", {}, async () => ({
content: [{ text: String(PHONE_NUMBER), type: "text" }],
}));
this.server.tool("about", "Description about this mcp server", {}, async () => ({
content: [],
structuredContent: {
name: "Github OAuth Demo",
description: "Github OAUTH DEMO",
},
}));
// Use the upstream access token to facilitate tools
this.server.tool("userInfoOctokit", "Get user info from GitHub, via Octokit", {}, async () => {
const octokit = new Octokit({ auth: this.props.accessToken });
return {
content: [
{
text: JSON.stringify(await octokit.rest.users.getAuthenticated()),
type: "text",
},
],
};
});
// Dynamically add tools based on the user's login. In this case, I want to limit
// access to my Image Generation tool to just me
if (ALLOWED_USERNAMES.has(this.props.login)) {
this.server.tool(
"generateImage",
"Generate an image using the `flux-1-schnell` model. Works best with 8 steps.",
{
prompt: z.string().describe("A text description of the image you want to generate."),
steps: z
.number()
.min(4)
.max(8)
.default(4)
.describe(
"The number of diffusion steps; higher values can improve quality but take longer. Must be between 4 and 8, inclusive.",
),
},
async ({ prompt, steps }) => {
const response = await this.env.AI.run("@cf/black-forest-labs/flux-1-schnell", {
prompt,
steps,
});
return {
content: [{ data: response.image!, mimeType: "image/jpeg", type: "image" }],
};
},
);
}
}
}
export default new OAuthProvider({
apiHandler: MyMCP.mount("/sse") as any,
apiRoute: "/sse",
authorizeEndpoint: "/authorize",
clientRegistrationEndpoint: "/register",
defaultHandler: GitHubHandler as any,
tokenEndpoint: "/token",
});