@librechat/agents 3.1.95 → 3.1.97

package/src/run.ts

@@ -1,5 +1,5 @@
 // src/run.ts
-import './instrumentation';
+import { initializeLangfuseTracing } from './instrumentation';
 import { PromptTemplate } from '@langchain/core/prompts';
 import { RunnableLambda } from '@langchain/core/runnables';
 import { AzureChatOpenAI, ChatOpenAI } from '@langchain/openai';
@@ -36,15 +36,16 @@ import {
   type CallbackEntry,
 } from '@/utils/callbacks';
 import {
-  createLegacyLangfuseHandler,
   createLangfuseTraceMetadata,
   createLangfuseHandler,
   disposeLangfuseHandler,
   getLangfuseTraceName,
-  hasExplicitLangfuseConfig,
-  hasLangfuseEnvConfig,
   isLangfuseCallbackHandler,
 } from '@/langfuse';
+import {
+  resolveLangfuseConfig,
+  withLangfuseToolOutputTracingConfig,
+} from '@/langfuseToolOutputTracing';
 import type { HookRegistry } from '@/hooks';
 
 export const defaultOmitOptions = new Set([
@@ -120,6 +121,7 @@ export class Run<_T extends t.BaseGraphState> {
   private handlerRegistry?: HandlerRegistry;
   private hookRegistry?: HookRegistry;
   private humanInTheLoop?: t.HumanInTheLoopConfig;
+  private langfuse?: t.LangfuseConfig;
   private toolOutputReferences?: t.ToolOutputReferencesConfig;
   private eagerEventToolExecution?: t.EagerEventToolExecutionConfig;
   private toolExecution?: t.ToolExecutionConfig;
@@ -166,6 +168,7 @@ export class Run<_T extends t.BaseGraphState> {
     this.handlerRegistry = handlerRegistry;
     this.hookRegistry = config.hooks;
     this.humanInTheLoop = config.humanInTheLoop;
+    this.langfuse = config.langfuse;
     this.toolOutputReferences = config.toolOutputReferences;
     this.eagerEventToolExecution = config.eagerEventToolExecution;
     this.toolExecution = config.toolExecution;
@@ -238,6 +241,7 @@ export class Run<_T extends t.BaseGraphState> {
       signal,
       runId: this.id,
       agents: [agentConfig],
+      langfuse: this.langfuse,
       tokenCounter: this.tokenCounter,
       indexTokenCountMap: this.indexTokenCountMap,
       calibrationRatio: this.calibrationRatio,
@@ -264,6 +268,7 @@ export class Run<_T extends t.BaseGraphState> {
       runId: this.id,
       agents,
       edges,
+      langfuse: this.langfuse,
       tokenCounter: this.tokenCounter,
       indexTokenCountMap: this.indexTokenCountMap,
       calibrationRatio: this.calibrationRatio,
@@ -546,6 +551,97 @@ export class Run<_T extends t.BaseGraphState> {
     };
   }
 
+  private shouldClearHookSession(streamThrew: boolean): boolean {
+    return (
+      this._interrupt == null || this._haltedReason != null || streamThrew
+    );
+  }
+
+  private isAwaitingResume(streamThrew: boolean): boolean {
+    return (
+      this._interrupt != null && this._haltedReason == null && !streamThrew
+    );
+  }
+
+  private getStreamLangfuseConfig(
+    graph: StandardGraph | MultiAgentGraph
+  ): t.LangfuseConfig | undefined {
+    const primaryContext = graph.agentContexts.get(graph.defaultAgentId);
+    if (primaryContext != null) {
+      return resolveLangfuseConfig(this.langfuse, primaryContext.langfuse);
+    }
+
+    for (const context of graph.agentContexts.values()) {
+      const langfuse = resolveLangfuseConfig(this.langfuse, context.langfuse);
+      if (langfuse != null) {
+        return langfuse;
+      }
+    }
+
+    return this.langfuse;
+  }
+
+  private getStreamToolOutputTracingLangfuseConfig(
+    graph: StandardGraph | MultiAgentGraph
+  ): t.LangfuseConfig | undefined {
+    const toolOutputTracingConfigs = Array.from(
+      graph.agentContexts.values()
+    )
+      .map((context) => {
+        return resolveLangfuseConfig(this.langfuse, context.langfuse)
+          ?.toolOutputTracing;
+      })
+      .filter((config): config is t.LangfuseToolOutputTracingConfig => {
+        return config != null;
+      });
+
+    if (toolOutputTracingConfigs.length === 0) {
+      return this.langfuse?.toolOutputTracing != null
+        ? { toolOutputTracing: this.langfuse.toolOutputTracing }
+        : undefined;
+    }
+    if (toolOutputTracingConfigs.length === 1) {
+      return { toolOutputTracing: toolOutputTracingConfigs[0] };
+    }
+
+    let enabled: boolean | undefined;
+    let redactionText: string | undefined;
+    let redactedToolNameMatchMode: 'exact' | 'partial' | undefined;
+    const redactedToolNames = new Set<string>();
+
+    for (const config of toolOutputTracingConfigs) {
+      if (config.enabled === false) {
+        enabled = false;
+      } else if (enabled !== false && config.enabled != null) {
+        enabled = config.enabled;
+      }
+
+      redactionText ??= config.redactionText;
+      if (config.redactedToolNameMatchMode === 'partial') {
+        redactedToolNameMatchMode = 'partial';
+      } else {
+        redactedToolNameMatchMode ??= config.redactedToolNameMatchMode;
+      }
+
+      for (const toolName of config.redactedToolNames ?? []) {
+        redactedToolNames.add(toolName);
+      }
+    }
+
+    return {
+      toolOutputTracing: {
+        ...(enabled != null ? { enabled } : {}),
+        ...(redactedToolNames.size > 0
+          ? { redactedToolNames: Array.from(redactedToolNames) }
+          : {}),
+        ...(redactedToolNameMatchMode != null
+          ? { redactedToolNameMatchMode }
+          : {}),
+        ...(redactionText != null ? { redactionText } : {}),
+      },
+    };
+  }
+
   async processStream(
     inputs: t.IState | Command,
     callerConfig: Partial<RunnableConfig> & {
@@ -564,6 +660,8 @@ export class Run<_T extends t.BaseGraphState> {
         'Graph not initialized. Make sure to use Run.create() to instantiate the Run.'
       );
     }
+    const graphRunnable = this.graphRunnable;
+    const graph = this.Graph;
 
     /**
      * `Command` inputs (currently only `Command({ resume })`) are
@@ -596,7 +694,7 @@ export class Run<_T extends t.BaseGraphState> {
      * boundary.
      */
     if (!isResume) {
-      this.Graph.resetValues(streamOptions?.keepContent);
+      graph.resetValues(streamOptions?.keepContent);
     }
     this._interrupt = undefined;
     this._haltedReason = undefined;
@@ -619,34 +717,33 @@ export class Run<_T extends t.BaseGraphState> {
       streamCallbacks ? [streamCallbacks, customHandler] : [customHandler]
     );
 
-    if (
-      hasLangfuseEnvConfig() &&
-      !hasExplicitLangfuseConfig(this.Graph.agentContexts.values())
-    ) {
-      const userId =
-        typeof config.configurable?.user_id === 'string'
-          ? config.configurable.user_id
-          : undefined;
-      const sessionId =
-        typeof config.configurable?.thread_id === 'string'
-          ? config.configurable.thread_id
-          : undefined;
-      const primaryContext = this.Graph.agentContexts.get(
-        this.Graph.defaultAgentId
-      );
-      const traceMetadata = createLangfuseTraceMetadata({
-        messageId: this.id,
-        parentMessageId: config.configurable?.requestBody?.parentMessageId,
-        agentName: primaryContext?.name,
-      });
-      const handler = createLegacyLangfuseHandler({
-        userId,
-        sessionId,
-        traceMetadata,
-        tags: ['librechat', 'agent'],
-      });
+    const primaryContext = graph.agentContexts.get(graph.defaultAgentId);
+    const userId =
+      typeof config.configurable?.user_id === 'string'
+        ? config.configurable.user_id
+        : undefined;
+    const sessionId =
+      typeof config.configurable?.thread_id === 'string'
+        ? config.configurable.thread_id
+        : undefined;
+    const traceMetadata = createLangfuseTraceMetadata({
+      messageId: this.id,
+      parentMessageId: config.configurable?.requestBody?.parentMessageId,
+      agentId: graph.defaultAgentId,
+      agentName: primaryContext?.name,
+    });
+    const streamLangfuseConfig = this.getStreamLangfuseConfig(graph);
+    initializeLangfuseTracing(streamLangfuseConfig);
+    const langfuseHandler = createLangfuseHandler({
+      langfuse: streamLangfuseConfig,
+      userId,
+      sessionId,
+      traceMetadata,
+      tags: ['librechat', 'agent'],
+    });
+    if (langfuseHandler != null) {
       config.runName = config.runName ?? getLangfuseTraceName(traceMetadata);
-      config.callbacks = appendCallbacks(config.callbacks, [handler]);
+      config.callbacks = appendCallbacks(config.callbacks, [langfuseHandler]);
     }
 
     if (!this.id) {
@@ -671,25 +768,6 @@ export class Run<_T extends t.BaseGraphState> {
       }
     }
 
-    /**
-     * `streamEvents` accepts both state inputs and `Command` (resume) at
-     * runtime, but our `CompiledStateWorkflow` type narrows the first
-     * arg to `BaseGraphState`. Cast on the call so the resume path
-     * type-checks without widening the wrapper for every caller.
-     */
-    const stream = this.graphRunnable.streamEvents(inputs as t.IState, config, {
-      raiseError: true,
-      /**
-       * Prevent EventStreamCallbackHandler from processing custom events.
-       * Custom events are already handled via our createCustomEventCallback()
-       * which routes them through the handlerRegistry.
-       * Without this flag, EventStreamCallbackHandler throws errors when
-       * custom events are dispatched for run IDs not in its internal map
-       * (due to timing issues in parallel execution or after run cleanup).
-       */
-      ignoreCustomEvent: true,
-    });
-
     /**
      * Tracks whether the stream loop threw. Used by the `finally`
      * block to decide whether to honor the interrupt-preservation
@@ -701,7 +779,26 @@ export class Run<_T extends t.BaseGraphState> {
      */
     let streamThrew = false;
 
-    try {
+    const consumeStream = async (): Promise<void> => {
+      /**
+       * `streamEvents` accepts both state inputs and `Command` (resume) at
+       * runtime, but our `CompiledStateWorkflow` type narrows the first
+       * arg to `BaseGraphState`. Cast on the call so the resume path
+       * type-checks without widening the wrapper for every caller.
+       */
+      const stream = graphRunnable.streamEvents(inputs as t.IState, config, {
+        raiseError: true,
+        /**
+         * Prevent EventStreamCallbackHandler from processing custom events.
+         * Custom events are already handled via our createCustomEventCallback()
+         * which routes them through the handlerRegistry.
+         * Without this flag, EventStreamCallbackHandler throws errors when
+         * custom events are dispatched for run IDs not in its internal map
+         * (due to timing issues in parallel execution or after run cleanup).
+         */
+        ignoreCustomEvent: true,
+      });
+
       for await (const event of stream) {
         const { data, metadata, ...info } = event;
 
@@ -797,9 +894,8 @@ export class Run<_T extends t.BaseGraphState> {
             hook_event_name: 'Stop',
             runId: this.id,
             threadId,
-            agentId: this.Graph.defaultAgentId,
-            messages:
-              this.Graph.getRunMessages() ?? stateInputs?.messages ?? [],
+            agentId: graph.defaultAgentId,
+            messages: graph.getRunMessages() ?? stateInputs?.messages ?? [],
             stopHookActive: false, // will be true when stop is triggered by a hook (Phase 2)
           },
           sessionId: this.id,
@@ -807,6 +903,14 @@ export class Run<_T extends t.BaseGraphState> {
           /* Stop hook errors must not masquerade as stream failures */
         });
       }
+    };
+
+    try {
+      await withLangfuseToolOutputTracingConfig(
+        streamLangfuseConfig,
+        consumeStream,
+        this.getStreamToolOutputTracingLangfuseConfig(graph)
+      );
     } catch (err) {
       streamThrew = true;
       if (this.hookRegistry?.hasHookFor('StopFailure', this.id) === true) {
@@ -842,11 +946,7 @@ export class Run<_T extends t.BaseGraphState> {
        * expected, sessions must drop). Every state where no resume
        * is expected clears.
        */
-      if (
-        this._interrupt == null ||
-        this._haltedReason != null ||
-        streamThrew
-      ) {
+      if (this.shouldClearHookSession(streamThrew)) {
         this.hookRegistry?.clearSession(this.id);
       }
       /**
@@ -858,6 +958,7 @@ export class Run<_T extends t.BaseGraphState> {
        * unaffected — their entries live under their own session ids.
        */
       this.hookRegistry?.clearHaltSignal(this.id);
+      await disposeLangfuseHandler(langfuseHandler);
 
       /**
        * Break the reference chain that keeps heavy data alive via
@@ -915,8 +1016,7 @@ export class Run<_T extends t.BaseGraphState> {
        * Run from scratch) is a separate concern; see
        * `HumanInTheLoopConfig` JSDoc.
        */
-      const awaitingResume =
-        this._interrupt != null && this._haltedReason == null && !streamThrew;
+      const awaitingResume = this.isAwaitingResume(streamThrew);
       if (!this.skipCleanup && !awaitingResume) {
         this.Graph.clearHeavyState();
       }
@@ -1172,25 +1272,18 @@ export class Run<_T extends t.BaseGraphState> {
         typeof chainOptions.configurable?.thread_id === 'string'
           ? chainOptions.configurable.thread_id
           : undefined;
-      const hasExplicitLangfuse =
-        this.Graph != null &&
-        hasExplicitLangfuseConfig(this.Graph.agentContexts.values());
-      if (titleContext?.langfuse != null) {
-        titleLangfuseHandler = createLangfuseHandler({
-          langfuse: titleContext.langfuse,
-          userId,
-          sessionId,
-          traceMetadata,
-          tags: ['librechat', 'title'],
-        });
-      } else if (hasLangfuseEnvConfig() && !hasExplicitLangfuse) {
-        titleLangfuseHandler = createLegacyLangfuseHandler({
-          userId,
-          sessionId,
-          traceMetadata,
-          tags: ['librechat', 'title'],
-        });
-      }
+      const titleLangfuseConfig = resolveLangfuseConfig(
+        this.langfuse,
+        titleContext?.langfuse
+      );
+      initializeLangfuseTracing(titleLangfuseConfig);
+      titleLangfuseHandler = createLangfuseHandler({
+        langfuse: titleLangfuseConfig,
+        userId,
+        sessionId,
+        traceMetadata,
+        tags: ['librechat', 'title'],
+      });
 
       if (titleLangfuseHandler != null) {
         chainOptions.callbacks = appendCallbacks(chainOptions.callbacks, [
@@ -1263,9 +1356,14 @@ export class Run<_T extends t.BaseGraphState> {
 
     try {
       try {
-        return await fullChain.invoke(
-          { input: inputText, output: response },
-          invokeConfig
+        return await withLangfuseToolOutputTracingConfig(
+          this.langfuse,
+          () =>
+            fullChain.invoke(
+              { input: inputText, output: response },
+              invokeConfig
+            ),
+          titleContext?.langfuse
         );
       } catch (_e) {
         // Fallback: strip callbacks to avoid EventStream tracer errors in certain environments
@@ -1278,9 +1376,14 @@ export class Run<_T extends t.BaseGraphState> {
         const safeConfig = Object.assign({}, rest, {
           callbacks: langfuseHandler ? [langfuseHandler] : [],
         });
-        return await fullChain.invoke(
-          { input: inputText, output: response },
-          safeConfig as Partial<RunnableConfig>
+        return await withLangfuseToolOutputTracingConfig(
+          this.langfuse,
+          () =>
+            fullChain.invoke(
+              { input: inputText, output: response },
+              safeConfig as Partial<RunnableConfig>
+            ),
+          titleContext?.langfuse
         );
       }
     } finally {

package/src/specs/langfuse-callbacks.test.ts

@@ -6,10 +6,23 @@ import type * as t from '@/types';
 
 const mockSpan = {
   end: jest.fn(),
+  spanContext: jest.fn(() => ({
+    traceId: 'trace-id',
+    spanId: 'span-id',
+    traceFlags: 1,
+  })),
   setAttributes: jest.fn(),
   setStatus: jest.fn(),
 };
 const mockStartSpan = jest.fn(() => mockSpan);
+const mockStartActiveSpan = jest.fn(
+  (
+    _name: string,
+    _options: unknown,
+    _context: unknown,
+    callback: (span: typeof mockSpan) => unknown
+  ) => callback(mockSpan)
+);
 const mockForceFlush = jest.fn();
 const mockShutdown = jest.fn();
 
@@ -22,6 +35,7 @@ jest.mock('@opentelemetry/sdk-trace-base', () => ({
   BasicTracerProvider: jest.fn().mockImplementation(() => ({
     forceFlush: mockForceFlush,
     getTracer: jest.fn(() => ({
+      startActiveSpan: mockStartActiveSpan,
       startSpan: mockStartSpan,
     })),
     shutdown: mockShutdown,
@@ -70,6 +84,169 @@ describe('Langfuse callback composition', () => {
 
     await run.processStream({ messages: [new HumanMessage('hello')] }, config);
 
-    expect(mockStartSpan).toHaveBeenCalled();
+    expect(mockStartActiveSpan).toHaveBeenCalled();
+    expect(mockForceFlush).toHaveBeenCalled();
+  });
+
+  it('attaches Langfuse callbacks for direct graph invocations', async () => {
+    const run = await Run.create<t.IState>({
+      runId: 'test-langfuse-direct-graph',
+      graphConfig: {
+        type: 'standard',
+        agents: [
+          {
+            agentId: 'agent_abc123',
+            name: 'DWAINE',
+            provider: Providers.OPENAI,
+            clientOptions: { model: 'gpt-4' },
+            tools: [],
+            langfuse: {
+              enabled: true,
+              publicKey: 'pk-test',
+              secretKey: 'sk-test',
+            },
+          },
+        ],
+      },
+      skipCleanup: true,
+    });
+
+    run.Graph?.overrideTestModel(['hello']);
+    const workflow = run.Graph?.createWorkflow();
+    await workflow?.invoke(
+      { messages: [new HumanMessage('hello')] },
+      {
+        callbacks: [],
+        configurable: { thread_id: 'thread-1', user_id: 'user-1' },
+      }
+    );
+
+    expect(mockStartActiveSpan).toHaveBeenCalled();
+  });
+
+  it('preserves per-agent Langfuse config when a stream callback already exists', async () => {
+    const { LangfuseSpanProcessor } = await import('@langfuse/otel');
+    const { initializeLangfuseTracing } = await import('@/instrumentation');
+    const { createLangfuseHandler } = await import('@/langfuse');
+    initializeLangfuseTracing({
+      publicKey: 'pk-run',
+      secretKey: 'sk-run',
+      baseUrl: 'https://langfuse.run',
+    });
+    const streamHandler = createLangfuseHandler({
+      langfuse: {
+        publicKey: 'pk-run',
+        secretKey: 'sk-run',
+        baseUrl: 'https://langfuse.run',
+      },
+    });
+    const run = await Run.create<t.IState>({
+      runId: 'test-langfuse-agent-callback-override',
+      graphConfig: {
+        type: 'standard',
+        agents: [
+          {
+            agentId: 'agent_abc123',
+            name: 'DWAINE',
+            provider: Providers.OPENAI,
+            clientOptions: { model: 'gpt-4' },
+            tools: [],
+            langfuse: {
+              enabled: true,
+              publicKey: 'pk-agent',
+              secretKey: 'sk-agent',
+              baseUrl: 'https://langfuse.agent',
+            },
+          },
+        ],
+      },
+      skipCleanup: true,
+    });
+
+    run.Graph?.overrideTestModel(['hello']);
+    const workflow = run.Graph?.createWorkflow();
+    await workflow?.invoke(
+      { messages: [new HumanMessage('hello')] },
+      {
+        callbacks: streamHandler != null ? [streamHandler] : [],
+        configurable: { thread_id: 'thread-1', user_id: 'user-1' },
+      }
+    );
+
+    expect(LangfuseSpanProcessor).toHaveBeenCalledWith(
+      expect.objectContaining({
+        publicKey: 'pk-agent',
+        secretKey: 'sk-agent',
+        baseUrl: 'https://langfuse.agent',
+      })
+    );
+  });
+
+  it('adds current agent metadata when a stream Langfuse callback already exists', async () => {
+    const metadataSpy = jest.fn();
+    const { createLangfuseHandler } = await import('@/langfuse');
+    const streamHandler = createLangfuseHandler({
+      langfuse: {
+        publicKey: 'pk-run',
+        secretKey: 'sk-run',
+        baseUrl: 'https://langfuse.run',
+      },
+    });
+    const run = await Run.create<t.IState>({
+      runId: 'test-langfuse-agent-metadata-with-stream-callback',
+      graphConfig: {
+        type: 'multi-agent',
+        agents: [
+          {
+            agentId: 'agent_default',
+            name: 'Default Agent',
+            provider: Providers.OPENAI,
+            clientOptions: { model: 'gpt-4' },
+            tools: [],
+          },
+          {
+            agentId: 'agent_specialist',
+            name: 'Specialist Agent',
+            provider: Providers.OPENAI,
+            clientOptions: { model: 'gpt-4' },
+            tools: [],
+          },
+        ],
+        edges: [],
+      },
+      skipCleanup: true,
+    });
+
+    run.Graph?.overrideTestModel(['hello from specialist']);
+    const agentNode = run.Graph?.createAgentNode('agent_specialist');
+    await agentNode?.invoke(
+      { messages: [new HumanMessage('hello')] },
+      {
+        callbacks: [
+          ...(streamHandler != null ? [streamHandler] : []),
+          {
+            handleChatModelStart: async (
+              _llm: unknown,
+              _messages: unknown,
+              _runId: string,
+              _parentRunId?: string,
+              _extraParams?: unknown,
+              _tags?: string[],
+              metadata?: Record<string, unknown>
+            ): Promise<void> => {
+              metadataSpy(metadata);
+            },
+          },
+        ],
+        configurable: { thread_id: 'thread-1', user_id: 'user-1' },
+      }
+    );
+
+    expect(metadataSpy).toHaveBeenCalledWith(
+      expect.objectContaining({
+        agentId: 'agent_specialist',
+        agentName: 'Specialist Agent',
+      })
+    );
   });
 });