diff --git a/packages/agent/src/framework-mounter.ts b/packages/agent/src/framework-mounter.ts index 9aa2da42ad..e3d6ad44aa 100644 --- a/packages/agent/src/framework-mounter.ts +++ b/packages/agent/src/framework-mounter.ts @@ -10,7 +10,8 @@ import path from 'path'; import FastifyAdapter from './fastify-adapter'; import InProcessDispatcher from './mcp-in-process-dispatcher'; -import McpMiddleware from './mcp-middleware'; +import isMcpRoute from './mcp-routes'; +import RootMiddleware from './root-middleware'; export default class FrameworkMounter { public standaloneServerPort: number; @@ -24,7 +25,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,7 +38,7 @@ export default class FrameworkMounter { this.prefix = prefix; this.logger = logger; this.fastifyAdapter = new FastifyAdapter(logger); - this.mcpMiddleware = new McpMiddleware(); + this.rootMiddleware = new RootMiddleware(); this.inProcessDispatcher = new InProcessDispatcher(logger); } @@ -45,7 +46,7 @@ export default class FrameworkMounter { * Set the MCP HTTP callback. Call this before mount() or remount(). */ protected setMcpCallback(callback: HttpCallback | null, routeMatcher?: McpRouteMatcher): void { - this.mcpMiddleware.setCallback(callback, routeMatcher); + this.rootMiddleware.set('mcp', callback, routeMatcher ?? isMcpRoute); } /** @@ -129,7 +130,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 +146,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 +178,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 +195,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 +225,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/mcp-routes.ts b/packages/agent/src/mcp-routes.ts new file mode 100644 index 0000000000..5fb9eb8640 --- /dev/null +++ b/packages/agent/src/mcp-routes.ts @@ -0,0 +1,9 @@ +/** + * MCP route patterns claimed at the root of the host application. The MCP server filters on them + * itself too; this is what tells the agent's root middleware which requests to offer it. + */ +const MCP_ROUTE_PATTERNS = ['/.well-known/', '/oauth/', '/mcp']; + +export default function isMcpRoute(url: string): boolean { + return MCP_ROUTE_PATTERNS.some(pattern => url.startsWith(pattern)); +} diff --git a/packages/agent/src/root-middleware.ts b/packages/agent/src/root-middleware.ts new file mode 100644 index 0000000000..15935b4e79 --- /dev/null +++ b/packages/agent/src/root-middleware.ts @@ -0,0 +1,65 @@ +import type { HttpCallback } from './types'; +import type Koa from 'koa'; + +import expressToKoa from './utils/express-to-koa'; + +export type RouteMatcher = (url: string) => boolean; + +interface Handler { + callback: HttpCallback; + matches: RouteMatcher; +} + +/** + * 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 (a restart, a `stop()`) without the + * mount points having to know what is registered. + */ +export default class RootMiddleware { + private readonly handlers = new Map(); + + set(name: string, callback: HttpCallback | null, matches: RouteMatcher): void { + if (callback) this.handlers.set(name, { callback, matches }); + else this.handlers.delete(name); + } + + private handlerFor(url: string): HttpCallback | null { + for (const { callback, matches } of this.handlers.values()) { + if (matches(url)) return callback; + } + + return null; + } + + /** + * 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.handlerFor(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/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..ed1c56cbac --- /dev/null +++ b/packages/agent/test/root-middleware.test.ts @@ -0,0 +1,87 @@ +/* eslint-disable @typescript-eslint/no-explicit-any */ +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(); + const mcp = jest.fn(); + const bff = jest.fn(); + middleware.set('mcp', mcp, url => url.startsWith('/mcp')); + middleware.set('bff', bff, url => url.startsWith('/bff')); + + middleware.getExpressMiddleware()({ url: '/bff/agent/v1' } as any, {} as any, jest.fn()); + + expect(bff).toHaveBeenCalled(); + expect(mcp).not.toHaveBeenCalled(); + }); + + it('should pass an unclaimed url straight to the host', () => { + const middleware = new RootMiddleware(); + const claimed = jest.fn(); + const next = jest.fn(); + middleware.set('mcp', claimed, url => url.startsWith('/mcp')); + + middleware.getExpressMiddleware()({ url: '/api/v1/forest' } as any, {} as any, next); + + expect(claimed).not.toHaveBeenCalled(); + expect(next).toHaveBeenCalled(); + }); + }); + + describe('getKoaMiddleware', () => { + it('should leave an unclaimed url to the next Koa middleware', async () => { + const middleware = new RootMiddleware(); + const claimed = jest.fn(); + const next = jest.fn().mockResolvedValue(undefined); + middleware.set('mcp', claimed, 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', async () => { + const middleware = new RootMiddleware(); + const claimed = jest.fn((_req, _res, done) => done()); + middleware.set('mcp', claimed as any, url => url.startsWith('/mcp')); + + await middleware.getKoaMiddleware()( + makeCtx('/mcp'), + jest.fn().mockResolvedValue(undefined) as any, + ); + + expect(claimed).toHaveBeenCalled(); + }); + }); + + describe('getCallback', () => { + it('should be null while nothing is registered, so the host needs no detour', () => { + expect(new RootMiddleware().getCallback()).toBeNull(); + }); + + it('should be null again once the last handler is removed', () => { + const middleware = new RootMiddleware(); + middleware.set('mcp', jest.fn(), () => true); + + middleware.set('mcp', null, () => true); + + expect(middleware.getCallback()).toBeNull(); + }); + + it('should dispatch to a registered handler', () => { + const middleware = new RootMiddleware(); + const claimed = jest.fn(); + middleware.set('mcp', claimed, url => url.startsWith('/mcp')); + + middleware.getCallback()?.({ url: '/mcp' } as any, {} as any, jest.fn()); + + expect(claimed).toHaveBeenCalled(); + }); + }); +});