From ad2b7c053fb4fb376d727d0f36c3c68f542e4458 Mon Sep 17 00:00:00 2001
From: Rushabh Lalwani <rlalwani@google.com>
Date: Tue, 8 Sep 2026 02:23:50 -0700
Subject: [PATCH] Pass through raw_trace_viewer arguments to the frontend by
 piping them natively from WASM via EventData maps. This bypasses the
 maybeFetchEventArgs fallback for raw trace viewer. Setup types in TS to
 handle natively populated args.

PiperOrigin-RevId: 977785875
---
 .../trace_viewer_container.scss                 | 17 ++++++++++++-----
 .../trace_viewer_container.ts                   |  1 +
 .../app/components/trace_viewer_v2/event_data.h |  1 +
 .../trace_viewer_v2/timeline/timeline.cc        |  7 +++++++
 .../trace_viewer_v2/timeline/timeline_test.cc   | 17 +++++++++++++++++
 5 files changed, 38 insertions(+), 5 deletions(-)

diff --git a/frontend/app/components/trace_viewer_container/trace_viewer_container.scss b/frontend/app/components/trace_viewer_container/trace_viewer_container.scss
index cc94c8110..7465acef2 100644
--- a/frontend/app/components/trace_viewer_container/trace_viewer_container.scss
+++ b/frontend/app/components/trace_viewer_container/trace_viewer_container.scss
@@ -410,6 +410,7 @@ canvas.emscripten {
       .detail-value {
         color: var(--mat-sys-on-surface);
         word-break: break-word;
+        white-space: pre-wrap;
       }
     }
 
@@ -579,8 +580,10 @@ canvas.emscripten {
   border-radius: 24px;
   font-size: 13px;
   font-family: 'Google Sans Text', Roboto, sans-serif;
-  box-shadow: 0 4px 12px rgba(0, 0, 0, 0.1), 0 1px 3px rgba(0,0,0,0.08);
-  border: 1px solid var(--md-sys-color-outline-variant, rgba(0,0,0,0.05));
+  box-shadow:
+    0 4px 12px rgba(0, 0, 0, 0.1),
+    0 1px 3px rgba(0, 0, 0, 0.08);
+  border: 1px solid var(--md-sys-color-outline-variant, rgba(0, 0, 0, 0.05));
   pointer-events: none;
   z-index: 10;
   white-space: nowrap;
@@ -614,8 +617,12 @@ canvas.emscripten {
 }
 
 @keyframes fadeIn {
-  0% { opacity: 0; }
-  100% { opacity: 1; }
+  0% {
+    opacity: 0;
+  }
+  100% {
+    opacity: 1;
+  }
 }
 
 .timeline-status-bar mat-icon {
@@ -646,5 +653,5 @@ canvas.emscripten {
   font-weight: 600;
   padding: 2px 6px;
   line-height: 14px;
-  box-shadow: 0 1px 1px rgba(0,0,0,0.05);
+  box-shadow: 0 1px 1px rgba(0, 0, 0, 0.05);
 }
diff --git a/frontend/app/components/trace_viewer_container/trace_viewer_container.ts b/frontend/app/components/trace_viewer_container/trace_viewer_container.ts
index 904547e9c..345ab06c8 100644
--- a/frontend/app/components/trace_viewer_container/trace_viewer_container.ts
+++ b/frontend/app/components/trace_viewer_container/trace_viewer_container.ts
@@ -121,6 +121,7 @@ export declare interface EntrySelectedEventDetail {
   uid?: string;
   hloModuleName?: string;
   hloOpName?: string;
+  args?: Record<string, string>;
 }
 
 // Type guard for the 'EntrySelected' custom event.
diff --git a/frontend/app/components/trace_viewer_v2/event_data.h b/frontend/app/components/trace_viewer_v2/event_data.h
index 0c01c5ee2..8d0656012 100644
--- a/frontend/app/components/trace_viewer_v2/event_data.h
+++ b/frontend/app/components/trace_viewer_v2/event_data.h
@@ -38,6 +38,7 @@ inline constexpr absl::string_view kEventSelectedUid = "uid";
 inline constexpr absl::string_view kEventSelectedHloModuleName =
     "hloModuleName";
 inline constexpr absl::string_view kEventSelectedHloOpName = "hloOpName";
+inline constexpr absl::string_view kEventSelectedArgs = "args";
 
 // Constants for fetch data event.
 inline constexpr absl::string_view kFetchData = "fetch_data";
diff --git a/frontend/app/components/trace_viewer_v2/timeline/timeline.cc b/frontend/app/components/trace_viewer_v2/timeline/timeline.cc
index 6cf462b93..117bb1af8 100644
--- a/frontend/app/components/trace_viewer_v2/timeline/timeline.cc
+++ b/frontend/app/components/trace_viewer_v2/timeline/timeline.cc
@@ -1537,6 +1537,13 @@ EventData Timeline::CreateBaseEventData(int event_index, bool is_hover) const {
       if (const auto it = args.find(kHloOp); it != args.end()) {
         event_data.try_emplace(kEventSelectedHloOpName, it->second);
       }
+
+      EventData js_args;
+      js_args.reserve(args.size());
+      for (const auto& [key, value] : args) {
+        js_args.try_emplace(key, value);
+      }
+      event_data.try_emplace(kEventSelectedArgs, std::move(js_args));
     }
   }
   return event_data;
diff --git a/frontend/app/components/trace_viewer_v2/timeline/timeline_test.cc b/frontend/app/components/trace_viewer_v2/timeline/timeline_test.cc
index cac5e7843..62be53094 100644
--- a/frontend/app/components/trace_viewer_v2/timeline/timeline_test.cc
+++ b/frontend/app/components/trace_viewer_v2/timeline/timeline_test.cc
@@ -4667,6 +4667,23 @@ TEST_F(RealTimelineImGuiFixture, ClickEventWithArgsSelectsEvent) {
   EXPECT_EQ(
       std::any_cast<std::string>(event_detail.at(kEventSelectedHloOpName)),
       "test_op");
+
+  ASSERT_TRUE(event_detail.contains(kEventSelectedArgs));
+  const auto* js_args =
+      std::any_cast<EventData>(&event_detail.find(kEventSelectedArgs)->second);
+  ASSERT_NE(js_args, nullptr);
+
+  const auto uid_it = js_args->find("uid");
+  ASSERT_NE(uid_it, js_args->end());
+  EXPECT_EQ(std::any_cast<std::string>(uid_it->second), "12345");
+
+  const auto hlo_module_it = js_args->find(kHloModule);
+  ASSERT_NE(hlo_module_it, js_args->end());
+  EXPECT_EQ(std::any_cast<std::string>(hlo_module_it->second), "test_module");
+
+  const auto hlo_op_it = js_args->find(kHloOp);
+  ASSERT_NE(hlo_op_it, js_args->end());
+  EXPECT_EQ(std::any_cast<std::string>(hlo_op_it->second), "test_op");
 }
 
 TEST_F(RealTimelineImGuiFixture, ClickEventSetsSelectionIndices) {