@librechat/agents 3.1.39 → 3.1.40

package/src/graphs/MultiAgentGraph.ts

@@ -15,6 +15,7 @@ import {
   getCurrentTaskInput,
   messagesStateReducer,
 } from '@langchain/langgraph';
+import type { LangGraphRunnableConfig } from '@langchain/langgraph';
 import type { BaseMessage, AIMessageChunk } from '@langchain/core/messages';
 import type { ToolRunnableConfig } from '@langchain/core/tools';
 import type * as t from '@/types';
@@ -745,7 +746,8 @@ export class MultiAgentGraph extends StandardGraph {
 
       /** Wrapper function that handles agentMessages channel, handoff reception, and conditional routing */
       const agentWrapper = async (
-        state: t.MultiAgentGraphState
+        state: t.MultiAgentGraphState,
+        config?: LangGraphRunnableConfig
       ): Promise<t.MultiAgentGraphState | Command> => {
         let result: t.MultiAgentGraphState;
 
@@ -817,7 +819,7 @@ export class MultiAgentGraph extends StandardGraph {
             ...state,
             messages: messagesForAgent,
           };
-          result = await agentSubgraph.invoke(transformedState);
+          result = await agentSubgraph.invoke(transformedState, config);
           result = {
             ...result,
             agentMessages: [],
@@ -863,14 +865,14 @@ export class MultiAgentGraph extends StandardGraph {
             ...state,
             messages: state.agentMessages,
           };
-          result = await agentSubgraph.invoke(transformedState);
+          result = await agentSubgraph.invoke(transformedState, config);
           result = {
             ...result,
             /** Clear agentMessages for next agent */
             agentMessages: [],
           };
         } else {
-          result = await agentSubgraph.invoke(state);
+          result = await agentSubgraph.invoke(state, config);
         }
 
         /** If agent has both handoff and direct edges, use Command for exclusive routing */

package/src/run.ts

@@ -4,6 +4,7 @@ import { CallbackHandler } from '@langfuse/langchain';
 import { PromptTemplate } from '@langchain/core/prompts';
 import { RunnableLambda } from '@langchain/core/runnables';
 import { AzureChatOpenAI, ChatOpenAI } from '@langchain/openai';
+import { BaseCallbackHandler } from '@langchain/core/callbacks/base';
 import type {
   MessageContentComplex,
   BaseMessage,
@@ -240,9 +241,14 @@ export class Run<_T extends t.BaseGraphState> {
       ? this.getCallbacks(streamOptions.callbacks)
       : [];
 
-    config.callbacks = baseCallbacks.concat(streamCallbacks).concat({
+    const customHandler = BaseCallbackHandler.fromMethods({
       [Callback.CUSTOM_EVENT]: customEventCallback,
     });
+    customHandler.awaitHandlers = true;
+
+    config.callbacks = baseCallbacks
+      .concat(streamCallbacks)
+      .concat(customHandler);
 
     if (
       isPresent(process.env.LANGFUSE_SECRET_KEY) &&

package/src/specs/custom-event-await.test.ts

@@ -0,0 +1,215 @@
+import { HumanMessage } from '@langchain/core/messages';
+import type * as t from '@/types';
+import { ToolEndHandler, ModelEndHandler } from '@/events';
+import { ContentTypes, GraphEvents, Providers } from '@/common';
+import { ChatModelStreamHandler, createContentAggregator } from '@/stream';
+import { Run } from '@/run';
+
+describe('Custom event handler awaitHandlers behavior', () => {
+  jest.setTimeout(15000);
+
+  const llmConfig: t.LLMConfig = {
+    provider: Providers.OPENAI,
+    streaming: true,
+    streamUsage: false,
+  };
+
+  const config = {
+    configurable: {
+      thread_id: 'test-thread',
+    },
+    streamMode: 'values' as const,
+    version: 'v2' as const,
+  };
+
+  it('should fully aggregate all content before processStream returns', async () => {
+    const longResponse =
+      'The quick brown fox jumps over the lazy dog and then runs across the field to find shelter from the rain';
+
+    let aggregateCallCount = 0;
+    const { contentParts, aggregateContent } = createContentAggregator();
+
+    const wrappedAggregate: t.ContentAggregator = (params) => {
+      aggregateCallCount++;
+      aggregateContent(params);
+    };
+
+    let messageDeltaCount = 0;
+
+    const customHandlers: Record<string | GraphEvents, t.EventHandler> = {
+      [GraphEvents.TOOL_END]: new ToolEndHandler(),
+      [GraphEvents.CHAT_MODEL_END]: new ModelEndHandler(),
+      [GraphEvents.CHAT_MODEL_STREAM]: new ChatModelStreamHandler(),
+      [GraphEvents.ON_RUN_STEP_COMPLETED]: {
+        handle: (
+          event: GraphEvents.ON_RUN_STEP_COMPLETED,
+          data: t.StreamEventData
+        ) => {
+          wrappedAggregate({
+            event,
+            data: data as unknown as { result: t.ToolEndEvent },
+          });
+        },
+      },
+      [GraphEvents.ON_RUN_STEP]: {
+        handle: (event: GraphEvents.ON_RUN_STEP, data: t.StreamEventData) => {
+          wrappedAggregate({ event, data: data as t.RunStep });
+        },
+      },
+      [GraphEvents.ON_RUN_STEP_DELTA]: {
+        handle: (
+          event: GraphEvents.ON_RUN_STEP_DELTA,
+          data: t.StreamEventData
+        ) => {
+          wrappedAggregate({ event, data: data as t.RunStepDeltaEvent });
+        },
+      },
+      [GraphEvents.ON_MESSAGE_DELTA]: {
+        handle: async (
+          event: GraphEvents.ON_MESSAGE_DELTA,
+          data: t.StreamEventData
+        ) => {
+          messageDeltaCount++;
+          wrappedAggregate({ event, data: data as t.MessageDeltaEvent });
+        },
+      },
+      [GraphEvents.ON_REASONING_DELTA]: {
+        handle: (
+          event: GraphEvents.ON_REASONING_DELTA,
+          data: t.StreamEventData
+        ) => {
+          wrappedAggregate({ event, data: data as t.ReasoningDeltaEvent });
+        },
+      },
+    };
+
+    const run = await Run.create<t.IState>({
+      runId: 'test-await-handlers',
+      graphConfig: {
+        type: 'standard',
+        llmConfig,
+      },
+      returnContent: true,
+      customHandlers,
+    });
+
+    run.Graph!.overrideTestModel([longResponse]);
+
+    const inputs = { messages: [new HumanMessage('hello')] };
+    const finalContentParts = await run.processStream(inputs, config);
+
+    expect(finalContentParts).toBeDefined();
+    expect(finalContentParts!.length).toBeGreaterThan(0);
+
+    expect(messageDeltaCount).toBeGreaterThan(0);
+    expect(aggregateCallCount).toBeGreaterThan(0);
+
+    const typedParts = contentParts as t.MessageContentComplex[];
+    const textParts = typedParts.filter(
+      (p: t.MessageContentComplex | undefined) =>
+        p !== undefined && p.type === ContentTypes.TEXT
+    );
+    expect(textParts.length).toBeGreaterThan(0);
+
+    const aggregatedText = textParts
+      .map(
+        (p) =>
+          (p as { type: string; [ContentTypes.TEXT]: string })[
+            ContentTypes.TEXT
+          ]
+      )
+      .join('');
+    expect(aggregatedText).toBe(longResponse);
+  });
+
+  it('should aggregate content from async handlers before processStream returns', async () => {
+    const response =
+      'This is a test of async handler aggregation with multiple tokens';
+
+    const { contentParts, aggregateContent } = createContentAggregator();
+
+    let asyncHandlerCompletions = 0;
+
+    const customHandlers: Record<string | GraphEvents, t.EventHandler> = {
+      [GraphEvents.TOOL_END]: new ToolEndHandler(),
+      [GraphEvents.CHAT_MODEL_END]: new ModelEndHandler(),
+      [GraphEvents.CHAT_MODEL_STREAM]: new ChatModelStreamHandler(),
+      [GraphEvents.ON_RUN_STEP_COMPLETED]: {
+        handle: (
+          event: GraphEvents.ON_RUN_STEP_COMPLETED,
+          data: t.StreamEventData
+        ) => {
+          aggregateContent({
+            event,
+            data: data as unknown as { result: t.ToolEndEvent },
+          });
+        },
+      },
+      [GraphEvents.ON_RUN_STEP]: {
+        handle: (event: GraphEvents.ON_RUN_STEP, data: t.StreamEventData) => {
+          aggregateContent({ event, data: data as t.RunStep });
+        },
+      },
+      [GraphEvents.ON_RUN_STEP_DELTA]: {
+        handle: (
+          event: GraphEvents.ON_RUN_STEP_DELTA,
+          data: t.StreamEventData
+        ) => {
+          aggregateContent({ event, data: data as t.RunStepDeltaEvent });
+        },
+      },
+      [GraphEvents.ON_MESSAGE_DELTA]: {
+        handle: async (
+          event: GraphEvents.ON_MESSAGE_DELTA,
+          data: t.StreamEventData
+        ) => {
+          await new Promise<void>((resolve) => setTimeout(resolve, 5));
+          aggregateContent({ event, data: data as t.MessageDeltaEvent });
+          asyncHandlerCompletions++;
+        },
+      },
+      [GraphEvents.ON_REASONING_DELTA]: {
+        handle: (
+          event: GraphEvents.ON_REASONING_DELTA,
+          data: t.StreamEventData
+        ) => {
+          aggregateContent({ event, data: data as t.ReasoningDeltaEvent });
+        },
+      },
+    };
+
+    const run = await Run.create<t.IState>({
+      runId: 'test-async-handlers',
+      graphConfig: {
+        type: 'standard',
+        llmConfig,
+      },
+      returnContent: true,
+      customHandlers,
+    });
+
+    run.Graph!.overrideTestModel([response]);
+
+    const inputs = { messages: [new HumanMessage('hello')] };
+    await run.processStream(inputs, config);
+
+    expect(asyncHandlerCompletions).toBeGreaterThan(0);
+
+    const typedParts = contentParts as t.MessageContentComplex[];
+    const textParts = typedParts.filter(
+      (p: t.MessageContentComplex | undefined) =>
+        p !== undefined && p.type === ContentTypes.TEXT
+    );
+    expect(textParts.length).toBeGreaterThan(0);
+
+    const aggregatedText = textParts
+      .map(
+        (p) =>
+          (p as { type: string; [ContentTypes.TEXT]: string })[
+            ContentTypes.TEXT
+          ]
+      )
+      .join('');
+    expect(aggregatedText).toBe(response);
+  });
+});