[Agent Builder] improve agent overall resilience (#241607)

## Summary

Fix elastic/search-team#11720
Fix elastic/search-team#11700

Improve the agent execution's resilience by catching all potential
errors returned from LLM calls and identifying the ones that we think we
can recover from.

At the moment, the "recoverable" errors we identify are:
- calling a tool which is not available
- calling a tool with invalid parameters
- empty text responses

For each type of error, we have specific logic to "represent" them to
the agent so that it can try to work around it. For example for
non-available tools, we return a tool response with an error stating
that the tool isn't available.

We retry a maximum of 3 times before giving up. Note that a successful
call will reset that counter (so that if we face an error again in a
later cycle, we get our error budget back)

This PR also introduce a new type of error in the inference plugin,
`ContextLengthExceededError`, which is thrown when the LLM call fails
for, context length / too many tokens errors. I did that in the
inference plugin because it makes sense to have that centralized and
available to everyone.

Note that I didn't implement connectivity-related recovery (e.g timeout,
token budget and so on), because this is already implemented in the
`inference` plugin.

Author: pgayvallet

x-pack/platform/packages/shared/ai-infra/inference-common/index.ts

@@ -55,7 +55,7 @@ export {
   type ChatCompletionToolValidationError,
   type ChatCompletionTokenLimitReachedError,
   isToolValidationError,
-  isTokenLimitReachedError,
+  isOutputTokenLimitReachedError,
   isToolNotFoundError,
   type ChatCompleteMetadata,
   type ConnectorTelemetryMetadata,

x-pack/platform/packages/shared/ai-infra/inference-common/src/chat_complete/errors.ts

@@ -12,17 +12,26 @@ import type { UnvalidatedToolCall } from './tools';
  * List of code of error that are specific to the {@link ChatCompleteAPI}
  */
 export enum ChatCompletionErrorCode {
-  TokenLimitReachedError = 'tokenLimitReachedError',
+  ContextLengthExceededError = 'contextLengthExceededError',
+  OutputTokenLimitReachedError = 'outputTokenLimitReachedError',
   ToolNotFoundError = 'toolNotFoundError',
   ToolValidationError = 'toolValidationError',
 }
 
 /**
- * Error thrown if the completion call fails because of a token limit
- * error, e.g. when the context window is higher than the limit
+ * Error thrown if the completion call fails because of a context length error,
+ * e.g. when too many input token or tool definitions are sent.
+ */
+export type ChatCompletionContextLengthExceededError = InferenceTaskError<
+  ChatCompletionErrorCode.ContextLengthExceededError,
+  {}
+>;
+
+/**
+ * Error thrown if the completion call fails because of an output token limit error
  */
 export type ChatCompletionTokenLimitReachedError = InferenceTaskError<
-  ChatCompletionErrorCode.TokenLimitReachedError,
+  ChatCompletionErrorCode.OutputTokenLimitReachedError,
   {
     tokenLimit?: number;
     tokenCount?: number;
@@ -38,6 +47,8 @@ export type ChatCompletionToolNotFoundError = InferenceTaskError<
   {
     /** The name of the tool that got called */
     name: string;
+    /** (unparsed) arguments the tool was called with*/
+    arguments: string;
   }
 >;
 
@@ -58,6 +69,18 @@ export type ChatCompletionToolValidationError = InferenceTaskError<
   }
 >;
 
+/**
+ * Check if an error is a {@link ChatCompletionContextLengthExceededError}
+ */
+export function isContextLengthExceededError(
+  error: Error
+): error is ChatCompletionContextLengthExceededError {
+  return (
+    error instanceof InferenceTaskError &&
+    error.code === ChatCompletionErrorCode.ContextLengthExceededError
+  );
+}
+
 /**
  * Check if an error is a {@link ChatCompletionToolValidationError}
  */
@@ -71,12 +94,12 @@ export function isToolValidationError(error?: Error): error is ChatCompletionToo
 /**
  * Check if an error is a {@link ChatCompletionTokenLimitReachedError}
  */
-export function isTokenLimitReachedError(
+export function isOutputTokenLimitReachedError(
   error: Error
 ): error is ChatCompletionTokenLimitReachedError {
   return (
     error instanceof InferenceTaskError &&
-    error.code === ChatCompletionErrorCode.TokenLimitReachedError
+    error.code === ChatCompletionErrorCode.OutputTokenLimitReachedError
   );
 }
 

x-pack/platform/packages/shared/ai-infra/inference-common/src/chat_complete/index.ts

@@ -77,7 +77,7 @@ export {
   type ChatCompletionToolValidationError,
   type ChatCompletionTokenLimitReachedError,
   isToolValidationError,
-  isTokenLimitReachedError,
+  isOutputTokenLimitReachedError,
   isToolNotFoundError,
 } from './errors';
 

x-pack/platform/packages/shared/ai-infra/inference-langchain/src/chat_model/inference_chat_model.test.ts

@@ -689,16 +689,15 @@ describe('InferenceChatModel', () => {
       });
       chatComplete.mockReturnValue(response);
 
-      const output = await chatModel.stream('Some question');
-
       const allChunks: AIMessageChunk[] = [];
       await expect(async () => {
+        const output = await chatModel.stream('Some question');
         for await (const chunk of output) {
           allChunks.push(chunk);
         }
       }).rejects.toThrowErrorMatchingInlineSnapshot(`"something went wrong"`);
 
-      expect(allChunks.length).toBe(2);
+      expect(allChunks.length).toBe(0);
     });
   });
 

