feat(ai-proxy): add UnauthorizedError/TooManyRequestsError and typed AI provider errors

packages/agent-toolkit/src/errors.ts

@@ -3,13 +3,15 @@ export class BusinessError extends Error {
   // INTERNAL USAGES
   public readonly isBusinessError = true;
   public baseBusinessErrorName: string;
+  public httpCode: number;
 
   public readonly data: Record<string, unknown> | undefined;
 
-  constructor(message?: string, data?: Record<string, unknown>, name?: string) {
+  constructor(message?: string, data?: Record<string, unknown>, name?: string, httpCode = 422) {
     super(message);
     this.name = name ?? this.constructor.name;
     this.data = data;
+    this.httpCode = httpCode;
   }
 
   /**
@@ -27,31 +29,43 @@ export class BusinessError extends Error {
 
 export class ValidationError extends BusinessError {
   constructor(message?: string, data?: Record<string, unknown>, name?: string) {
-    super(message, data, name);
+    super(message, data, name, 400);
     this.baseBusinessErrorName = 'ValidationError';
   }
 }
 export class BadRequestError extends BusinessError {
   constructor(message?: string, data?: Record<string, unknown>, name?: string) {
-    super(message, data, name);
+    super(message, data, name, 400);
     this.baseBusinessErrorName = 'BadRequestError';
   }
 }
 export class UnprocessableError extends BusinessError {
   constructor(message?: string, data?: Record<string, unknown>, name?: string) {
-    super(message, data, name);
+    super(message, data, name, 422);
     this.baseBusinessErrorName = 'UnprocessableError';
   }
 }
 export class ForbiddenError extends BusinessError {
   constructor(message?: string, data?: Record<string, unknown>, name?: string) {
-    super(message, data, name);
+    super(message, data, name, 403);
     this.baseBusinessErrorName = 'ForbiddenError';
   }
 }
 export class NotFoundError extends BusinessError {
   constructor(message?: string, data?: Record<string, unknown>, name?: string) {
-    super(message, data, name);
+    super(message, data, name, 404);
     this.baseBusinessErrorName = 'NotFoundError';
   }
 }
+export class UnauthorizedError extends BusinessError {
+  constructor(message?: string, data?: Record<string, unknown>, name?: string) {
+    super(message, data, name, 401);
+    this.baseBusinessErrorName = 'UnauthorizedError';
+  }
+}
+export class TooManyRequestsError extends BusinessError {
+  constructor(message?: string, data?: Record<string, unknown>, name?: string) {
+    super(message, data, name, 429);
+    this.baseBusinessErrorName = 'TooManyRequestsError';
+  }
+}

packages/agent-toolkit/test/errors.test.ts

@@ -4,6 +4,8 @@ import {
   BusinessError,
   ForbiddenError,
   NotFoundError,
+  TooManyRequestsError,
+  UnauthorizedError,
   UnprocessableError,
   ValidationError,
 } from '../src/errors';
@@ -67,6 +69,8 @@ describe('errors', () => {
       { ErrorClass: UnprocessableError, errorName: 'UnprocessableError' },
       { ErrorClass: ForbiddenError, errorName: 'ForbiddenError' },
       { ErrorClass: NotFoundError, errorName: 'NotFoundError' },
+      { ErrorClass: UnauthorizedError, errorName: 'UnauthorizedError' },
+      { ErrorClass: TooManyRequestsError, errorName: 'TooManyRequestsError' },
     ])('$errorName should have the correct baseBusinessErrorName', ({ ErrorClass, errorName }) => {
       const error = new ErrorClass('test');
       expect(error.baseBusinessErrorName).toEqual(errorName);

packages/agent/src/routes/system/error-handling.ts

@@ -6,6 +6,8 @@ import {
   BusinessError,
   ForbiddenError,
   NotFoundError,
+  TooManyRequestsError,
+  UnauthorizedError,
   UnprocessableError,
   ValidationError,
 } from '@forestadmin/datasource-toolkit';
@@ -60,6 +62,10 @@ export default class ErrorHandling extends BaseRoute {
       case BusinessError.isOfType(error, BadRequestError):
         return HttpCode.BadRequest;
 
+      case error instanceof UnauthorizedError:
+      case BusinessError.isOfType(error, UnauthorizedError):
+        return HttpCode.Unauthorized;
+
       case error instanceof ForbiddenError:
       case BusinessError.isOfType(error, ForbiddenError):
         return HttpCode.Forbidden;
@@ -68,6 +74,10 @@ export default class ErrorHandling extends BaseRoute {
       case BusinessError.isOfType(error, NotFoundError):
         return HttpCode.NotFound;
 
+      case error instanceof TooManyRequestsError:
+      case BusinessError.isOfType(error, TooManyRequestsError):
+        return HttpCode.TooManyRequests;
+
       case error instanceof UnprocessableError:
       case BusinessError.isOfType(error, UnprocessableError):
       case error instanceof BusinessError:

packages/agent/src/types.ts

@@ -52,12 +52,14 @@ export type HttpCallback = (req: IncomingMessage, res: ServerResponse, next?: ()
 
 export enum HttpCode {
   BadRequest = 400,
+  Unauthorized = 401,
   Forbidden = 403,
+  NotFound = 404,
+  Unprocessable = 422,
+  TooManyRequests = 429,
   InternalServerError = 500,
   NoContent = 204,
-  NotFound = 404,
   Ok = 200,
-  Unprocessable = 422,
 }
 
 export enum RouteType {

packages/agent/test/routes/system/error-handling.test.ts

@@ -4,6 +4,8 @@ import {
   BadRequestError,
   ForbiddenError,
   NotFoundError,
+  TooManyRequestsError,
+  UnauthorizedError,
   UnprocessableError,
   ValidationError,
 } from '@forestadmin/datasource-toolkit';
@@ -133,6 +135,32 @@ describe('ErrorHandling', () => {
       expect(console.error).not.toHaveBeenCalled();
     });
 
+    test('it should set the status and body for unauthorized errors', async () => {
+      const context = createMockContext();
+      const next = jest.fn().mockRejectedValue(new UnauthorizedError('unauthorized'));
+
+      await expect(handleError.call(route, context, next)).rejects.toThrow();
+
+      expect(context.response.status).toStrictEqual(HttpCode.Unauthorized);
+      expect(context.response.body).toStrictEqual({
+        errors: [{ detail: 'unauthorized', name: 'UnauthorizedError', status: 401 }],
+      });
+      expect(console.error).not.toHaveBeenCalled();
+    });
+
+    test('it should set the status and body for too many requests errors', async () => {
+      const context = createMockContext();
+      const next = jest.fn().mockRejectedValue(new TooManyRequestsError('rate limited'));
+
+      await expect(handleError.call(route, context, next)).rejects.toThrow();
+
+      expect(context.response.status).toStrictEqual(HttpCode.TooManyRequests);
+      expect(context.response.body).toStrictEqual({
+        errors: [{ detail: 'rate limited', name: 'TooManyRequestsError', status: 429 }],
+      });
+      expect(console.error).not.toHaveBeenCalled();
+    });
+
     test('it should set the status and body for other errors and prevent information leak', async () => {
       const context = createMockContext();
       const next = jest
@@ -258,6 +286,8 @@ describe('ErrorHandling', () => {
       { Error: UnprocessableError, errorName: 'UnprocessableError' },
       { Error: ForbiddenError, errorName: 'ForbiddenError' },
       { Error: NotFoundError, errorName: 'NotFoundError' },
+      { Error: UnauthorizedError, errorName: 'UnauthorizedError' },
+      { Error: TooManyRequestsError, errorName: 'TooManyRequestsError' },
     ])('should have the right baseBusinessErrorName', ({ Error, errorName }) => {
       const extendedError = new Error('message');
 

packages/ai-proxy/src/errors.ts

@@ -1,23 +1,17 @@
 /**
- * All custom AI errors extend HTTP-status error classes (BadRequestError, NotFoundError,
- * UnprocessableError) from datasource-toolkit. This allows the agent's error middleware
- * to map them to their natural HTTP status codes automatically.
+ * All custom AI errors extend their matching HTTP-status error class from datasource-toolkit.
+ * This allows the agent's error middleware to map them to their natural HTTP status codes.
  */
 
 // eslint-disable-next-line max-classes-per-file
 import {
   BadRequestError,
   NotFoundError,
+  TooManyRequestsError,
+  UnauthorizedError,
   UnprocessableError,
 } from '@forestadmin/datasource-toolkit';
 
-export class AIError extends UnprocessableError {
-  constructor(message: string) {
-    super(message);
-    this.name = 'AIError';
-  }
-}
-
 export class AIBadRequestError extends BadRequestError {
   constructor(message: string) {
     super(message);
@@ -41,24 +35,50 @@ export class AINotFoundError extends NotFoundError {
   }
 }
 
-export class AIUnprocessableError extends UnprocessableError {
+export class AIProviderError extends UnprocessableError {
+  readonly provider: string;
   readonly cause?: Error;
 
-  constructor(message: string, options?: { cause?: Error }) {
-    super(message);
-    this.name = 'AIUnprocessableError';
+  constructor(provider: string, options?: { cause?: Error }) {
+    super(`Error while calling ${provider}: ${options?.cause?.message ?? 'unknown'}`);
+    this.name = 'AIProviderError';
+    this.provider = provider;
+    if (options?.cause) this.cause = options.cause;
+  }
+}
+
+export class AITooManyRequestsError extends TooManyRequestsError {
+  readonly provider: string;
+  readonly cause?: Error;
+
+  constructor(provider: string, options?: { cause?: Error }) {
+    super(`${provider} rate limit exceeded`);
+    this.name = 'AITooManyRequestsError';
+    this.provider = provider;
+    if (options?.cause) this.cause = options.cause;
+  }
+}
+
+export class AIUnauthorizedError extends UnauthorizedError {
+  readonly provider: string;
+  readonly cause?: Error;
+
+  constructor(provider: string, options?: { cause?: Error }) {
+    super(`${provider} authentication failed: check your API key configuration`);
+    this.name = 'AIUnauthorizedError';
+    this.provider = provider;
     if (options?.cause) this.cause = options.cause;
   }
 }
 
-export class AINotConfiguredError extends AIError {
+export class AINotConfiguredError extends UnprocessableError {
   constructor(message = 'AI is not configured') {
     super(message);
     this.name = 'AINotConfiguredError';
   }
 }
 
-export class AIToolUnprocessableError extends AIUnprocessableError {
+export class AIToolUnprocessableError extends UnprocessableError {
   constructor(message: string) {
     super(message);
     this.name = 'AIToolUnprocessableError';
@@ -72,7 +92,7 @@ export class AIToolNotFoundError extends AINotFoundError {
   }
 }
 
-export class McpError extends AIError {
+export class McpError extends UnprocessableError {
   constructor(message: string) {
     super(message);
     this.name = 'McpError';

packages/ai-proxy/src/provider-dispatcher.ts

@@ -4,12 +4,19 @@ import type { RemoteTools } from './remote-tools';
 import type { DispatchBody } from './schemas/route';
 import type { AIMessage, BaseMessageLike } from '@langchain/core/messages';
 
+import { BusinessError, UnprocessableError } from '@forestadmin/datasource-toolkit';
 import { ChatAnthropic } from '@langchain/anthropic';
 import { convertToOpenAIFunction } from '@langchain/core/utils/function_calling';
 import { ChatOpenAI } from '@langchain/openai';
 
 import AnthropicAdapter from './anthropic-adapter';
-import { AIBadRequestError, AINotConfiguredError, AIUnprocessableError } from './errors';
+import {
+  AIBadRequestError,
+  AINotConfiguredError,
+  AIProviderError,
+  AITooManyRequestsError,
+  AIUnauthorizedError,
+} from './errors';
 import { LangChainAdapter } from './langchain-adapter';
 
 // Re-export types for consumers
@@ -100,7 +107,7 @@ export default class ProviderDispatcher {
     const rawResponse = response.additional_kwargs.__raw_response as ChatCompletionResponse;
 
     if (!rawResponse) {
-      throw new AIUnprocessableError(
+      throw new UnprocessableError(
         'OpenAI response missing raw response data. This may indicate an API change.',
       );
     }
@@ -138,42 +145,19 @@ export default class ProviderDispatcher {
     return LangChainAdapter.convertResponse(response, this.modelName);
   }
 
-  /**
-   * Wraps provider errors into AI-specific error types.
-   *
-   * TODO: Currently all provider errors are wrapped as AIUnprocessableError,
-   * losing the original HTTP semantics (429 rate limit, 401 auth failure).
-   * To fix this properly we need to:
-   * 1. Add UnauthorizedError and TooManyRequestsError to datasource-toolkit
-   * 2. Add corresponding cases in the agent's error-handling middleware
-   * 3. Create AIProviderError, AIRateLimitError, AIAuthenticationError in ai-proxy
-   *    with baseBusinessErrorName overrides for correct HTTP status mapping
-   */
   private static wrapProviderError(error: unknown, providerName: string): Error {
-    if (error instanceof AIUnprocessableError) return error;
-    if (error instanceof AIBadRequestError) return error;
+    if (error instanceof BusinessError) return error;
 
     if (!(error instanceof Error)) {
-      return new AIUnprocessableError(`Error while calling ${providerName}: ${String(error)}`);
+      return new AIProviderError(providerName, { cause: new Error(String(error)) });
     }
 
     const { status } = error as Error & { status?: number };
 
-    if (status === 429) {
-      return new AIUnprocessableError(`${providerName} rate limit exceeded: ${error.message}`, {
-        cause: error,
-      });
-    }
-
-    if (status === 401) {
-      return new AIUnprocessableError(`${providerName} authentication failed: ${error.message}`, {
-        cause: error,
-      });
-    }
+    if (status === 429) return new AITooManyRequestsError(providerName, { cause: error });
+    if (status === 401) return new AIUnauthorizedError(providerName, { cause: error });
 
-    return new AIUnprocessableError(`Error while calling ${providerName}: ${error.message}`, {
-      cause: error,
-    });
+    return new AIProviderError(providerName, { cause: error });
   }
 
   private enrichToolDefinitions(tools?: ChatCompletionTool[]): ChatCompletionTool[] | undefined {