fix(server-utils): Instrument LangGraph stream executions

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

@@ -5,7 +5,10 @@ import {
   GEN_AI_INPUT_MESSAGES,
   GEN_AI_OPERATION_NAME,
   GEN_AI_PIPELINE_NAME,
+  GEN_AI_RESPONSE_FINISH_REASONS,
   GEN_AI_RESPONSE_MODEL,
+  GEN_AI_RESPONSE_STREAMING,
+  GEN_AI_RESPONSE_TEXT,
   GEN_AI_USAGE_INPUT_TOKENS,
   GEN_AI_USAGE_OUTPUT_TOKENS,
   GEN_AI_USAGE_TOTAL_TOKENS,
@@ -17,7 +20,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 +32,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 +77,40 @@ 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"}]',
+      });
+      expect(streamSpan!.attributes[GEN_AI_RESPONSE_TEXT]).toEqual({
+        type: 'string',
+        value: '[{"role":"assistant","content":"Mock response from LangGraph agent"}]',
+      });
+      expect(streamSpan!.attributes[GEN_AI_RESPONSE_MODEL]).toEqual({
+        type: 'string',
+        value: 'mock-model',
+      });
+      expect(streamSpan!.attributes[GEN_AI_RESPONSE_FINISH_REASONS]).toEqual({
+        type: 'array',
+        value: ['stop'],
+      });
+      expect(streamSpan!.attributes[GEN_AI_USAGE_INPUT_TOKENS]).toEqual({
+        type: 'integer',
+        value: 20,
+      });
+      expect(streamSpan!.attributes[GEN_AI_USAGE_OUTPUT_TOKENS]).toEqual({
+        type: 'integer',
+        value: 10,
+      });
+      expect(streamSpan!.attributes[GEN_AI_USAGE_TOTAL_TOKENS]).toEqual({
+        type: 'integer',
+        value: 30,
+      });
     })
     .start(signal);
   await runner.makeRequest('get', '/');

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

@@ -44,6 +44,16 @@ async function run() {
         { role: 'user', content: 'Tell me about the weather' },
       ],
     });
+
+    const stream = await graph.stream({
+      messages: [{ role: 'user', content: 'Stream the weather forecast' }],
+    });
+    await stream.pipeTo(new WritableStream());
+
+    const canceledStream = await graph.stream({
+      messages: [{ role: 'user', content: 'Cancel the weather forecast' }],
+    });
+    await canceledStream.cancel('no longer needed');
   });
 
   await Sentry.flush(2000);

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

@@ -6,7 +6,9 @@ import {
   GEN_AI_INPUT_MESSAGES,
   GEN_AI_OPERATION_NAME,
   GEN_AI_PIPELINE_NAME,
+  GEN_AI_RESPONSE_FINISH_REASONS,
   GEN_AI_RESPONSE_MODEL,
+  GEN_AI_RESPONSE_STREAMING,
   GEN_AI_RESPONSE_TEXT,
   GEN_AI_RESPONSE_TOOL_CALLS,
   GEN_AI_SYSTEM_INSTRUCTIONS,
@@ -31,14 +33,16 @@ describe('LangGraph integration', () => {
         .expect({ transaction: { transaction: 'langgraph-test' } })
         .expect({
           span: container => {
-            expect(container.items).toHaveLength(2);
+            expect(container.items).toHaveLength(4);
             expect(container.items.map(span => span.name).sort()).toEqual([
               'invoke_agent weather_assistant',
               '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(4);
             for (const span of invokeAgentSpans) {
               expect(span.status).toBe('ok');
               expect(span.attributes['sentry.op'].value).toBe('gen_ai.invoke_agent');
@@ -47,6 +51,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 streamSpans = invokeAgentSpans.filter(
+              span => span.attributes[GEN_AI_RESPONSE_STREAMING]?.value === true,
+            );
+            expect(streamSpans).toHaveLength(2);
           },
         })
         .start()
@@ -61,7 +70,7 @@ describe('LangGraph integration', () => {
         .expect({ transaction: { transaction: 'langgraph-test' } })
         .expect({
           span: container => {
-            expect(container.items).toHaveLength(2);
+            expect(container.items).toHaveLength(4);
 
             const weatherTodaySpan = container.items.find(span =>
               getStringAttributeValue(span.attributes[GEN_AI_INPUT_MESSAGES]?.value)?.includes(
@@ -83,6 +92,30 @@ 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);
+            expect(weatherStreamSpan!.attributes[GEN_AI_RESPONSE_TEXT].value).toBe(
+              '[{"role":"assistant","content":"Mock LLM response"}]',
+            );
+            expect(weatherStreamSpan!.attributes[GEN_AI_RESPONSE_MODEL].value).toBe('mock-model');
+            expect(weatherStreamSpan!.attributes[GEN_AI_RESPONSE_FINISH_REASONS].value).toEqual(['stop']);
+            expect(weatherStreamSpan!.attributes[GEN_AI_USAGE_INPUT_TOKENS].value).toBe(20);
+            expect(weatherStreamSpan!.attributes[GEN_AI_USAGE_OUTPUT_TOKENS].value).toBe(10);
+            expect(weatherStreamSpan!.attributes[GEN_AI_USAGE_TOTAL_TOKENS].value).toBe(30);
+
+            const canceledStreamSpan = container.items.find(span =>
+              getStringAttributeValue(span.attributes[GEN_AI_INPUT_MESSAGES]?.value)?.includes(
+                'Cancel the weather forecast',
+              ),
+            );
+            expect(canceledStreamSpan).toBeDefined();
+            expect(canceledStreamSpan!.attributes[GEN_AI_RESPONSE_STREAMING].value).toBe(true);
           },
         })
         .start()