diff --git a/packages/agent/src/agent.ts b/packages/agent/src/agent.ts index b13b25c9cb..f46aa15828 100644 --- a/packages/agent/src/agent.ts +++ b/packages/agent/src/agent.ts @@ -3,8 +3,7 @@ import type { ForestAdminHttpDriverServices } from './services'; import type { AgentOptions, AgentOptionsWithDefaults, - HttpCallback, - McpRouteMatcher, + RootHandler, WorkflowExecutorEmbedOptions, } from './types'; import type { @@ -101,12 +100,12 @@ export default class Agent extends FrameworkMounter let mounted = false; try { - const { router, mcpHttpCallback, mcpIsMcpRoute } = await this.buildRouterAndSendSchema(); + const { router, mcp } = await this.buildRouterAndSendSchema(); await this.options.forestAdminClient.subscribeToServerEvents(); this.options.forestAdminClient.onRefreshCustomizations(this.restart.bind(this)); - this.setMcpCallback(mcpHttpCallback ?? null, mcpIsMcpRoute); + this.setMcpCallback(mcp ?? null); await this.mount(router); mounted = true; @@ -180,9 +179,9 @@ export default class Agent extends FrameworkMounter try { // We force sending schema when restarting - const { router, mcpHttpCallback, mcpIsMcpRoute } = await this.buildRouterAndSendSchema(); + const { router, mcp } = await this.buildRouterAndSendSchema(); - this.setMcpCallback(mcpHttpCallback ?? null, mcpIsMcpRoute); + this.setMcpCallback(mcp ?? null); await this.remount(router); } finally { this.isRestarting = false; @@ -373,8 +372,7 @@ export default class Agent extends FrameworkMounter */ private async getRouter(dataSource: DataSource): Promise<{ router: Router; - mcpHttpCallback?: HttpCallback; - mcpIsMcpRoute?: McpRouteMatcher; + mcp?: RootHandler; }> { // Bootstrap app const services = makeServices(this.options); @@ -383,12 +381,10 @@ export default class Agent extends FrameworkMounter await Promise.all(routes.map(route => route.bootstrap())); // Initialize MCP server if enabled via mountAiMcpServer() - let mcpHttpCallback: HttpCallback | undefined; - let mcpIsMcpRoute: McpRouteMatcher | undefined; + let mcp: RootHandler | undefined; if (this.mcpEnabled) { - ({ httpCallback: mcpHttpCallback, isMcpRoute: mcpIsMcpRoute } = - await this.initializeMcpServer()); + mcp = await this.initializeMcpServer(); } // Build main router @@ -413,7 +409,7 @@ export default class Agent extends FrameworkMounter router.use(correlationIdMiddleware); routes.forEach(route => route.setupRoutes(router)); - return { router, mcpHttpCallback, mcpIsMcpRoute }; + return { router, mcp }; } /** @@ -421,10 +417,7 @@ export default class Agent extends FrameworkMounter * Uses dynamic import to defer loading until mountAiMcpServer() is actually used. * This avoids loading the mcp-server dependency at startup for users who don't use MCP. */ - private async initializeMcpServer(): Promise<{ - httpCallback: HttpCallback; - isMcpRoute: McpRouteMatcher; - }> { + private async initializeMcpServer(): Promise { const mcpLogger = (level, message) => this.options.logger(level, `[MCP] ${message}`); try { @@ -455,8 +448,8 @@ export default class Agent extends FrameworkMounter agentDispatcher: this.getInProcessDispatcher(), }); - const httpCallback = await mcpServer.getHttpCallback(); - const isMcpRoute = makeIsMcpRoute(this.mcpBasePath); + const callback = await mcpServer.getHttpCallback(); + const matches = makeIsMcpRoute(this.mcpBasePath); mcpLogger('Info', 'Server initialized successfully'); mcpLogger( @@ -464,7 +457,7 @@ export default class Agent extends FrameworkMounter 'Tool calls dispatch in-process and skip any middleware mounted in front of the agent', ); - return { httpCallback, isMcpRoute }; + return { callback, matches }; } catch (error) { const { message } = error as Error; mcpLogger('Error', `Failed to initialize MCP server: ${message}`); @@ -503,8 +496,7 @@ export default class Agent extends FrameworkMounter private async buildRouterAndSendSchema(): Promise<{ router: Router; - mcpHttpCallback?: HttpCallback; - mcpIsMcpRoute?: McpRouteMatcher; + mcp?: RootHandler; }> { const { isProduction, logger, typingsPath, typingsMaxDepth } = this.options; diff --git a/packages/agent/src/framework-mounter.ts b/packages/agent/src/framework-mounter.ts index 9aa2da42ad..4fa10994fd 100644 --- a/packages/agent/src/framework-mounter.ts +++ b/packages/agent/src/framework-mounter.ts @@ -1,5 +1,5 @@ /* eslint-disable @typescript-eslint/no-explicit-any */ -import type { HttpCallback, McpRouteMatcher } from './types'; +import type { HttpCallback, RootHandler } from './types'; import type { Logger } from '@forestadmin/datasource-toolkit'; import type net from 'net'; @@ -10,7 +10,7 @@ import path from 'path'; import FastifyAdapter from './fastify-adapter'; import InProcessDispatcher from './mcp-in-process-dispatcher'; -import McpMiddleware from './mcp-middleware'; +import RootMiddleware from './root-middleware'; export default class FrameworkMounter { public standaloneServerPort: number; @@ -24,7 +24,7 @@ export default class FrameworkMounter { private readonly logger: Logger; private readonly fastifyAdapter: FastifyAdapter; - private readonly mcpMiddleware: McpMiddleware; + private readonly rootMiddleware: RootMiddleware; private readonly inProcessDispatcher: InProcessDispatcher; private inProcessHookRegistered = false; @@ -37,15 +37,17 @@ export default class FrameworkMounter { this.prefix = prefix; this.logger = logger; this.fastifyAdapter = new FastifyAdapter(logger); - this.mcpMiddleware = new McpMiddleware(); + this.rootMiddleware = new RootMiddleware(logger); this.inProcessDispatcher = new InProcessDispatcher(logger); } /** - * Set the MCP HTTP callback. Call this before mount() or remount(). + * Register the MCP server at the root of the host application, or `null` to unregister it. The + * matcher comes with the callback: there is no default pattern set, so nothing can claim host + * urls by accident. Call this before mount() or remount(). */ - protected setMcpCallback(callback: HttpCallback | null, routeMatcher?: McpRouteMatcher): void { - this.mcpMiddleware.setCallback(callback, routeMatcher); + protected setMcpCallback(handler: RootHandler | null): void { + this.rootMiddleware.set('mcp', handler); } /** @@ -129,7 +131,7 @@ export default class FrameworkMounter { */ mountOnExpress(express: any): this { // MCP middleware - the callback handles its own path filtering and calls next() for non-MCP routes - express.use(this.mcpMiddleware.getExpressMiddleware()); + express.use(this.rootMiddleware.getExpressMiddleware()); // Mount main forest routes at /{prefix}/forest express.use(this.completeMountPrefix, this.getConnectCallback(false)); @@ -145,7 +147,7 @@ export default class FrameworkMounter { */ mountOnFastify(fastify: any): this { // MCP middleware at root - the callback handles its own path filtering - this.fastifyAdapter.useCallback(fastify, this.mcpMiddleware.getExpressMiddleware(), '/'); + this.fastifyAdapter.useCallback(fastify, this.rootMiddleware.getExpressMiddleware(), '/'); // Mount main forest routes const callback = this.getConnectCallback(false); @@ -177,7 +179,7 @@ export default class FrameworkMounter { }); // MCP middleware - intercepts MCP routes before they reach Koa's body parser - koa.use(this.mcpMiddleware.getKoaMiddleware()); + koa.use(this.rootMiddleware.getKoaMiddleware()); koa.use(parentRouter.routes()); this.logger('Info', `Successfully mounted on Koa`); @@ -194,12 +196,12 @@ export default class FrameworkMounter { if (adapter.constructor.name === 'ExpressAdapter') { // MCP middleware at root - the callback handles its own path filtering - nestJs.use(this.mcpMiddleware.getExpressMiddleware()); + nestJs.use(this.rootMiddleware.getExpressMiddleware()); // Mount main forest routes nestJs.use(this.completeMountPrefix, callback); } else { // Fastify adapter - MCP middleware at root - this.fastifyAdapter.useCallback(nestJs, this.mcpMiddleware.getExpressMiddleware(), '/'); + this.fastifyAdapter.useCallback(nestJs, this.rootMiddleware.getExpressMiddleware(), '/'); this.fastifyAdapter.useCallback(nestJs, callback, this.completeMountPrefix); } @@ -224,10 +226,10 @@ export default class FrameworkMounter { return (req, res) => { // For standalone server (nested), check MCP callback first // The MCP callback handles its own path filtering - const mcpCallback = this.mcpMiddleware.getCallback(); + const rootCallback = this.rootMiddleware.getCallback(); - if (nested && mcpCallback) { - mcpCallback(req, res, () => { + if (nested && rootCallback) { + rootCallback(req, res, () => { // next() called means not an MCP route - forward to main handler if (handler) { handler(req, res); diff --git a/packages/agent/src/mcp-middleware.ts b/packages/agent/src/mcp-middleware.ts deleted file mode 100644 index 64f79b31bc..0000000000 --- a/packages/agent/src/mcp-middleware.ts +++ /dev/null @@ -1,69 +0,0 @@ -import type { HttpCallback, McpRouteMatcher } from './types'; -import type Koa from 'koa'; - -import expressToKoa from './utils/express-to-koa'; - -/** - * MCP route patterns that should be intercepted by the MCP middleware. - */ -const MCP_ROUTE_PATTERNS = ['/.well-known/', '/oauth/', '/mcp']; - -/** - * Check if a URL matches any MCP route pattern. - */ -function isMcpRoute(url: string): boolean { - return MCP_ROUTE_PATTERNS.some(pattern => url.startsWith(pattern)); -} - -/** - * Factory functions for creating MCP middleware for different frameworks. - * MCP middleware intercepts MCP-specific routes (/.well-known/*, /oauth/*, /mcp) - * and forwards them to the MCP callback, while passing through other routes. - */ -export default class McpMiddleware { - private mcpHttpCallback: HttpCallback | null = null; - private routeMatcher: McpRouteMatcher = isMcpRoute; - - setCallback(callback: HttpCallback | null, routeMatcher?: McpRouteMatcher): void { - this.mcpHttpCallback = callback; - this.routeMatcher = routeMatcher ?? isMcpRoute; - } - - /** - * Get the current MCP HTTP callback. - */ - getCallback(): HttpCallback | null { - return this.mcpHttpCallback; - } - - /** - * Get an Express/Connect-style middleware that forwards requests to the MCP callback. - * The MCP callback handles path filtering and calls next() for non-MCP routes. - */ - getExpressMiddleware(): HttpCallback { - return (req, res, next) => { - if (this.mcpHttpCallback) { - this.mcpHttpCallback(req, res, next); - } else if (next) { - next(); - } - }; - } - - /** - * Get a Koa middleware that forwards requests to the MCP callback. - * Uses expressToKoa wrapper for proper Koa integration. - */ - getKoaMiddleware(): Koa.Middleware { - return expressToKoa( - (req, res, next) => { - if (this.mcpHttpCallback) { - this.mcpHttpCallback(req, res, next); - } else if (next) { - next(); - } - }, - url => this.routeMatcher(url), - ); - } -} diff --git a/packages/agent/src/root-middleware.ts b/packages/agent/src/root-middleware.ts new file mode 100644 index 0000000000..525a2336bd --- /dev/null +++ b/packages/agent/src/root-middleware.ts @@ -0,0 +1,79 @@ +import type { HttpCallback, RootHandler } from './types'; +import type { Logger } from '@forestadmin/datasource-toolkit'; +import type Koa from 'koa'; + +import expressToKoa from './utils/express-to-koa'; + +/** + * Handlers the agent serves at the root of the host application rather than under its own router + * prefix — the MCP server, the embedded BFF. Each one claims a set of paths through its matcher; + * anything else falls through to the host, untouched. + * + * Registered by name so a handler can be replaced or removed on a restart without the mount points + * having to know what is registered. + */ +export default class RootMiddleware { + private readonly handlers = new Map(); + private readonly logger: Logger; + + constructor(logger: Logger) { + this.logger = logger; + } + + set(name: string, handler: RootHandler | null): void { + if (handler) this.handlers.set(name, handler); + else this.handlers.delete(name); + } + + private entryFor(url: string): [string, RootHandler] | null { + for (const entry of this.handlers) { + if (entry[1].matches(url)) return entry; + } + + return null; + } + + /** + * First registered matcher wins. Overlapping matchers cannot be caught at registration — the + * predicates are opaque — so the winner is named here instead: a request answered by the wrong + * handler is otherwise indistinguishable from one whose route was never registered. + */ + private handlerFor(url: string): HttpCallback | null { + const entry = this.entryFor(url); + + if (!entry) return null; + + this.logger('Debug', `Root middleware: '${entry[0]}' claims ${url}`); + + return entry[1].callback; + } + + /** + * Connect-style middleware. A handler that decides not to answer calls `next()`, which passes the + * request on to the host exactly as if nothing had been registered. + */ + getExpressMiddleware(): HttpCallback { + return (req, res, next) => { + const handler = this.handlerFor(req.url ?? '/'); + + if (handler) { + handler(req, res, next); + + return; + } + + next?.(); + }; + } + + getKoaMiddleware(): Koa.Middleware { + return expressToKoa(this.getExpressMiddleware(), url => this.entryFor(url) !== null); + } + + /** The combined callback, or null when nothing is registered and the host needs no detour. */ + getCallback(): HttpCallback | null { + if (this.handlers.size === 0) return null; + + return this.getExpressMiddleware(); + } +} diff --git a/packages/agent/src/types.ts b/packages/agent/src/types.ts index 77bbd6e8bd..c66feda2d9 100644 --- a/packages/agent/src/types.ts +++ b/packages/agent/src/types.ts @@ -162,7 +162,10 @@ export type AgentOptionsWithDefaults = Readonly< export type HttpCallback = (req: IncomingMessage, res: ServerResponse, next?: () => void) => void; -export type McpRouteMatcher = (url: string) => boolean; +export type RouteMatcher = (url: string) => boolean; + +/** A handler served at the root of the host application, with the urls it claims. */ +export type RootHandler = { callback: HttpCallback; matches: RouteMatcher }; export enum HttpCode { BadRequest = 400, diff --git a/packages/agent/test/agent.test.ts b/packages/agent/test/agent.test.ts index 94f6542965..58cac2670c 100644 --- a/packages/agent/test/agent.test.ts +++ b/packages/agent/test/agent.test.ts @@ -762,7 +762,7 @@ describe('Agent', () => { test('threads a basePath-scoped route matcher to the MCP middleware', async () => { const options = factories.forestAdminHttpDriverOptions.build(); const agent = new Agent(options); - type SetMcpCallback = (cb: unknown, matcher?: (url: string) => boolean) => void; + type SetMcpCallback = (handler: { matches: (url: string) => boolean } | null) => void; const setMcpCallbackSpy = jest.spyOn( agent as unknown as { setMcpCallback: SetMcpCallback }, 'setMcpCallback', @@ -771,7 +771,7 @@ describe('Agent', () => { agent.mountAiMcpServer({ basePath: '/ai' }); await agent.start(); - const matcher = setMcpCallbackSpy.mock.calls.at(-1)?.[1]; + const matcher = setMcpCallbackSpy.mock.calls.at(-1)?.[0]?.matches; expect(matcher?.('/ai/mcp')).toBe(true); expect(matcher?.('/oauth/token')).toBe(false); }); diff --git a/packages/agent/test/mcp-middleware.test.ts b/packages/agent/test/mcp-middleware.test.ts deleted file mode 100644 index 1a9251d968..0000000000 --- a/packages/agent/test/mcp-middleware.test.ts +++ /dev/null @@ -1,35 +0,0 @@ -import McpMiddleware from '../src/mcp-middleware'; - -describe('McpMiddleware', () => { - function makeCtx(url: string) { - return { url, req: {}, res: { once: jest.fn() }, respond: true } as any; - } - - test('getKoaMiddleware gates requests with the custom route matcher', async () => { - const middleware = new McpMiddleware(); - const callback = jest.fn((req, res, next) => next()); - middleware.setCallback(callback, url => url.startsWith('/custom')); - - const koaMiddleware = middleware.getKoaMiddleware(); - - await koaMiddleware(makeCtx('/mcp'), jest.fn().mockResolvedValue(undefined) as any); - expect(callback).not.toHaveBeenCalled(); - - await koaMiddleware(makeCtx('/custom/x'), jest.fn().mockResolvedValue(undefined) as any); - expect(callback).toHaveBeenCalledTimes(1); - }); - - test('falls back to the default root matcher when none is provided', async () => { - const middleware = new McpMiddleware(); - const callback = jest.fn((req, res, next) => next()); - middleware.setCallback(callback); - - const koaMiddleware = middleware.getKoaMiddleware(); - - await koaMiddleware(makeCtx('/api/other'), jest.fn().mockResolvedValue(undefined) as any); - expect(callback).not.toHaveBeenCalled(); - - await koaMiddleware(makeCtx('/mcp'), jest.fn().mockResolvedValue(undefined) as any); - expect(callback).toHaveBeenCalledTimes(1); - }); -}); diff --git a/packages/agent/test/root-middleware.test.ts b/packages/agent/test/root-middleware.test.ts new file mode 100644 index 0000000000..80e322c090 --- /dev/null +++ b/packages/agent/test/root-middleware.test.ts @@ -0,0 +1,112 @@ +import RootMiddleware from '../src/root-middleware'; + +describe('RootMiddleware', () => { + function makeCtx(url: string) { + return { url, req: { url }, res: { once: jest.fn() }, respond: true } as any; + } + + describe('getExpressMiddleware', () => { + it('should hand the request to the handler whose matcher claims the url', () => { + const middleware = new RootMiddleware(jest.fn()); + const mcp = jest.fn(); + const bff = jest.fn(); + const req = { url: '/bff/agent/v1' } as any; + const res = {} as any; + const next = jest.fn(); + middleware.set('mcp', { callback: mcp, matches: url => url.startsWith('/mcp') }); + middleware.set('bff', { callback: bff, matches: url => url.startsWith('/bff') }); + + middleware.getExpressMiddleware()(req, res, next); + + expect(bff).toHaveBeenCalledWith(req, res, next); + expect(mcp).not.toHaveBeenCalled(); + }); + + it('should pass an unclaimed url straight to the host', () => { + const middleware = new RootMiddleware(jest.fn()); + const claimed = jest.fn(); + const next = jest.fn(); + middleware.set('mcp', { callback: claimed, matches: url => url.startsWith('/mcp') }); + + middleware.getExpressMiddleware()({ url: '/api/v1/forest' } as any, {} as any, next); + + expect(claimed).not.toHaveBeenCalled(); + expect(next).toHaveBeenCalled(); + }); + + it('should claim nothing beyond what the registered matcher declares', () => { + const middleware = new RootMiddleware(jest.fn()); + const claimed = jest.fn(); + const next = jest.fn(); + middleware.set('mcp', { callback: claimed, matches: url => url.startsWith('/custom') }); + + middleware.getExpressMiddleware()({ url: '/mcp' } as any, {} as any, next); + + expect(claimed).not.toHaveBeenCalled(); + expect(next).toHaveBeenCalled(); + }); + + it('should name the handler that claimed the url', () => { + const logger = jest.fn(); + const middleware = new RootMiddleware(logger); + middleware.set('mcp', { callback: jest.fn(), matches: url => url.startsWith('/mcp') }); + + middleware.getExpressMiddleware()({ url: '/mcp' } as any, {} as any, jest.fn()); + + expect(logger).toHaveBeenCalledWith('Debug', "Root middleware: 'mcp' claims /mcp"); + }); + }); + + describe('getKoaMiddleware', () => { + it('should leave an unclaimed url to the next Koa middleware', async () => { + const middleware = new RootMiddleware(jest.fn()); + const claimed = jest.fn(); + const next = jest.fn().mockResolvedValue(undefined); + middleware.set('mcp', { callback: claimed, matches: url => url.startsWith('/mcp') }); + + await middleware.getKoaMiddleware()(makeCtx('/api/v1/forest'), next as any); + + expect(claimed).not.toHaveBeenCalled(); + expect(next).toHaveBeenCalled(); + }); + + it('should offer a claimed url to its handler with the node request, response and a next', async () => { + const middleware = new RootMiddleware(jest.fn()); + const claimed = jest.fn((_req, _res, done) => done()); + const ctx = makeCtx('/mcp'); + middleware.set('mcp', { callback: claimed as any, matches: url => url.startsWith('/mcp') }); + + await middleware.getKoaMiddleware()(ctx, jest.fn().mockResolvedValue(undefined) as any); + + expect(claimed).toHaveBeenCalledWith(ctx.req, ctx.res, expect.any(Function)); + }); + }); + + describe('getCallback', () => { + it('should be null while nothing is registered, so the host needs no detour', () => { + expect(new RootMiddleware(jest.fn()).getCallback()).toBeNull(); + }); + + it('should be null again once the last handler is removed', () => { + const middleware = new RootMiddleware(jest.fn()); + middleware.set('mcp', { callback: jest.fn(), matches: () => true }); + + middleware.set('mcp', null); + + expect(middleware.getCallback()).toBeNull(); + }); + + it('should dispatch to a registered handler', () => { + const middleware = new RootMiddleware(jest.fn()); + const claimed = jest.fn(); + const req = { url: '/mcp' } as any; + const res = {} as any; + const next = jest.fn(); + middleware.set('mcp', { callback: claimed, matches: url => url.startsWith('/mcp') }); + + middleware.getCallback()?.(req, res, next); + + expect(claimed).toHaveBeenCalledWith(req, res, next); + }); + }); +});