fix: Refine user message extraction from invocation groups.

The findUserMsgFromInvocGroup function now requires both 'gcp.vertex.agent.invocation_id' and 'gcp.vertex.agent.llm_request' attributes to be present on a span to extract the user message.

PiperOrigin-RevId: 841953660

Author: google-genai-bot

src/app/components/trace-tab/trace-tab.component.spec.ts

@@ -146,4 +146,76 @@ describe('TraceTabComponent', () => {
       // mocking TraceTreeComponent
     });
   });
+
+  describe('findUserMsgFromInvocGroup', () => {
+    it('should find user message from span with both invocation_id and llm_request',
+       () => {
+         // First span has only invocation_id, second span has both
+         const group: Span[] = [
+           {
+             name: 'invocation',
+             start_time: 1733084700000000000,
+             end_time: 1733084760000000000,
+             span_id: 'span-1',
+             trace_id: 'trace-1',
+             attributes: {
+               'gcp.vertex.agent.invocation_id': 'invoc-1',
+             },
+           },
+           {
+             name: 'call_llm',
+             start_time: 1733084710000000000,
+             end_time: 1733084750000000000,
+             span_id: 'span-2',
+             parent_span_id: 'span-1',
+             trace_id: 'trace-1',
+             attributes: {
+               'gcp.vertex.agent.invocation_id': 'invoc-1',
+               'gcp.vertex.agent.llm_request':
+                   '{"contents":[{"role":"user","parts":[{"text":"hi"}]}]}',
+             },
+           },
+         ];
+
+         const result = component.findUserMsgFromInvocGroup(group);
+         expect(result).toBe('hi');
+       });
+
+    it('should return fallback when no span has llm_request', () => {
+      const group: Span[] = [
+        {
+          name: 'invocation',
+          start_time: 1733084700000000000,
+          end_time: 1733084760000000000,
+          span_id: 'span-1',
+          trace_id: 'trace-1',
+          attributes: {
+            'gcp.vertex.agent.invocation_id': 'invoc-1',
+          },
+        },
+      ];
+
+      const result = component.findUserMsgFromInvocGroup(group);
+      expect(result).toBe('[no invocation id found]');
+    });
+
+    it('should return error message on invalid JSON', () => {
+      const group: Span[] = [
+        {
+          name: 'call_llm',
+          start_time: 1733084700000000000,
+          end_time: 1733084760000000000,
+          span_id: 'span-1',
+          trace_id: 'trace-1',
+          attributes: {
+            'gcp.vertex.agent.invocation_id': 'invoc-1',
+            'gcp.vertex.agent.llm_request': 'invalid json{',
+          },
+        },
+      ];
+
+      const result = component.findUserMsgFromInvocGroup(group);
+      expect(result).toBe('[error parsing request]');
+    });
+  });
 });

src/app/components/trace-tab/trace-tab.component.ts

@@ -70,19 +70,28 @@ export class TraceTabComponent implements OnInit, OnChanges {
   }
 
   findUserMsgFromInvocGroup(group: any[]) {
+    // Find a span that has both invocation_id and llm_request
+    // The invocation_id is present on multiple spans, but llm_request
+    // is only on the call_llm span
     const eventItem = group?.find(
         item => item.attributes !== undefined &&
-            'gcp.vertex.agent.invocation_id' in item.attributes)
+            'gcp.vertex.agent.invocation_id' in item.attributes &&
+            'gcp.vertex.agent.llm_request' in item.attributes)
 
     if (!eventItem) {
       return '[no invocation id found]';
     }
 
-    const requestJson =
-        JSON.parse(eventItem.attributes['gcp.vertex.agent.llm_request']) as LlmRequest
-    const userContent =
-        requestJson.contents.filter((c: any) => c.role == 'user').at(-1)
-    return userContent?.parts[0]?.text ?? '[attachment]';
+    try {
+      const requestJson =
+          JSON.parse(eventItem.attributes['gcp.vertex.agent.llm_request']) as
+          LlmRequest
+      const userContent =
+          requestJson.contents.filter((c: any) => c.role == 'user').at(-1)
+      return userContent?.parts[0]?.text ?? '[attachment]';
+    } catch {
+      return '[error parsing request]';
+    }
   }
 
   findInvocIdFromTraceId(traceId: string) {