@librechat/agents 2.4.321 → 3.0.0-rc1

package/dist/types/utils/events.d.ts

@@ -0,0 +1,6 @@
+import type { RunnableConfig } from '@langchain/core/runnables';
+/**
+ * Safely dispatches a custom event and properly awaits it to avoid
+ * race conditions where events are dispatched after run cleanup.
+ */
+export declare function safeDispatchCustomEvent(event: string, payload: unknown, config?: RunnableConfig): Promise<void>;

package/dist/types/utils/title.d.ts

@@ -1,3 +1,4 @@
 import type { Runnable } from '@langchain/core/runnables';
-import * as t from '@/types';
+import type * as t from '@/types';
 export declare const createTitleRunnable: (model: t.ChatModelInstance, _titlePrompt?: string) => Promise<Runnable>;
+export declare const createCompletionTitleRunnable: (model: t.ChatModelInstance, titlePrompt?: string) => Promise<Runnable>;

package/dist/types/utils/tokens.d.ts

@@ -1,3 +1,27 @@
 import type { BaseMessage } from '@langchain/core/messages';
 export declare function getTokenCountForMessage(message: BaseMessage, getTokenCount: (text: string) => number): number;
+/**
+ * Creates a singleton token counter function that reuses the same encoder instance.
+ * This avoids creating multiple function closures and prevents potential memory issues.
+ */
 export declare const createTokenCounter: () => Promise<(message: BaseMessage) => number>;
+/**
+ * Utility to manage the token encoder lifecycle explicitly.
+ * Useful for applications that need fine-grained control over resource management.
+ */
+export declare const TokenEncoderManager: {
+    /**
+     * Pre-initializes the encoder. This can be called during app startup
+     * to avoid lazy loading delays later.
+     */
+    initialize(): Promise<void>;
+    /**
+     * Clears the cached encoder and token counter.
+     * Useful for testing or when you need to force a fresh reload.
+     */
+    reset(): void;
+    /**
+     * Checks if the encoder has been initialized.
+     */
+    isInitialized(): boolean;
+};

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@librechat/agents",
-  "version": "2.4.321",
+  "version": "3.0.00-rc1",
   "main": "./dist/cjs/main.cjs",
   "module": "./dist/esm/main.mjs",
   "types": "./dist/types/index.d.ts",
@@ -47,14 +47,19 @@
     "image": "node -r dotenv/config --loader ./tsconfig-paths-bootstrap.mjs --experimental-specifier-resolution=node ./src/scripts/image.ts --provider 'google' --name 'Jo' --location 'New York, NY'",
     "code_exec_files": "node -r dotenv/config --loader ./tsconfig-paths-bootstrap.mjs --experimental-specifier-resolution=node ./src/scripts/code_exec_files.ts --provider 'openAI' --name 'Jo' --location 'New York, NY'",
     "code_exec_simple": "node -r dotenv/config --loader ./tsconfig-paths-bootstrap.mjs --experimental-specifier-resolution=node ./src/scripts/code_exec_simple.ts --provider 'google' --name 'Jo' --location 'New York, NY'",
-    "simple": "node -r dotenv/config --loader ./tsconfig-paths-bootstrap.mjs --experimental-specifier-resolution=node ./src/scripts/simple.ts --provider 'xai' --name 'Jo' --location 'New York, NY'",
+    "simple": "node -r dotenv/config --loader ./tsconfig-paths-bootstrap.mjs --experimental-specifier-resolution=node ./src/scripts/simple.ts --provider 'anthropic' --name 'Jo' --location 'New York, NY'",
     "caching": "node -r dotenv/config --loader ./tsconfig-paths-bootstrap.mjs --experimental-specifier-resolution=node ./src/scripts/caching.ts --name 'Jo' --location 'New York, NY'",
     "thinking": "node -r dotenv/config --loader ./tsconfig-paths-bootstrap.mjs --experimental-specifier-resolution=node ./src/scripts/thinking.ts --name 'Jo' --location 'New York, NY'",
     "memory": "node -r dotenv/config --loader ./tsconfig-paths-bootstrap.mjs --experimental-specifier-resolution=node ./src/scripts/memory.ts --provider 'openAI' --name 'Jo' --location 'New York, NY'",
