Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
36 changes: 14 additions & 22 deletions packages/agent/src/agent.ts
Original file line number Diff line number Diff line change
Expand Up @@ -3,8 +3,7 @@ import type { ForestAdminHttpDriverServices } from './services';
import type {
AgentOptions,
AgentOptionsWithDefaults,
HttpCallback,
McpRouteMatcher,
RootHandler,
WorkflowExecutorEmbedOptions,
} from './types';
import type {
Expand Down Expand Up @@ -101,12 +100,12 @@ export default class Agent<S extends TSchema = TSchema> 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;

Expand Down Expand Up @@ -180,9 +179,9 @@ export default class Agent<S extends TSchema = TSchema> 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;
Expand Down Expand Up @@ -373,8 +372,7 @@ export default class Agent<S extends TSchema = TSchema> extends FrameworkMounter
*/
private async getRouter(dataSource: DataSource): Promise<{
router: Router;
mcpHttpCallback?: HttpCallback;
mcpIsMcpRoute?: McpRouteMatcher;
mcp?: RootHandler;
}> {
// Bootstrap app
const services = makeServices(this.options);
Expand All @@ -383,12 +381,10 @@ export default class Agent<S extends TSchema = TSchema> 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
Expand All @@ -413,18 +409,15 @@ export default class Agent<S extends TSchema = TSchema> extends FrameworkMounter
router.use(correlationIdMiddleware);
routes.forEach(route => route.setupRoutes(router));

return { router, mcpHttpCallback, mcpIsMcpRoute };
return { router, mcp };
}

/**
* Initialize the MCP server.
* 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<RootHandler> {
const mcpLogger = (level, message) => this.options.logger(level, `[MCP] ${message}`);

try {
Expand Down Expand Up @@ -455,16 +448,16 @@ export default class Agent<S extends TSchema = TSchema> 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(
'Info',
'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}`);
Expand Down Expand Up @@ -503,8 +496,7 @@ export default class Agent<S extends TSchema = TSchema> extends FrameworkMounter

private async buildRouterAndSendSchema(): Promise<{
router: Router;
mcpHttpCallback?: HttpCallback;
mcpIsMcpRoute?: McpRouteMatcher;
mcp?: RootHandler;
}> {
const { isProduction, logger, typingsPath, typingsMaxDepth } = this.options;

Expand Down
32 changes: 17 additions & 15 deletions packages/agent/src/framework-mounter.ts
Original file line number Diff line number Diff line change
@@ -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';

Expand All @@ -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;
Expand All @@ -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;

Expand All @@ -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);
}

/**
Expand Down Expand Up @@ -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));
Expand All @@ -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);
Expand Down Expand Up @@ -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`);

Expand All @@ -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);
}

Expand All @@ -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);
Expand Down
69 changes: 0 additions & 69 deletions packages/agent/src/mcp-middleware.ts

This file was deleted.

79 changes: 79 additions & 0 deletions packages/agent/src/root-middleware.ts
Original file line number Diff line number Diff line change
@@ -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<string, RootHandler>();
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?.();

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Claude Opus 5 (claude-opus-5): Preferential

Two small contract-versus-code mismatches in this file, neither reachable today.

The defensive 404 disappeared. At the merge base an unclaimed request still reached the MCP callback, which owned the decision and had an explicit last resort when next was missing — res.writeHead(404, …) under a "this shouldn't happen in normal usage" comment. next?.() replaces that with a no-op: no status, no log, the socket open until the client times out. I enumerated all five mount points plus the nested standalone path and every one supplies a next, so this is not observable — it is worth a line only because the optional-call is now the only thing between a future next-less caller and a hung connection.

The class comment claims a removal path no code takes. It says a handler can be removed on "a restart, a stop()". restart() only ever replaces, and FrameworkMounter.stop() runs its onStop tasks and nothing else — no path calls set(name, null, …) on stop, so after agent.stop() a still-serving Express or NestJS host keeps routing /mcp into a stopped MCP callback. The behaviour is unchanged from the merge base; only the comment asserting otherwise is new, and it sits in the one file a maintainer will read while diagnosing exactly that. Either wire stop() to deregister, or cut the parenthetical.

Worth a sentence in the PR body either way: if a handler is ever deregistered while the host keeps the middleware, its requests fall through to the host's 404 — a reasonable contract, but indistinguishable from "the route was never registered", so the deregistration wants a log.

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Comment fixed — the parenthetical now says "on a restart" only, since stop() runs its onStop tasks and never deregisters. Not restoring the 404: the middleware's contract is that an unclaimed url is not its business, and the old 404 belonged to the MCP callback, not to the dispatcher — having RootMiddleware answer a request no handler claimed is exactly the behaviour this PR removes from the connect path. All five mount points plus the nested standalone path supply a next; if a terminal, next-less caller is ever added it needs to say so at construction rather than have the middleware guess per request.

};
}

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();
}
}
5 changes: 4 additions & 1 deletion packages/agent/src/types.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down
4 changes: 2 additions & 2 deletions packages/agent/test/agent.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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',
Expand All @@ -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);
});
Expand Down
Loading
Loading