feat(agent): add addRouter method for custom HTTP routes

packages/agent/src/agent.ts

@@ -4,6 +4,8 @@ import type {
   AgentOptions,
   AgentOptionsWithDefaults,
   AiConfiguration,
+  CustomRouterCallback,
+  CustomRouterOptions,
   HttpCallback,
 } from './types';
 import type {
@@ -17,8 +19,21 @@ import type {
 import type { DataSource, DataSourceFactory } from '@forestadmin/datasource-toolkit';
 import type { ForestSchema } from '@forestadmin/forestadmin-client';
 
-import { isModelSupportingTools } from '@forestadmin/ai-proxy';
+import {
+  AIBadRequestError,
+  AIError,
+  AINotFoundError,
+  Router as AiProxyRouter,
+  extractMcpOauthTokensFromHeaders,
+  injectOauthTokens,
+  isModelSupportingTools,
+} from '@forestadmin/ai-proxy';
 import { DataSourceCustomizer } from '@forestadmin/datasource-customizer';
+import {
+  BadRequestError,
+  NotFoundError,
+  UnprocessableError,
+} from '@forestadmin/datasource-toolkit';
 import bodyParser from '@koa/bodyparser';
 import cors from '@koa/cors';
 import Router from '@koa/router';
@@ -27,6 +42,7 @@ import stringify from 'json-stringify-pretty-compact';
 
 import FrameworkMounter from './framework-mounter';
 import makeRoutes from './routes';
+import CustomRoute from './routes/custom/custom-route';
 import makeServices from './services';
 import CustomizationService from './services/model-customizations/customization';
 import SchemaGenerator from './utils/forest-schema/generator';
@@ -50,6 +66,12 @@ export default class Agent<S extends TSchema = TSchema> extends FrameworkMounter
   protected schemaGenerator: SchemaGenerator;
   protected aiConfigurations: AiConfiguration[] = [];
 
+  /** Custom router callbacks registered via addRouter() */
+  private customRouterCallbacks: Array<{
+    callback: CustomRouterCallback;
+    options?: CustomRouterOptions;
+  }> = [];
+
   /** Whether MCP server should be mounted */
   private mcpEnabled = false;
 
@@ -259,13 +281,97 @@ export default class Agent<S extends TSchema = TSchema> extends FrameworkMounter
         'Make sure to test Forest Admin AI features thoroughly to ensure compatibility.',
     );
 
+    // Store for schema metadata
     this.aiConfigurations.push(configuration);
 
+    // Register AI proxy route via addRouter
+    this.addRouter((router, context) => {
+      const aiProxyRouter = new AiProxyRouter({
+        aiConfigurations: [configuration],
+        logger: context.logger,
+      });
+
+      router.post('/_internal/ai-proxy/:route', async ctx => {
+        try {
+          const tokensByMcpServerName = extractMcpOauthTokensFromHeaders(ctx.request.headers);
+          const mcpConfigs =
+            await context.options.forestAdminClient.mcpServerConfigService.getConfiguration();
+
+          // eslint-disable-next-line @typescript-eslint/no-explicit-any
+          ctx.response.body = await aiProxyRouter.route({
+            route: ctx.params.route,
+            body: ctx.request.body,
+            query: ctx.query,
+            mcpConfigs: injectOauthTokens({ mcpConfigs, tokensByMcpServerName }),
+          } as any);
+          ctx.response.status = 200;
+        } catch (error) {
+          if (error instanceof AIError) {
+            context.logger('Error', `AI proxy error: ${(error as Error).message}`, error);
+
+            if (error instanceof AIBadRequestError) throw new BadRequestError(error.message);
+            if (error instanceof AINotFoundError) throw new NotFoundError(error.message);
+            throw new UnprocessableError(error.message);
+          }
+
+          throw error;
+        }
+      });
+    });
+
+    return this;
+  }
+
+  /**
+   * Add custom HTTP routes to the agent.
+   *
+   * Routes are authenticated by default (require Forest Admin authentication).
+   * Use `{ authenticated: false }` to create public routes.
+   *
+   * @param callback - Function that receives a Koa router and context to define routes
+   * @param options - Configuration options for the routes
+   * @param options.authenticated - Whether routes require authentication (default: true)
+   * @param options.prefix - URL prefix for all routes in this router (default: '')
+   * @returns The agent instance for chaining
+   *
+   * @example
+   * // Simple authenticated route
+   * agent.addRouter((router, context) => {
+   *   router.get('/stats', async (ctx) => {
+   *     const users = await context.dataSource
+   *       .getCollection('users')
+   *       .list(caller, new PaginatedFilter({}), new Projection('id'));
+   *     ctx.body = { count: users.length };
+   *   });
+   * });
+   *
+   * @example
+   * // Public route (no authentication required)
+   * agent.addRouter(
+   *   (router) => {
+   *     router.get('/health', (ctx) => {
+   *       ctx.body = { status: 'ok' };
+   *     });
+   *   },
+   *   { authenticated: false }
+   * );
+   */
+  addRouter(callback: CustomRouterCallback, options?: CustomRouterOptions): this {
+    this.customRouterCallbacks.push({ callback, options });
+
     return this;
   }
 
   protected getRoutes(dataSource: DataSource, services: ForestAdminHttpDriverServices) {
-    return makeRoutes(dataSource, this.options, services, this.aiConfigurations);
+    const routes = makeRoutes(dataSource, this.options, services);
+
+    // Add custom routes (includes AI proxy route if addAi was called)
+    for (const { callback, options } of this.customRouterCallbacks) {
+      routes.push(new CustomRoute(services, this.options, dataSource, callback, options));
+    }
+
+    // Re-sort routes to ensure custom routes are in the correct order
+    return routes.sort((a, b) => a.type - b.type);
   }
 
   /**

packages/agent/src/index.ts

@@ -8,7 +8,12 @@ export function createAgent<S extends TSchema = TSchema>(options: AgentOptions):
 }
 
 export { Agent };
-export { AgentOptions } from './types';
+export {
+  AgentOptions,
+  CustomRouterCallback,
+  CustomRouterContext,
+  CustomRouterOptions,
+} from './types';
 export * from '@forestadmin/datasource-customizer';
 
 // export is necessary for the agent-generator package

packages/agent/src/routes/custom/custom-route.ts

@@ -0,0 +1,50 @@
+import type { ForestAdminHttpDriverServices } from '../../services';
+import type {
+  AgentOptionsWithDefaults,
+  CustomRouterCallback,
+  CustomRouterOptions,
+} from '../../types';
+import type { DataSource } from '@forestadmin/datasource-toolkit';
+import type Router from '@koa/router';
+
+import KoaRouter from '@koa/router';
+
+import { RouteType } from '../../types';
+import BaseRoute from '../base-route';
+import CustomRouterContextImpl from './custom-router-context';
+
+export default class CustomRoute extends BaseRoute {
+  private readonly dataSource: DataSource;
+  private readonly callback: CustomRouterCallback;
+  private readonly customOptions: Required<CustomRouterOptions>;
+
+  get type(): RouteType {
+    return this.customOptions.authenticated ? RouteType.CustomRoute : RouteType.PublicCustomRoute;
+  }
+
+  constructor(
+    services: ForestAdminHttpDriverServices,
+    options: AgentOptionsWithDefaults,
+    dataSource: DataSource,
+    callback: CustomRouterCallback,
+    customOptions?: CustomRouterOptions,
+  ) {
+    super(services, options);
+    this.dataSource = dataSource;
+    this.callback = callback;
+    this.customOptions = {
+      authenticated: customOptions?.authenticated ?? true,
+      prefix: customOptions?.prefix ?? '',
+    };
+  }
+
+  setupRoutes(router: Router): void {
+    const context = new CustomRouterContextImpl(this.dataSource, this.services, this.options);
+    const customRouter = new KoaRouter({ prefix: this.customOptions.prefix || undefined });
+
+    this.callback(customRouter, context);
+
+    router.use(customRouter.routes());
+    router.use(customRouter.allowedMethods());
+  }
+}

packages/agent/src/routes/custom/custom-router-context.ts

@@ -0,0 +1,23 @@
+import type { ForestAdminHttpDriverServices } from '../../services';
+import type { AgentOptionsWithDefaults, CustomRouterContext } from '../../types';
+import type { DataSource, Logger } from '@forestadmin/datasource-toolkit';
+
+export default class CustomRouterContextImpl implements CustomRouterContext {
+  readonly dataSource: DataSource;
+  readonly services: ForestAdminHttpDriverServices;
+  readonly options: AgentOptionsWithDefaults;
+
+  constructor(
+    dataSource: DataSource,
+    services: ForestAdminHttpDriverServices,
+    options: AgentOptionsWithDefaults,
+  ) {
+    this.dataSource = dataSource;
+    this.services = services;
+    this.options = options;
+  }
+
+  get logger(): Logger {
+    return this.options.logger;
+  }
+}

packages/agent/src/routes/index.ts

@@ -1,5 +1,5 @@
 import type { ForestAdminHttpDriverServices as Services } from '../services';
-import type { AiConfiguration, AgentOptionsWithDefaults as Options } from '../types';
+import type { AgentOptionsWithDefaults as Options } from '../types';
 import type BaseRoute from './base-route';
 import type { DataSource } from '@forestadmin/datasource-toolkit';
 
@@ -14,7 +14,6 @@ import Get from './access/get';
 import List from './access/list';
 import ListRelated from './access/list-related';
 import NativeQueryDatasource from './access/native-query-datasource';
-import AiProxyRoute from './ai/ai-proxy';
 import Capabilities from './capabilities';
 import ActionRoute from './modification/action/action';
 import AssociateRelated from './modification/associate-related';
@@ -165,21 +164,10 @@ function getActionRoutes(
   return routes;
 }
 
-function getAiRoutes(
-  options: Options,
-  services: Services,
-  aiConfigurations: AiConfiguration[],
-): BaseRoute[] {
-  if (aiConfigurations.length === 0) return [];
-
-  return [new AiProxyRoute(services, options, aiConfigurations)];
-}
-
 export default function makeRoutes(
   dataSource: DataSource,
   options: Options,
   services: Services,
-  aiConfigurations: AiConfiguration[] = [],
 ): BaseRoute[] {
   const routes = [
     ...getRootRoutes(options, services),
@@ -189,7 +177,6 @@ export default function makeRoutes(
     ...getApiChartRoutes(dataSource, options, services),
     ...getRelatedRoutes(dataSource, options, services),
     ...getActionRoutes(dataSource, options, services),
-    ...getAiRoutes(options, services, aiConfigurations),
   ];
 
   // Ensure routes and middlewares are loaded in the right order.