-    "tool-test": "node -r dotenv/config --loader ./tsconfig-paths-bootstrap.mjs --experimental-specifier-resolution=node ./src/scripts/tools.ts --provider 'alibaba' --name 'Jo' --location 'New York, NY'",
+    "tool-test": "node -r dotenv/config --loader ./tsconfig-paths-bootstrap.mjs --experimental-specifier-resolution=node ./src/scripts/tools.ts --provider 'bedrock' --name 'Jo' --location 'New York, NY'",
     "search": "node -r dotenv/config --loader ./tsconfig-paths-bootstrap.mjs --experimental-specifier-resolution=node ./src/scripts/search.ts --provider 'anthropic' --name 'Jo' --location 'New York, NY'",
+    "ant_web_search": "node -r dotenv/config --loader ./tsconfig-paths-bootstrap.mjs --experimental-specifier-resolution=node ./src/scripts/ant_web_search.ts --name 'Jo' --location 'New York, NY'",
     "abort": "node -r dotenv/config --loader ./tsconfig-paths-bootstrap.mjs --experimental-specifier-resolution=node ./src/scripts/abort.ts --provider 'openAI' --name 'Jo' --location 'New York, NY'",
     "start:cli2": "node -r dotenv/config --loader ./tsconfig-paths-bootstrap.mjs --experimental-specifier-resolution=node ./src/scripts/cli2.ts --provider 'anthropic' --name 'Jo' --location 'New York, NY'",
+    "multi-agent-test": "node -r dotenv/config --loader ./tsconfig-paths-bootstrap.mjs --experimental-specifier-resolution=node ./src/scripts/multi-agent-test.ts",
+    "multi-agent-parallel": "node -r dotenv/config --loader ./tsconfig-paths-bootstrap.mjs --experimental-specifier-resolution=node ./src/scripts/multi-agent-parallel.ts",
+    "multi-agent-sequence": "node -r dotenv/config --loader ./tsconfig-paths-bootstrap.mjs --experimental-specifier-resolution=node ./src/scripts/multi-agent-sequence.ts",
+    "multi-agent-conditional": "node -r dotenv/config --loader ./tsconfig-paths-bootstrap.mjs --experimental-specifier-resolution=node ./src/scripts/multi-agent-conditional.ts",
     "script2": "node -r dotenv/config --loader ./tsconfig-paths-bootstrap.mjs --experimental-specifier-resolution=node ./src/proto/example_test.ts",
     "script3": "node -r dotenv/config --loader ./tsconfig-paths-bootstrap.mjs --experimental-specifier-resolution=node ./src/proto/example_test_anthropic.ts",
     "script4": "node -r dotenv/config --loader ./tsconfig-paths-bootstrap.mjs --experimental-specifier-resolution=node ./src/scripts/cli4.ts --name 'Jo' --location 'New York, NY'",
@@ -65,35 +70,46 @@
     "start:collab": "node --loader ./tsconfig-paths-bootstrap.mjs --experimental-specifier-resolution=node ./src/scripts/main.ts",
     "start:collab5": "node --loader ./tsconfig-paths-bootstrap.mjs --experimental-specifier-resolution=node ./src/scripts/collab_design_v5.ts",
     "start:dev": "node --loader ./tsconfig-paths-bootstrap.mjs --experimental-specifier-resolution=node ./src/main.ts",
+    "supervised": "node -r dotenv/config --loader ./tsconfig-paths-bootstrap.mjs --experimental-specifier-resolution=node ./src/scripts/supervised.ts --provider anthropic --name Jo --location \"New York, NY\"",
     "test": "jest",
+    "test:memory": "NODE_OPTIONS='--expose-gc' npx jest src/specs/title.memory-leak.test.ts",
+    "test:all": "npm test -- --testPathIgnorePatterns=title.memory-leak.test.ts && npm run test:memory",
     "reinstall": "npm run clean && npm ci && rm -rf ./dist && npm run build",
     "re": "bun run clean && bun install && rm -rf ./dist && turbo build",
     "lint": "eslint \"{,!(node_modules|venv)/**/}*.{js,jsx,ts,tsx}\" --fix",
     "format": "prettier --write ."
   },
