openxla · copybara-service · Feb 9, 2026
diff --git a/frontend/app/components/trace_viewer/trace_viewer.ng.html b/frontend/app/components/trace_viewer/trace_viewer.ng.html
@@ -1 +1,9 @@
-<iframe [src]="url | safe"></iframe>
+<trace-viewer-container [url]="url"
+                        [traceViewerModule]="traceViewerModule"
+                        [useTraceViewerV2]="useTraceViewerV2"
+                        [selectedEvent]="selectedEvent"
+                        [selectedEventProperties]="selectedEventProperties"
+                        [eventDetailColumns]="eventDetailColumns"
+                        (eventSelected)="onEventSelected($event)"
+                        (searchEvents)="onSearchEvents($event)">
+</trace-viewer-container>
diff --git a/frontend/app/components/trace_viewer/trace_viewer.ts b/frontend/app/components/trace_viewer/trace_viewer.ts
@@ -1,34 +1,61 @@
 import {PlatformLocation} from '@angular/common';
-import {Component, inject, Injector, OnDestroy} from '@angular/core';
+import {HttpClient} from '@angular/common/http';
+import {AfterViewInit, Component, inject, Injector, OnDestroy, OnInit, ViewChild} from '@angular/core';
 import {ActivatedRoute} from '@angular/router';
 import {Store} from '@ngrx/store';
 import {API_PREFIX, DATA_API, PLUGIN_NAME} from 'org_xprof/frontend/app/common/constants/constants';
 import {HostMetadata} from 'org_xprof/frontend/app/common/interfaces/hosts';
 import {NavigationEvent} from 'org_xprof/frontend/app/common/interfaces/navigation_event';
+import {EntrySelectedEventDetail, SelectedEvent, SelectedEventProperty, TraceViewerContainer, } from 'org_xprof/frontend/app/components/trace_viewer_container/trace_viewer_container';
+import {TraceData as MainTraceData, traceViewerV2Main, TraceViewerV2Module} from 'org_xprof/frontend/app/components/trace_viewer_v2/main';
 import {SOURCE_CODE_SERVICE_INTERFACE_TOKEN} from 'org_xprof/frontend/app/services/source_code_service/source_code_service_interface';
 import {getHostsState} from 'org_xprof/frontend/app/store/selectors';
 import {combineLatest, ReplaySubject} from 'rxjs';
 import {takeUntil} from 'rxjs/operators';
 
+interface TraceData {
+  traceEvents?: Array<{[key: string]: unknown}>;
+  [key: string]: unknown;
+}
+
+/**
+ * The name of the event selected custom event, dispatched from WASM in Trace
+ * Viewer v2.
+ */
+export const EVENT_SELECTED_EVENT_NAME = 'eventselected';
+
 /** A trace viewer component. */
 @Component({
   standalone: false,
   selector: 'trace-viewer',
   templateUrl: './trace_viewer.ng.html',
   styleUrls: ['./trace_viewer.css']
 })
-export class TraceViewer implements OnDestroy {
+export class TraceViewer implements OnInit, AfterViewInit, OnDestroy {
   /** Handles on-destroy Subject, used to unsubscribe. */
   private readonly destroyed = new ReplaySubject<void>(1);
+  private navigationEvent: NavigationEvent = {};
   private readonly injector = inject(Injector);
   private readonly store = inject(Store<{}>);
 
   url = '';
   pathPrefix = '';
   sourceCodeServiceIsAvailable = false;
   hostList: string[] = [];
+  useTraceViewerV2 = true;
+  traceViewerModule: TraceViewerV2Module|null = null;
+  selectedEvent: SelectedEvent|null = null;
+  selectedEventProperties: SelectedEventProperty[] = [];
+  eventDetailColumns: string[] = ['property', 'value'];
+  private readonly eventArgsCache = new Map<string, {[key: string]: string}>();
+  private queryString = '';
+  searching = false;
+
+  @ViewChild(TraceViewerContainer, {static: false})
+  container?: TraceViewerContainer;
 
   constructor(
+      private readonly httpClient: HttpClient,
       platformLocation: PlatformLocation,
       route: ActivatedRoute,
   ) {
@@ -44,7 +71,8 @@ export class TraceViewer implements OnDestroy {
             this.hostList =
                 hostsMetadata.map((host: HostMetadata) => host.hostname);
           }
-          this.update(params as NavigationEvent);
+          this.navigationEvent = {...params, ...queryParams};
+          this.update(this.navigationEvent);
         });
 
     // We don't need the source code service to be persistently available.
@@ -59,41 +87,169 @@ export class TraceViewer implements OnDestroy {
         });
   }
 
