refactor(agent): decouple ai-proxy via factory injection pattern

packages/_example/src/forest/agent.ts

@@ -2,6 +2,7 @@ import type { Schema } from './typings';
 import type { AgentOptions } from '@forestadmin/agent';
 
 import { createAgent } from '@forestadmin/agent';
+import { createAiProvider } from '@forestadmin/ai-proxy';
 import { createMongoDataSource } from '@forestadmin/datasource-mongo';
 import { createMongooseDataSource } from '@forestadmin/datasource-mongoose';
 import { createSequelizeDataSource } from '@forestadmin/datasource-sequelize';
@@ -94,5 +95,13 @@ export default function makeAgent() {
     .customizeCollection('post', customizePost)
     .customizeCollection('comment', customizeComment)
     .customizeCollection('review', customizeReview)
-    .customizeCollection('sales', customizeSales);
+    .customizeCollection('sales', customizeSales)
+    .addAi(
+      createAiProvider({
+        model: 'gpt-4o',
+        provider: 'openai',
+        name: 'test',
+        apiKey: process.env.OPENAI_API_KEY,
+      }),
+    );
 }

packages/agent-toolkit/src/index.ts

@@ -1 +1 @@
-export {};
+export type { AiProviderDefinition, AiProviderMeta, AiRouter } from './interfaces/ai';

packages/agent-toolkit/src/interfaces/ai.ts

@@ -0,0 +1,35 @@
+/**
+ * Inlined from datasource-toolkit/factory.ts to avoid a dependency on a lower-level package.
+ * Must stay structurally compatible with datasource-toolkit's Logger & LoggerLevel.
+ * If those types change, update these copies accordingly.
+ */
+export type LoggerLevel = 'Debug' | 'Info' | 'Warn' | 'Error';
+export type Logger = (level: LoggerLevel, message: string, error?: Error) => void;
+
+/** Metadata describing a configured AI provider, used in schema reporting and logging. */
+export interface AiProviderMeta {
+  name: string;
+  provider: string;
+  model: string;
+}
+
+export interface AiRouter {
+  /**
+   * Route a request to the AI proxy.
+   *
+   * Implementations should throw BusinessError subclasses (BadRequestError, NotFoundError,
+   * UnprocessableError) for proper HTTP status mapping by the agent's error middleware.
+   */
+  route(args: {
+    route: string;
+    body?: unknown;
+    query?: Record<string, string | string[] | undefined>;
+    mcpServerConfigs?: unknown;
+    headers?: Record<string, string | string[] | undefined>;
+  }): Promise<unknown>;
+}
+
+export interface AiProviderDefinition {
+  providers: AiProviderMeta[];
+  init(logger: Logger): AiRouter;
+}

packages/agent/package.json

