diff --git a/packages/trace-viewer/src/ui/sourceTab.tsx b/packages/trace-viewer/src/ui/sourceTab.tsx
index 63f6b2032d..7a5b7d6a61 100644
--- a/packages/trace-viewer/src/ui/sourceTab.tsx
+++ b/packages/trace-viewer/src/ui/sourceTab.tsx
@@ -23,12 +23,14 @@ import { StackTraceView } from './stackTrace';
 import { CodeMirrorWrapper } from '@web/components/codeMirrorWrapper';
 import type { SourceHighlight } from '@web/components/codeMirrorWrapper';
 import type { SourceModel } from './modelUtil';
+import type { StackFrame } from '@protocol/channels';
 
 export const SourceTab: React.FunctionComponent<{
   action: ActionTraceEvent | undefined,
   sources: Map<string, SourceModel>,
   hideStackFrames?: boolean,
-}> = ({ action, sources, hideStackFrames }) => {
+  fallbackLocation?: StackFrame,
+}> = ({ action, sources, hideStackFrames, fallbackLocation }) => {
   const [lastAction, setLastAction] = React.useState<ActionTraceEvent | undefined>();
   const [selectedFrame, setSelectedFrame] = React.useState<number>(0);
 
@@ -39,28 +41,35 @@ export const SourceTab: React.FunctionComponent<{
     }
   }, [action, lastAction, setLastAction, setSelectedFrame]);
 
-  const source = useAsyncMemo<SourceModel>(async () => {
-    const file = action?.stack?.[selectedFrame].file;
-    if (!file)
-      return { errors: [], content: undefined };
-    const source = sources.get(file)!;
+  const { source, targetLine, highlight } = useAsyncMemo<{ source: SourceModel, targetLine: number, highlight: SourceHighlight[] }>(async () => {
+    const location = action?.stack?.[selectedFrame] || fallbackLocation;
+    if (!location?.file)
+      return { source: { errors: [], content: undefined }, targetLine: 0, highlight: [] };
+
+    let source = sources.get(location.file);
+    // Fallback location can fall outside the sources model.
+    if (!source) {
+      source = { errors: [], content: undefined };
+      sources.set(location.file, source);
+    }
+
+    const targetLine = location.line || 0;
+    const highlight: SourceHighlight[] = source.errors.map(e => ({ type: 'error', line: e.location.line, message: e.error!.message }));
+    highlight.push({ line: targetLine, type: 'running' });
+
     if (source.content === undefined) {
-      const sha1 = await calculateSha1(file);
+      const sha1 = await calculateSha1(location.file);
       try {
         let response = await fetch(`sha1/src@${sha1}.txt`);
         if (response.status === 404)
-          response = await fetch(`file?path=${file}`);
+          response = await fetch(`file?path=${location.file}`);
         source.content = await response.text();
       } catch {
-        source.content = `<Unable to read "${file}">`;
+        source.content = `<Unable to read "${location.file}">`;
       }
     }
-    return source;
-  }, [action, selectedFrame], { errors: [], content: 'Loading\u2026' });
-
-  const targetLine = action?.stack?.[selectedFrame]?.line || 0;
-  const highlight: SourceHighlight[] = source.errors.map(e => ({ type: 'error', line: e.location.line, message: e.error!.message }));
-  highlight.push({ line: targetLine, type: 'running' });
+    return { source, targetLine, highlight };
+  }, [action, selectedFrame, fallbackLocation], { source: { errors: [], content: 'Loading\u2026' }, targetLine: 0, highlight: [] });
 
   return <SplitView sidebarSize={200} orientation='horizontal' sidebarHidden={hideStackFrames}>
     <CodeMirrorWrapper text={source.content || ''} language='javascript' highlight={highlight} revealLine={targetLine} readOnly={true} lineNumbers={true} />
diff --git a/packages/trace-viewer/src/ui/watchMode.tsx b/packages/trace-viewer/src/ui/watchMode.tsx
index 77bed628c1..5e3d7083e4 100644
--- a/packages/trace-viewer/src/ui/watchMode.tsx
+++ b/packages/trace-viewer/src/ui/watchMode.tsx
@@ -445,7 +445,7 @@ const TraceView: React.FC<{
     };
   }, [result, outputDir, testCase, setModel, counter, setCounter]);
 
