diff --git a/packages/trace-viewer/src/ui/attachmentsTab.tsx b/packages/trace-viewer/src/ui/attachmentsTab.tsx
index 69bfcd68cc..1aa272aa02 100644
--- a/packages/trace-viewer/src/ui/attachmentsTab.tsx
+++ b/packages/trace-viewer/src/ui/attachmentsTab.tsx
@@ -147,6 +147,7 @@ export const AttachmentsTab: React.FunctionComponent<{
 
 function attachmentURL(attachment: Attachment, queryParams: Record<string, string> = {}) {
   const params = new URLSearchParams(queryParams);
+  // TODO: put in testServerURL here
   if (attachment.sha1) {
     params.set('trace', attachment.traceUrl);
     return 'sha1/' + attachment.sha1 + '?' + params.toString();
diff --git a/packages/trace-viewer/src/ui/uiModeTraceView.tsx b/packages/trace-viewer/src/ui/uiModeTraceView.tsx
index eee20439e4..de6fde259d 100644
--- a/packages/trace-viewer/src/ui/uiModeTraceView.tsx
+++ b/packages/trace-viewer/src/ui/uiModeTraceView.tsx
@@ -24,6 +24,7 @@ import type { ContextEntry } from '../types/entries';
 import type { SourceLocation } from './modelUtil';
 import { MultiTraceModel } from './modelUtil';
 import { Workbench } from './workbench';
+import { testServerBaseURL } from '../searchParams';
 
 export const TraceView: React.FC<{
   item: { treeItem?: TreeItem, testFile?: SourceLocation, testCase?: reporterTypes.TestCase },
diff --git a/packages/trace-viewer/src/ui/uiModeView.tsx b/packages/trace-viewer/src/ui/uiModeView.tsx
index 69a5988641..cad9794d8f 100644
--- a/packages/trace-viewer/src/ui/uiModeView.tsx
+++ b/packages/trace-viewer/src/ui/uiModeView.tsx
@@ -37,6 +37,7 @@ import { FiltersView } from './uiModeFiltersView';
 import { TestListView } from './uiModeTestListView';
 import { TraceView } from './uiModeTraceView';
 import { SettingsView } from './settingsView';
+import { testServerWebSocketURL, searchParams } from '../searchParams';
 
 let xtermSize = { cols: 80, rows: 24 };
 const xtermDataSource: XtermDataSource = {
@@ -46,10 +47,6 @@ const xtermDataSource: XtermDataSource = {
   resize: () => {},
 };
 
-const searchParams = new URLSearchParams(window.location.search);
-const guid = searchParams.get('ws');
-const wsURL = new URL(`../${guid}`, window.location.toString());
-wsURL.protocol = (window.location.protocol === 'https:' ? 'wss:' : 'ws:');
 const queryParams = {
   args: searchParams.getAll('arg'),
   grep: searchParams.get('grep') || undefined,
@@ -109,7 +106,7 @@ export const UIModeView: React.FC<{}> = ({
   const inputRef = React.useRef<HTMLInputElement>(null);
 
   const reloadTests = React.useCallback(() => {
-    setTestServerConnection(new TestServerConnection(new WebSocketTestServerTransport(wsURL)));
+    setTestServerConnection(new TestServerConnection(new WebSocketTestServerTransport(testServerWebSocketURL)));
   }, []);
 
   // Load tests on startup.