From ef9c29e0b38785bf2ea9b9b5844717ff94d6d6f5 Mon Sep 17 00:00:00 2001 From: Nicolas Bouliol Date: Tue, 1 Sep 2026 23:40:08 +0200 Subject: [PATCH 1/2] refactor(agent): serve root-mounted handlers through one registry `McpMiddleware` held a single callback, so the MCP server was the only thing the agent could serve at the root of a host application. The embedded BFF needs the same treatment, on paths of its own. `RootMiddleware` keeps named handlers, each with the matcher that says which urls it claims; anything unclaimed falls through to the host untouched. The five mount points are unchanged apart from the type they call, and the MCP matcher is now applied on the connect path too, where the callback used to be trusted to filter itself. Co-Authored-By: Claude Opus 5 (1M context) --- packages/agent/src/framework-mounter.ts | 25 +++--- packages/agent/src/mcp-middleware.ts | 69 ---------------- packages/agent/src/mcp-routes.ts | 9 +++ packages/agent/src/root-middleware.ts | 65 +++++++++++++++ packages/agent/test/mcp-middleware.test.ts | 35 --------- packages/agent/test/root-middleware.test.ts | 87 +++++++++++++++++++++ 6 files changed, 174 insertions(+), 116 deletions(-) delete mode 100644 packages/agent/src/mcp-middleware.ts create mode 100644 packages/agent/src/mcp-routes.ts create mode 100644 packages/agent/src/root-middleware.ts delete mode 100644 packages/agent/test/mcp-middleware.test.ts create mode 100644 packages/agent/test/root-middleware.test.ts 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(); + }); + }); +}); From f73332faad83746ced434e429433bd36ddee7da4 Mon Sep 17 00:00:00 2001 From: Nicolas Bouliol Date: Wed, 2 Sep 2026 16:49:56 +0200 Subject: [PATCH 2/2] refactor(agent): pair every root handler with the matcher it is served under MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The default route matcher was a crude `startsWith` over `/.well-known/`, `/oauth/` and `/mcp`: it claimed a host's ACME challenges, its own `/oauth/callback` and `/mcp-dashboard`. It was dead — both callers pass the basePath-scoped matcher `makeIsMcpRoute` builds — but promoting it to the default of all five mount points meant the first caller written without a matcher would silently swallow host traffic. A handler and its matcher now travel together as `RootHandler`, from `initializeMcpServer` down to `RootMiddleware.set`, so registering one without the other no longer compiles. `McpRouteMatcher` and `RouteMatcher`, the same signature declared twice, collapse into one. Overlapping matchers still cannot be caught at registration — the predicates are opaque — so dispatch names the winner in a debug line: a request answered by the wrong handler is otherwise indistinguishable from a route that was never registered. Tests assert which arguments a handler was called with, `next` included: dropping it made the Express path hang a Koa request forever and nothing failed. The deleted suite's negative case comes back — a matcher claiming `/custom` must leave `/mcp` to the host. Co-Authored-By: Claude Opus 5 (1M context) --- packages/agent/src/agent.ts | 36 ++++------ packages/agent/src/framework-mounter.ts | 13 ++-- packages/agent/src/mcp-routes.ts | 9 --- packages/agent/src/root-middleware.ts | 48 ++++++++----- packages/agent/src/types.ts | 5 +- packages/agent/test/agent.test.ts | 4 +- packages/agent/test/root-middleware.test.ts | 77 ++++++++++++++------- 7 files changed, 109 insertions(+), 83 deletions(-) delete mode 100644 packages/agent/src/mcp-routes.ts 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 e3d6ad44aa..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,6 @@ import path from 'path'; import FastifyAdapter from './fastify-adapter'; import InProcessDispatcher from './mcp-in-process-dispatcher'; -import isMcpRoute from './mcp-routes'; import RootMiddleware from './root-middleware'; export default class FrameworkMounter { @@ -38,15 +37,17 @@ export default class FrameworkMounter { this.prefix = prefix; this.logger = logger; this.fastifyAdapter = new FastifyAdapter(logger); - this.rootMiddleware = new RootMiddleware(); + 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.rootMiddleware.set('mcp', callback, routeMatcher ?? isMcpRoute); + protected setMcpCallback(handler: RootHandler | null): void { + this.rootMiddleware.set('mcp', handler); } /** diff --git a/packages/agent/src/mcp-routes.ts b/packages/agent/src/mcp-routes.ts deleted file mode 100644 index 5fb9eb8640..0000000000 --- a/packages/agent/src/mcp-routes.ts +++ /dev/null @@ -1,9 +0,0 @@ -/** - * 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 index 15935b4e79..525a2336bd 100644 --- a/packages/agent/src/root-middleware.ts +++ b/packages/agent/src/root-middleware.ts @@ -1,39 +1,53 @@ -import type { HttpCallback } from './types'; +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'; -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. + * 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 handlers = new Map(); + private readonly logger: Logger; - set(name: string, callback: HttpCallback | null, matches: RouteMatcher): void { - if (callback) this.handlers.set(name, { callback, matches }); + 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 handlerFor(url: string): HttpCallback | null { - for (const { callback, matches } of this.handlers.values()) { - if (matches(url)) return callback; + 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. @@ -53,7 +67,7 @@ export default class RootMiddleware { } getKoaMiddleware(): Koa.Middleware { - return expressToKoa(this.getExpressMiddleware(), url => this.handlerFor(url) !== null); + return expressToKoa(this.getExpressMiddleware(), url => this.entryFor(url) !== null); } /** The combined callback, or null when nothing is registered and the host needs no detour. */ diff --git a/packages/agent/src/types.ts b/packages/agent/src/types.ts index b49b9be255..3643b84186 100644 --- a/packages/agent/src/types.ts +++ b/packages/agent/src/types.ts @@ -153,7 +153,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/root-middleware.test.ts b/packages/agent/test/root-middleware.test.ts index ed1c56cbac..80e322c090 100644 --- a/packages/agent/test/root-middleware.test.ts +++ b/packages/agent/test/root-middleware.test.ts @@ -1,4 +1,3 @@ -/* eslint-disable @typescript-eslint/no-explicit-any */ import RootMiddleware from '../src/root-middleware'; describe('RootMiddleware', () => { @@ -8,37 +7,62 @@ describe('RootMiddleware', () => { describe('getExpressMiddleware', () => { it('should hand the request to the handler whose matcher claims the url', () => { - const middleware = new RootMiddleware(); + const middleware = new RootMiddleware(jest.fn()); const mcp = jest.fn(); const bff = jest.fn(); - middleware.set('mcp', mcp, url => url.startsWith('/mcp')); - middleware.set('bff', bff, url => url.startsWith('/bff')); + 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()({ url: '/bff/agent/v1' } as any, {} as any, jest.fn()); + middleware.getExpressMiddleware()(req, res, next); - expect(bff).toHaveBeenCalled(); + expect(bff).toHaveBeenCalledWith(req, res, next); expect(mcp).not.toHaveBeenCalled(); }); it('should pass an unclaimed url straight to the host', () => { - const middleware = new RootMiddleware(); + const middleware = new RootMiddleware(jest.fn()); const claimed = jest.fn(); const next = jest.fn(); - middleware.set('mcp', claimed, url => url.startsWith('/mcp')); + 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(); + const middleware = new RootMiddleware(jest.fn()); const claimed = jest.fn(); const next = jest.fn().mockResolvedValue(undefined); - middleware.set('mcp', claimed, url => url.startsWith('/mcp')); + middleware.set('mcp', { callback: claimed, matches: url => url.startsWith('/mcp') }); await middleware.getKoaMiddleware()(makeCtx('/api/v1/forest'), next as any); @@ -46,42 +70,43 @@ describe('RootMiddleware', () => { expect(next).toHaveBeenCalled(); }); - it('should offer a claimed url to its handler', async () => { - const middleware = new RootMiddleware(); + 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()); - middleware.set('mcp', claimed as any, url => url.startsWith('/mcp')); + const ctx = makeCtx('/mcp'); + middleware.set('mcp', { callback: claimed as any, matches: url => url.startsWith('/mcp') }); - await middleware.getKoaMiddleware()( - makeCtx('/mcp'), - jest.fn().mockResolvedValue(undefined) as any, - ); + await middleware.getKoaMiddleware()(ctx, jest.fn().mockResolvedValue(undefined) as any); - expect(claimed).toHaveBeenCalled(); + 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().getCallback()).toBeNull(); + expect(new RootMiddleware(jest.fn()).getCallback()).toBeNull(); }); it('should be null again once the last handler is removed', () => { - const middleware = new RootMiddleware(); - middleware.set('mcp', jest.fn(), () => true); + const middleware = new RootMiddleware(jest.fn()); + middleware.set('mcp', { callback: jest.fn(), matches: () => true }); - middleware.set('mcp', null, () => true); + middleware.set('mcp', null); expect(middleware.getCallback()).toBeNull(); }); it('should dispatch to a registered handler', () => { - const middleware = new RootMiddleware(); + const middleware = new RootMiddleware(jest.fn()); const claimed = jest.fn(); - middleware.set('mcp', claimed, url => url.startsWith('/mcp')); + 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()?.({ url: '/mcp' } as any, {} as any, jest.fn()); + middleware.getCallback()?.(req, res, next); - expect(claimed).toHaveBeenCalled(); + expect(claimed).toHaveBeenCalledWith(req, res, next); }); }); });