-  return <Workbench key='workbench' model={model} hideTimelineBars={true} hideStackFrames={true} showSourcesFirst={true} />;
+  return <Workbench key='workbench' model={model} hideTimelineBars={true} hideStackFrames={true} showSourcesFirst={true} defaultSourceLocation={testCase?.location} />;
 };
 
 declare global {
diff --git a/packages/trace-viewer/src/ui/workbench.tsx b/packages/trace-viewer/src/ui/workbench.tsx
index 4020a411f0..f5caccf209 100644
--- a/packages/trace-viewer/src/ui/workbench.tsx
+++ b/packages/trace-viewer/src/ui/workbench.tsx
@@ -30,13 +30,15 @@ import type { TabbedPaneTabModel } from '@web/components/tabbedPane';
 import { Timeline } from './timeline';
 import './workbench.css';
 import { MetadataView } from './metadataView';
+import type { Location } from '../../../playwright-test/types/testReporter';
 
 export const Workbench: React.FunctionComponent<{
   model?: MultiTraceModel,
   hideTimelineBars?: boolean,
   hideStackFrames?: boolean,
   showSourcesFirst?: boolean,
-}> = ({ model, hideTimelineBars, hideStackFrames, showSourcesFirst }) => {
+  defaultSourceLocation?: Location,
+}> = ({ model, hideTimelineBars, hideStackFrames, showSourcesFirst, defaultSourceLocation }) => {
   const [selectedAction, setSelectedAction] = React.useState<ActionTraceEvent | undefined>(undefined);
   const [highlightedAction, setHighlightedAction] = React.useState<ActionTraceEvent | undefined>();
   const [selectedNavigatorTab, setSelectedNavigatorTab] = React.useState<string>('actions');
@@ -66,7 +68,7 @@ export const Workbench: React.FunctionComponent<{
   const sourceTab: TabbedPaneTabModel = {
     id: 'source',
     title: 'Source',
-    render: () => <SourceTab action={activeAction} sources={model?.sources || new Map()} hideStackFrames={hideStackFrames}/>
+    render: () => <SourceTab action={activeAction} sources={model?.sources || new Map()} hideStackFrames={hideStackFrames} fallbackLocation={defaultSourceLocation}/>
   };
   const consoleTab: TabbedPaneTabModel = {
     id: 'console',
diff --git a/tests/playwright-test/ui-mode-test-source.spec.ts b/tests/playwright-test/ui-mode-test-source.spec.ts
new file mode 100644
index 0000000000..2772f7eb04
--- /dev/null
+++ b/tests/playwright-test/ui-mode-test-source.spec.ts
@@ -0,0 +1,57 @@
+/**
+ * Copyright (c) Microsoft Corporation.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+import { test, expect, dumpTestTree } from './ui-mode-fixtures';
+
+test.describe.configure({ mode: 'parallel' });
+
+const basicTestTree = {
+  'a.test.ts': `
+    import { test } from '@playwright/test';
+    test('first', () => {});
+    test('second', () => {});
+  `,
+  'b.test.ts': `
+    import { test } from '@playwright/test';
+    test('third', () => {});
+  `,
+};
+
+test('should show selected test in sources', async ({ runUITest }) => {
+  const page = await runUITest(basicTestTree);
+  await expect.poll(dumpTestTree(page), { timeout: 15000 }).toBe(`
+    ▼ ◯ a.test.ts
+        ◯ first
+        ◯ second
+    ▼ ◯ b.test.ts
+        ◯ third
+  `);
+
+  await page.getByTestId('test-tree').getByText('first').click();
+  await expect(
+      page.locator('.CodeMirror .source-line-running'),
+  ).toHaveText(`3    test('first', () => {});`);
+
+  await page.getByTestId('test-tree').getByText('second').click();
+  await expect(
+      page.locator('.CodeMirror .source-line-running'),
+  ).toHaveText(`4    test('second', () => {});`);
+
+  await page.getByTestId('test-tree').getByText('third').click();
+  await expect(
+      page.locator('.CodeMirror .source-line-running'),
+  ).toHaveText(`3    test('third', () => {});`);
+});