ADK changes

PiperOrigin-RevId: 858348685

Author: reniowood

src/app/core/services/trace.service.spec.ts

@@ -122,6 +122,25 @@ describe('TraceService', () => {
          const indices = await firstValueFrom(service.hoveredMessageIndices$);
          expect(indices).toEqual([1]);
        });
+
+    it(
+        'should safely handle messages with missing event data', async () => {
+          const messages = [
+            {role: 'bot', eventId: 'e1'},
+            {role: 'bot', eventId: 'e2'},  // e2 is missing in eventData
+          ];
+          const eventData = new Map<string, any>([
+            ['e1', {invocationId: 'inv1'}],
+          ]);
+          const span: Partial<Span> = {
+            attributes: {'gcp.vertex.agent.event_id': 'e1'},
+          };
+          service.setMessages(messages);
+          service.setEventData(eventData);
+          service.setHoveredMessages(span as Span, 'inv1');
+          const indices = await firstValueFrom(service.hoveredMessageIndices$);
+          expect(indices).toEqual([0]);
+        });
   });
 
   describe('resetTraceService', () => {

src/app/core/services/trace.service.ts

@@ -14,79 +14,73 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
+
 import {Injectable} from '@angular/core';
-import {BehaviorSubject, Subject} from 'rxjs';
+import {BehaviorSubject} from 'rxjs';
+
 import {Span} from '../models/Trace';
+
 import {TraceService as TraceServiceInterface} from './interfaces/trace';
 
 @Injectable({providedIn: 'root'})
 export class TraceService implements TraceServiceInterface {
-    private selectedTraceRowSource = new BehaviorSubject<Span | undefined>(undefined);
-    selectedTraceRow$ = this.selectedTraceRowSource.asObservable();
+  private selectedTraceRowSource =
+      new BehaviorSubject<Span|undefined>(undefined);
+  selectedTraceRow$ = this.selectedTraceRowSource.asObservable();
 
-    private eventDataSource = new BehaviorSubject<Map<string, any> | undefined>(undefined);
-    eventData$ = this.eventDataSource.asObservable();
+  private eventDataSource =
+      new BehaviorSubject<Map<string, any>|undefined>(undefined);
+  eventData$ = this.eventDataSource.asObservable();
 
-    private hoveredMessageIndicesSource = new BehaviorSubject<number[]>([]);
-    hoveredMessageIndices$ = this.hoveredMessageIndicesSource.asObservable();
+  private hoveredMessageIndicesSource = new BehaviorSubject<number[]>([]);
+  hoveredMessageIndices$ = this.hoveredMessageIndicesSource.asObservable();
 
-    private messagesSource = new BehaviorSubject<any[]>([]);
-    messages$ = this.messagesSource.asObservable();
+  private messagesSource = new BehaviorSubject<any[]>([]);
+  messages$ = this.messagesSource.asObservable();
 
-    selectedRow(span: Span | undefined) {
-        this.selectedTraceRowSource.next(span);
-    }
+  selectedRow(span: Span|undefined) {
+    this.selectedTraceRowSource.next(span);
+  }
 
-    setEventData(data: Map<string, any> | undefined) {
-        this.eventDataSource.next(data);
-    }
+  setEventData(data: Map<string, any>|undefined) {
+    this.eventDataSource.next(data);
+  }
 
-    setMessages(messages: any[]) {
-        this.messagesSource.next(messages)
-    }
+  setMessages(messages: any[]) {
+    this.messagesSource.next(messages);
+  }
 
-    setHoveredMessages(span: Span | undefined, invocationId: string) {
-        if (!span) {
-            this.hoveredMessageIndicesSource.next([]);
-            return;
-        }
+  setHoveredMessages(span: Span|undefined, invocationId: string) {
+    if (!span) {
+      this.hoveredMessageIndicesSource.next([]);
+      return;
+    }
 
-        const attributes = span.attributes
-        const hasEvent: boolean = attributes && attributes['gcp.vertex.agent.event_id']
-        let index = 0;
-        const messageIndices = [];
-        for (const msg of this.messagesSource.value) {
-            if (msg.role == 'user') {
-                index++
-                continue
-            }
+    const attributes = span.attributes;
+    const hasEvent: boolean =
+        attributes && attributes['gcp.vertex.agent.event_id'];
+    const messageIndices = [];
+    for (const [index, msg] of this.messagesSource.value.entries()) {
+      if (msg.role === 'user') {
+        continue;
+      }
 
-            if (this.eventDataSource.value?.get(msg.eventId).invocationId != invocationId) {
-                index++
-                continue
-            }
+      if (this.eventDataSource.value?.get(msg.eventId)?.invocationId !==
+          invocationId) {
+        continue;
+      }
 
-            if (!hasEvent) {
-                messageIndices.push(index)
-                index++
-                continue
-            } else {
-                if (attributes['gcp.vertex.agent.event_id'] == msg.eventId) {
-                    messageIndices.push(index)
-                    index++
-                    continue
-                } else {
-                    index++
-                    continue
-                }
-            }
-        }
-        this.hoveredMessageIndicesSource.next(messageIndices);
+      if (!hasEvent ||
+          attributes['gcp.vertex.agent.event_id'] === msg.eventId) {
+        messageIndices.push(index);
+      }
     }
+    this.hoveredMessageIndicesSource.next(messageIndices);
+  }
 
-    resetTraceService() {
-        this.eventDataSource.next(undefined);
-        this.messagesSource.next([]);
-        this.hoveredMessageIndicesSource.next([])
-    }
+  resetTraceService() {
+    this.eventDataSource.next(undefined);
+    this.messagesSource.next([]);
+    this.hoveredMessageIndicesSource.next([]);
+  }
 }