import { McpServer } from "@modelcontextprotocol/sdk/server/mcp.js";
import { McpAgent } from "agents/mcp";
import { z } from "zod";
import { ALLOWED_USERNAMES } from "./config";
import { Props } from "./types";
export class MyMCP extends McpAgent<Env, Record<string, never>, Props> {
server = new McpServer({
name: "Google OAuth Proxy Demo",
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: [{ text: String(a + b), type: "text" }],
}),
);
// Use the upstream access token to facilitate tools
this.server.tool(
"userInfoGoogle",
"Get user info from Google",
{},
async () => {
const response = await fetch("https://openidconnect.googleapis.com/v1/userinfo", {
headers: {
Authorization: `Bearer ${this.props.accessToken}`,
},
});
const userInfo = await response.json();
return {
content: [
{
text: JSON.stringify(userInfo),
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 (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" }],
};
},
);
}
}
}
