refactor(mcp-server): extract framework-agnostic handlers

packages/agent/package.json

@@ -14,6 +14,7 @@
   "dependencies": {
     "@fast-csv/format": "^4.3.5",
     "@fastify/express": "^1.1.0",
+    "@forestadmin/ai-proxy": "0.1.0",
     "@forestadmin/datasource-customizer": "1.67.1",
     "@forestadmin/datasource-toolkit": "1.50.0",
     "@forestadmin/forestadmin-client": "1.36.14",

packages/agent/src/agent.ts

@@ -10,6 +10,13 @@ import {
 } from '@forestadmin/datasource-customizer';
 import { DataSource, DataSourceFactory } from '@forestadmin/datasource-toolkit';
 import { ForestSchema } from '@forestadmin/forestadmin-client';
+import {
+  McpHandlers,
+  isMcpRoute,
+  type AuthorizeParams,
+  type McpHandlersOptions,
+  type TokenParams,
+} from '@forestadmin/mcp-server';
 import bodyParser from '@koa/bodyparser';
 import cors from '@koa/cors';
 import Router from '@koa/router';
@@ -20,7 +27,7 @@ import FrameworkMounter from './framework-mounter';
 import makeRoutes from './routes';
 import makeServices, { ForestAdminHttpDriverServices } from './services';
 import CustomizationService from './services/model-customizations/customization';
-import { AgentOptions, AgentOptionsWithDefaults } from './types';
+import { AgentOptions, AgentOptionsWithDefaults, McpOptions } from './types';
 import SchemaGenerator from './utils/forest-schema/generator';
 import OptionsValidator from './utils/options-validator';
 
@@ -41,6 +48,12 @@ export default class Agent<S extends TSchema = TSchema> extends FrameworkMounter
   protected customizationService: CustomizationService;
   protected schemaGenerator: SchemaGenerator;
 
+  /** MCP options registered via useMcp() */
+  private mcpOptions: McpOptions | null = null;
+
+  /** MCP handlers instance */
+  private mcpHandlers: McpHandlers | null = null;
+
   /**
    * Create a new Agent Builder.
    * If any options are missing, the default will be applied:
@@ -73,11 +86,12 @@ export default class Agent<S extends TSchema = TSchema> extends FrameworkMounter
    * Start the agent.
    */
   async start(): Promise<void> {
-    const router = await this.buildRouterAndSendSchema();
+    const { router, mcpRouter } = await this.buildRouterAndSendSchema();
 
     await this.options.forestAdminClient.subscribeToServerEvents();
     this.options.forestAdminClient.onRefreshCustomizations(this.restart.bind(this));
 
+    this.setMcpRouter(mcpRouter ?? null);
     await this.mount(router);
   }
 
@@ -96,9 +110,10 @@ export default class Agent<S extends TSchema = TSchema> extends FrameworkMounter
    */
   async restart(): Promise<void> {
     // We force sending schema when restarting
-    const updatedRouter = await this.buildRouterAndSendSchema();
+    const { router, mcpRouter } = await this.buildRouterAndSendSchema();
 
-    await this.remount(updatedRouter);
+    this.setMcpRouter(mcpRouter ?? null);
+    await this.remount(router);
   }
 
   /**
@@ -190,20 +205,45 @@ export default class Agent<S extends TSchema = TSchema> extends FrameworkMounter
     return this;
   }
 
+  /**
+   * Enable MCP (Model Context Protocol) server support.
+   * This allows AI assistants to interact with your Forest Admin data.
+   *
+   * @param options Optional configuration for the MCP server
+   * @example
+   * agent.useMcp({ baseUrl: 'https://my-app.example.com' });
+   */
+  useMcp(options?: McpOptions): this {
+    this.mcpOptions = options || {};
+
+    return this;
+  }
+
   protected getRoutes(dataSource: DataSource, services: ForestAdminHttpDriverServices) {
     return makeRoutes(dataSource, this.options, services);
   }
 
   /**
    * Create an http handler which can respond to all queries which are expected from an agent.
+   * Returns the main router and optional MCP router.
    */
-  private async getRouter(dataSource: DataSource): Promise<Router> {
+  private async getRouter(
+    dataSource: DataSource,
+  ): Promise<{ router: Router; mcpRouter?: Router }> {
     // Bootstrap app
     const services = makeServices(this.options);
     const routes = this.getRoutes(dataSource, services);
+
     await Promise.all(routes.map(route => route.bootstrap()));
 
-    // Build router
+    // Initialize MCP router if configured via useMcp()
+    let mcpRouter: Router | undefined;
+
+    if (this.mcpOptions !== null) {
+      mcpRouter = await this.createMcpRouter();
+    }
+
+    // Build main router
     const router = new Router();
     router.all('(.*)', cors({ credentials: true, maxAge: 24 * 3600, privateNetworkAccess: true }));
     router.use(
@@ -216,10 +256,180 @@ export default class Agent<S extends TSchema = TSchema> extends FrameworkMounter
     );
     routes.forEach(route => route.setupRoutes(router));
 
+    return { router, mcpRouter };
+  }
+
+  /**
+   * Create a Koa router with MCP routes using the framework-agnostic handlers.
+   */
+  private async createMcpRouter(): Promise<Router> {
+    const handlerOptions: McpHandlersOptions = {
+      forestServerUrl: this.options.forestServerUrl,
+      envSecret: this.options.envSecret,
+      authSecret: this.options.authSecret,
+      logger: this.options.logger,
+    };
+
+    // Reuse existing handlers instance or create new one
+    if (!this.mcpHandlers) {
+      this.mcpHandlers = new McpHandlers(handlerOptions);
+      await this.mcpHandlers.initialize();
+    }
+
+    const handlers = this.mcpHandlers;
+
+    // Determine base URL
+    const baseUrl = this.mcpOptions?.baseUrl
+      ? new URL('/', this.mcpOptions.baseUrl)
+      : handlers.getBaseUrl();
+
+    if (!baseUrl) {
+      throw new Error(
+        'Could not determine base URL for MCP server. ' +
+          'Either provide a baseUrl option to useMcp() or ensure the Forest Admin environment has an api_endpoint configured.',
+      );
+    }
+
+    // Create router with MCP routes
+    const router = new Router();
+
+    // CORS for MCP routes
+    router.use(cors({ credentials: true }));
+
+    // Body parser for MCP routes
+    router.use(
+      bodyParser({
+        encoding: 'utf-8',
+        parsedMethods: ['POST'],
+      }),
+    );
+
+    // OAuth metadata endpoint (RFC 8414)
+    router.get('/.well-known/oauth-authorization-server', ctx => {
+      ctx.body = handlers.getOAuthMetadata(baseUrl);
+    });
+
+    // OAuth protected resource metadata (RFC 9728)
+    router.get('/.well-known/oauth-protected-resource', ctx => {
+      ctx.body = handlers.getProtectedResourceMetadata(baseUrl);
+    });
+
+    // Authorization endpoint
+    router.get('/oauth/authorize', async ctx => {
+      await this.handleAuthorize(ctx, handlers);
+    });
+
+    router.post('/oauth/authorize', async ctx => {
+      await this.handleAuthorize(ctx, handlers);
+    });
+
+    // Token endpoint
+    router.post('/oauth/token', async ctx => {
+      await this.handleToken(ctx, handlers);
+    });
+
+    // MCP endpoint (protected)
+    router.post('/mcp', async ctx => {
+      await this.handleMcp(ctx, handlers);
+    });
+
+    this.options.logger('Info', 'MCP server initialized successfully');
+
     return router;
   }
 
-  private async buildRouterAndSendSchema(): Promise<Router> {
+  private async handleAuthorize(ctx: any, handlers: McpHandlers): Promise<void> {
+    ctx.set('Cache-Control', 'no-store');
+
+    const params = ctx.method === 'POST' ? ctx.request.body : ctx.query;
+    const authorizeParams: AuthorizeParams = {
+      clientId: params.client_id,
+      redirectUri: params.redirect_uri,
+      codeChallenge: params.code_challenge,
+      state: params.state,
+      scope: params.scope,
+    };
+
+    const result = await handlers.handleAuthorize(authorizeParams);
+
+    if (result.type === 'redirect') {
+      ctx.redirect(result.url);
+    } else {
+      ctx.status = result.status;
+      ctx.body = { error: result.error, error_description: result.errorDescription };
+    }
+  }
+
+  private async handleToken(ctx: any, handlers: McpHandlers): Promise<void> {
+    const body = ctx.request.body as Record<string, string>;
+    const tokenParams: TokenParams = {
+      grantType: body.grant_type,
+      clientId: body.client_id,
+      code: body.code,
+      codeVerifier: body.code_verifier,
+      redirectUri: body.redirect_uri,
+      refreshToken: body.refresh_token,
+    };
+
+    const result = await handlers.handleToken(tokenParams);
+
+    if (result.type === 'success') {
+      ctx.body = result.tokens;
+    } else {
+      ctx.status = result.status;
+      ctx.body = { error: result.error, error_description: result.errorDescription };
+    }
+  }
+
+  private async handleMcp(ctx: any, handlers: McpHandlers): Promise<void> {
+    // Extract bearer token
+    const authHeader = ctx.get('Authorization');
+
+    if (!authHeader?.startsWith('Bearer ')) {
+      ctx.status = 401;
+      ctx.body = { error: 'unauthorized', error_description: 'Bearer token required' };
+
+      return;
+    }
+
+    const token = authHeader.slice(7);
+
+    try {
+      // Verify token
+      const authInfo = await handlers.verifyAccessToken(token);
+
+      // Check scopes
+      if (!authInfo.scopes.includes('mcp:read')) {
+        ctx.status = 403;
+        ctx.body = { error: 'insufficient_scope', error_description: 'mcp:read scope required' };
+
+        return;
+      }
+
+      // Handle MCP request - this writes directly to the response for streaming
+      const result = await handlers.handleMcpRequest(ctx.req, ctx.res, ctx.request.body);
+
+      if (result.type === 'handled') {
+        ctx.respond = false; // Let the transport handle the response
+      } else {
+        ctx.status = result.status;
+        ctx.body = result.error;
+      }
+    } catch (error: any) {
+      this.options.logger('Error', `MCP Error: ${error}`);
+      ctx.status = error.message?.includes('token') ? 401 : 500;
+      ctx.body = {
+        jsonrpc: '2.0',
+        error: { code: -32603, message: error.message || 'Internal error' },
+        id: null,
+      };
+    }
+  }
+
+  private async buildRouterAndSendSchema(): Promise<{
+    router: Router;
+    mcpRouter?: Router;
+  }> {
     const { isProduction, logger, typingsPath, typingsMaxDepth } = this.options;
 
     // It allows to rebuild the full customization stack with no code customizations
@@ -231,15 +441,15 @@ export default class Agent<S extends TSchema = TSchema> extends FrameworkMounter
     this.nocodeCustomizer.use(this.customizationService.addCustomizations);
 
     const dataSource = await this.nocodeCustomizer.getDataSource(logger);
-    const [router] = await Promise.all([
+    const [routers] = await Promise.all([
       this.getRouter(dataSource),
       this.sendSchema(dataSource),
       !isProduction && typingsPath
         ? this.customizer.updateTypesOnFileSystem(typingsPath, typingsMaxDepth, logger)
         : Promise.resolve(),
     ]);
 
-    return router;
+    return routers;
   }
 
   /**