fix(server-utils): Instrument LangGraph stream executions

Fixes getsentry#19626

Author: psh4607

dev-packages/cloudflare-integration-tests/suites/tracing/langgraph/index.ts

@@ -61,6 +61,13 @@ export default Sentry.withSentry(
         messages: [{ role: 'user', content: 'What is the weather in SF?' }],
       });
 
+      const stream = await compiled.stream({
+        messages: [{ role: 'user', content: 'Stream the weather in SF' }],
+      });
+      for await (const _chunk of stream) {
+        // Consuming the iterator is what runs the graph and completes the agent span.
+      }
+
       return new Response(JSON.stringify({ success: true }));
     },
   },

dev-packages/cloudflare-integration-tests/suites/tracing/langgraph/test.ts

@@ -6,6 +6,7 @@ import {
   GEN_AI_OPERATION_NAME,
   GEN_AI_PIPELINE_NAME,
   GEN_AI_RESPONSE_MODEL,
+  GEN_AI_RESPONSE_STREAMING,
   GEN_AI_USAGE_INPUT_TOKENS,
   GEN_AI_USAGE_OUTPUT_TOKENS,
   GEN_AI_USAGE_TOTAL_TOKENS,
@@ -17,7 +18,7 @@ import { createRunner } from '../../../runner';
 // want to test that the instrumentation does not break in our
 // cloudflare SDK.
 
-it('traces langgraph compile and invoke operations', async ({ signal }) => {
+it('traces langgraph invoke and stream operations', async ({ signal }) => {
   const runner = createRunner(__dirname)
     .ignore('event')
     .expect(envelope => {
@@ -29,13 +30,14 @@ it('traces langgraph compile and invoke operations', async ({ signal }) => {
       const container = envelope[1]?.[1]?.[1] as any;
       expect(container).toBeDefined();
 
-      expect(container.items).toHaveLength(1);
+      expect(container.items).toHaveLength(2);
       expect(container.items.map((span: SerializedStreamedSpan) => span.name).sort()).toEqual([
         'invoke_agent weather_assistant',
+        'invoke_agent weather_assistant',
       ]);
 
       const invokeAgentSpan = container.items.find(
-        (span: SerializedStreamedSpan) => span.name === 'invoke_agent weather_assistant',
+        (span: SerializedStreamedSpan) => span.attributes[GEN_AI_RESPONSE_STREAMING] === undefined,
       );
       expect(invokeAgentSpan).toBeDefined();
       expect(invokeAgentSpan!.status).toBe('ok');
@@ -73,6 +75,16 @@ it('traces langgraph compile and invoke operations', async ({ signal }) => {
         type: 'integer',
         value: 30,
       });
+
+      const streamSpan = container.items.find(
+        (span: SerializedStreamedSpan) => span.attributes[GEN_AI_RESPONSE_STREAMING]?.value === true,
+      );
+      expect(streamSpan).toBeDefined();
+      expect(streamSpan!.status).toBe('ok');
+      expect(streamSpan!.attributes[GEN_AI_INPUT_MESSAGES]).toEqual({
+        type: 'string',
+        value: '[{"role":"user","content":"Stream the weather in SF"}]',
+      });
     })
     .start(signal);
   await runner.makeRequest('get', '/');

dev-packages/node-integration-tests/suites/tracing/langgraph/scenario.mjs

@@ -44,6 +44,13 @@ async function run() {
         { role: 'user', content: 'Tell me about the weather' },
       ],
     });
+
+    const stream = await graph.stream({
+      messages: [{ role: 'user', content: 'Stream the weather forecast' }],
+    });
+    for await (const _chunk of stream) {
+      // Consuming the iterator is what runs the graph and completes the agent span.
+    }
   });
 
   await Sentry.flush(2000);

dev-packages/node-integration-tests/suites/tracing/langgraph/test.ts

@@ -7,6 +7,7 @@ import {
   GEN_AI_OPERATION_NAME,
   GEN_AI_PIPELINE_NAME,
   GEN_AI_RESPONSE_MODEL,
+  GEN_AI_RESPONSE_STREAMING,
   GEN_AI_RESPONSE_TEXT,
   GEN_AI_RESPONSE_TOOL_CALLS,
   GEN_AI_SYSTEM_INSTRUCTIONS,
@@ -31,14 +32,15 @@ describe('LangGraph integration', () => {
         .expect({ transaction: { transaction: 'langgraph-test' } })
         .expect({
           span: container => {
-            expect(container.items).toHaveLength(2);
+            expect(container.items).toHaveLength(3);
             expect(container.items.map(span => span.name).sort()).toEqual([
               'invoke_agent weather_assistant',
               'invoke_agent weather_assistant',
+              'invoke_agent weather_assistant',
             ]);
 
             const invokeAgentSpans = container.items.filter(span => span.name === 'invoke_agent weather_assistant');
-            expect(invokeAgentSpans).toHaveLength(2);
+            expect(invokeAgentSpans).toHaveLength(3);
             for (const span of invokeAgentSpans) {
               expect(span.status).toBe('ok');
               expect(span.attributes['sentry.op'].value).toBe('gen_ai.invoke_agent');
@@ -47,6 +49,11 @@ describe('LangGraph integration', () => {
               expect(span.attributes[GEN_AI_AGENT_NAME].value).toBe('weather_assistant');
               expect(span.attributes[GEN_AI_PIPELINE_NAME].value).toBe('weather_assistant');
             }
+
+            const streamSpan = invokeAgentSpans.find(
+              span => span.attributes[GEN_AI_RESPONSE_STREAMING]?.value === true,
+            );
+            expect(streamSpan).toBeDefined();
           },
         })
         .start()
@@ -61,7 +68,7 @@ describe('LangGraph integration', () => {
         .expect({ transaction: { transaction: 'langgraph-test' } })
         .expect({
           span: container => {
-            expect(container.items).toHaveLength(2);
+            expect(container.items).toHaveLength(3);
 
             const weatherTodaySpan = container.items.find(span =>
               getStringAttributeValue(span.attributes[GEN_AI_INPUT_MESSAGES]?.value)?.includes(
@@ -83,6 +90,14 @@ describe('LangGraph integration', () => {
             expect(weatherDetailsSpan!.name).toBe('invoke_agent weather_assistant');
             expect(weatherDetailsSpan!.status).toBe('ok');
             expect(weatherDetailsSpan!.attributes['sentry.op'].value).toBe('gen_ai.invoke_agent');
+
+            const weatherStreamSpan = container.items.find(span =>
+              getStringAttributeValue(span.attributes[GEN_AI_INPUT_MESSAGES]?.value)?.includes(
+                'Stream the weather forecast',
+              ),
+            );
+            expect(weatherStreamSpan).toBeDefined();
+            expect(weatherStreamSpan!.attributes[GEN_AI_RESPONSE_STREAMING].value).toBe(true);
           },
         })
         .start()