@librechat/agents 2.4.322 → 3.0.0-rc10

package/src/graphs/Graph.ts

@@ -4,37 +4,44 @@ import { nanoid } from 'nanoid';
 import { concat } from '@langchain/core/utils/stream';
 import { ToolNode } from '@langchain/langgraph/prebuilt';
 import { ChatVertexAI } from '@langchain/google-vertexai';
-import { START, END, StateGraph } from '@langchain/langgraph';
-import { Runnable, RunnableConfig } from '@langchain/core/runnables';
-import { dispatchCustomEvent } from '@langchain/core/callbacks/dispatch';
 import {
-  AIMessageChunk,
+  START,
+  END,
+  Command,
+  StateGraph,
+  Annotation,
+  messagesStateReducer,
+} from '@langchain/langgraph';
+import {
+  Runnable,
+  RunnableConfig,
+  RunnableLambda,
+} from '@langchain/core/runnables';
+import {
   ToolMessage,
   SystemMessage,
+  AIMessageChunk,
 } from '@langchain/core/messages';
 import type {
-  BaseMessage,
   BaseMessageFields,
   UsageMetadata,
+  BaseMessage,
 } from '@langchain/core/messages';
+import type { ToolCall } from '@langchain/core/messages/tool';
 import type * as t from '@/types';
 import {
-  Providers,
-  GraphEvents,
   GraphNodeKeys,
-  StepTypes,
-  Callback,
   ContentTypes,
+  GraphEvents,
+  Providers,
+  StepTypes,
 } from '@/common';
-import type { ToolCall } from '@langchain/core/messages/tool';
-import { getChatModelClass, manualToolStreamProviders } from '@/llm/providers';
-import { ToolNode as CustomToolNode, toolsCondition } from '@/tools/ToolNode';
 import {
-  createPruneMessages,
+  formatAnthropicArtifactContent,
+  convertMessagesToContent,
   modifyDeltaProperties,
   formatArtifactPayload,
-  convertMessagesToContent,
-  formatAnthropicArtifactContent,
+  createPruneMessages,
 } from '@/messages';
 import {
   resetIfNotEmpty,
@@ -43,38 +50,37 @@ import {
   joinKeys,
   sleep,
 } from '@/utils';
+import { getChatModelClass, manualToolStreamProviders } from '@/llm/providers';
+import { ToolNode as CustomToolNode, toolsCondition } from '@/tools/ToolNode';
 import { ChatOpenAI, AzureChatOpenAI } from '@/llm/openai';
+import { safeDispatchCustomEvent } from '@/utils/events';
+import { AgentContext } from '@/agents/AgentContext';
 import { createFakeStreamingLLM } from '@/llm/fake';
 import { HandlerRegistry } from '@/events';
 
 const { AGENT, TOOLS } = GraphNodeKeys;
-export type GraphNode = GraphNodeKeys | typeof START;
-export type ClientCallback<T extends unknown[]> = (
-  graph: StandardGraph,
-  ...args: T
-) => void;
-export type ClientCallbacks = {
-  [Callback.TOOL_ERROR]?: ClientCallback<[Error, string]>;
-  [Callback.TOOL_START]?: ClientCallback<unknown[]>;
-  [Callback.TOOL_END]?: ClientCallback<unknown[]>;
-};
-export type SystemCallbacks = {
-  [K in keyof ClientCallbacks]: ClientCallbacks[K] extends ClientCallback<
-    infer Args
-  >
-    ? (...args: Args) => void
-    : never;
-};
 
 export abstract class Graph<
   T extends t.BaseGraphState = t.BaseGraphState,
-  // eslint-disable-next-line @typescript-eslint/no-unused-vars
-  TNodeName extends string = string,
+  _TNodeName extends string = string,
 > {
   abstract resetValues(): void;
-  abstract createGraphState(): t.GraphStateChannels<T>;
-  abstract initializeTools(): CustomToolNode<T> | ToolNode<T>;
-  abstract initializeModel(): Runnable;
+  abstract initializeTools({
+    currentTools,
+    currentToolMap,
+  }: {
+    currentTools?: t.GraphTools;
+    currentToolMap?: t.ToolMap;
+  }): CustomToolNode<T> | ToolNode<T>;
+  abstract initializeModel({
+    currentModel,
+    tools,
+    clientOptions,
+  }: {
+    currentModel?: t.ChatModel;
+    tools?: t.GraphTools;
+    clientOptions?: t.ClientOptions;
+  }): Runnable;
   abstract getRunMessages(): BaseMessage[] | undefined;
   abstract getContentParts(): t.MessageContentComplex[] | undefined;
   abstract generateStepId(stepKey: string): [string, number];
@@ -85,28 +91,32 @@ export abstract class Graph<
   abstract checkKeyList(keyList: (string | number | undefined)[]): boolean;
   abstract getStepIdByKey(stepKey: string, index?: number): string;
   abstract getRunStep(stepId: string): t.RunStep | undefined;
-  abstract dispatchRunStep(stepKey: string, stepDetails: t.StepDetails): string;
-  abstract dispatchRunStepDelta(id: string, delta: t.ToolCallDelta): void;
-  abstract dispatchMessageDelta(id: string, delta: t.MessageDelta): void;
+  abstract dispatchRunStep(
+    stepKey: string,
+    stepDetails: t.StepDetails
+  ): Promise<string>;
+  abstract dispatchRunStepDelta(
+    id: string,
+    delta: t.ToolCallDelta
+  ): Promise<void>;
+  abstract dispatchMessageDelta(
+    id: string,
+    delta: t.MessageDelta
+  ): Promise<void>;
   abstract dispatchReasoningDelta(
     stepId: string,
     delta: t.ReasoningDelta
-  ): void;
+  ): Promise<void>;
   abstract handleToolCallCompleted(
     data: t.ToolEndData,
-    metadata?: Record<string, unknown>
-  ): void;
-
-  abstract createCallModel(): (
-    state: T,
-    config?: RunnableConfig
-  ) => Promise<Partial<T>>;
-  abstract createWorkflow(): t.CompiledWorkflow<T>;
-  lastToken?: string;
-  tokenTypeSwitch?: 'reasoning' | 'content';
-  reasoningKey: 'reasoning_content' | 'reasoning' = 'reasoning_content';
-  currentTokenType: ContentTypes.TEXT | ContentTypes.THINK | 'think_and_text' =
-    ContentTypes.TEXT;
+    metadata?: Record<string, unknown>,
+    omitOutput?: boolean
+  ): Promise<void>;
+
+  abstract createCallModel(
+    agentId?: string,
+    currentModel?: t.ChatModel
+  ): (state: T, config?: RunnableConfig) => Promise<Partial<T>>;
   messageStepHasToolCalls: Map<string, boolean> = new Map();
   messageIdsByStepKey: Map<string, string> = new Map();
   prelimMessageIdsByStepKey: Map<string, string> = new Map();