+  ngOnInit() {
+    // Event listeners are handled by TraceViewerContainer.
+  }
+
+  ngAfterViewInit() {
+    if (this.useTraceViewerV2) {
+      // Use setTimeout to defer initialization to the next tick to avoid
+      // ExpressionChangedAfterItHasBeenCheckedError when setting 'loading'
+      // state.
+      setTimeout(() => {
+        this.initializeWasmApp();
+      });
+    }
+  }
+
+  async initializeWasmApp() {
+    this.traceViewerModule = await traceViewerV2Main();
+    this.update(this.navigationEvent);
+  }
+
   update(event: NavigationEvent) {
     const isStreaming = (event.tag === 'trace_viewer@');
     const run = event.run || '';
     const tag = event.tag || '';
     const runPath = event.run_path || '';
     const sessionPath = event.session_path || '';
-    let queryString = `run=${run}&tag=${tag}`;
+    this.queryString = `run=${run}&tag=${tag}`;
 
     if (sessionPath) {
-      queryString += `&session_path=${sessionPath}`;
+      this.queryString += `&session_path=${sessionPath}`;
     } else if (runPath) {
-      queryString += `&run_path=${runPath}`;
+      this.queryString += `&run_path=${runPath}`;
     }
 
     if (event.hosts && typeof event.hosts === 'string') {
       // Since event.hosts is a comma-separated string, we can use it directly.
-      queryString += `&hosts=${event.hosts}`;
+      this.queryString += `&hosts=${event.hosts}`;
     } else if (event.host) {
-      queryString += `&host=${event.host}`;
+      this.queryString += `&host=${event.host}`;
     } else {
-      queryString +=
+      this.queryString +=
           `&host=${this.hostList.length > 0 ? this.hostList[0] : ''}`;
     }
 
-    const traceDataUrl = `${this.pathPrefix}${DATA_API}?${queryString}`;
-    this.url = `${this.pathPrefix}${API_PREFIX}${
-        PLUGIN_NAME}/trace_viewer_index.html?is_streaming=${
-        isStreaming}&is_oss=true&trace_data_url=${
-        encodeURIComponent(traceDataUrl)}&source_code_service=${
-        this.sourceCodeServiceIsAvailable}`;
+    const traceDataUrl = `${this.pathPrefix}${DATA_API}?${this.queryString}`;
+
+    if (this.useTraceViewerV2) {
+      if (this.traceViewerModule && this.traceViewerModule.loadJsonData) {
+        this.traceViewerModule.loadJsonData(traceDataUrl);
+      }
+    } else {
+      this.url = `${this.pathPrefix}${API_PREFIX}${
+          PLUGIN_NAME}/trace_viewer_index.html?is_streaming=${
+          isStreaming}&is_oss=true&trace_data_url=${
+          encodeURIComponent(traceDataUrl)}&source_code_service=${
+          this.sourceCodeServiceIsAvailable}`;
+    }
   }
 
   ngOnDestroy() {
     // Unsubscribes all pending subscriptions.
     this.destroyed.next();
     this.destroyed.complete();
   }