+  "overrides": {
+    "@langchain/openai": "0.5.18",
+    "@browserbasehq/stagehand": {
+      "openai": "$openai"
+    }
+  },
   "dependencies": {
-    "@langchain/anthropic": "^0.3.20",
-    "@langchain/aws": "0.1.8",
-    "@langchain/community": "^0.3.42",
-    "@langchain/core": "^0.3.55",
-    "@langchain/deepseek": "^0.0.1",
-    "@langchain/google-genai": "^0.2.8",
-    "@langchain/google-vertexai": "^0.2.8",
-    "@langchain/langgraph": "^0.2.72",
-    "@langchain/mistralai": "^0.2.0",
-    "@langchain/ollama": "^0.2.0",
-    "@langchain/openai": "^0.5.10",
-    "@langchain/xai": "^0.0.2",
+    "@langchain/anthropic": "^0.3.26",
+    "@langchain/aws": "^0.1.12",
+    "@langchain/community": "^0.3.47",
+    "@langchain/core": "^0.3.72",
+    "@langchain/deepseek": "^0.0.2",
+    "@langchain/google-genai": "^0.2.13",
+    "@langchain/google-vertexai": "^0.2.13",
+    "@langchain/langgraph": "^0.4.9",
+    "@langchain/mistralai": "^0.2.1",
+    "@langchain/ollama": "^0.2.3",
+    "@langchain/openai": "0.5.18",
+    "@langchain/xai": "^0.0.3",
     "cheerio": "^1.0.0",
     "dotenv": "^16.4.7",
     "https-proxy-agent": "^7.0.6",
-    "nanoid": "^3.3.7"
+    "nanoid": "^3.3.7",
+    "openai": "^5.10.1"
   },
   "imports": {
     "@/*": "./src/*",
     "~/*": "./*"
   },
   "devDependencies": {
+    "@anthropic-ai/vertex-sdk": "^0.12.0",
     "@eslint/compat": "^1.2.7",
     "@rollup/plugin-alias": "^5.1.0",
     "@rollup/plugin-commonjs": "^28.0.3",
@@ -107,7 +123,7 @@
     "@types/yargs-parser": "^21.0.3",
     "@typescript-eslint/eslint-plugin": "^8.24.0",
     "@typescript-eslint/parser": "^8.24.0",
-    "eslint": "^9.22.0",
+    "eslint": "^9.31.0",
     "eslint-import-resolver-typescript": "^3.7.0",
     "eslint-plugin-import": "^2.31.0",
     "husky": "^9.1.7",
@@ -122,7 +138,8 @@
     "tsc-alias": "^1.8.10",
     "tsconfig-paths": "^4.2.0",
     "tslib": "^2.6.3",
-    "typescript": "^5.5.3"
+    "typescript": "^5.5.3",
+    "winston": "^3.17.0"
   },
   "lint-staged": {
     "*.{js,ts}": [

package/src/agents/AgentContext.ts

@@ -0,0 +1,315 @@
+/* eslint-disable no-console */
+// src/agents/AgentContext.ts
+import { zodToJsonSchema } from 'zod-to-json-schema';
+import { SystemMessage } from '@langchain/core/messages';
+import { RunnableLambda } from '@langchain/core/runnables';
+import type {
+  UsageMetadata,
+  BaseMessage,
+  BaseMessageFields,
+} from '@langchain/core/messages';
+import type { RunnableConfig, Runnable } from '@langchain/core/runnables';
+import type * as t from '@/types';
+import type { createPruneMessages } from '@/messages';
+import { ContentTypes, Providers } from '@/common';
+
+/**
+ * Encapsulates agent-specific state that can vary between agents in a multi-agent system
+ */
+export class AgentContext {
+  /**
+   * Create an AgentContext from configuration with token accounting initialization
+   */
+  static fromConfig(
+    agentConfig: t.AgentInputs,
+    tokenCounter?: t.TokenCounter,
+    indexTokenCountMap?: Record<string, number>
+  ): AgentContext {
+    const {
+      agentId,
+      provider,
+      clientOptions,
+      tools,
+      toolMap,
+      toolEnd,
+      instructions,
+      additional_instructions,
+      streamBuffer,
+      maxContextTokens,
+      reasoningKey,
+    } = agentConfig;
+
+    const agentContext = new AgentContext({
+      agentId,
+      provider,
+      clientOptions,
+      maxContextTokens,
+      streamBuffer,
+      tools,
+      toolMap,
+      instructions,
+      additionalInstructions: additional_instructions,
+      reasoningKey,
+      toolEnd,
+      instructionTokens: 0,
+      tokenCounter,
+    });
+
+    if (tokenCounter) {
+      const tokenMap = indexTokenCountMap || {};
+      agentContext.indexTokenCountMap = tokenMap;
+      agentContext.tokenCalculationPromise = agentContext
+        .calculateInstructionTokens(tokenCounter)
+        .then(() => {
+          // Update token map with instruction tokens
+          agentContext.updateTokenMapWithInstructions(tokenMap);
+        })
+        .catch((err) => {
+          console.error('Error calculating instruction tokens:', err);
+        });
+    } else if (indexTokenCountMap) {
+      agentContext.indexTokenCountMap = indexTokenCountMap;
+    }
+
+    return agentContext;
+  }
+
+  /** Agent identifier */
+  agentId: string;
+  /** Provider for this specific agent */
+  provider: Providers;
+  /** Client options for this agent */
+  clientOptions?: t.ClientOptions;
+  /** Token count map indexed by message position */
+  indexTokenCountMap: Record<string, number | undefined> = {};
+  /** Maximum context tokens for this agent */
+  maxContextTokens?: number;
+  /** Current usage metadata for this agent */
+  currentUsage?: Partial<UsageMetadata>;
+  /** Prune messages function configured for this agent */
+  pruneMessages?: ReturnType<typeof createPruneMessages>;
+  /** Token counter function for this agent */
+  tokenCounter?: t.TokenCounter;
+  /** Instructions/system message token count */
+  instructionTokens: number = 0;
+  /** The amount of time that should pass before another consecutive API call */
+  streamBuffer?: number;
+  /** Last stream call timestamp for rate limiting */
+  lastStreamCall?: number;
+  /** Tools available to this agent */
+  tools?: t.GraphTools;
+  /** Tool map for this agent */
+  toolMap?: t.ToolMap;
+  /** Instructions for this agent */
+  instructions?: string;
+  /** Additional instructions for this agent */
+  additionalInstructions?: string;
+  /** Reasoning key for this agent */
+  reasoningKey: 'reasoning_content' | 'reasoning' = 'reasoning_content';
+  /** Last token for reasoning detection */
+  lastToken?: string;
+  /** Token type switch state */
+  tokenTypeSwitch?: 'reasoning' | 'content';
+  /** Current token type being processed */
+  currentTokenType: ContentTypes.TEXT | ContentTypes.THINK | 'think_and_text' =
+    ContentTypes.TEXT;
+  /** Whether tools should end the workflow */
+  toolEnd: boolean = false;
+  /** System runnable for this agent */
+  systemRunnable?: Runnable<
+    BaseMessage[],
+    (BaseMessage | SystemMessage)[],
+    RunnableConfig<Record<string, unknown>>
+  >;
+  /** Promise for token calculation initialization */
+  tokenCalculationPromise?: Promise<void>;
+
+  constructor({
+    agentId,
+    provider,
+    clientOptions,
+    maxContextTokens,
+    streamBuffer,
+    tokenCounter,
+    tools,
+    toolMap,
+    instructions,
+    additionalInstructions,
+    reasoningKey,
+    toolEnd,
+    instructionTokens,
+  }: {
+    agentId: string;
+    provider: Providers;
+    clientOptions?: t.ClientOptions;
+    maxContextTokens?: number;
+    streamBuffer?: number;
+    tokenCounter?: t.TokenCounter;
+    tools?: t.GraphTools;
+    toolMap?: t.ToolMap;
+    instructions?: string;
+    additionalInstructions?: string;
+    reasoningKey?: 'reasoning_content' | 'reasoning';
+    toolEnd?: boolean;
+    instructionTokens?: number;
+  }) {
+    this.agentId = agentId;
+    this.provider = provider;
+    this.clientOptions = clientOptions;
+    this.maxContextTokens = maxContextTokens;
+    this.streamBuffer = streamBuffer;
+    this.tokenCounter = tokenCounter;
+    this.tools = tools;
+    this.toolMap = toolMap;
+    this.instructions = instructions;
+    this.additionalInstructions = additionalInstructions;
+    if (reasoningKey) {
+      this.reasoningKey = reasoningKey;
+    }
+    if (toolEnd !== undefined) {
+      this.toolEnd = toolEnd;
+    }
+    if (instructionTokens !== undefined) {
+      this.instructionTokens = instructionTokens;
+    }
+
+    this.systemRunnable = this.createSystemRunnable();
+  }
+
+  /**
+   * Create system runnable from instructions and calculate tokens if tokenCounter is available
+   */
+  private createSystemRunnable():
+    | Runnable<
+        BaseMessage[],
+        (BaseMessage | SystemMessage)[],
+        RunnableConfig<Record<string, unknown>>
+      >
+    | undefined {
+    let finalInstructions: string | BaseMessageFields | undefined =
+      this.instructions;
+
+    if (
+      this.additionalInstructions != null &&
+      this.additionalInstructions !== ''
+    ) {
+      finalInstructions =
+        finalInstructions != null && finalInstructions
+          ? `${finalInstructions}\n\n${this.additionalInstructions}`
+          : this.additionalInstructions;
+    }
+
+    // Handle Anthropic prompt caching
+    if (
+      finalInstructions != null &&
+      finalInstructions !== '' &&
+      this.provider === Providers.ANTHROPIC
+    ) {
+      const anthropicOptions = this.clientOptions as
+        | t.AnthropicClientOptions
+        | undefined;
+      const defaultHeaders = anthropicOptions?.clientOptions?.defaultHeaders as
+        | Record<string, string>
+        | undefined;
+      const anthropicBeta = defaultHeaders?.['anthropic-beta'];
+      if (
+        typeof anthropicBeta === 'string' &&
+        anthropicBeta.includes('prompt-caching')
+      ) {
+        finalInstructions = {
+          content: [
+            {
+              type: 'text',
+              text: this.instructions,
+              cache_control: { type: 'ephemeral' },
+            },
+          ],
+        };
+      }
+    }
+
+    if (finalInstructions != null && finalInstructions !== '') {
+      const systemMessage = new SystemMessage(finalInstructions);
+
+      if (this.tokenCounter) {
+        this.instructionTokens += this.tokenCounter(systemMessage);
+      }
+
+      return RunnableLambda.from((messages: BaseMessage[]) => {
+        return [systemMessage, ...messages];
+      }).withConfig({ runName: 'prompt' });
+    }
+
+    return undefined;
+  }
+
+  /**
+   * Reset context for a new run
+   */
+  reset(): void {
+    this.instructionTokens = 0;
+    this.lastToken = undefined;
+    this.indexTokenCountMap = {};
+    this.currentUsage = undefined;
+    this.pruneMessages = undefined;
+    this.lastStreamCall = undefined;
+    this.tokenTypeSwitch = undefined;
+    this.currentTokenType = ContentTypes.TEXT;
+  }
+
+  /**
+   * Update the token count map with instruction tokens
+   */
+  updateTokenMapWithInstructions(baseTokenMap: Record<string, number>): void {
+    if (this.instructionTokens > 0) {
+      // Shift all indices by the instruction token count
+      const shiftedMap: Record<string, number> = {};
+      for (const [key, value] of Object.entries(baseTokenMap)) {
+        const index = parseInt(key, 10);
+        if (!isNaN(index)) {
+          shiftedMap[String(index)] =
+            value + (index === 0 ? this.instructionTokens : 0);
+        }
+      }
+      this.indexTokenCountMap = shiftedMap;
+    } else {
+      this.indexTokenCountMap = { ...baseTokenMap };
+    }
+  }
+
+  /**
+   * Calculate tool tokens and add to instruction tokens
+   * Note: System message tokens are calculated during systemRunnable creation
+   */
+  async calculateInstructionTokens(
+    tokenCounter: t.TokenCounter
+  ): Promise<void> {
+    let toolTokens = 0;
+    if (this.tools && this.tools.length > 0) {
+      for (const tool of this.tools) {
+        const genericTool = tool as Record<string, unknown>;
+        if (
+          genericTool.schema != null &&
+          typeof genericTool.schema === 'object'
+        ) {
+          const schema = genericTool.schema as {
+            describe: (desc: string) => unknown;
+          };
+          const describedSchema = schema.describe(
+            (genericTool.description as string) || ''
+          );
+          const jsonSchema = zodToJsonSchema(
+            describedSchema as Parameters<typeof zodToJsonSchema>[0],
+            (genericTool.name as string) || ''
+          );
+          toolTokens += tokenCounter(
+            new SystemMessage(JSON.stringify(jsonSchema))
+          );
+        }
+      }
+    }
+
+    // Add tool tokens to existing instruction tokens (which may already include system message tokens)
+    this.instructionTokens += toolTokens;
+  }
+}

package/src/common/enum.ts

@@ -87,8 +87,9 @@ export enum Providers {
 }
 
 export enum GraphNodeKeys {
-  TOOLS = 'tools',
-  AGENT = 'agent',
+  TOOLS = 'tools=',
+  AGENT = 'agent=',
+  ROUTER = 'router',
   PRE_TOOLS = 'pre_tools',
   POST_TOOLS = 'post_tools',
 }
@@ -117,10 +118,12 @@ export enum ContentTypes {
   IMAGE_FILE = 'image_file',
   /** Anthropic */
   THINKING = 'thinking',
-  /** Bedrock */
-  REASONING_CONTENT = 'reasoning_content',
+  /** Vertex AI / Google Common */
+  REASONING = 'reasoning',
   /** Multi-Agent Switch */
   AGENT_UPDATE = 'agent_update',
+  /** Bedrock */
+  REASONING_CONTENT = 'reasoning_content',
 }
 
 export enum ToolCallTypes {
@@ -136,6 +139,7 @@ export enum Callback {
   TOOL_ERROR = 'handleToolError',
   TOOL_START = 'handleToolStart',
   TOOL_END = 'handleToolEnd',
+  CUSTOM_EVENT = 'handleCustomEvent',
   /*
   LLM_START = 'handleLLMStart',
   LLM_NEW_TOKEN = 'handleLLMNewToken',
@@ -151,7 +155,6 @@ export enum Callback {
   RETRIEVER_START = 'handleRetrieverStart',
   RETRIEVER_END = 'handleRetrieverEnd',
   RETRIEVER_ERROR = 'handleRetrieverError',
-  CUSTOM_EVENT = 'handleCustomEvent'
   */
 }
 
@@ -162,6 +165,12 @@ export enum Constants {
   CONTENT_AND_ARTIFACT = 'content_and_artifact',
 }
 
+export enum TitleMethod {
+  STRUCTURED = 'structured',
+  FUNCTIONS = 'functions',
+  COMPLETION = 'completion',
+}
+
 export enum EnvVar {
   CODE_API_KEY = 'LIBRECHAT_CODE_API_KEY',
   CODE_BASEURL = 'LIBRECHAT_CODE_BASEURL',

package/src/events.ts

@@ -1,10 +1,11 @@
 /* eslint-disable no-console */
 // src/events.ts
 import type {
+  ToolMessage,
   UsageMetadata,
   BaseMessageFields,
 } from '@langchain/core/messages';
-import type { Graph } from '@/graphs';
+import type { MultiAgentGraph, StandardGraph } from '@/graphs';
 import type * as t from '@/types';
 import { handleToolCalls } from '@/tools/handlers';
 import { Providers } from '@/common';
@@ -30,12 +31,12 @@ export class ModelEndHandler implements t.EventHandler {
     this.collectedUsage = collectedUsage;
   }
 
-  handle(
+  async handle(
     event: string,
     data: t.ModelEndData,
     metadata?: Record<string, unknown>,
-    graph?: Graph
-  ): void {
+    graph?: StandardGraph | MultiAgentGraph
+  ): Promise<void> {
     if (!graph || !metadata) {
       console.warn(`Graph or metadata not found in ${event} event`);
       return;
@@ -58,25 +59,35 @@ export class ModelEndHandler implements t.EventHandler {
       { depth: null }
     );
 
-    if (metadata.provider !== Providers.GOOGLE) {
+    const agentContext = graph.getAgentContext(metadata);
+
+    if (
+      agentContext.provider !== Providers.GOOGLE &&
+      agentContext.provider !== Providers.BEDROCK
+    ) {
       return;
     }
 
-    handleToolCalls(data?.output?.tool_calls, metadata, graph);
+    await handleToolCalls(data?.output?.tool_calls, metadata, graph);
   }
 }
 
 export class ToolEndHandler implements t.EventHandler {
   private callback?: t.ToolEndCallback;
-  constructor(callback?: t.ToolEndCallback) {
+  private omitOutput?: (name?: string) => boolean;
+  constructor(
+    callback?: t.ToolEndCallback,
+    omitOutput?: (name?: string) => boolean
+  ) {
     this.callback = callback;
+    this.omitOutput = omitOutput;
   }
-  handle(
+  async handle(
     event: string,
     data: t.StreamEventData | undefined,
     metadata?: Record<string, unknown>,
-    graph?: Graph
-  ): void {
+    graph?: StandardGraph | MultiAgentGraph
+  ): Promise<void> {
     if (!graph || !metadata) {
       console.warn(`Graph or metadata not found in ${event} event`);
       return;
@@ -89,10 +100,10 @@ export class ToolEndHandler implements t.EventHandler {
     }
 
     this.callback?.(toolEndData, metadata);
-
-    graph.handleToolCallCompleted(
+    await graph.handleToolCallCompleted(
       { input: toolEndData.input, output: toolEndData.output },
-      metadata
+      metadata,
+      this.omitOutput?.((toolEndData.output as ToolMessage | undefined)?.name)
     );
   }
 }