x-pack/platform/packages/shared/ai-infra/inference-langchain/src/chat_model/utils/observable_to_generator.test.ts

@@ -39,6 +39,30 @@ describe('toAsyncIterator', () => {
       }
     }).rejects.toThrowErrorMatchingInlineSnapshot(`"something went wrong"`);
 
-    expect(output).toEqual([1, 2, 3]);
+    // Fail-fast behavior: queued values are discarded when error occurs
+    expect(output).toEqual([]);
+  });
+
+  it('throws an error when the source observable errors while iterator is waiting', async () => {
+    const obs$ = new Observable<number>((subscriber) => {
+      subscriber.next(1);
+      subscriber.next(2);
+
+      // Delay before erroring, so the iterator will be waiting for the next value
+      setTimeout(() => {
+        subscriber.error(new Error('delayed error'));
+      }, 10);
+    });
+
+    const output: number[] = [];
+    const iterator = toAsyncIterator(obs$);
+
+    await expect(async () => {
+      for await (const event of iterator) {
+        output.push(event);
+      }
+    }).rejects.toThrowErrorMatchingInlineSnapshot(`"delayed error"`);
+
+    expect(output).toEqual([1, 2]);
   });
 });

x-pack/platform/packages/shared/ai-infra/inference-langchain/src/chat_model/utils/observable_to_generator.ts

@@ -17,6 +17,7 @@ export function toAsyncIterator<T>(observable: Observable<T>): AsyncIterableIter
 
   const queue: Array<IteratorResult<T>> = [];
   let done = false;
+  let error: any = null;
 
   const subscription = observable.subscribe({
     next(value) {
@@ -28,11 +29,14 @@ export function toAsyncIterator<T>(observable: Observable<T>): AsyncIterableIter
       }
     },
     error(err) {
+      done = true;
+      error = err;
+      // Clear any queued values - we fail fast
+      queue.length = 0;
       if (reject) {
         reject(err);
         reject = null;
-      } else {
-        queue.push(Promise.reject(err) as any); // Queue an error
+        resolve = null;
       }
     },
     complete() {
@@ -49,6 +53,11 @@ export function toAsyncIterator<T>(observable: Observable<T>): AsyncIterableIter
       return this;
     },
     next() {
+      // Check for error first - fail fast
+      if (error !== null) {
+        return Promise.reject(error);
+      }
+
       if (queue.length > 0) {
         return Promise.resolve(queue.shift()!);
       }
@@ -66,9 +75,9 @@ export function toAsyncIterator<T>(observable: Observable<T>): AsyncIterableIter
       subscription.unsubscribe();
       return Promise.resolve({ value: undefined, done: true });
     },
-    throw(error?: any) {
+    throw(err?: any) {
       subscription.unsubscribe();
-      return Promise.reject(error);
+      return Promise.reject(err);
     },
   };
 }

x-pack/platform/packages/shared/onechat/onechat-common/agents/execution_errors.ts

