diff --git a/packages/trace-viewer/src/ui/inspectorTab.tsx b/packages/trace-viewer/src/ui/inspectorTab.tsx
index 3e096032df..7b3a90b709 100644
--- a/packages/trace-viewer/src/ui/inspectorTab.tsx
+++ b/packages/trace-viewer/src/ui/inspectorTab.tsx
@@ -27,13 +27,16 @@ export const InspectorTab: React.FunctionComponent<{
   highlightedLocator: string,
   setHighlightedLocator: (locator: string) => void,
 }> = ({ sdkLanguage, setIsInspecting, highlightedLocator, setHighlightedLocator }) => {
-  return <div className='vbox'>
-    <CodeMirrorWrapper text={highlightedLocator} language={sdkLanguage} focusOnChange={true} wrapLines={true} onChange={text => {
-      // Updating text needs to go first - react can squeeze a render between the state updates.
-      setHighlightedLocator(text);
-      setIsInspecting(false);
-    }}></CodeMirrorWrapper>
-    <div style={{ position: 'absolute', right: '0', top: '0' }}>
+  return <div className='vbox' style={{ backgroundColor: 'var(--vscode-sideBar-background)' }}>
+    <div style={{ margin: '10px 0px 10px 10px', color: 'var(--vscode-editorCodeLens-foreground)', flex: 'none' }}>Locator</div>
+    <div style={{ margin: '0 10px 10px', flex: 'auto' }}>
+      <CodeMirrorWrapper text={highlightedLocator} language={sdkLanguage} focusOnChange={true} isFocused={true} wrapLines={true} onChange={text => {
+        // Updating text needs to go first - react can squeeze a render between the state updates.
+        setHighlightedLocator(text);
+        setIsInspecting(false);
+      }}></CodeMirrorWrapper>
+    </div>
+    <div style={{ position: 'absolute', right: 5, top: 5 }}>
       <ToolbarButton icon='files' title='Copy locator' onClick={() => {
         copy(highlightedLocator);
       }}></ToolbarButton>
diff --git a/packages/trace-viewer/src/ui/workbench.tsx b/packages/trace-viewer/src/ui/workbench.tsx
index c2bdb0f320..5e7639fcab 100644
--- a/packages/trace-viewer/src/ui/workbench.tsx
+++ b/packages/trace-viewer/src/ui/workbench.tsx
@@ -68,17 +68,23 @@ export const Workbench: React.FunctionComponent<{
       setSelectedAction(failedAction);
     else if (model?.actions.length)
       setSelectedAction(model.actions[model.actions.length - 1]);
-  }, [model, selectedAction, setSelectedAction, setSelectedPropertiesTab, initialSelection]);
+  }, [model, selectedAction, setSelectedAction, initialSelection]);
 
   const onActionSelected = React.useCallback((action: ActionTraceEventInContext) => {
     setSelectedAction(action);
     onSelectionChanged?.(action);
   }, [setSelectedAction, onSelectionChanged]);
 
+  const selectPropertiesTab = React.useCallback((tab: string) => {
+    setSelectedPropertiesTab(tab);
+    if (tab !== 'inspector')
+      setIsInspecting(false);
+  }, []);
+
   const locatorPicked = React.useCallback((locator: string) => {
     setHighlightedLocator(locator);
-    setSelectedPropertiesTab('inspector');
-  }, []);
+    selectPropertiesTab('inspector');
+  }, [selectPropertiesTab]);
 
   const sdkLanguage = model?.sdkLanguage || 'javascript';
 
@@ -171,11 +177,16 @@ export const Workbench: React.FunctionComponent<{
         <TabbedPane
           tabs={tabs}
           selectedTab={selectedPropertiesTab}
-          setSelectedTab={setSelectedPropertiesTab}
+          setSelectedTab={selectPropertiesTab}
           leftToolbar={[
-            <ToolbarButton icon='microscope' title='Pick locator' toggled={isInspecting} onClick={() => {
+            <ToolbarButton title='Pick locator' toggled={isInspecting} onClick={() => {
+              if (!isInspecting)
+                selectPropertiesTab('inspector');
               setIsInspecting(!isInspecting);
-            }}></ToolbarButton>
+            }}><svg xmlns="http://www.w3.org/2000/svg" height="24" viewBox="0 -960 960 960" width="24">
+                <path d="M450-42v-75q-137-14-228-105T117-450H42v-60h75q14-137 105-228t228-105v-75h60v75q137 14 228 105t105 228h75v60h-75q-14 137-105 228T510-117v75h-60Zm30-134q125 0 214.5-89.5T784-480q0-125-89.5-214.5T480-784q-125 0-214.5 89.5T176-480q0 125 89.5 214.5T480-176Zm0-154q-63 0-106.5-43.5T330-480q0-63 43.5-106.5T480-630q63 0 106.5 43.5T630-480q0 63-43.5 106.5T480-330Zm0-60q38 0 64-26t26-64q0-38-26-64t-64-26q-38 0-64 26t-26 64q0 38 26 64t64 26Zm0-90Z"/>
+              </svg>
+            </ToolbarButton>
           ]}
         />
       </SplitView>
@@ -191,7 +202,7 @@ export const Workbench: React.FunctionComponent<{
               selectedTime={selectedTime}
               onSelected={onActionSelected}
               onHighlighted={setHighlightedAction}
-              revealConsole={() => setSelectedPropertiesTab('console')}
+              revealConsole={() => selectPropertiesTab('console')}
               isLive={isLive}
             />
           },
diff --git a/packages/web/src/components/codeMirrorWrapper.tsx b/packages/web/src/components/codeMirrorWrapper.tsx
index c8c96b2fd5..ae42dd7ce8 100644
--- a/packages/web/src/components/codeMirrorWrapper.tsx
+++ b/packages/web/src/components/codeMirrorWrapper.tsx
@@ -36,6 +36,7 @@ export interface SourceProps {
   highlight?: SourceHighlight[];
   revealLine?: number;
   lineNumbers?: boolean;
+  isFocused?: boolean;
   focusOnChange?: boolean;
   wrapLines?: boolean;
   onChange?: (text: string) => void;
@@ -48,6 +49,7 @@ export const CodeMirrorWrapper: React.FC<SourceProps> = ({
   highlight,
   revealLine,
   lineNumbers,
+  isFocused,
   focusOnChange,
   wrapLines,
   onChange,
@@ -94,10 +96,12 @@ export const CodeMirrorWrapper: React.FC<SourceProps> = ({
         lineWrapping: wrapLines,
       });
       codemirrorRef.current = { cm };
+      if (isFocused)
+        cm.focus();
       setCodemirror(cm);
       return cm;
     })();
-  }, [modulePromise, codemirror, codemirrorElement, language, lineNumbers, wrapLines, readOnly]);
+  }, [modulePromise, codemirror, codemirrorElement, language, lineNumbers, wrapLines, readOnly, isFocused]);
 
   React.useEffect(() => {
     if (codemirrorRef.current)