+
+  onEventSelected(event: EntrySelectedEventDetail|null) {
+    if (!event) {
+      this.selectedEvent = null;
+      this.selectedEventProperties = [];
+      return;
+    }
+
+    this.selectedEvent = {
+      name: event.name,
+      startUsFormatted: event.startUsFormatted,
+      durationUsFormatted: event.durationUsFormatted,
+    };
+
+    const properties: SelectedEventProperty[] = [];
+    properties.push({property: 'Name', value: event.name});
+    properties.push({property: 'Start Time', value: event.startUsFormatted});
+    properties.push({property: 'Duration', value: event.durationUsFormatted});
+    if (event.hloModuleName) {
+      properties.push({property: 'HLO Module', value: event.hloModuleName});
+    }
+    if (event.hloOpName) {
+      properties.push({property: 'HLO Op', value: event.hloOpName});
+    }
+    this.selectedEventProperties = properties;
+
+    if (event.uid) {
+      this.maybeFetchEventArgs(
+          event.name, event.startUs, event.durationUs, event.uid);
+    }
+  }
+
+  onSearchEvents(query: string) {
+    if (!query) {
+      this.traceViewerModule?.setSearchResultsInWasm({traceEvents: []});
+      this.container?.updateSearchResultCountText();
+      return;
+    }
+    this.searching = true;
+    const searchParams = new URLSearchParams(this.queryString);
+    searchParams.set('search_prefix', query);
+    this.httpClient
+        .get<TraceData>(
+            `${this.pathPrefix}${DATA_API}?${searchParams.toString()}`)
+        .pipe(takeUntil(this.destroyed))
+        .subscribe((data) => {
+          this.searching = false;
+          if (this.traceViewerModule && data) {
+            this.traceViewerModule.setSearchResultsInWasm({
+              ...data,
+              traceEvents: data.traceEvents || [],
+            } as MainTraceData);
+            this.container?.updateSearchResultCountText();
+          }
+        });
+  }
+
+  private maybeFetchEventArgs(
+      name: string, startUs: number, durationUs: number, uid: string) {
+    const key = `${name}-${startUs}-${durationUs}`;
+    if (this.eventArgsCache.has(key)) {
+      if (this.selectedEvent) {
+        this.addArgsToSelectedEvent(this.eventArgsCache.get(key)!);
+      }
+      return;
+    }
+
+    const params = new URLSearchParams(this.queryString);
+    params.set('event_name', name);
+    params.set('start_time_ms', (startUs / 1000).toString());
+    params.set('duration_ms', (durationUs / 1000).toString());
+    params.set('unique_id', Math.floor(Number(uid)).toString());
+
+    this.httpClient
+        .get<TraceData>(`${this.pathPrefix}${DATA_API}?${params.toString()}`)
+        .pipe(takeUntil(this.destroyed))
+        .subscribe((data) => {
+          const traceData = data as TraceData;
+          if (!traceData || !traceData.traceEvents ||
+              traceData.traceEvents.length === 0) {
+            return;
+          }
+          const lastEvent =
+              traceData.traceEvents[traceData.traceEvents.length - 1];
+          if (lastEvent['ph'] === 'X' && this.selectedEvent &&
+              lastEvent['args']) {
+            const args = lastEvent['args'] as {[key: string]: string};
+            this.eventArgsCache.set(key, args);
+            this.addArgsToSelectedEvent(args);
+          }
+        });
+  }
+
+  private addArgsToSelectedEvent(args: {[key: string]: string}) {
+    if (!this.selectedEvent) return;
+    const properties = [...this.selectedEventProperties];
+    for (const key of Object.keys(args)) {
+      properties.push({property: key, value: args[key]});
+    }
+    this.selectedEventProperties = properties;
+  }
 }
diff --git a/frontend/app/components/trace_viewer/trace_viewer_module.ts b/frontend/app/components/trace_viewer/trace_viewer_module.ts
@@ -1,7 +1,9 @@
 import {CommonModule} from '@angular/common';
+import {HttpClientModule} from '@angular/common/http';
 import {NgModule} from '@angular/core';
 import {MatIconModule} from '@angular/material/icon';
 import {MatProgressBarModule} from '@angular/material/progress-bar';
+import {TraceViewerContainer} from 'org_xprof/frontend/app/components/trace_viewer_container/trace_viewer_container';
 import {PipesModule} from 'org_xprof/frontend/app/pipes/pipes_module';
 
 import {TraceViewer} from './trace_viewer';
@@ -11,9 +13,11 @@ import {TraceViewer} from './trace_viewer';
   declarations: [TraceViewer],
   imports: [
     CommonModule,
+    HttpClientModule,
     MatIconModule,
     MatProgressBarModule,
     PipesModule,
+    TraceViewerContainer,
   ],
   exports: [TraceViewer]
 })

diff --git a/frontend/app/components/trace_viewer_container/trace_viewer_container.ts b/frontend/app/components/trace_viewer_container/trace_viewer_container.ts
@@ -8,7 +8,7 @@ import {MatInputModule} from '@angular/material/input';
 import {MatProgressBarModule} from '@angular/material/progress-bar';
 import {MatProgressSpinnerModule} from '@angular/material/progress-spinner';
 import {MatTableModule} from '@angular/material/table';
