feat(query-builder): Add support for shift+arrow for selecting tokens

static/app/components/searchQueryBuilder/hooks/useQueryBuilderGridItem.tsx

@@ -6,6 +6,7 @@ import type {ListState} from '@react-stately/list';
 import type {FocusableElement, Node} from '@react-types/shared';
 
 import {useSearchQueryBuilder} from 'sentry/components/searchQueryBuilder/context';
+import {useKeyboardSelection} from 'sentry/components/searchQueryBuilder/hooks/useVirtualSelectionCursor';
 import {findNearestFreeTextKey} from 'sentry/components/searchQueryBuilder/utils';
 import type {ParseResultToken} from 'sentry/components/searchSyntax/parser';
 
@@ -95,6 +96,7 @@ export function useQueryBuilderGridItem(
 ) {
   const {wrapperRef} = useSearchQueryBuilder();
   const {rowProps, gridCellProps} = useGridListItem({node: item}, state, ref);
+  const {selectInDirection} = useKeyboardSelection();
 
   // When focus is inside the input, we want to handle some things differently.
   // Returns true if the default behavior should be used, false if not.
@@ -160,6 +162,15 @@ export function useQueryBuilderGridItem(
           return;
         }
 
+        if (e.shiftKey) {
+          selectInDirection({
+            state,
+            beginNewSelectionFromKey: item.key,
+            direction: 'right',
+          });
+          return;
+        }
+
         // Option/Ctrl + ArrowRight should skip focus over to the next free text token
         if (isMac() ? e.altKey : e.ctrlKey) {
           const nextKey = state.collection.getKeyAfter(item.key);
@@ -186,6 +197,15 @@ export function useQueryBuilderGridItem(
           return;
         }
 
+        if (e.shiftKey) {
+          selectInDirection({
+            state,
+            beginNewSelectionFromKey: item.key,
+            direction: 'left',
+          });
+          return;
+        }
+
         // Option/Ctrl + ArrowLeft should skip focus over to the next free text token
         if (isMac() ? e.altKey : e.ctrlKey) {
           const previousKey = state.collection.getKeyBefore(item.key);
@@ -200,7 +220,7 @@ export function useQueryBuilderGridItem(
         focusPreviousGridCell({walker, wrapperRef, state, item});
       }
     },
-    [handleInputKeyDown, item, state, wrapperRef]
+    [handleInputKeyDown, item, selectInDirection, state, wrapperRef]
   );
 
   const onKeyDown = useCallback(

static/app/components/searchQueryBuilder/hooks/useQueryBuilderState.tsx

@@ -323,7 +323,7 @@ function replaceTokensWithText(
 ): QueryBuilderState {
   const newQuery = replaceTokensWithPadding(state.query, action.tokens, action.text);
   const cursorPosition =
-    (action.tokens[0]?.location.start.offset ?? 0) + action.text.length;
+    (action.tokens[0]?.location.start.offset ?? 0) + action.text.length; // TODO: Ensure this is sorted
   const newParsedQuery = parseQueryBuilderValue(newQuery, getFieldDefinition);
   const focusedToken = newParsedQuery?.find(
     token => token.type === Token.FREE_TEXT && token.location.end.offset >= cursorPosition

static/app/components/searchQueryBuilder/hooks/useVirtualSelectionCursor.tsx

@@ -0,0 +1,121 @@
+import {
+  createContext,
+  type ReactNode,
+  useCallback,
+  useContext,
+  useMemo,
+  useRef,
+} from 'react';
+import type {ListState} from '@react-stately/list';
+import type {Key} from '@react-types/shared';
+
+import {findNearestFreeTextKey} from 'sentry/components/searchQueryBuilder/utils';
+import {type ParseResultToken, Token} from 'sentry/components/searchSyntax/parser';
+
+type SelectFunc = (params: {
+  direction: 'left' | 'right';
+  state: ListState<ParseResultToken>;
+  beginNewSelectionFromKey?: Key;
+  toEnd?: boolean;
+}) => void;
+
+export interface KeyboardSelectionData {
+  selectInDirection: SelectFunc;
+}
+
+export function useKeyboardSelection() {
+  return useContext(KeyboardSelectionContext);
+}
+
+function getKeysBetween(state: ListState<ParseResultToken>, key1: Key, key2: Key) {
+  const keys = [...state.collection.getKeys()];
+
+  const keyIndex1 = keys.indexOf(key1);
+  const keyIndex2 = keys.indexOf(key2);
+
+  if (keyIndex1 < keyIndex2) {
+    return keys.slice(keyIndex1, keyIndex2 + 1);
+  }
+
+  return keys.slice(keyIndex2, keyIndex1 + 1);
+}
+
+function combineSelection(state: ListState<ParseResultToken>, newSelection: Key[]) {
+  const currentSelection = new Set(state.selectionManager.selectedKeys);
+
+  for (const key of newSelection) {
+    if (currentSelection.has(key)) {
+      currentSelection.delete(key);
+    } else {
+      currentSelection.add(key);
+    }
+  }
+
+  return currentSelection;
+}
+
+function useKeyboardSelectionState() {
+  const cursorKeyPositionRef = useRef<Key | null>();
+
+  const selectInDirection = useCallback<SelectFunc>(
+    ({state, beginNewSelectionFromKey, direction}) => {
+      const fromKey =
+        beginNewSelectionFromKey ??
+        cursorKeyPositionRef.current ??
+        (direction === 'left'
+          ? state.selectionManager.firstSelectedKey
+          : state.selectionManager.lastSelectedKey);
+
+      if (!fromKey) {
+        return;
+      }
+
+      // Get the start key to make the new selection from.
+      // If we are alre
+      const nextKeyInDirection =
+        direction === 'left'
+          ? state.collection.getKeyBefore(fromKey)
+          : state.collection.getKeyAfter(fromKey);
+      const fromItem = state.collection.getItem(fromKey);
+      const startKey =
+        fromItem?.value?.type === Token.FREE_TEXT ? nextKeyInDirection : fromKey;
+
+      if (!startKey) {
+        return;
+      }
+
+      const endKey = findNearestFreeTextKey(state, startKey, direction);
+
+      if (!endKey) {
+        return;
+      }
+
+      const newSelection = getKeysBetween(state, startKey, endKey);
+
+      state.selectionManager.setSelectedKeys(combineSelection(state, newSelection));
+      cursorKeyPositionRef.current = endKey;
+    },
+    []
+  );
+
+  return useMemo(
+    () => ({
+      selectInDirection,
+    }),
+    [selectInDirection]
+  );
+}
+
+const KeyboardSelectionContext = createContext<KeyboardSelectionData>({
+  selectInDirection: () => {},
+});
+
+export function KeyboardSelection({children}: {children: ReactNode}) {
+  const state = useKeyboardSelectionState();
+
+  return (
+    <KeyboardSelectionContext.Provider value={state}>
+      {children}
+    </KeyboardSelectionContext.Provider>
+  );
+}

static/app/components/searchQueryBuilder/index.spec.tsx

@@ -1011,6 +1011,53 @@ describe('SearchQueryBuilder', function () {
       expect(getLastInput()).toHaveFocus();
     });
 
+    it('extends selection with shift+arrow keys', async function () {
+      render(
+        <SearchQueryBuilder
+          {...defaultProps}
+          initialQuery="browser.name:firefox assigned:me"
+        />
+      );
+
+      await userEvent.click(getLastInput());
+
+      // Shift+ArrowLeft should select assigned:me
+      await userEvent.keyboard('{Shift>}{ArrowLeft}{/Shift}');
+      await waitFor(() => {
+        expect(screen.getByRole('row', {name: 'assigned:me'})).toHaveAttribute(
+          'aria-selected',
+          'true'
+        );
+      });
+
+      // Shift+ArrowLeft again should select browser.name
+      await userEvent.keyboard('{Shift>}{ArrowLeft}{/Shift}');
+      await waitFor(() => {
+        expect(screen.getByRole('row', {name: 'browser.name:firefox'})).toHaveAttribute(
+          'aria-selected',
+          'true'
+        );
+      });
+      // assigned:me should still be selected
+      expect(screen.getByRole('row', {name: 'assigned:me'})).toHaveAttribute(
+        'aria-selected',
+        'true'
+      );
+
+      // Shift+ArrowRight should unselect browser.name:firefox
+      await userEvent.keyboard('{Shift>}{ArrowRight}{/Shift}');
+      await waitFor(() => {
+        expect(
+          screen.getByRole('row', {name: 'browser.name:firefox'})
+        ).not.toHaveAttribute('aria-selected', 'true');
+      });
+      // assigned:me should still be selected
+      expect(screen.getByRole('row', {name: 'assigned:me'})).toHaveAttribute(
+        'aria-selected',
+        'true'
+      );
+    });
+
     it('when focus is in a filter segment, backspace first focuses the filter then deletes it', async function () {
       render(
         <SearchQueryBuilder {...defaultProps} initialQuery="browser.name:firefox" />

static/app/components/searchQueryBuilder/selectionKeyHandler.tsx

@@ -3,6 +3,7 @@ import {VisuallyHidden} from '@react-aria/visually-hidden';
 import type {ListState} from '@react-stately/list';
 
 import {useSearchQueryBuilder} from 'sentry/components/searchQueryBuilder/context';
+import {useKeyboardSelection} from 'sentry/components/searchQueryBuilder/hooks/useVirtualSelectionCursor';
 import {findNearestFreeTextKey} from 'sentry/components/searchQueryBuilder/utils';
 import type {ParseResultToken} from 'sentry/components/searchSyntax/parser';
 import {defined} from 'sentry/utils';
@@ -24,8 +25,10 @@ type SelectionKeyHandlerProps = {
 export const SelectionKeyHandler = forwardRef(
   ({state, undo}: SelectionKeyHandlerProps, ref: ForwardedRef<HTMLInputElement>) => {
     const {dispatch, disabled} = useSearchQueryBuilder();
+    const {selectInDirection} = useKeyboardSelection();
 
-    const selectedTokens = Array.from(state.selectionManager.selectedKeys)
+    const selectedTokens: ParseResultToken[] = [...state.collection.getKeys()]
+      .filter(key => state.selectionManager.selectedKeys.has(key))
       .map(key => state.collection.getItem(key)?.value)
       .filter(defined);
 
@@ -70,6 +73,12 @@ export const SelectionKeyHandler = forwardRef(
           case 'ArrowRight':
             e.preventDefault();
             e.stopPropagation();
+
+            if (e.shiftKey) {
+              selectInDirection({state, direction: 'right'});
+              return;
+            }
+
             state.selectionManager.clearSelection();
             state.selectionManager.setFocusedKey(
               findNearestFreeTextKey(
@@ -82,6 +91,12 @@ export const SelectionKeyHandler = forwardRef(
           case 'ArrowLeft':
             e.preventDefault();
             e.stopPropagation();
+
+            if (e.shiftKey) {
+              selectInDirection({state, direction: 'left'});
+              return;
+            }
+
             state.selectionManager.clearSelection();
             state.selectionManager.setFocusedKey(
               findNearestFreeTextKey(
@@ -142,14 +157,13 @@ export const SelectionKeyHandler = forwardRef(
             return;
         }
       },
-      [dispatch, selectedTokens, state, undo]
+      [dispatch, selectInDirection, selectedTokens, state, undo]
     );
 
     // Using VisuallyHidden because display: none will not allow the input to be focused
     return (
       <VisuallyHidden>
         <input
-          aria-hidden
           data-test-id="selection-key-handler"
           ref={ref}
           tabIndex={-1}

static/app/components/searchQueryBuilder/tokenizedQueryGrid.tsx

@@ -10,6 +10,7 @@ import {useSearchQueryBuilder} from 'sentry/components/searchQueryBuilder/contex
 import {useQueryBuilderGrid} from 'sentry/components/searchQueryBuilder/hooks/useQueryBuilderGrid';
 import {useSelectOnDrag} from 'sentry/components/searchQueryBuilder/hooks/useSelectOnDrag';
 import {useUndoStack} from 'sentry/components/searchQueryBuilder/hooks/useUndoStack';
+import {KeyboardSelection} from 'sentry/components/searchQueryBuilder/hooks/useVirtualSelectionCursor';
 import {SelectionKeyHandler} from 'sentry/components/searchQueryBuilder/selectionKeyHandler';
 import {SearchQueryBuilderBoolean} from 'sentry/components/searchQueryBuilder/tokens/boolean';
 import {SearchQueryBuilderFilter} from 'sentry/components/searchQueryBuilder/tokens/filter/filter';
@@ -143,18 +144,20 @@ export function TokenizedQueryGrid({label, actionBarWidth}: TokenizedQueryGridPr
   }
 
   return (
-    <Grid
-      aria-label={label ?? t('Create a search query')}
-      items={parsedQuery}
-      selectionMode="multiple"
-      actionBarWidth={actionBarWidth}
-    >
-      {item => (
-        <Item key={makeTokenKey(item, parsedQuery)}>
-          {item.text.trim() ? item.text : t('Space')}
-        </Item>
-      )}
-    </Grid>
+    <KeyboardSelection>
+      <Grid
+        aria-label={label ?? t('Create a search query')}
+        items={parsedQuery}
+        selectionMode="multiple"
+        actionBarWidth={actionBarWidth}
+      >
+        {item => (
+          <Item key={makeTokenKey(item, parsedQuery)}>
+            {item.text.trim() ? item.text : t('Space')}
+          </Item>
+        )}
+      </Grid>
+    </KeyboardSelection>
   );
 }