@@ -115,95 +125,52 @@ export abstract class Graph<
   stepKeyIds: Map<string, string[]> = new Map<string, string[]>();
   contentIndexMap: Map<string, number> = new Map();
   toolCallStepIds: Map<string, string> = new Map();
-  currentUsage: Partial<UsageMetadata> | undefined;
-  indexTokenCountMap: Record<string, number | undefined> = {};
-  maxContextTokens: number | undefined;
-  pruneMessages?: ReturnType<typeof createPruneMessages>;
-  /** The amount of time that should pass before another consecutive API call */
-  streamBuffer: number | undefined;
-  tokenCounter?: t.TokenCounter;
   signal?: AbortSignal;
+  /** Set of invoked tool call IDs from non-message run steps completed mid-run, if any */
+  invokedToolIds?: Set<string>;
+  handlerRegistry: HandlerRegistry | undefined;
 }
 
-export class StandardGraph extends Graph<t.BaseGraphState, GraphNode> {
-  private graphState: t.GraphStateChannels<t.BaseGraphState>;
-  clientOptions: t.ClientOptions;
-  boundModel?: Runnable;
-  /** The last recorded timestamp that a stream API call was invoked */
-  lastStreamCall: number | undefined;
-  handlerRegistry: HandlerRegistry | undefined;
-  systemMessage: SystemMessage | undefined;
+export class StandardGraph extends Graph<t.BaseGraphState, t.GraphNode> {
+  overrideModel?: t.ChatModel;
+  /** Optional compile options passed into workflow.compile() */
+  compileOptions?: t.CompileOptions | undefined;
   messages: BaseMessage[] = [];
   runId: string | undefined;
-  tools?: t.GenericTool[];
-  toolMap?: t.ToolMap;
   startIndex: number = 0;
-  provider: Providers;
-  toolEnd: boolean;
   signal?: AbortSignal;
+  /** Map of agent contexts by agent ID */
+  agentContexts: Map<string, AgentContext> = new Map();
+  /** Default agent ID to use */
+  defaultAgentId: string;
 
   constructor({
+    // parent-level graph inputs
     runId,
-    tools,
     signal,
-    toolMap,
-    provider,
-    streamBuffer,
-    instructions,
-    reasoningKey,
-    clientOptions,
-    toolEnd = false,
-    additional_instructions = '',
+    agents,
+    tokenCounter,
+    indexTokenCountMap,
   }: t.StandardGraphInput) {
     super();
     this.runId = runId;
-    this.tools = tools;
     this.signal = signal;
-    this.toolEnd = toolEnd;
-    this.toolMap = toolMap;
-    this.provider = provider;
-    this.streamBuffer = streamBuffer;
-    this.clientOptions = clientOptions;
-    this.graphState = this.createGraphState();
-    this.boundModel = this.initializeModel();
-    if (reasoningKey) {
-      this.reasoningKey = reasoningKey;
-    }
 
-    let finalInstructions: string | BaseMessageFields | undefined =
-      instructions;
-    if (additional_instructions) {
-      finalInstructions =
-        finalInstructions != null && finalInstructions
-          ? `${finalInstructions}\n\n${additional_instructions}`
-          : additional_instructions;
+    if (agents.length === 0) {
+      throw new Error('At least one agent configuration is required');
     }
 
-    if (
-      finalInstructions != null &&
-      finalInstructions &&
-      provider === Providers.ANTHROPIC &&
-      ((
-        clientOptions as t.AnthropicClientOptions
-      ).clientOptions?.defaultHeaders?.['anthropic-beta']?.includes(
-        'prompt-caching'
-      ) ??
-        false)
-    ) {
-      finalInstructions = {
-        content: [
-          {
-            type: 'text',
-            text: instructions,
-            cache_control: { type: 'ephemeral' },
-          },
-        ],
-      };
-    }
+    for (const agentConfig of agents) {
+      const agentContext = AgentContext.fromConfig(
+        agentConfig,
+        tokenCounter,
+        indexTokenCountMap
+      );
 
-    if (finalInstructions != null && finalInstructions !== '') {
-      this.systemMessage = new SystemMessage(finalInstructions);
+      this.agentContexts.set(agentConfig.agentId, agentContext);
     }
+
+    this.defaultAgentId = agents[0].agentId;
   }
 
   /* Init */
@@ -222,23 +189,17 @@ export class StandardGraph extends Graph<t.BaseGraphState, GraphNode> {
       new Map()
     );
     this.messageStepHasToolCalls = resetIfNotEmpty(
-      this.prelimMessageIdsByStepKey,
+      this.messageStepHasToolCalls,
       new Map()
     );
     this.prelimMessageIdsByStepKey = resetIfNotEmpty(
       this.prelimMessageIdsByStepKey,
       new Map()
     );
-    this.currentTokenType = resetIfNotEmpty(
-      this.currentTokenType,
-      ContentTypes.TEXT
-    );
-    this.lastToken = resetIfNotEmpty(this.lastToken, undefined);
-    this.tokenTypeSwitch = resetIfNotEmpty(this.tokenTypeSwitch, undefined);
-    this.indexTokenCountMap = resetIfNotEmpty(this.indexTokenCountMap, {});
-    this.currentUsage = resetIfNotEmpty(this.currentUsage, undefined);
-    this.tokenCounter = resetIfNotEmpty(this.tokenCounter, undefined);
-    this.maxContextTokens = resetIfNotEmpty(this.maxContextTokens, undefined);
+    this.invokedToolIds = resetIfNotEmpty(this.invokedToolIds, undefined);
+    for (const context of this.agentContexts.values()) {
+      context.reset();
+    }
   }
 
   /* Run Step Processing */
@@ -251,6 +212,33 @@ export class StandardGraph extends Graph<t.BaseGraphState, GraphNode> {
     return undefined;
   }
 
+  getAgentContext(metadata: Record<string, unknown> | undefined): AgentContext {
+    if (!metadata) {
+      throw new Error('No metadata provided to retrieve agent context');
+    }
+
+    const currentNode = metadata.langgraph_node as string;
+    if (!currentNode) {
+      throw new Error(
+        'No langgraph_node in metadata to retrieve agent context'
+      );
+    }
+
+    let agentId: string | undefined;
+    if (currentNode.startsWith(AGENT)) {
+      agentId = currentNode.substring(AGENT.length);
+    } else if (currentNode.startsWith(TOOLS)) {
+      agentId = currentNode.substring(TOOLS.length);
+    }
+
+    const agentContext = this.agentContexts.get(agentId ?? '');
+    if (!agentContext) {
+      throw new Error(`No agent context found for agent ID ${agentId}`);
+    }
+
+    return agentContext;
+  }
+
   getStepKey(metadata: Record<string, unknown> | undefined): string {
     if (!metadata) return '';
 
@@ -304,13 +292,19 @@ export class StandardGraph extends Graph<t.BaseGraphState, GraphNode> {
       metadata.langgraph_step as number,
       metadata.checkpoint_ns as string,
     ];
+
+    const agentContext = this.getAgentContext(metadata);
     if (
-      this.currentTokenType === ContentTypes.THINK ||
-      this.currentTokenType === 'think_and_text'
+      agentContext.currentTokenType === ContentTypes.THINK ||
+      agentContext.currentTokenType === 'think_and_text'
     ) {
       keyList.push('reasoning');
     }
 
+    if (this.invokedToolIds != null && this.invokedToolIds.size > 0) {
+      keyList.push(this.invokedToolIds.size + '');
+    }
+
     return keyList;
   }
 
@@ -330,87 +324,126 @@ export class StandardGraph extends Graph<t.BaseGraphState, GraphNode> {
 
   /* Graph */
 
-  createGraphState(): t.GraphStateChannels<t.BaseGraphState> {
-    return {
-      messages: {
-        value: (x: BaseMessage[], y: BaseMessage[]): BaseMessage[] => {
-          if (!x.length) {
-            if (this.systemMessage) {
-              x.push(this.systemMessage);
-            }
+  createSystemRunnable({
+    provider,
+    clientOptions,
+    instructions,
+    additional_instructions,
+  }: {
+    provider?: Providers;
+    clientOptions?: t.ClientOptions;
+    instructions?: string;
+    additional_instructions?: string;
+  }): t.SystemRunnable | undefined {
+    let finalInstructions: string | BaseMessageFields | undefined =
+      instructions;
+    if (additional_instructions != null && additional_instructions !== '') {
+      finalInstructions =
+        finalInstructions != null && finalInstructions
+          ? `${finalInstructions}\n\n${additional_instructions}`
+          : additional_instructions;
+    }
 
-            this.startIndex = x.length + y.length;
-          }
-          const current = x.concat(y);
-          this.messages = current;
-          return current;
-        },
-        default: () => [],
-      },
-    };
+    if (
+      finalInstructions != null &&
+      finalInstructions &&
+      provider === Providers.ANTHROPIC &&
+      ((
+        (clientOptions as t.AnthropicClientOptions).clientOptions
+          ?.defaultHeaders as Record<string, string> | undefined
+      )?.['anthropic-beta']?.includes('prompt-caching') ??
+        false)
+    ) {
+      finalInstructions = {
+        content: [
+          {
+            type: 'text',
+            text: instructions,
+            cache_control: { type: 'ephemeral' },
+          },
+        ],
+      };
+    }
+
+    if (finalInstructions != null && finalInstructions !== '') {
+      const systemMessage = new SystemMessage(finalInstructions);
+      return RunnableLambda.from((messages: BaseMessage[]) => {
+        return [systemMessage, ...messages];
+      }).withConfig({ runName: 'prompt' });
+    }
   }
 
-  initializeTools():
-    | CustomToolNode<t.BaseGraphState>
-    | ToolNode<t.BaseGraphState> {
-    // return new ToolNode<t.BaseGraphState>(this.tools);
+  initializeTools({
+    currentTools,
+    currentToolMap,
+  }: {
+    currentTools?: t.GraphTools;
+    currentToolMap?: t.ToolMap;
+  }): CustomToolNode<t.BaseGraphState> | ToolNode<t.BaseGraphState> {
     return new CustomToolNode<t.BaseGraphState>({
-      tools: this.tools || [],
-      toolMap: this.toolMap,
+      tools: (currentTools as t.GenericTool[] | undefined) ?? [],
+      toolMap: currentToolMap,
       toolCallStepIds: this.toolCallStepIds,
       errorHandler: (data, metadata) =>
         StandardGraph.handleToolCallErrorStatic(this, data, metadata),
     });
   }
 
-  initializeModel(): Runnable {
-    const ChatModelClass = getChatModelClass(this.provider);
-    const model = new ChatModelClass(this.clientOptions);
+  initializeModel({
+    provider,
+    tools,
+    clientOptions,
+  }: {
+    provider: Providers;
+    tools?: t.GraphTools;
+    clientOptions?: t.ClientOptions;
+  }): Runnable {
+    const ChatModelClass = getChatModelClass(provider);
+    const model = new ChatModelClass(clientOptions ?? {});
 
     if (
-      isOpenAILike(this.provider) &&
+      isOpenAILike(provider) &&
       (model instanceof ChatOpenAI || model instanceof AzureChatOpenAI)
     ) {
-      model.temperature = (this.clientOptions as t.OpenAIClientOptions)
+      model.temperature = (clientOptions as t.OpenAIClientOptions)
         .temperature as number;
-      model.topP = (this.clientOptions as t.OpenAIClientOptions).topP as number;
-      model.frequencyPenalty = (this.clientOptions as t.OpenAIClientOptions)
+      model.topP = (clientOptions as t.OpenAIClientOptions).topP as number;
+      model.frequencyPenalty = (clientOptions as t.OpenAIClientOptions)
         .frequencyPenalty as number;
-      model.presencePenalty = (this.clientOptions as t.OpenAIClientOptions)
+      model.presencePenalty = (clientOptions as t.OpenAIClientOptions)
         .presencePenalty as number;
-      model.n = (this.clientOptions as t.OpenAIClientOptions).n as number;
+      model.n = (clientOptions as t.OpenAIClientOptions).n as number;
     } else if (
-      this.provider === Providers.VERTEXAI &&
+      provider === Providers.VERTEXAI &&
       model instanceof ChatVertexAI
     ) {
-      model.temperature = (this.clientOptions as t.VertexAIClientOptions)
+      model.temperature = (clientOptions as t.VertexAIClientOptions)
         .temperature as number;
-      model.topP = (this.clientOptions as t.VertexAIClientOptions)
-        .topP as number;
-      model.topK = (this.clientOptions as t.VertexAIClientOptions)
-        .topK as number;
-      model.topLogprobs = (this.clientOptions as t.VertexAIClientOptions)
+      model.topP = (clientOptions as t.VertexAIClientOptions).topP as number;
+      model.topK = (clientOptions as t.VertexAIClientOptions).topK as number;
+      model.topLogprobs = (clientOptions as t.VertexAIClientOptions)
         .topLogprobs as number;
-      model.frequencyPenalty = (this.clientOptions as t.VertexAIClientOptions)
+      model.frequencyPenalty = (clientOptions as t.VertexAIClientOptions)
         .frequencyPenalty as number;
-      model.presencePenalty = (this.clientOptions as t.VertexAIClientOptions)
+      model.presencePenalty = (clientOptions as t.VertexAIClientOptions)
         .presencePenalty as number;
-      model.maxOutputTokens = (this.clientOptions as t.VertexAIClientOptions)
+      model.maxOutputTokens = (clientOptions as t.VertexAIClientOptions)
         .maxOutputTokens as number;
     }
 
-    if (!this.tools || this.tools.length === 0) {
+    if (!tools || tools.length === 0) {
       return model as unknown as Runnable;
     }
 
-    return (model as t.ModelWithTools).bindTools(this.tools);
+    return (model as t.ModelWithTools).bindTools(tools);
   }
+
   overrideTestModel(
     responses: string[],
     sleep?: number,
     toolCalls?: ToolCall[]
   ): void {
-    this.boundModel = createFakeStreamingLLM({
+    this.overrideModel = createFakeStreamingLLM({
       responses,
       sleep,
       toolCalls,
@@ -418,42 +451,84 @@ export class StandardGraph extends Graph<t.BaseGraphState, GraphNode> {
   }
 
   getNewModel({
-    clientOptions = {},
-    omitOriginalOptions,
+    provider,
+    clientOptions,
   }: {
+    provider: Providers;
     clientOptions?: t.ClientOptions;
-    omitOriginalOptions?: Set<string>;
   }): t.ChatModelInstance {
-    const ChatModelClass = getChatModelClass(this.provider);
-    const _options = omitOriginalOptions
-      ? Object.fromEntries(
-        Object.entries(this.clientOptions).filter(
-          ([key]) => !omitOriginalOptions.has(key)
-        )
-      )
-      : this.clientOptions;
-    const options = Object.assign(_options, clientOptions);
-    return new ChatModelClass(options);
+    const ChatModelClass = getChatModelClass(provider);
+    return new ChatModelClass(clientOptions ?? {});
   }
 
-  storeUsageMetadata(finalMessage?: BaseMessage): void {
+  getUsageMetadata(
+    finalMessage?: BaseMessage
+  ): Partial<UsageMetadata> | undefined {
     if (
       finalMessage &&
       'usage_metadata' in finalMessage &&
       finalMessage.usage_metadata != null
     ) {
-      this.currentUsage = finalMessage.usage_metadata as Partial<UsageMetadata>;
+      return finalMessage.usage_metadata as Partial<UsageMetadata>;
     }
   }
 
-  cleanupSignalListener(): void {
+  /** Execute model invocation with streaming support */
+  private async attemptInvoke(
+    {
+      currentModel,
+      finalMessages,
+      provider,
+      tools,
+    }: {
+      currentModel?: t.ChatModel;
+      finalMessages: BaseMessage[];
+      provider: Providers;
+      tools?: t.GraphTools;
+    },
+    config?: RunnableConfig
+  ): Promise<Partial<t.BaseGraphState>> {
+    const model = this.overrideModel ?? currentModel;
+    if (!model) {
+      throw new Error('No model found');
+    }
+
+    if ((tools?.length ?? 0) > 0 && manualToolStreamProviders.has(provider)) {
+      if (!model.stream) {
+        throw new Error('Model does not support stream');
+      }
+      const stream = await model.stream(finalMessages, config);
+      let finalChunk: AIMessageChunk | undefined;
+      for await (const chunk of stream) {
+        await safeDispatchCustomEvent(
+          GraphEvents.CHAT_MODEL_STREAM,
+          { chunk, emitted: true },
+          config
+        );
+        finalChunk = finalChunk ? concat(finalChunk, chunk) : chunk;
+      }
+      finalChunk = modifyDeltaProperties(provider, finalChunk);
+      return { messages: [finalChunk as AIMessageChunk] };
+    } else {
+      const finalMessage = await model.invoke(finalMessages, config);
+      if ((finalMessage.tool_calls?.length ?? 0) > 0) {
+        finalMessage.tool_calls = finalMessage.tool_calls?.filter(
+          (tool_call: ToolCall) => !!tool_call.name
+        );
+      }
+      return { messages: [finalMessage] };
+    }
+  }
+
+  cleanupSignalListener(currentModel?: t.ChatModel): void {
     if (!this.signal) {
       return;
     }
-    if (!this.boundModel) {
+    const model = this.overrideModel ?? currentModel;
+    if (!model) {
       return;
     }
-    const client = (this.boundModel as ChatOpenAI | undefined)?.exposedClient;
+    const client = (model as ChatOpenAI | undefined)?.exposedClient;
     if (!client?.abortHandler) {
       return;
     }
@@ -461,18 +536,30 @@ export class StandardGraph extends Graph<t.BaseGraphState, GraphNode> {
     client.abortHandler = undefined;
   }
 
-  createCallModel() {
+  createCallModel(agentId = 'default', currentModel?: t.ChatModel) {
     return async (
       state: t.BaseGraphState,
       config?: RunnableConfig
     ): Promise<Partial<t.BaseGraphState>> => {
-      const { provider = '' } =
-        (config?.configurable as t.GraphConfig | undefined) ?? {};
-      if (this.boundModel == null) {
+      /**
+       * Get agent context - it must exist by this point
+       */
+      const agentContext = this.agentContexts.get(agentId);
+      if (!agentContext) {
+        throw new Error(`Agent context not found for agentId: ${agentId}`);
+      }
+
+      const model = this.overrideModel ?? currentModel;
+      if (!model) {
         throw new Error('No Graph model found');
       }
-      if (!config || !provider) {
-        throw new Error(`No ${config ? 'provider' : 'config'} provided`);
+      if (!config) {
+        throw new Error('No config provided');
+      }
+
+      // Ensure token calculations are complete before proceeding
+      if (agentContext.tokenCalculationPromise) {
+        await agentContext.tokenCalculationPromise;
       }
       if (!config.signal) {
         config.signal = this.signal;
@@ -482,35 +569,35 @@ export class StandardGraph extends Graph<t.BaseGraphState, GraphNode> {
 
       let messagesToUse = messages;
       if (
-        !this.pruneMessages &&
-        this.tokenCounter &&
-        this.maxContextTokens != null &&
-        this.indexTokenCountMap[0] != null
+        !agentContext.pruneMessages &&
+        agentContext.tokenCounter &&
+        agentContext.maxContextTokens != null &&
+        agentContext.indexTokenCountMap[0] != null
       ) {
         const isAnthropicWithThinking =
-          (this.provider === Providers.ANTHROPIC &&
-            (this.clientOptions as t.AnthropicClientOptions).thinking !=
+          (agentContext.provider === Providers.ANTHROPIC &&
+            (agentContext.clientOptions as t.AnthropicClientOptions).thinking !=
               null) ||
-          (this.provider === Providers.BEDROCK &&
-            (this.clientOptions as t.BedrockAnthropicInput)
+          (agentContext.provider === Providers.BEDROCK &&
+            (agentContext.clientOptions as t.BedrockAnthropicInput)
               .additionalModelRequestFields?.['thinking'] != null);
 
-        this.pruneMessages = createPruneMessages({
-          provider: this.provider,
-          indexTokenCountMap: this.indexTokenCountMap,
-          maxTokens: this.maxContextTokens,
-          tokenCounter: this.tokenCounter,
+        agentContext.pruneMessages = createPruneMessages({
           startIndex: this.startIndex,
+          provider: agentContext.provider,
+          tokenCounter: agentContext.tokenCounter,
+          maxTokens: agentContext.maxContextTokens,
           thinkingEnabled: isAnthropicWithThinking,
+          indexTokenCountMap: agentContext.indexTokenCountMap,
         });
       }
-      if (this.pruneMessages) {
-        const { context, indexTokenCountMap } = this.pruneMessages({
+      if (agentContext.pruneMessages) {
+        const { context, indexTokenCountMap } = agentContext.pruneMessages({
           messages,
-          usageMetadata: this.currentUsage,
+          usageMetadata: agentContext.currentUsage,
           // startOnMessageType: 'human',
         });
-        this.indexTokenCountMap = indexTokenCountMap;
+        agentContext.indexTokenCountMap = indexTokenCountMap;
         messagesToUse = context;
       }
 
@@ -525,7 +612,7 @@ export class StandardGraph extends Graph<t.BaseGraphState, GraphNode> {
           : null;
 
       if (
-        provider === Providers.BEDROCK &&
+        agentContext.provider === Providers.BEDROCK &&
         lastMessageX instanceof AIMessageChunk &&
         lastMessageY instanceof ToolMessage &&
         typeof lastMessageX.content === 'string'
@@ -535,92 +622,165 @@ export class StandardGraph extends Graph<t.BaseGraphState, GraphNode> {
 
       const isLatestToolMessage = lastMessageY instanceof ToolMessage;
 
-      if (isLatestToolMessage && provider === Providers.ANTHROPIC) {
+      if (
+        isLatestToolMessage &&
+        agentContext.provider === Providers.ANTHROPIC
+      ) {
         formatAnthropicArtifactContent(finalMessages);
       } else if (
         isLatestToolMessage &&
-        (isOpenAILike(provider) || isGoogleLike(provider))
+        (isOpenAILike(agentContext.provider) ||
+          isGoogleLike(agentContext.provider))
       ) {
         formatArtifactPayload(finalMessages);
       }
 
-      if (this.lastStreamCall != null && this.streamBuffer != null) {
-        const timeSinceLastCall = Date.now() - this.lastStreamCall;
-        if (timeSinceLastCall < this.streamBuffer) {
+      if (
+        agentContext.lastStreamCall != null &&
+        agentContext.streamBuffer != null
+      ) {
+        const timeSinceLastCall = Date.now() - agentContext.lastStreamCall;
+        if (timeSinceLastCall < agentContext.streamBuffer) {
           const timeToWait =
-            Math.ceil((this.streamBuffer - timeSinceLastCall) / 1000) * 1000;
+            Math.ceil((agentContext.streamBuffer - timeSinceLastCall) / 1000) *
+            1000;
           await sleep(timeToWait);
         }
       }
 
-      this.lastStreamCall = Date.now();
+      agentContext.lastStreamCall = Date.now();
 
-      let result: Partial<t.BaseGraphState>;
-      if (
-        (this.tools?.length ?? 0) > 0 &&
-        manualToolStreamProviders.has(provider)
-      ) {
-        const stream = await this.boundModel.stream(finalMessages, config);
-        let finalChunk: AIMessageChunk | undefined;
-        for await (const chunk of stream) {
-          dispatchCustomEvent(GraphEvents.CHAT_MODEL_STREAM, { chunk }, config);
-          if (!finalChunk) {
-            finalChunk = chunk;
-          } else {
-            finalChunk = concat(finalChunk, chunk);
+      let result: Partial<t.BaseGraphState> | undefined;
+      const fallbacks =
+        (agentContext.clientOptions as t.LLMConfig | undefined)?.fallbacks ??
+        [];
+      try {
+        result = await this.attemptInvoke(
+          {
+            currentModel: model,
+            finalMessages,
+            provider: agentContext.provider,
+            tools: agentContext.tools,
+          },
+          config
+        );
+      } catch (primaryError) {
+        let lastError: unknown = primaryError;
+        for (const fb of fallbacks) {
+          try {
+            let model = this.getNewModel({
+              provider: fb.provider,
+              clientOptions: fb.clientOptions,
+            });
+            const bindableTools = agentContext.tools;
+            model = (
+              !bindableTools || bindableTools.length === 0
+                ? model
+                : model.bindTools(bindableTools)
+            ) as t.ChatModelInstance;
+            result = await this.attemptInvoke(
+              {
+                currentModel: model,
+                finalMessages,
+                provider: fb.provider,
+                tools: agentContext.tools,
+              },
+              config
+            );
+            lastError = undefined;
+            break;
+          } catch (e) {
+            lastError = e;
+            continue;
           }
         }
-
-        finalChunk = modifyDeltaProperties(this.provider, finalChunk);
-        result = { messages: [finalChunk as AIMessageChunk] };
-      } else {
-        const finalMessage = (await this.boundModel.invoke(
-          finalMessages,
-          config
-        )) as AIMessageChunk;
-        if ((finalMessage.tool_calls?.length ?? 0) > 0) {
-          finalMessage.tool_calls = finalMessage.tool_calls?.filter(
-            (tool_call) => {
-              if (!tool_call.name) {
-                return false;
-              }
-              return true;
-            }
-          );
+        if (lastError !== undefined) {
+          throw lastError;
         }
-        result = { messages: [finalMessage] };
       }
 
-      this.storeUsageMetadata(result.messages?.[0]);
+      if (!result) {
+        throw new Error('No result after model invocation');
+      }
+      agentContext.currentUsage = this.getUsageMetadata(result.messages?.[0]);
       this.cleanupSignalListener();
       return result;
     };
   }
 
-  createWorkflow(): t.CompiledWorkflow<t.BaseGraphState> {
+  createAgentNode(agentId: string): t.CompiledAgentWorfklow {
+    const agentContext = this.agentContexts.get(agentId);
+    if (!agentContext) {
+      throw new Error(`Agent context not found for agentId: ${agentId}`);
+    }
+
+    let currentModel = this.initializeModel({
+      tools: agentContext.tools,
+      provider: agentContext.provider,
+      clientOptions: agentContext.clientOptions,
+    });
+
+    if (agentContext.systemRunnable) {
+      currentModel = agentContext.systemRunnable.pipe(currentModel);
+    }
+
+    const agentNode = `${AGENT}${agentId}` as const;
+    const toolNode = `${TOOLS}${agentId}` as const;
+
     const routeMessage = (
       state: t.BaseGraphState,
       config?: RunnableConfig
     ): string => {
       this.config = config;
-      // const lastMessage = state.messages[state.messages.length - 1] as AIMessage;
-      // if (!lastMessage?.tool_calls?.length) {
-      //   return END;
-      // }
-      // return TOOLS;
-      return toolsCondition(state);
+      return toolsCondition(state, toolNode, this.invokedToolIds);
     };
 
-    const workflow = new StateGraph<t.BaseGraphState>({
-      channels: this.graphState,
-    })
-      .addNode(AGENT, this.createCallModel())
-      .addNode(TOOLS, this.initializeTools())
-      .addEdge(START, AGENT)
-      .addConditionalEdges(AGENT, routeMessage)
-      .addEdge(TOOLS, this.toolEnd ? END : AGENT);
+    const StateAnnotation = Annotation.Root({
+      messages: Annotation<BaseMessage[]>({
+        reducer: messagesStateReducer,
+        default: () => [],
+      }),
+    });
+
+    const workflow = new StateGraph(StateAnnotation)
+      .addNode(agentNode, this.createCallModel(agentId, currentModel))
+      .addNode(
+        toolNode,
+        this.initializeTools({
+          currentTools: agentContext.tools,
+          currentToolMap: agentContext.toolMap,
+        })
+      )
+      .addEdge(START, agentNode)
+      .addConditionalEdges(agentNode, routeMessage)
+      .addEdge(toolNode, agentContext.toolEnd ? END : agentNode);
+
+    // Cast to unknown to avoid tight coupling to external types; options are opt-in
+    return workflow.compile(this.compileOptions as unknown as never);
+  }
 
-    return workflow.compile();
+  createWorkflow(): t.CompiledStateWorkflow {
+    /** Use the default (first) agent for now */
+    const agentNode = this.createAgentNode(this.defaultAgentId);
+    const StateAnnotation = Annotation.Root({
+      messages: Annotation<BaseMessage[]>({
+        reducer: (a, b) => {
+          if (!a.length) {
+            this.startIndex = a.length + b.length;
+          }
+          const result = messagesStateReducer(a, b);
+          this.messages = result;
+          return result;
+        },
+        default: () => [],
+      }),
+    });
+    const workflow = new StateGraph(StateAnnotation)
+      .addNode(this.defaultAgentId, agentNode, { ends: [END] })
+      .addEdge(START, this.defaultAgentId)
+      .compile();
+
+    return workflow;
   }
 
   /* Dispatchers */
@@ -628,7 +788,10 @@ export class StandardGraph extends Graph<t.BaseGraphState, GraphNode> {
   /**
    * Dispatches a run step to the client, returns the step ID
    */
-  dispatchRunStep(stepKey: string, stepDetails: t.StepDetails): string {
+  async dispatchRunStep(
+    stepKey: string,
+    stepDetails: t.StepDetails
+  ): Promise<string> {
     if (!this.config) {
       throw new Error('No config provided');
     }
@@ -660,14 +823,19 @@ export class StandardGraph extends Graph<t.BaseGraphState, GraphNode> {
 
     this.contentData.push(runStep);
     this.contentIndexMap.set(stepId, runStep.index);
-    dispatchCustomEvent(GraphEvents.ON_RUN_STEP, runStep, this.config);
+    await safeDispatchCustomEvent(
+      GraphEvents.ON_RUN_STEP,
+      runStep,
+      this.config
+    );
     return stepId;
   }
 
-  handleToolCallCompleted(
+  async handleToolCallCompleted(
     data: t.ToolEndData,
-    metadata?: Record<string, unknown>
-  ): void {
+    metadata?: Record<string, unknown>,
+    omitOutput?: boolean
+  ): Promise<void> {
     if (!this.config) {
       throw new Error('No config provided');
     }
@@ -676,7 +844,11 @@ export class StandardGraph extends Graph<t.BaseGraphState, GraphNode> {
       return;
     }
 
-    const { input, output } = data;
+    const { input, output: _output } = data;
+    if ((_output as Command | undefined)?.lg_name === 'Command') {
+      return;
+    }
+    const output = _output as ToolMessage;
     const { tool_call_id } = output;
     const stepId = this.toolCallStepIds.get(tool_call_id) ?? '';
     if (!stepId) {
@@ -688,41 +860,45 @@ export class StandardGraph extends Graph<t.BaseGraphState, GraphNode> {
       throw new Error(`No run step found for stepId ${stepId}`);
     }
 
+    const dispatchedOutput =
+      typeof output.content === 'string'
+        ? output.content
+        : JSON.stringify(output.content);
+
     const args = typeof input === 'string' ? input : input.input;
     const tool_call = {
       args: typeof args === 'string' ? args : JSON.stringify(args),
       name: output.name ?? '',
       id: output.tool_call_id,
-      output:
-        typeof output.content === 'string'
-          ? output.content
-          : JSON.stringify(output.content),
+      output: omitOutput === true ? '' : dispatchedOutput,
       progress: 1,
     };
 
-    this.handlerRegistry?.getHandler(GraphEvents.ON_RUN_STEP_COMPLETED)?.handle(
-      GraphEvents.ON_RUN_STEP_COMPLETED,
-      {
-        result: {
-          id: stepId,
-          index: runStep.index,
-          type: 'tool_call',
-          tool_call,
-        } as t.ToolCompleteEvent,
-      },
-      metadata,
-      this
-    );
+    await this.handlerRegistry
+      ?.getHandler(GraphEvents.ON_RUN_STEP_COMPLETED)
+      ?.handle(
+        GraphEvents.ON_RUN_STEP_COMPLETED,
+        {
+          result: {
+            id: stepId,
+            index: runStep.index,
+            type: 'tool_call',
+            tool_call,
+          } as t.ToolCompleteEvent,
+        },
+        metadata,
+        this
+      );
   }
   /**
    * Static version of handleToolCallError to avoid creating strong references
    * that prevent garbage collection
    */
-  static handleToolCallErrorStatic(
+  static async handleToolCallErrorStatic(
     graph: StandardGraph,
     data: t.ToolErrorData,
     metadata?: Record<string, unknown>
-  ): void {
+  ): Promise<void> {
     if (!graph.config) {
       throw new Error('No config provided');
     }
@@ -752,7 +928,7 @@ export class StandardGraph extends Graph<t.BaseGraphState, GraphNode> {
       progress: 1,
     };
 
-    graph.handlerRegistry
+    await graph.handlerRegistry
       ?.getHandler(GraphEvents.ON_RUN_STEP_COMPLETED)
       ?.handle(
         GraphEvents.ON_RUN_STEP_COMPLETED,
@@ -773,14 +949,17 @@ export class StandardGraph extends Graph<t.BaseGraphState, GraphNode> {
    * Instance method that delegates to the static method
    * Kept for backward compatibility
    */
-  handleToolCallError(
+  async handleToolCallError(
     data: t.ToolErrorData,
     metadata?: Record<string, unknown>
-  ): void {
-    StandardGraph.handleToolCallErrorStatic(this, data, metadata);
+  ): Promise<void> {
+    await StandardGraph.handleToolCallErrorStatic(this, data, metadata);
   }
 
-  dispatchRunStepDelta(id: string, delta: t.ToolCallDelta): void {
+  async dispatchRunStepDelta(
+    id: string,
+    delta: t.ToolCallDelta
+  ): Promise<void> {
     if (!this.config) {
       throw new Error('No config provided');
     } else if (!id) {
@@ -790,14 +969,14 @@ export class StandardGraph extends Graph<t.BaseGraphState, GraphNode> {
       id,
       delta,
     };
-    dispatchCustomEvent(
+    await safeDispatchCustomEvent(
       GraphEvents.ON_RUN_STEP_DELTA,
       runStepDelta,
       this.config
     );
   }
 
-  dispatchMessageDelta(id: string, delta: t.MessageDelta): void {
+  async dispatchMessageDelta(id: string, delta: t.MessageDelta): Promise<void> {
     if (!this.config) {
       throw new Error('No config provided');
     }
@@ -805,14 +984,17 @@ export class StandardGraph extends Graph<t.BaseGraphState, GraphNode> {
       id,
       delta,
     };
-    dispatchCustomEvent(
+    await safeDispatchCustomEvent(
       GraphEvents.ON_MESSAGE_DELTA,
       messageDelta,
       this.config
     );
   }
 
-  dispatchReasoningDelta = (stepId: string, delta: t.ReasoningDelta): void => {
+  dispatchReasoningDelta = async (
+    stepId: string,
+    delta: t.ReasoningDelta
+  ): Promise<void> => {
     if (!this.config) {
       throw new Error('No config provided');
     }
@@ -820,7 +1002,7 @@ export class StandardGraph extends Graph<t.BaseGraphState, GraphNode> {
       id: stepId,
       delta,
     };
-    dispatchCustomEvent(
+    await safeDispatchCustomEvent(
       GraphEvents.ON_REASONING_DELTA,
       reasoningDelta,
       this.config