Refactor stream concatenation, add streaming disabled metric

amychisholm03 · amychisholm03 · commit 2e336d37b2db · 2026-01-27T14:06:18.000-08:00
diff --git a/lib/subscribers/langchain/runnable-stream.js b/lib/subscribers/langchain/runnable-stream.js
@@ -12,7 +12,6 @@ class LangchainRunnableStreamSubscriber extends LangchainRunnableSubscriber {
   }
 
   asyncEnd(data) {
-    // Exit early if disabled.
     if (!this.enabled) {
       this.logger.debug('`ai_monitoring.enabled` is set to false, stream will not be instrumented.')
       return
@@ -23,15 +22,16 @@ class LangchainRunnableStreamSubscriber extends LangchainRunnableSubscriber {
       return
     }
 
-    // Get context.
     const ctx = this.agent.tracer.getContext()
     const { transaction } = ctx
     if (transaction?.isActive() !== true) {
       return
     }
 
-    // Extract data.
     const request = data?.arguments?.[0]
+    // Requests via LangGraph API have the `messages` property with the
+    // information we need, otherwise it just lives on the `request`
+    // object directly.
     const userRequest = request?.messages ? request.messages?.[0] : request
     const params = data?.arguments?.[1] || {}
     const metadata = params?.metadata ?? {}
@@ -41,7 +41,7 @@ class LangchainRunnableStreamSubscriber extends LangchainRunnableSubscriber {
     // Note: as of 18.x `ReadableStream` is a global
     // eslint-disable-next-line n/no-unsupported-features/node-builtins
     if (response instanceof ReadableStream) {
-      this.wrapNextHandler({ response, ctx, request: userRequest, metadata, tags })
+      this.instrumentStream({ response, ctx, request: userRequest, metadata, tags })
     } else {
       // Input error occurred which means a stream was not created.
       // Skip instrumenting streaming and create Llm Events from
@@ -67,22 +67,22 @@ class LangchainRunnableStreamSubscriber extends LangchainRunnableSubscriber {
    * @param {object} params.metadata metadata for the call
    * @param {Array} params.tags tags for the call
    */
-  wrapNextHandler({ ctx, response, request, metadata, tags }) {
+  instrumentStream({ ctx, response, request, metadata, tags }) {
     const self = this
     const orig = response.getReader
     response.getReader = function wrapedGetReader() {
       const reader = orig.apply(this, arguments)
       const origRead = reader.read
-      let content = ''
-      let langgraphMessages = []
+      const accumulator = { content: '', langgraphMessages: [] }
       reader.read = async function wrappedRead(...args) {
         try {
           const result = await origRead.apply(this, args)
           if (result?.done) {
             // only create Llm events when stream iteration is done
-            const responseMsgs = langgraphMessages.length > 0
-              ? langgraphMessages.filter((msg) => msg.constructor?.name !== 'HumanMessage')
-              : content
+            const responseMsgs = self.getResponseMessages(
+              accumulator.langgraphMessages,
+              accumulator.content
+            )
             self.recordChatCompletionEvents({
               ctx,
               response: responseMsgs,
@@ -92,29 +92,14 @@ class LangchainRunnableStreamSubscriber extends LangchainRunnableSubscriber {
             })
           } else {
             // Concat the streamed content
-            if (result?.value?.messages || result?.value?.agent?.messages) {
-              // LangGraph case:
-              // The result.value.messages field contains all messages,
-              // request and response, and adds new events for the length
-              // of the stream. The last iteration will contain all messages
-              // in the stream so we can just re-assign it.
-              langgraphMessages = result?.value?.messages ?? result?.value?.agent?.messages
-            } else if (typeof result?.value?.content === 'string') {
-              // LangChain MessageChunk case
-              content += result.value.content
-            } else if (typeof result?.value === 'string') {
-              // Base LangChain case
-              content += result.value
-            } else if (typeof result?.value?.[0] === 'string') {
-              // Array parser case
-              content += result.value[0]
-            }
+            self.accumulateStreamContent(result, accumulator)
           }
           return result
         } catch (error) {
-          const responseMsgs = langgraphMessages.length > 0
-            ? langgraphMessages.filter((msg) => msg.constructor?.name !== 'HumanMessage')
-            : content
+          const responseMsgs = self.getResponseMessages(
+            accumulator.langgraphMessages,
+            accumulator.content
+          )
           self.recordChatCompletionEvents({
             ctx,
             request,
@@ -125,14 +110,55 @@ class LangchainRunnableStreamSubscriber extends LangchainRunnableSubscriber {
           })
           throw error
         } finally {
-          // update segment duration on every stream iteration to extend
-          // the timer
+          // update segment duration on every stream
+          // iteration to extend the timer
           ctx.segment.touch()
         }
       }
       return reader
     }
   }
+
+  /**
+   * Accumulates streamed content from various LangChain/LangGraph result formats.
+   *
+   * @param {object} result the stream result chunk
+   * @param {object} accumulator object containing content string and langgraphMessages array
+   */
+  accumulateStreamContent(result, accumulator) {
+    if (result?.value?.messages || result?.value?.agent?.messages) {
+      // LangGraph case:
+      // The result.value.messages field contains all messages,
+      // request and response, and adds new events for the length
+      // of the stream. The last iteration will contain all messages
+      // in the stream, so we can just re-assign `langgraphMessages`.
+      accumulator.langgraphMessages = result?.value?.messages ?? result?.value?.agent?.messages
+    } else if (typeof result?.value?.content === 'string') {
+      // LangChain MessageChunk case
+      accumulator.content += result.value.content
+    } else if (typeof result?.value === 'string') {
+      // Base LangChain case
+      accumulator.content += result.value
+    } else if (typeof result?.value?.[0] === 'string') {
+      // LangChain array parser case
+      accumulator.content += result.value[0]
+    }
+  }
+
+  /**
+   * Extracts response messages from accumulated stream data.
+   * For LangGraph, filters out HumanMessages.
+   * For LangChain, returns concatenated content.
+   *
+   * @param {object[]} langgraphMessages accumulated LangGraph messages
+   * @param {string} content accumulated LangChain content
+   * @returns {object[]|string} the response messages
+   */
+  getResponseMessages(langgraphMessages, content) {
+    return langgraphMessages.length > 0
+      ? langgraphMessages.filter((msg) => msg.constructor?.name !== 'HumanMessage')
+      : content
+  }
 }
 
 module.exports = LangchainRunnableStreamSubscriber
diff --git a/lib/subscribers/langchain/runnable.js b/lib/subscribers/langchain/runnable.js
@@ -20,6 +20,9 @@ class LangchainRunnableSubscriber extends AiMonitoringChatSubscriber {
   asyncEnd(data) {
     const ctx = this.agent.tracer.getContext()
     const request = data?.arguments?.[0]
+    // Requests via LangGraph API have the `messages` property with the
+    // information we need, otherwise it just lives on the `request`
+    // object directly.
     const userRequest = request?.messages ? request.messages?.[0] : request
     const params = data?.arguments?.[1] || {}
     const { result: response, error: err } = data
diff --git a/lib/subscribers/langgraph/graph-stream.js b/lib/subscribers/langgraph/graph-stream.js
@@ -4,38 +4,44 @@
  */
 
 const AiMonitoringSubscriber = require('../ai-monitoring/base')
-const { DESTINATIONS } = require('#agentlib/config/attribute-filter.js')
 const LlmErrorMessage = require('#agentlib/llm-events/error-message.js')
 const LangGraphAgentEvent = require('#agentlib/llm-events/langgraph/agent.js')
-const { AI: { LANGGRAPH } } = require('#agentlib/metrics/names.js')
+const { AI: { LANGGRAPH, STREAMING_DISABLED } } = require('#agentlib/metrics/names.js')
 
 class LangGraphStreamSubscriber extends AiMonitoringSubscriber {
   constructor({ agent, logger }) {
     super({ agent,
       logger,
       packageName: '@langchain/langgraph',
       channelName: 'nr_stream',
-      name: `${LANGGRAPH.AGENT}/stream`,
+      // 'agent' is the default name for an unknown AI
+      // agent or one with no name
+      name: `${LANGGRAPH.AGENT}/stream/agent`,
       trackingPrefix: LANGGRAPH.TRACKING_PREFIX })
     this.events = ['asyncEnd']
   }
 
-  get enabled() {
-    return super.enabled && this.agent.config.ai_monitoring.streaming.enabled
+  get streamingEnabled() {
+    return this.agent.config.ai_monitoring.streaming.enabled
   }
 
   handler(data, ctx) {
-    // Update segment name with LangGraph agent name.
-    const aiAgentName = data?.self?.name ?? 'agent'
-    this.name = `${LANGGRAPH.AGENT}/stream/${aiAgentName}`
+    // Store LangGraph AI agent name and update
+    // segment name to use it.
+    this.aiAgentName = data?.self?.name ?? 'agent'
+    this.name = `${LANGGRAPH.AGENT}/stream/${this.aiAgentName}`
     return super.handler(data, ctx)
   }
 
   asyncEnd(data) {
-    // Get constants and exit early if need be
-    const { agent, logger } = this
+    const { agent, logger, aiAgentName } = this
     if (!this.enabled) {
-      logger.debug('LangGraph streaming instrumentation is disabled, not recording Llm events.')
+      logger.debug('LangGraph instrumentation is disabled, not instrumenting stream.')
+      return
+    }
+    if (!this.streamingEnabled) {
+      logger.debug('LangGraph streaming instrumentation is disabled, not instrumenting stream.')
+      this.agent.metrics.getOrCreateMetric(STREAMING_DISABLED).incrementCallCount()
       return
     }
     const ctx = agent.tracer.getContext()
@@ -44,17 +50,9 @@ class LangGraphStreamSubscriber extends AiMonitoringSubscriber {
       return
     }
 
-    // Extract data
     const { error: initialErr, result: stream } = data
-    const aiAgentName = data?.self?.name ?? 'agent'
-
-    // Mark as LLM event
-    transaction.trace.attributes.addAttribute(DESTINATIONS.TRANS_EVENT, 'llm', true)
-
-    // Add AI Agent component attribute
     segment.addSpanAttribute('subcomponent', `{"type": "APM-AI_AGENT", "name": ${aiAgentName}}`)
 
-    // If error already:
     if (initialErr) {
       this.recordAiAgentEvent({ aiAgentName, transaction, segment, error: initialErr })
       return
@@ -63,11 +61,11 @@ class LangGraphStreamSubscriber extends AiMonitoringSubscriber {
     // Note: as of 18.x `ReadableStream` is a global
     // eslint-disable-next-line n/no-unsupported-features/node-builtins
     if (stream instanceof ReadableStream) {
-      this.wrapNextHandler({ stream, segment, transaction, aiAgentName })
+      this.instrumentStream({ stream, segment, transaction, aiAgentName })
     }
   }
 
-  wrapNextHandler({ stream, segment, transaction, aiAgentName }) {
+  instrumentStream({ stream, segment, transaction, aiAgentName }) {
     const self = this
     const orig = stream.getReader
     stream.getReader = function wrappedGetReaderLangGraph() {
@@ -84,8 +82,8 @@ class LangGraphStreamSubscriber extends AiMonitoringSubscriber {
           self.recordAiAgentEvent({ aiAgentName, transaction, segment, error: err })
           throw err
         } finally {
-          // update segment duration on every stream iteration to extend
-          // the timer
+          // update segment duration on every stream
+          // iteration to extend the timer
           segment.touch()
         }
       }
diff --git a/test/versioned/langgraph/graph-stream.test.js b/test/versioned/langgraph/graph-stream.test.js
