AI Assistant: Add possibility to select/deselect rows on any page depending on scope (#33865)

Co-authored-by: anna.shakhova <anna.shakhova@devexpress.com>

Author: dmirgaev

packages/devextreme/js/__internal/grids/grid_core/ai_assistant/commands/__tests__/selection.test.ts

@@ -6,8 +6,12 @@ import {
 import { z } from 'zod';
 
 import {
+  isKeyShapeValid,
+  normalizeKey,
   // eslint-disable-next-line spellcheck/spell-checker
-  isKeyShapeValid, normalizeKey, optionalNullish, resolveFilterValue,
+  optionalNullish,
+  resolveFilterValue,
+  splitIntoLoadWindows,
 } from '../utils';
 
 describe('normalizeKey', () => {
@@ -168,3 +172,33 @@ describe('resolveFilterValue', () => {
     expect(resolveFilterValue('date', true)).toBe(true);
   });
 });
+
+describe('splitIntoLoadWindows', () => {
+  it('returns an empty array for an empty input', () => {
+    expect(splitIntoLoadWindows([], 10)).toEqual([]);
+  });
+
+  it('wraps a single index into a single window', () => {
+    expect(splitIntoLoadWindows([5], 10)).toEqual([[5]]);
+  });
+
+  it('merges across gaps while the span stays within the window', () => {
+    expect(splitIntoLoadWindows([1, 3, 5], 10)).toEqual([[1, 3, 5]]);
+  });
+
+  it('starts a new window when the span would exceed the limit', () => {
+    expect(splitIntoLoadWindows([1, 2, 4, 5, 6, 10], 3)).toEqual([
+      [1, 2], [4, 5, 6], [10],
+    ]);
+  });
+
+  it('sorts unsorted input before windowing', () => {
+    expect(splitIntoLoadWindows([5, 1, 6, 2, 10, 7], 3)).toEqual([
+      [1, 2], [5, 6, 7], [10],
+    ]);
+  });
+
+  it('deduplicates repeated indexes', () => {
+    expect(splitIntoLoadWindows([1, 1, 2, 2, 3], 10)).toEqual([[1, 2, 3]]);
+  });
+});

packages/devextreme/js/__internal/grids/grid_core/ai_assistant/commands/__tests__/utils.test.ts

@@ -25,8 +25,8 @@ import {
   clearSelectionCommand,
   deselectAllCommand,
   selectAllCommand,
-  selectByIndexesCommand,
   selectByKeysCommand,
+  selectionByIndexesCommand,
 } from './selection';
 import {
   clearSortingCommand,
@@ -49,7 +49,7 @@ export const coreCommands = [
   clearSelectionCommand,
   deselectAllCommand,
   selectAllCommand,
-  selectByIndexesCommand,
+  selectionByIndexesCommand,
   selectByKeysCommand,
   clearSortingCommand,
   sortingCommand,

packages/devextreme/js/__internal/grids/grid_core/ai_assistant/commands/index.ts

@@ -1,11 +1,17 @@
 import type { CommandResult } from '@ts/grids/grid_core/ai_assistant/types';
+import type { InternalGrid, RowKey } from '@ts/grids/grid_core/m_types';
 import { z } from 'zod';
 
 import { defineGridCommand } from './defineGridCommand';
 import {
-  compositeKeyPairSchema, isKeyShapeValid, normalizeKey,
+  compositeKeyPairSchema,
+  isKeyShapeValid,
+  normalizeKey,
+  splitIntoLoadWindows,
 } from './utils';
 
+const MAX_LOAD_WINDOW_PAGES = 5;
+
 const selectByKeysCommandSchema = z.object({
   keys: z.array(z.union([
     z.string(),
@@ -45,51 +51,167 @@ export const selectByKeysCommand = defineGridCommand({
   },
 });
 
-const selectByIndexesCommandSchema = z.object({
+const selectionByIndexesCommandSchema = z.object({
   indexes: z.array(z.number().int().min(1)).min(1),
   mode: z.enum(['select', 'deselect']),
+  scope: z.enum(['allPages', 'page']),
 }).strict();
 
-export const selectByIndexesCommand = defineGridCommand({
-  name: 'selectByIndexes',
-  description: 'Select or deselect specific rows by their 1-based indexes within the current page. '
-    + 'Index 1 is the first row on the visible page; group/header rows are not addressable. '
-    + 'Set mode to "deselect" to remove the listed rows from the current selection (e.g. "unselect row 1"); set it to "select" to select them. '
-    + 'When mode is "select", the listed rows replace the current selection. '
-    + 'To target rows that are not on the current page, use selectByKeys, or call pageIndex first to switch the page. '
-    + 'To clear selection only within the current selectAll scope, use deselectAll; to clear selection across all pages regardless of selectAllMode, use clearSelection.',
-  schema: selectByIndexesCommandSchema,
+// Maps 1-based indexes to the keys at those positions;
+// an index past the last entry rejects the whole set.
+const pickKeysByIndex = (
+  keys: RowKey[],
+  indexes: number[],
+): RowKey[] | null => {
+  const normalizedRowIndexes = indexes.map((index) => index - 1);
+  const allIndexesValid = normalizedRowIndexes.every(
+    (index) => index < keys.length,
+  );
+
+  if (!allIndexesValid) {
+    return null;
+  }
+
+  return normalizedRowIndexes.map((index) => keys[index]);
+};
+
+const resolveKeysFromCurrentPage = (
+  component: InternalGrid,
+  indexes: number[],
+): RowKey[] | null => {
+  // Group/footer rows are not counted, so indexes address the Nth data row.
+  const items = component.getController('data').items();
+  const dataItems = items.filter((item) => item.rowType === 'data');
+  const dataKeys = dataItems.map((item) => item.key);
+
+  return pickKeysByIndex(dataKeys, indexes);
+};
+
+const resolveKeysFromAllPagesRemote = async (
+  component: InternalGrid,
+  indexes: number[],
+): Promise<RowKey[] | null> => {
+  const dataSource = component.getDataSource();
+  const store = dataSource?.store();
+
+  if (!dataSource || !store) {
+    return null;
+  }
+
+  const keyExpr = component.option('keyExpr') ?? store.key();
+
+  if (!keyExpr) {
+    return null;
+  }
+
+  // Under grouping store.load returns group structures rather than flat rows,
+  // so an index no longer maps to a single data row. Fail instead of resolving meaningless keys
+  const grouping = dataSource.group();
+  const isGrouped = Array.isArray(grouping) ? grouping.length > 0 : !!grouping;
+
+  if (isGrouped) {
+    return null;
+  }
+
+  const dataController = component.getController('data');
+  const filter = dataController.getCombinedFilter(true);
+  const baseLoadOptions = {
+    ...dataSource.loadOptions(),
+    filter,
+  };
+
+  const windows = splitIntoLoadWindows(indexes, dataSource.pageSize() * MAX_LOAD_WINDOW_PAGES);
+
+  const loadedWindows = await Promise.all(windows.map((window) => {
+    const skip = window[0] - 1;
+    const take = window[window.length - 1] - window[0] + 1;
+    return store.load({ ...baseLoadOptions, skip, take })
+      .then((result) => {
+        const rows = Array.isArray(result) ? result : (result as { data: unknown[] }).data;
+        return { window, rows };
+      });
+  }));
+
+  const keys: RowKey[] = [];
+
+  for (const { window, rows } of loadedWindows) {
+    if (!Array.isArray(rows)) {
+      return null;
+    }
+
+    for (const index of window) {
+      // The requested index maps to `window[0]` offset within the loaded rows
+      const row = rows[index - window[0]];
+
+      if (row === undefined) {
+        return null;
+      }
+
+      keys.push(store.keyOf(row));
+    }
+  }
+
+  return keys;
+};
+
+const resolveKeysFromAllPagesLocal = async (
+  component: InternalGrid,
+  indexes: number[],
+): Promise<RowKey[] | null> => {
+  const keys = await component.getController('data').getAllDataRowKeys();
+
+  return pickKeysByIndex(keys, indexes);
+};
+
+// Picks the "allPages" strategy by paging mode:
+// with local paging the full dataset is already on the client, so read the cache;
+// with remote paging rows are fetched by position.
+const resolveKeysFromAllPages = (
+  component: InternalGrid,
+  indexes: number[],
+): Promise<RowKey[] | null> => {
+  const dataController = component.getController('data');
+  const isRemotePaging = !!dataController.dataSource()?.remoteOperations()?.paging;
+
+  return isRemotePaging
+    ? resolveKeysFromAllPagesRemote(component, indexes)
+    : resolveKeysFromAllPagesLocal(component, indexes);
+};
+
+export const selectionByIndexesCommand = defineGridCommand({
+  name: 'selectionByIndexes',
+  description: 'Select or deselect rows by their 1-based indexes. '
+    + 'Indexes start at 1: "the first row" is index 1 and "the 5th row" is index 5. Do NOT use 0-based counting here, this command is 1-based. '
+    + 'Always set scope to choose how indexes are interpreted: '
+    + '"allPages" — indexes are positions within the currently filtered and sorted dataset, NOT limited to the current page; index 1 is the first row of the dataset, regardless of pageIndex/pageSize. Use this when the user does NOT explicitly refer to the visible page (e.g. "select rows 1 to 100"). '
+    + '"page" — indexes are positions within the currently rendered page; index 1 is the first data row on the visible page and group/header rows are not counted. Use this ONLY when the user explicitly mentions the current/visible page (e.g. "select the first 3 rows on the current page", "deselect row 2 on this page"). '
+    + 'Set mode to "select" to add the listed rows to the current selection (multiple calls accumulate, so previously selected ranges are kept); set mode to "deselect" to remove the listed rows from the current selection (e.g. "unselect row 1"). '
+    + 'To clear selection only within the current selectAll scope, use deselectAll; to clear selection across all pages regardless of selectAllMode, use clearSelection. '
+    + 'To target rows by key value rather than by index, use selectByKeys.',
+  schema: selectionByIndexesCommandSchema,
   execute: (component, { success, failure }) => async (args): Promise<CommandResult> => {
     const rowIndexes = args.indexes.join(', ');
     const action = args.mode === 'deselect' ? 'Deselect' : 'Select';
-    const defaultMessage = `${action} row(s) number ${rowIndexes} on the current page.`;
+    const scopeSuffix = args.scope === 'page' ? ' on the current page' : '';
+    const defaultMessage = `${action} row(s) number ${rowIndexes}${scopeSuffix}.`;
 
     if (component.option('selection.mode') === 'none') {
       return failure(defaultMessage);
     }
 
-    const items = component.getController('data').items();
-    const normalizedRowIndexes = args.indexes.map((index) => index - 1);
-    const allIndexesValid = normalizedRowIndexes.every(
-      (index) => items[index]?.rowType === 'data',
-    );
+    try {
+      const keys = args.scope === 'page'
+        ? resolveKeysFromCurrentPage(component, args.indexes)
+        : await resolveKeysFromAllPages(component, args.indexes);
 
-    if (!allIndexesValid) {
-      return failure(defaultMessage);
-    }
+      if (keys === null) {
+        return failure(defaultMessage);
+      }
 
-    try {
-      switch (args.mode) {
-        case 'deselect': {
-          const itemKeys = normalizedRowIndexes.map((index) => items[index].key);
-          await component.deselectRows(itemKeys);
-          break;
-        }
-        case 'select':
-          await component.selectRowsByIndexes(normalizedRowIndexes);
-          break;
-        default:
-          return failure(defaultMessage);
+      if (args.mode === 'deselect') {
+        await component.deselectRows(keys);
+      } else {
+        await component.selectRows(keys, true);
       }
 
       return success(defaultMessage);

packages/devextreme/js/__internal/grids/grid_core/ai_assistant/commands/selection.ts

@@ -67,6 +67,27 @@ export const isKeyShapeValid = (
   return keyExpr.every((field) => field in key);
 };
 
+export const splitIntoLoadWindows = (
+  indexes: number[],
+  maxWindowSize: number,
+): number[][] => {
+  const sorted = [...new Set(indexes)].sort((a, b) => a - b);
+  const windows: number[][] = [];
+
+  sorted.forEach((index) => {
+    const current = windows.at(-1);
+
+    // Sorted indexes are merged into the same window within maxWindowSize
+    if (current && index - current[0] + 1 <= maxWindowSize) {
+      current.push(index);
+    } else {
+      windows.push([index]);
+    }
+  });
+
+  return windows;
+};
+
 type FilterExprValue = BasicFilterExpr['value'];
 
 export function resolveFilterValue(

packages/devextreme/js/__internal/grids/grid_core/ai_assistant/commands/utils.ts

@@ -23,7 +23,7 @@ import type { ValidatingController } from '@ts/grids/grid_core/validating/m_vali
 import { AI_COLUMN_NAME } from '../ai_column/const';
 import modules from '../m_modules';
 import type {
-  Controllers, Module,
+  Controllers, Module, RowKey,
 } from '../m_types';
 import gridCoreUtils from '../m_utils';
 import type { VirtualScrollController } from '../virtual_scrolling/m_virtual_scrolling_core';
@@ -1542,6 +1542,13 @@ export class DataController extends DataHelperMixin(modules.Controller) {
     return d;
   }
 
+  public getAllDataRowKeys(): Promise<RowKey[]> {
+    return Promise.resolve(this.loadAll(undefined) as unknown as Promise<Item[]>)
+      .then((items) => items
+        .filter((item) => item.rowType === 'data')
+        .map((item) => item.key));
+  }
+
   public getKeyByRowIndex(rowIndex, byLoaded?) {
     const item = this.items(byLoaded)[rowIndex];
     if (item) {