@@ -13,7 +13,7 @@
   },
   "dependencies": {
     "@fast-csv/format": "^4.3.5",
-    "@forestadmin/ai-proxy": "1.4.2",
+    "@forestadmin/agent-toolkit": "1.0.0",
     "@forestadmin/datasource-customizer": "1.67.3",
     "@forestadmin/datasource-toolkit": "1.50.1",
     "@forestadmin/forestadmin-client": "1.37.11",

packages/agent/src/agent.ts

@@ -1,11 +1,7 @@
 /* eslint-disable @typescript-eslint/no-explicit-any */
 import type { ForestAdminHttpDriverServices } from './services';
-import type {
-  AgentOptions,
-  AgentOptionsWithDefaults,
-  AiConfiguration,
-  HttpCallback,
-} from './types';
+import type { AgentOptions, AgentOptionsWithDefaults, HttpCallback } from './types';
+import type { AiProviderDefinition } from '@forestadmin/agent-toolkit';
 import type {
   CollectionCustomizer,
   DataSourceChartDefinition,
@@ -47,7 +43,7 @@ export default class Agent<S extends TSchema = TSchema> extends FrameworkMounter
   protected nocodeCustomizer: DataSourceCustomizer<S>;
   protected customizationService: CustomizationService;
   protected schemaGenerator: SchemaGenerator;
-  protected aiConfigurations: AiConfiguration[] = [];
+  protected aiProvider: AiProviderDefinition | null = null;
 
   /** Whether MCP server should be mounted */
   private mcpEnabled = false;
@@ -207,6 +203,7 @@ export default class Agent<S extends TSchema = TSchema> extends FrameworkMounter
    * Enable MCP (Model Context Protocol) server support.
    * This allows AI assistants to interact with your Forest Admin data.
    *
+   * @see {@link https://docs.forestadmin.com/developer-guide-agents-nodejs/agent-customization/ai/mcp-server}
    * @example
    * agent.mountAiMcpServer();
    */
@@ -222,42 +219,51 @@ export default class Agent<S extends TSchema = TSchema> extends FrameworkMounter
    * All AI requests from Forest Admin are forwarded to your agent and processed locally.
    * Your data and API keys never transit through Forest Admin servers, ensuring full privacy.
    *
-   * @param configuration - The AI provider configuration
-   * @param configuration.name - A unique name to identify this AI configuration
-   * @param configuration.provider - The AI provider to use ('openai')
-   * @param configuration.apiKey - Your API key for the chosen provider
-   * @param configuration.model - The model to use (e.g., 'gpt-4o')
+   * Requires the `@forestadmin/ai-proxy` package to be installed:
+   * ```bash
+   * npm install @forestadmin/ai-proxy
+   * ```
+   *
+   * @see {@link https://docs.forestadmin.com/developer-guide-agents-nodejs/agent-customization/ai/self-hosted-ai}
+   * @param provider - An AI provider definition created via `createAiProvider` from `@forestadmin/ai-proxy`
    * @returns The agent instance for chaining
    * @throws Error if addAi is called more than once
    *
    * @example
-   * agent.addAi({
+   * import { createAiProvider } from '@forestadmin/ai-proxy';
+   *
+   * agent.addAi(createAiProvider({
    *   name: 'assistant',
    *   provider: 'openai',
    *   apiKey: process.env.OPENAI_API_KEY,
    *   model: 'gpt-4o',
-   * });
+   * }));
    */
-  addAi(configuration: AiConfiguration): this {
-    if (this.aiConfigurations.length > 0) {
+  addAi(provider: AiProviderDefinition): this {
+    if (this.aiProvider) {
       throw new Error(
         'addAi can only be called once. Multiple AI configurations are not supported yet.',
       );
     }
 
-    this.options.logger(
-      'Warn',
-      `AI configuration added with model '${configuration.model}'. ` +
-        'Make sure to test Forest Admin AI features thoroughly to ensure compatibility.',
-    );
+    this.aiProvider = provider;
 
-    this.aiConfigurations.push(configuration);
+    for (const p of provider.providers) {
+      this.options.logger(
+        'Warn',
+        `AI configuration added with model '${p.model}'. ` +
+          'Make sure to test Forest Admin AI features thoroughly to ensure compatibility.',
+      );
+    }
 
     return this;
   }
 
   protected getRoutes(dataSource: DataSource, services: ForestAdminHttpDriverServices) {
-    return makeRoutes(dataSource, this.options, services, this.aiConfigurations);
+    // init() is called on every start/restart to recreate routing state with a fresh Router.
+    const aiRouter = this.aiProvider?.init(this.options.logger) ?? null;
+
+    return makeRoutes(dataSource, this.options, services, aiRouter);
   }
 
   /**
@@ -380,9 +386,10 @@ export default class Agent<S extends TSchema = TSchema> extends FrameworkMounter
     let schema: Pick<ForestSchema, 'collections'>;
 
     // Get the AI configurations for schema metadata
+    const aiMeta = this.aiProvider?.providers ?? [];
     const { meta } = SchemaGenerator.buildMetadata(
       this.customizationService.buildFeatures(),
-      this.aiConfigurations,
+      aiMeta,
     );
 
     // When using experimental no-code features even in production we need to build a new schema

packages/agent/src/index.ts

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

packages/agent/src/routes/ai/ai-proxy.ts

@@ -1,74 +1,40 @@
 import type { ForestAdminHttpDriverServices } from '../../services';
-import type { AgentOptionsWithDefaults, AiConfiguration } from '../../types';
+import type { AgentOptionsWithDefaults } from '../../types';
+import type { AiRouter } from '@forestadmin/agent-toolkit';
 import type KoaRouter from '@koa/router';
 import type { Context } from 'koa';
 
-import {
-  AIBadRequestError,
-  AIError,
-  AINotConfiguredError,
-  AINotFoundError,
-  Router as AiProxyRouter,
-  extractMcpOauthTokensFromHeaders,
-  injectOauthTokens,
-} from '@forestadmin/ai-proxy';
-import {
-  BadRequestError,
-  NotFoundError,
-  UnprocessableError,
-} from '@forestadmin/datasource-toolkit';
-
 import { HttpCode, RouteType } from '../../types';
 import BaseRoute from '../base-route';
 
 export default class AiProxyRoute extends BaseRoute {
   readonly type = RouteType.PrivateRoute;
-  private readonly aiProxyRouter: AiProxyRouter;
+  private readonly aiRouter: AiRouter;
 
   constructor(
     services: ForestAdminHttpDriverServices,
     options: AgentOptionsWithDefaults,
-    aiConfigurations: AiConfiguration[],
+    aiRouter: AiRouter,
   ) {
     super(services, options);
-    this.aiProxyRouter = new AiProxyRouter({
-      aiConfigurations,
-      logger: this.options.logger,
-    });
+    this.aiRouter = aiRouter;
   }
 
   setupRoutes(router: KoaRouter): void {
     router.post('/_internal/ai-proxy/:route', this.handleAiProxy.bind(this));
   }
 
   private async handleAiProxy(context: Context): Promise<void> {
-    try {
-      const tokensByMcpServerName = extractMcpOauthTokensFromHeaders(context.request.headers);
-
-      const mcpConfigs =
-        await this.options.forestAdminClient.mcpServerConfigService.getConfiguration();
-
-      context.response.body = await this.aiProxyRouter.route({
-        route: context.params.route,
-        body: context.request.body,
-        query: context.query,
-        mcpConfigs: injectOauthTokens({ mcpConfigs, tokensByMcpServerName }),
-      });
-      context.response.status = HttpCode.Ok;
-    } catch (error) {
-      if (error instanceof AIError) {
-        this.options.logger('Error', `AI proxy error: ${error.message}`, error);
-
-        if (error instanceof AINotConfiguredError) {
-          throw new UnprocessableError('AI is not configured. Please call addAi() on your agent.');
-        }
-
-        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;
-    }
+    const mcpServerConfigs =
+      await this.options.forestAdminClient.mcpServerConfigService.getConfiguration();
+
+    context.response.body = await this.aiRouter.route({
+      route: context.params.route,
+      body: context.request.body,
+      query: context.query,
+      headers: context.request.headers,
+      mcpServerConfigs,
+    });
+    context.response.status = HttpCode.Ok;
   }
 }

packages/agent/src/routes/index.ts

@@ -1,6 +1,7 @@
 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 { AiRouter } from '@forestadmin/agent-toolkit';
 import type { DataSource } from '@forestadmin/datasource-toolkit';
 
 import CollectionApiChartRoute from './access/api-chart-collection';
@@ -165,21 +166,17 @@ function getActionRoutes(
   return routes;
 }
 
-function getAiRoutes(
-  options: Options,
-  services: Services,
-  aiConfigurations: AiConfiguration[],
-): BaseRoute[] {
-  if (aiConfigurations.length === 0) return [];
+function getAiRoutes(options: Options, services: Services, aiRouter: AiRouter | null): BaseRoute[] {
+  if (!aiRouter) return [];
 
-  return [new AiProxyRoute(services, options, aiConfigurations)];
+  return [new AiProxyRoute(services, options, aiRouter)];
 }
 
 export default function makeRoutes(
   dataSource: DataSource,
   options: Options,
   services: Services,
-  aiConfigurations: AiConfiguration[] = [],
+  aiRouter: AiRouter | null = null,
 ): BaseRoute[] {
   const routes = [
     ...getRootRoutes(options, services),
@@ -189,7 +186,7 @@ export default function makeRoutes(
     ...getApiChartRoutes(dataSource, options, services),
     ...getRelatedRoutes(dataSource, options, services),
     ...getActionRoutes(dataSource, options, services),
-    ...getAiRoutes(options, services, aiConfigurations),
+    ...getAiRoutes(options, services, aiRouter),
   ];
 
   // Ensure routes and middlewares are loaded in the right order.

packages/agent/src/types.ts

@@ -1,9 +1,9 @@
-import type { AiConfiguration, AiProvider } from '@forestadmin/ai-proxy';
+import type { AiProviderDefinition } from '@forestadmin/agent-toolkit';
 import type { CompositeId, Logger, LoggerLevel } from '@forestadmin/datasource-toolkit';
 import type { ForestAdminClient } from '@forestadmin/forestadmin-client';
 import type { IncomingMessage, ServerResponse } from 'http';
 
-export type { AiConfiguration, AiProvider };
+export type { AiProviderDefinition };
 
 /** Options to configure behavior of an agent's forestadmin driver */
 export type AgentOptions = {

packages/agent/src/utils/forest-schema/generator.ts

@@ -1,4 +1,5 @@
-import type { AgentOptionsWithDefaults, AiConfiguration } from '../../types';
+import type { AgentOptionsWithDefaults } from '../../types';
+import type { AiProviderMeta } from '@forestadmin/agent-toolkit';
 import type { DataSource } from '@forestadmin/datasource-toolkit';
 import type { ForestSchema } from '@forestadmin/forestadmin-client';
 
@@ -23,7 +24,7 @@ export default class SchemaGenerator {
 
   static buildMetadata(
     features: Record<string, string> | null,
-    aiConfigurations: AiConfiguration[] = [],
+    aiProviders: AiProviderMeta[] = [],
   ): Pick<ForestSchema, 'meta'> {
     const { version } = require('../../../package.json'); // eslint-disable-line @typescript-eslint/no-var-requires,global-require
 
@@ -33,8 +34,8 @@ export default class SchemaGenerator {
         liana_version: version,
         liana_features: features,
         ai_llms:
-          aiConfigurations.length > 0
-            ? aiConfigurations.map(c => ({ name: c.name, provider: c.provider }))
+          aiProviders.length > 0
+            ? aiProviders.map(p => ({ name: p.name, provider: p.provider, model: p.model }))
             : null,
         stack: {
           engine: 'nodejs',