index.ts•2.21 kB
import OAuthProvider from "@cloudflare/workers-oauth-provider";
import { McpAgent } from "agents/mcp";
import { McpServer } from "@modelcontextprotocol/sdk/server/mcp.js";
import { z } from "zod";
import { AuthkitHandler } from "./authkit-handler";
import type { Props } from "./props";
export class MyMCP extends McpAgent<Env, unknown, Props> {
server = new McpServer({
name: "MCP server demo using AuthKit",
version: "1.0.0",
});
async init() {
// Hello, world!
this.server.tool(
"add",
"Add two numbers the way only MCP can",
{ a: z.number(), b: z.number() },
async ({ a, b }) => ({
content: [{ type: "text", text: String(a + b) }],
})
);
// Dynamically add tools based on the user's permissions. They must have the
// `image_generation` permission to use this tool.
if (this.props.permissions.includes("image_generation")) {
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 }) => {
// TODO: Update the `McpAgent` type to pass its `Env` generic parameter
// down to the `DurableObject` type it extends to avoid this cast.
const env = this.env as Env;
const response = await env.AI.run(
"@cf/black-forest-labs/flux-1-schnell",
{
prompt,
steps,
}
);
return {
content: [
{
type: "image",
data: response.image!,
mimeType: "image/jpeg",
},
],
};
}
);
}
}
}
export default new OAuthProvider({
apiRoute: "/sse",
apiHandler: MyMCP.mount("/sse") as any, // Use 'any' for maximum flexibility
defaultHandler: AuthkitHandler as any, // Use 'any' for maximum flexibility
authorizeEndpoint: "/authorize",
tokenEndpoint: "/token",
clientRegistrationEndpoint: "/register",
});