From 0e1aeaaecf1b0ce0e0d7dde1d51383b86c4c3f72 Mon Sep 17 00:00:00 2001
From: Pavel Feldman <pavel.feldman@gmail.com>
Date: Fri, 19 May 2023 15:18:18 -0700
Subject: [PATCH] chore: do not show stale source in the trace (#23163)

---
 packages/trace-viewer/src/sw.ts                 |  8 ++++++--
 packages/trace-viewer/src/ui/actionList.tsx     | 11 ++++++-----
 packages/trace-viewer/src/ui/attachmentsTab.tsx | 17 +++++++++--------
 packages/trace-viewer/src/ui/modelUtil.ts       | 12 ++++++++----
 packages/trace-viewer/src/ui/timeline.tsx       | 10 +++++-----
 packages/trace-viewer/src/ui/workbench.tsx      | 13 ++++++-------
 6 files changed, 40 insertions(+), 31 deletions(-)

diff --git a/packages/trace-viewer/src/sw.ts b/packages/trace-viewer/src/sw.ts
index 59c1340c95..79a8e2b7e3 100644
--- a/packages/trace-viewer/src/sw.ts
+++ b/packages/trace-viewer/src/sw.ts
@@ -116,8 +116,12 @@ async function doFetch(event: FetchEvent): Promise<Response> {
     }
 
     if (relativePath.startsWith('/sha1/')) {
-      // Sha1 is unique, load it from either of the models for simplicity.
-      for (const { traceModel } of loadedTraces.values()) {
+      // Sha1 for sources is based on the file path, can't load it of a random model.
+      const traceUrls = clientIdToTraceUrls.get(event.clientId);
+      for (const [trace, { traceModel }] of loadedTraces) {
+        // We will accept explicit ?trace= value as well as the clientId associated with the trace.
+        if (traceUrl !== trace && !traceUrls.includes(trace))
+          continue;
         const blob = await traceModel!.resourceForSha1(relativePath.slice('/sha1/'.length));
         if (blob)
           return new Response(blob, { status: 200 });
diff --git a/packages/trace-viewer/src/ui/actionList.tsx b/packages/trace-viewer/src/ui/actionList.tsx
index f883e55e5d..2fc95fe9ec 100644
--- a/packages/trace-viewer/src/ui/actionList.tsx
+++ b/packages/trace-viewer/src/ui/actionList.tsx
@@ -23,13 +23,14 @@ import { asLocator } from '@isomorphic/locatorGenerators';
 import type { Language } from '@isomorphic/locatorGenerators';
 import type { TreeState } from '@web/components/treeView';
 import { TreeView } from '@web/components/treeView';
+import type { ActionTraceEventInContext } from './modelUtil';
 
 export interface ActionListProps {
-  actions: ActionTraceEvent[],
-  selectedAction: ActionTraceEvent | undefined,
+  actions: ActionTraceEventInContext[],
+  selectedAction: ActionTraceEventInContext | undefined,
   sdkLanguage: Language | undefined;
-  onSelected: (action: ActionTraceEvent) => void,
-  onHighlighted: (action: ActionTraceEvent | undefined) => void,
+  onSelected: (action: ActionTraceEventInContext) => void,
+  onHighlighted: (action: ActionTraceEventInContext | undefined) => void,
   revealConsole: () => void,
   isLive?: boolean,
 }
@@ -38,7 +39,7 @@ type ActionTreeItem = {
   id: string;
   children: ActionTreeItem[];
   parent: ActionTreeItem | undefined;
-  action?: ActionTraceEvent;
+  action?: ActionTraceEventInContext;
 };
 
 const ActionTreeView = TreeView<ActionTreeItem>;
diff --git a/packages/trace-viewer/src/ui/attachmentsTab.tsx b/packages/trace-viewer/src/ui/attachmentsTab.tsx
index c74cea8b6f..f118751aa3 100644
--- a/packages/trace-viewer/src/ui/attachmentsTab.tsx
+++ b/packages/trace-viewer/src/ui/attachmentsTab.tsx
@@ -14,39 +14,40 @@
  * limitations under the License.
  */
 
-import type { ActionTraceEvent } from '@trace/trace';
 import * as React from 'react';
 import './attachmentsTab.css';
 import { ImageDiffView } from '@web/components/imageDiffView';
 import type { TestAttachment } from '@web/components/imageDiffView';
+import type { ActionTraceEventInContext } from './modelUtil';
 
 export const AttachmentsTab: React.FunctionComponent<{
-  action: ActionTraceEvent | undefined,
+  action: ActionTraceEventInContext | undefined,
 }> = ({ action }) => {
   if (!action)
     return null;
   const expected = action.attachments?.find(a => a.name.endsWith('-expected.png') && (a.path || a.sha1)) as TestAttachment | undefined;
   const actual = action.attachments?.find(a => a.name.endsWith('-actual.png') && (a.path || a.sha1)) as TestAttachment | undefined;
   const diff = action.attachments?.find(a => a.name.endsWith('-diff.png') && (a.path || a.sha1)) as TestAttachment | undefined;
+  const traceUrl = action.context.traceUrl;
 
   return <div className='attachments-tab'>
     {expected && actual && <div className='attachments-section'>Image diff</div>}
     {expected && actual && <ImageDiffView imageDiff={{
       name: 'Image diff',
-      expected: { attachment: { ...expected, path: attachmentURL(expected) }, title: 'Expected' },
-      actual: { attachment: { ...actual, path: attachmentURL(actual) } },
-      diff: diff ? { attachment: { ...diff, path: attachmentURL(diff) } } : undefined,
+      expected: { attachment: { ...expected, path: attachmentURL(traceUrl, expected) }, title: 'Expected' },
+      actual: { attachment: { ...actual, path: attachmentURL(traceUrl, actual) } },
+      diff: diff ? { attachment: { ...diff, path: attachmentURL(traceUrl, diff) } } : undefined,
     }} />}
     {<div className='attachments-section'>Attachments</div>}
     {action.attachments?.map(a => {
       return <div className='attachment-item'>
-        <a target='_blank' href={`sha1/${a.sha1}`}>{a.name}</a>
+        <a target='_blank' href={attachmentURL(traceUrl, a)}>{a.name}</a>
       </div>;
     })}
   </div>;
 };
 
-function attachmentURL(attachment: {
+function attachmentURL(traceUrl: string, attachment: {
   name: string;
   contentType: string;
   path?: string;
@@ -54,6 +55,6 @@ function attachmentURL(attachment: {
   body?: string;
 }) {
   if (attachment.sha1)
-    return 'sha1/' + attachment.sha1;
+    return 'sha1/' + attachment.sha1 + '?trace=' + traceUrl;
   return 'file?path=' + attachment.path;
 }
diff --git a/packages/trace-viewer/src/ui/modelUtil.ts b/packages/trace-viewer/src/ui/modelUtil.ts
index 30f0b11831..a86be35393 100644
--- a/packages/trace-viewer/src/ui/modelUtil.ts
+++ b/packages/trace-viewer/src/ui/modelUtil.ts
@@ -37,6 +37,10 @@ export type SourceModel = {
   content: string | undefined;
 };
 
+export type ActionTraceEventInContext = ActionTraceEvent & {
+  context: ContextEntry;
+};
+
 export class MultiTraceModel {
   readonly startTime: number;
   readonly endTime: number;
@@ -46,7 +50,7 @@ export class MultiTraceModel {
   readonly title?: string;
   readonly options: trace.BrowserContextEventOptions;
   readonly pages: PageEntry[];
-  readonly actions: trace.ActionTraceEvent[];
+  readonly actions: ActionTraceEventInContext[];
   readonly events: trace.EventTraceEvent[];
   readonly hasSource: boolean;
   readonly sdkLanguage: Language | undefined;
@@ -89,7 +93,7 @@ function indexModel(context: ContextEntry) {
 }
 
 function mergeActions(contexts: ContextEntry[]) {
-  const map = new Map<string, ActionTraceEvent>();
+  const map = new Map<string, ActionTraceEventInContext>();
 
   // Protocol call aka isPrimary contexts have startTime/endTime as server-side times.
   // Step aka non-isPrimary contexts have startTime/endTime are client-side times.
@@ -100,7 +104,7 @@ function mergeActions(contexts: ContextEntry[]) {
 
   for (const context of primaryContexts) {
     for (const action of context.actions)
-      map.set(`${action.apiName}@${action.wallTime}`, action);
+      map.set(`${action.apiName}@${action.wallTime}`, { ...action, context });
     if (!offset && context.actions.length)
       offset = context.actions[0].startTime - context.actions[0].wallTime;
   }
@@ -126,7 +130,7 @@ function mergeActions(contexts: ContextEntry[]) {
           existing.parentId = action.parentId;
         continue;
       }
-      map.set(key, action);
+      map.set(key, { ...action, context });
     }
   }
 
diff --git a/packages/trace-viewer/src/ui/timeline.tsx b/packages/trace-viewer/src/ui/timeline.tsx
index f756229533..4886aac38b 100644
--- a/packages/trace-viewer/src/ui/timeline.tsx
+++ b/packages/trace-viewer/src/ui/timeline.tsx
@@ -15,16 +15,16 @@
   limitations under the License.
 */
 
-import type { ActionTraceEvent, EventTraceEvent } from '@trace/trace';
+import type { EventTraceEvent } from '@trace/trace';
 import { msToString, useMeasure } from '@web/uiUtils';
 import * as React from 'react';
 import type { Boundaries } from '../geometry';
 import { FilmStrip } from './filmStrip';
-import type { MultiTraceModel } from './modelUtil';
+import type { ActionTraceEventInContext, MultiTraceModel } from './modelUtil';
 import './timeline.css';
 
 type TimelineBar = {
-  action?: ActionTraceEvent;
+  action?: ActionTraceEventInContext;
   event?: EventTraceEvent;
   leftPosition: number;
   rightPosition: number;
@@ -38,8 +38,8 @@ type TimelineBar = {
 
 export const Timeline: React.FunctionComponent<{
   model: MultiTraceModel | undefined,
-  selectedAction: ActionTraceEvent | undefined,
-  onSelected: (action: ActionTraceEvent) => void,
+  selectedAction: ActionTraceEventInContext | undefined,
+  onSelected: (action: ActionTraceEventInContext) => void,
   hideTimelineBars?: boolean,
 }> = ({ model, selectedAction, onSelected, hideTimelineBars }) => {
   const [measure, ref] = useMeasure<HTMLDivElement>();
diff --git a/packages/trace-viewer/src/ui/workbench.tsx b/packages/trace-viewer/src/ui/workbench.tsx
index c044d8d712..0f3e55cc66 100644
--- a/packages/trace-viewer/src/ui/workbench.tsx
+++ b/packages/trace-viewer/src/ui/workbench.tsx
@@ -14,14 +14,13 @@
   limitations under the License.
 */
 
-import type { ActionTraceEvent } from '@trace/trace';
 import { SplitView } from '@web/components/splitView';
 import * as React from 'react';
 import { ActionList } from './actionList';
 import { CallTab } from './callTab';
 import { ConsoleTab } from './consoleTab';
 import * as modelUtil from './modelUtil';
-import type { MultiTraceModel } from './modelUtil';
+import type { ActionTraceEventInContext, MultiTraceModel } from './modelUtil';
 import { NetworkTab } from './networkTab';
 import { SnapshotTab } from './snapshotTab';
 import { SourceTab } from './sourceTab';
@@ -39,12 +38,12 @@ export const Workbench: React.FunctionComponent<{
   showSourcesFirst?: boolean,
   rootDir?: string,
   fallbackLocation?: modelUtil.SourceLocation,
-  initialSelection?: ActionTraceEvent,
-  onSelectionChanged?: (action: ActionTraceEvent) => void,
+  initialSelection?: ActionTraceEventInContext,
+  onSelectionChanged?: (action: ActionTraceEventInContext) => void,
   isLive?: boolean,
 }> = ({ model, hideTimelineBars, hideStackFrames, showSourcesFirst, rootDir, fallbackLocation, initialSelection, onSelectionChanged, isLive }) => {
-  const [selectedAction, setSelectedAction] = React.useState<ActionTraceEvent | undefined>(undefined);
-  const [highlightedAction, setHighlightedAction] = React.useState<ActionTraceEvent | undefined>();
+  const [selectedAction, setSelectedAction] = React.useState<ActionTraceEventInContext | undefined>(undefined);
+  const [highlightedAction, setHighlightedAction] = React.useState<ActionTraceEventInContext | undefined>();
   const [selectedNavigatorTab, setSelectedNavigatorTab] = React.useState<string>('actions');
   const [selectedPropertiesTab, setSelectedPropertiesTab] = React.useState<string>(showSourcesFirst ? 'source' : 'call');
   const activeAction = model ? highlightedAction || selectedAction : undefined;
@@ -63,7 +62,7 @@ export const Workbench: React.FunctionComponent<{
       setSelectedAction(model.actions[model.actions.length - 1]);
   }, [model, selectedAction, setSelectedAction, setSelectedPropertiesTab, initialSelection]);
 
-  const onActionSelected = React.useCallback((action: ActionTraceEvent) => {
+  const onActionSelected = React.useCallback((action: ActionTraceEventInContext) => {
     setSelectedAction(action);
     onSelectionChanged?.(action);
   }, [setSelectedAction, onSelectionChanged]);