@@ -0,0 +1,49 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0; you may not use this file except in compliance with the Elastic License
+ * 2.0.
+ */
+
+export enum AgentExecutionErrorCode {
+  /** too many input tokens */
+  contextLengthExceeded = 'context_length_exceeded',
+  /** agent called a tool not currently available */
+  toolNotFound = 'tool_not_found',
+  /** agent called a tool with invalid arguments */
+  toolValidationError = 'tool_validation_error',
+  /** agent replied with an empty response */
+  emptyResponse = 'empty_response',
+  /** any uncategorized error */
+  unknownError = 'unknown_error',
+  /** invalid workflow state - should never be surfaced */
+  invalidState = 'invalid_state',
+}
+
+export interface ToolNotFoundErrorMeta {
+  /** name of the tool which was called */
+  toolName: string;
+  /** arguments the tool was called with */
+  toolArgs: string | Record<string, any>;
+}
+
+export interface TooValidationErrorMeta {
+  /** name of the tool which was called */
+  toolName: string;
+  /** arguments the tool was called with */
+  toolArgs: string | Record<string, any>;
+  /** schema validation error, if any */
+  validationError?: string;
+}
+
+interface ExecutionErrorMetaMap {
+  [AgentExecutionErrorCode.toolNotFound]: ToolNotFoundErrorMeta;
+  [AgentExecutionErrorCode.toolValidationError]: TooValidationErrorMeta;
+  [AgentExecutionErrorCode.contextLengthExceeded]: {};
+  [AgentExecutionErrorCode.unknownError]: {};
+  [AgentExecutionErrorCode.invalidState]: {};
+  [AgentExecutionErrorCode.emptyResponse]: {};
+}
+
+export type ExecutionErrorMetaOf<ErrCode extends AgentExecutionErrorCode> =
+  ExecutionErrorMetaMap[ErrCode];

x-pack/platform/packages/shared/onechat/onechat-common/agents/index.ts

@@ -19,3 +19,4 @@ export {
   type ResolvedAgentCapabilities,
   getKibanaDefaultAgentCapabilities,
 } from './capabilities';
+export { AgentExecutionErrorCode } from './execution_errors';

x-pack/platform/packages/shared/onechat/onechat-common/base/errors.ts

@@ -6,6 +6,7 @@
  */
 
 import { ServerSentEventError } from '@kbn/sse-utils';
+import type { AgentExecutionErrorCode, ExecutionErrorMetaOf } from '../agents/execution_errors';
 
 /**
  * Code to identify onechat errors
@@ -16,6 +17,7 @@ export enum OnechatErrorCode {
   toolNotFound = 'toolNotFound',
   agentNotFound = 'agentNotFound',
   conversationNotFound = 'conversationNotFound',
+  agentExecutionError = 'agentExecutionError',
   requestAborted = 'requestAborted',
 }
 
@@ -186,6 +188,34 @@ export const createRequestAbortedError = (
   return new OnechatError(OnechatErrorCode.requestAborted, message, meta ?? {});
 };
 
+/**
+ * Represents an error related to agent execution
+ */
+export type OnechatAgentExecutionError<
+  ErrCode extends AgentExecutionErrorCode = AgentExecutionErrorCode
+> = OnechatError<
+  OnechatErrorCode.agentExecutionError,
+  { errCode: ErrCode } & ExecutionErrorMetaOf<ErrCode>
+>;
+
+/**
+ * Checks if the given error is a {@link OnechatInternalError}
+ */
+export const isAgentExecutionError = (err: unknown): err is OnechatAgentExecutionError => {
+  return isOnechatError(err) && err.code === OnechatErrorCode.agentExecutionError;
+};
+
+export const createAgentExecutionError = <ErrCode extends AgentExecutionErrorCode>(
+  message: string,
+  code: ErrCode,
+  meta: ExecutionErrorMetaOf<ErrCode>
+): OnechatAgentExecutionError<ErrCode> => {
+  return new OnechatError(OnechatErrorCode.agentExecutionError, message, {
+    ...meta,
+    errCode: code,
+  });
+};
+
 /**
  * Global utility exposing all error utilities from a single export.
  */
@@ -195,8 +225,10 @@ export const OnechatErrorUtils = {
   isToolNotFoundError,
   isAgentNotFoundError,
   isConversationNotFoundError,
+  isAgentExecutionError,
   createInternalError,
   createToolNotFoundError,
   createAgentNotFoundError,
   createConversationNotFoundError,
+  createAgentExecutionError,
 };

x-pack/platform/packages/shared/onechat/onechat-genai-utils/langchain/messages.ts

@@ -111,15 +111,17 @@ export const createToolResultMessage = ({
   });
 };
 
-export const createToolCallMessage = (toolCall: ToolCall, message?: string): AIMessage => {
+export const createToolCallMessage = (
+  toolCallOrCalls: ToolCall | ToolCall[],
+  message?: string
+): AIMessage => {
+  const toolCalls = isArray(toolCallOrCalls) ? toolCallOrCalls : [toolCallOrCalls];
   return new AIMessage({
     content: message ?? '',
-    tool_calls: [
-      {
-        id: toolCall.toolCallId,
-        name: toolCall.toolName,
-        args: toolCall.args,
-      },
-    ],
+    tool_calls: toolCalls.map((toolCall) => ({
+      id: toolCall.toolCallId,
+      name: toolCall.toolName,
+      args: toolCall.args,
+    })),
   });
 };