-import {isSearchEventsEvent, LOADING_STATUS_UPDATE_EVENT_NAME, SEARCH_EVENTS_EVENT_NAME, type SearchEventsEventDetail, TraceViewerV2LoadingStatus, type TraceViewerV2Module} from 'org_xprof/frontend/app/components/trace_viewer_v2/main';
+import {isSearchEventsEvent, LOADING_STATUS_UPDATE_EVENT_NAME, SEARCH_EVENTS_EVENT_NAME, TraceViewerV2LoadingStatus, type TraceViewerV2Module} from 'org_xprof/frontend/app/components/trace_viewer_v2/main';
 import {PipesModule} from 'org_xprof/frontend/app/pipes/pipes_module';
 import {interval, ReplaySubject, Subject, Subscription} from 'rxjs';
 import {debounceTime, takeUntil} from 'rxjs/operators';
@@ -133,8 +133,8 @@ export class TraceViewerContainer implements OnInit, OnDestroy, AfterViewInit {
   @Input() selectedEventProperties: SelectedEventProperty[] = [];
   @Input() eventDetailColumns: string[] = [];
   @Output()
-  readonly eventSelected = new EventEmitter<EntrySelectedEventDetail>();
-  @Output() readonly searchEvents = new EventEmitter<SearchEventsEventDetail>();
+  readonly eventSelected = new EventEmitter<EntrySelectedEventDetail|null>();
+  @Output() readonly searchEvents = new EventEmitter<string>();
 
   @ViewChild('tvIframe') tvIframe?: ElementRef<HTMLIFrameElement>;
 
@@ -262,14 +262,18 @@ export class TraceViewerContainer implements OnInit, OnDestroy, AfterViewInit {
     if (!isEntrySelectedEvent(e)) {
       return;
     }
-    this.eventSelected.emit(e.detail);
+    if (e.detail.eventIndex === -1) {
+      this.eventSelected.emit(null);
+    } else {
+      this.eventSelected.emit(e.detail);
+    }
   };
 
   private readonly searchEventsEventListener = (e: Event) => {
     if (!isSearchEventsEvent(e)) {
       return;
     }
-    this.searchEvents.emit(e.detail);
+    this.searchEvents.emit(e.detail.events_query);
   };
 
   /**
@@ -332,6 +336,7 @@ export class TraceViewerContainer implements OnInit, OnDestroy, AfterViewInit {
 
   onSearchEvent(query: string) {
     this.search$.next(query);
+    this.searchEvents.emit(query);
   }
 
   nextSearchResult() {

diff --git a/frontend/app/components/trace_viewer_v2/BUILD b/frontend/app/components/trace_viewer_v2/BUILD
@@ -126,6 +126,7 @@ wasm_cc_library(
 BINARY_LINKOPTS = [
     "-sASYNCIFY",
     "-sMODULARIZE",
+    "-sMINIMAL_RUNTIME=0",
     "-sENVIRONMENT=web",
     "-sALLOW_MEMORY_GROWTH",
     "-sEXPORT_NAME=loadWasmTraceViewerModule",
@@ -188,7 +189,6 @@ filegroup(
     srcs = [
         "index.html",
         ":app_bundle.js",
-        ":trace_viewer_v2_wasm/trace_viewer_v2.js",
-        ":trace_viewer_v2_wasm/trace_viewer_v2.wasm",
+        ":trace_viewer_v2_wasm",
     ],
 )
diff --git a/frontend/app/components/trace_viewer_v2/main.ts b/frontend/app/components/trace_viewer_v2/main.ts
@@ -209,7 +209,38 @@ async function loadAndStartWasm(
   return traceviewerModule;
 }
 
+async function ensureWasmModuleIsLoaded(): Promise<void> {
+  // tslint:disable-next-line:no-any
+  if (typeof (window as any).loadWasmTraceViewerModule !== 'undefined') {
+    return;
+  }
+  return new Promise((resolve, reject) => {
+    const existingScript = document.querySelector(
+        'script[src*="trace_viewer_v2.js"]',
+    );
+    if (existingScript) {
+      existingScript.addEventListener('load', () => {
+        resolve();
+      });
+      existingScript.addEventListener('error', () => {
+        reject(new Error('Failed to load WASM module.'));
+      });
+      return;
+    }
+    const script = document.createElement('script');
+    script.src = 'trace_viewer_v2.js';
+    script.onload = () => {
+      resolve();
+    };
+    script.onerror = () => {
+      reject(new Error('Failed to load WASM module.'));
+    };
+    document.body.appendChild(script);
+  });
+}
+
 async function initGpuAndStartWasmApp(): Promise<TraceViewerV2Module> {
+  await ensureWasmModuleIsLoaded();
   const canvas = document.querySelector('#canvas') as HTMLCanvasElement;
   if (!canvas) {
     throw new Error('Could not find canvas element with id="canvas"');