/**
* My MCP Server
*
* A template MCP server with Google OAuth authentication.
*
* ===== CUSTOMIZE: Update the description above =====
*/
import OAuthProvider from '@cloudflare/workers-oauth-provider';
import { McpServer } from '@modelcontextprotocol/sdk/server/mcp.js';
import { McpAgent } from 'agents/mcp';
import { z, ZodTypeAny } from 'zod';
import { zodToJsonSchema } from 'zod-to-json-schema';
import { getIdentityHandler, getProviderName, supportsTokenRefresh } from './oauth/handler-factory';
import { refreshAccessToken as refreshGoogleToken } from './auth/identity/google';
import { refreshMicrosoftToken } from './auth/identity/microsoft';
import { adminApp } from './admin/routes';
import { timingSafeEqual } from './lib/crypto';
import { createAuth } from './lib/auth';
import { getLoginPage } from './pages/login';
import { getConsentPage } from './pages/consent';
import { createTokenManager, TokenNotFoundError, TokenExpiredError } from './lib/token-manager';
// Import modular tools, resources, prompts
import { registerAllTools, getToolsMetadata as getToolsMeta, type ToolContext } from './tools';
import { registerAllResources, type ResourceContext } from './resources';
import { registerAllPrompts } from './prompts';
import type { Env, Props, ToolMetadata, ToolCategory } from './types';
// Re-export types for wrangler
export type { Env };
/**
* My MCP Server
*
* ===== CUSTOMIZE: Update the class name and description =====
*/
export class MyMCP extends McpAgent<Env, Record<string, never>, Props> {
server = new McpServer({
name: 'my-mcp-server', // ===== CUSTOMIZE: Server name =====
version: '1.0.0',
});
// Tool registry for introspection and defer_loading support
private toolRegistry: Map<string, ToolMetadata> = new Map();
// Tool handlers for direct execution (used by admin chat)
// eslint-disable-next-line @typescript-eslint/no-explicit-any
private toolHandlers: Map<string, (args: Record<string, unknown>) => Promise<any>> = new Map();
/**
* Register a tool with metadata for introspection.
* This wraps server.tool() and stores metadata for the admin dashboard and AI chat.
*/
private registerTool<T extends Record<string, ZodTypeAny>>(
name: string,
description: string,
schema: T,
// eslint-disable-next-line @typescript-eslint/no-explicit-any
handler: (args: z.infer<z.ZodObject<T>>, extra?: any) => any,
options?: {
defer_loading?: boolean;
category?: ToolCategory;
tags?: string[];
requiresAuth?: string;
authScopes?: string[];
}
): void {
// Register with MCP server
// eslint-disable-next-line @typescript-eslint/no-explicit-any
(this.server.tool as any)(name, description, schema, handler);
// Store handler for direct execution (admin chat)
this.toolHandlers.set(name, handler as (args: Record<string, unknown>) => Promise<unknown>);
// Store metadata for introspection
const zodSchema = z.object(schema);
const jsonSchema = zodToJsonSchema(zodSchema, { target: 'openApi3' });
this.toolRegistry.set(name, {
name,
description,
inputSchema: jsonSchema as Record<string, unknown>,
defer_loading: options?.defer_loading,
category: options?.category,
tags: options?.tags,
requiresAuth: options?.requiresAuth,
authScopes: options?.authScopes,
});
}
/**
* Get metadata for all registered tools.
*/
getToolsMetadata(): ToolMetadata[] {
return Array.from(this.toolRegistry.values());
}
/**
* Execute a tool directly (for admin chat - bypasses MCP protocol).
*/
async executeTool(
toolName: string,
args: Record<string, unknown>
): Promise<{ success: boolean; result?: unknown; error?: string }> {
// Ensure tools are registered
if (this.toolHandlers.size === 0) {
await this.init();
}
const handler = this.toolHandlers.get(toolName);
if (!handler) {
return { success: false, error: `Tool '${toolName}' not found. Available: ${Array.from(this.toolHandlers.keys()).join(', ')}` };
}
try {
const result = await handler(args);
const content = result?.content;
if (Array.isArray(content)) {
const textContent = content
.filter((c: { type?: string; text?: string }) => c.type === 'text')
.map((c: { text?: string }) => c.text)
.join('\n');
// When tool returns isError, put the message in 'error' field for UI
if (result.isError) {
return { success: false, error: textContent };
}
return { success: true, result: textContent || result };
}
return { success: !result.isError, result };
} catch (error) {
return {
success: false,
error: error instanceof Error ? error.message : 'Tool execution failed',
};
}
}
/**
* Get tool context for modular tools that need agent access.
* Always returns context with env (for Layer 2 shared services).
* Props may be undefined for unauthenticated admin chat contexts.
*/
private getToolContext(): ToolContext {
return {
props: this.props,
authorizedFetch: this.props ? this.authorizedFetch.bind(this) : undefined,
env: this.env,
};
}
/**
* Get resource context for modular resources
*/
private getResourceContext(): ResourceContext | undefined {
return {
props: this.props,
serverName: 'my-mcp-server',
serverVersion: '1.0.0',
toolsCount: this.toolRegistry.size,
};
}
async init() {
// ===== REGISTER TOOLS FROM MODULES =====
// Tools are defined in src/tools/ - edit those files to customize
registerAllTools(
this.registerTool.bind(this),
this.getToolContext.bind(this)
);
// ===== REGISTER RESOURCES FROM MODULES =====
// Resources are defined in src/resources/ - edit those files to customize
registerAllResources(
// eslint-disable-next-line @typescript-eslint/no-explicit-any
this.server.resource.bind(this.server) as any,
this.getResourceContext.bind(this)
);
// ===== REGISTER PROMPTS FROM MODULES =====
// Prompts are defined in src/prompts/ - edit those files to customize
registerAllPrompts(
// eslint-disable-next-line @typescript-eslint/no-explicit-any
this.server.prompt.bind(this.server) as any
);
// ===== ADD CUSTOM TOOLS HERE =====
// You can still add tools directly if needed:
// this.registerTool(name, description, schema, handler, options);
// ===== INTERNAL AGENT (Optional) =====
// Enable with ENABLE_INTERNAL_AGENT=true in wrangler.jsonc
if (this.env.ENABLE_INTERNAL_AGENT === 'true') {
await this.registerInternalAgentTool();
}
}
/**
* Register the ask_agent tool (internal agent pattern)
*/
private async registerInternalAgentTool(): Promise<void> {
const { runAgent } = await import('./lib/agent');
const { getOrCreateConversation, getConversation, getMessages, addMessage, toChatMessages } = await import('./lib/memory');
this.registerTool(
'ask_agent',
'Ask the internal AI agent to help with a task. The agent has access to various tools and will use them as needed.',
{
query: z.string().describe('The task or question for the agent'),
conversation_id: z.string().optional().describe('Conversation ID for context continuity'),
},
async ({ query, conversation_id }) => {
try {
const userEmail = this.props?.email;
let history: import('./types').ChatMessage[] = [];
// Load existing conversation with ownership verification
if (conversation_id && this.env.DB && this.env.ENABLE_CONVERSATION_MEMORY === 'true') {
const existingConv = await getConversation(this.env.DB, conversation_id);
if (!existingConv) {
return {
content: [{ type: 'text', text: 'Conversation not found' }],
isError: true,
};
}
// Verify ownership - users can only access their own conversations
if (existingConv.metadata?.userEmail && existingConv.metadata.userEmail !== userEmail) {
return {
content: [{ type: 'text', text: 'Access denied to this conversation' }],
isError: true,
};
}
const messages = await getMessages(this.env.DB, conversation_id);
history = toChatMessages(messages);
}
const result = await runAgent(
{
env: this.env,
tools: this.getToolsMetadata().filter(t => t.name !== 'ask_agent'),
executeTool: this.executeTool.bind(this),
},
query,
history
);
let conversationId = conversation_id;
if (this.env.DB && this.env.ENABLE_CONVERSATION_MEMORY === 'true') {
// Store userEmail in metadata for ownership tracking
const conv = await getOrCreateConversation(this.env.DB, conversation_id, {
source: 'mcp',
userEmail: userEmail,
});
conversationId = conv.id;
await addMessage(this.env.DB, conv.id, {
conversationId: conv.id,
role: 'user',
content: query,
});
await addMessage(this.env.DB, conv.id, {
conversationId: conv.id,
role: 'assistant',
content: result.response,
});
}
const responseObj = {
conversation_id: conversationId,
response: result.response,
tools_used: result.toolsUsed,
};
return {
content: [{ type: 'text', text: JSON.stringify(responseObj, null, 2) }],
};
} catch (error) {
const message = error instanceof Error ? error.message : 'Agent error';
return {
content: [{ type: 'text', text: `Error: ${message}` }],
isError: true,
};
}
},
{
category: 'integration',
tags: ['agent', 'ai', 'assistant'],
}
);
}
// ===== HELPER METHODS =====
/**
* Ensure we have a valid access token, refreshing if expired.
* Uses TokenManager for encrypted storage and automatic refresh.
* Falls back to props storage for backward compatibility.
*/
private async ensureValidToken(): Promise<string> {
if (!this.props?.id) {
throw new Error('Not authenticated. Please reconnect your account.');
}
const provider = (this.props.identityProvider || getProviderName(this.env)) as 'google' | 'microsoft' | 'github';
// Try TokenManager first (encrypted storage with auto-refresh)
try {
const tokenManager = createTokenManager(this.env);
const tokenData = await tokenManager.get({
userId: this.props.id,
provider,
});
// Sync token back to props for DO storage (backward compatibility)
if (tokenData.accessToken !== this.props.accessToken) {
this.props = {
...this.props,
accessToken: tokenData.accessToken,
refreshToken: tokenData.refreshToken ?? this.props.refreshToken,
tokenExpiresAt: tokenData.expiresAt,
};
await this.ctx.storage.put('props', this.props);
console.log('TokenManager: Synced refreshed token to props');
}
return tokenData.accessToken;
} catch (error) {
// Handle TokenManager-specific errors
if (error instanceof TokenNotFoundError) {
console.log('TokenManager: Token not found, falling back to props');
// Fall through to legacy props-based flow
} else if (error instanceof TokenExpiredError) {
console.error('TokenManager: Token expired and refresh failed');
throw new Error('Session expired. Please reconnect your account.');
} else {
console.error('TokenManager: Unexpected error:', error);
// Fall through to legacy props-based flow
}
}
// ===== LEGACY FALLBACK: Use props-based token refresh =====
// This handles tokens stored before TokenManager was integrated
if (!this.props?.accessToken) {
throw new Error('Not authenticated. Please reconnect your account.');
}
// GitHub tokens don't expire - always return current token
if (!supportsTokenRefresh(provider)) {
return this.props.accessToken;
}
const FIVE_MINUTES_MS = 5 * 60 * 1000;
const isExpired = Date.now() >= (this.props.tokenExpiresAt ?? 0) - FIVE_MINUTES_MS;
if (!isExpired) return this.props.accessToken;
if (!this.props.refreshToken) {
console.warn('Token expired but no refresh token available');
throw new Error('Session expired. Please reconnect your account.');
}
console.log(`Legacy refresh: Token expired, refreshing (provider: ${provider})...`);
let newTokens;
switch (provider) {
case 'microsoft':
newTokens = await refreshMicrosoftToken({
clientId: this.env.MICROSOFT_CLIENT_ID!,
clientSecret: this.env.MICROSOFT_CLIENT_SECRET!,
refreshToken: this.props.refreshToken,
tenant: this.env.MICROSOFT_TENANT_ID,
});
break;
case 'google':
default:
newTokens = await refreshGoogleToken({
client_id: this.env.GOOGLE_CLIENT_ID,
client_secret: this.env.GOOGLE_CLIENT_SECRET,
refresh_token: this.props.refreshToken,
});
break;
}
if (!newTokens) {
console.error('Token refresh failed');
throw new Error('Session expired. Please reconnect your account.');
}
this.props = {
...this.props,
accessToken: newTokens.accessToken,
refreshToken: newTokens.refreshToken ?? this.props.refreshToken,
tokenExpiresAt: newTokens.expiresAt,
};
await this.ctx.storage.put('props', this.props);
// Store in TokenManager for future requests (migration path)
try {
const tokenManager = createTokenManager(this.env);
await tokenManager.store({
userId: this.props.id,
provider,
accessToken: newTokens.accessToken,
refreshToken: newTokens.refreshToken ?? this.props.refreshToken,
expiresAt: newTokens.expiresAt,
scopes: [], // Scopes not tracked in legacy flow
});
console.log('Legacy refresh: Migrated token to TokenManager');
} catch (storeError) {
console.error('Legacy refresh: Failed to store in TokenManager:', storeError);
}
console.log('Token refreshed successfully (legacy flow)');
return this.props.accessToken;
}
/**
* Make an authorized fetch request with automatic token refresh
*/
private async authorizedFetch(url: string, options: RequestInit = {}): Promise<Response> {
const accessToken = await this.ensureValidToken();
const response = await fetch(url, {
...options,
headers: {
...options.headers,
Authorization: `Bearer ${accessToken}`,
},
});
if (response.status === 401) {
this.props = { ...this.props!, accessToken: '', tokenExpiresAt: 0 };
await this.ctx.storage.put('props', this.props);
throw new Error('Access revoked. Please reconnect your account.');
}
return response;
}
}
// ===== OAUTH PROVIDER + HEADER AUTH =====
const sseHandler = MyMCP.serveSSE('/sse');
const mcpHandler = MyMCP.serve('/mcp');
/**
* Create OAuth provider with the appropriate identity handler.
* Handler is selected based on IDENTITY_PROVIDER env var.
*/
function createOAuthProvider(env: Env) {
return new OAuthProvider({
apiHandlers: {
'/sse': sseHandler,
'/mcp': mcpHandler,
},
authorizeEndpoint: '/authorize',
clientRegistrationEndpoint: '/register',
defaultHandler: getIdentityHandler(env) as unknown as ExportedHandler,
tokenEndpoint: '/token',
});
}
function handleMcpAuth(
request: Request,
env: Env,
ctx: ExecutionContext,
props: Props
): Promise<Response> {
const url = new URL(request.url);
const authCtx = { ...ctx, props };
if (url.pathname === '/sse') {
return sseHandler.fetch(request, env, authCtx);
} else {
return mcpHandler.fetch(request, env, authCtx);
}
}
export default {
async fetch(request: Request, env: Env, ctx: ExecutionContext): Promise<Response> {
const url = new URL(request.url);
const authHeader = request.headers.get('Authorization');
// ═══════════════════════════════════════════════════════════════════
// better-auth routes (social login, session management, OAuth provider)
// ═══════════════════════════════════════════════════════════════════
if (url.pathname.startsWith('/api/auth/')) {
try {
const auth = createAuth(env);
return auth.handler(request);
} catch (error) {
console.error('[better-auth] Error:', error);
return new Response(JSON.stringify({ error: 'Auth error' }), {
status: 500,
headers: { 'Content-Type': 'application/json' },
});
}
}
// Login page (for social login)
if (url.pathname === '/login') {
const callbackURL = url.searchParams.get('callbackURL') || undefined;
return new Response(getLoginPage(env, callbackURL), {
headers: { 'Content-Type': 'text/html' },
});
}
// Consent page (for OAuth authorization)
if (url.pathname === '/consent') {
return new Response(getConsentPage(url.searchParams), {
headers: { 'Content-Type': 'text/html' },
});
}
// MCP auth complete endpoint (handles better-auth callback for MCP OAuth)
// Must be handled before OAuthProvider to avoid token validation
if (url.pathname === '/mcp-auth-complete') {
const handler = getIdentityHandler(env);
return handler.fetch(request, env, ctx);
}
// Route admin pages and API
// Matches: /admin, /admin/services/:id, /api/admin/*
if (url.pathname.startsWith('/admin') || url.pathname.startsWith('/api/admin')) {
return adminApp.fetch(request, env, ctx);
}
// Check for Bearer token auth on MCP endpoints
if (authHeader?.startsWith('Bearer ') &&
(url.pathname === '/sse' || url.pathname === '/mcp')) {
const token = authHeader.slice(7);
// Check legacy AUTH_TOKEN (timing-safe comparison to prevent timing attacks)
if (env.AUTH_TOKEN && timingSafeEqual(token, env.AUTH_TOKEN)) {
return handleMcpAuth(request, env, ctx, {
id: 'legacy-auth',
email: 'legacy-auth@system',
name: 'Legacy Auth Token',
accessToken: '',
} as Props);
}
}
// Check for x-api-key header (better-auth apiKey plugin)
const apiKeyHeader = request.headers.get('x-api-key');
if (apiKeyHeader && (url.pathname === '/sse' || url.pathname === '/mcp')) {
try {
const auth = createAuth(env);
const result = await auth.api.verifyApiKey({
body: { key: apiKeyHeader },
});
if (result.valid && result.key) {
return handleMcpAuth(request, env, ctx, {
id: result.key.userId,
email: `apikey-${result.key.id}@system`,
name: result.key.name || 'API Key User',
accessToken: '',
} as Props);
}
} catch (error) {
console.error('[API Key] Validation error:', error);
// Fall through to OAuth provider
}
}
// Fall through to OAuth provider (handler selected based on IDENTITY_PROVIDER)
const oauthProvider = createOAuthProvider(env);
return oauthProvider.fetch(request, env, ctx);
},
};
