feat(askai): Aggregate MCP search tool calls (#2891)

* feat(askai): add Agent Studio memory support

* refactor(askai): remove Ask AI transport abstraction

* feat(askai): Feedback notes and tags

* fix: bump css bundle size limit

* move feedback actions to components

* feat(askai): Aggregate MCP search tool calls

AGENTS.md

@@ -27,13 +27,10 @@ bun install
 bun run build
 
 # Build specific package
-bun --filter @docsearch/react build
+bun run --filter @docsearch/react build
 
 # Watch mode (all packages)
 bun run watch
-
-# Clean builds
-bun run build:clean
 ```
 
 ## Test Commands
@@ -43,13 +40,7 @@ bun run build:clean
 bun run test
 
 # Run a single test file
-bun run test packages/docsearch-react/src/__tests__/utils.test.ts
-
-# Run tests matching a pattern
-bun run test --testNamePattern="extractLinksFromText"
-
-# Run tests in watch mode
-bun run test --watch
+bun run test --run packages/docsearch-react/src/__tests__/utils.test.ts
 
 # Type checking
 bun run test:types
@@ -58,6 +49,8 @@ bun run test:types
 bun run test:size
 ```
 
+When running tests, prefer to run specific files with the `--run` flag to prevent running with watch mode.
+
 ## Lint Commands
 
 ```bash
@@ -68,15 +61,16 @@ bun run lint
 bun run lint:css
 ```
 
-## E2E Testing (Cypress)
+## E2E Testing (Playwright)
 
 ```bash
 # Run Cypress tests
-bun run cy:run
+bun run pw:run
 
 # Run with specific browser
-bun run cy:run:chrome
-bun run cy:run:firefox
+bun run pw:run:chromium
+bun run pw:run:firefox
+bun run pw:run:webkit
 ```
 
 ## Code Style Guidelines

packages/docsearch-react/src/components/ToolCall.tsx

@@ -3,6 +3,7 @@ import React, { useMemo } from 'react';
 
 import { LoadingIcon, MemoryIcon, SearchIcon, ToolIcon } from '../icons';
 import type { AIToolPart, MemoryToolPart, SearchToolPart, ToolCalls, ToolDefinition } from '../types/AskiAi';
+import { isSearchToolPart } from '../utils/ai';
 
 import { ToolState } from './ui/ToolState';
 
@@ -161,10 +162,6 @@ function MemoryTool({ part, translations }: { part: MemoryToolPart; translations
   );
 }
 
-function isSearchToolPart(part: AIToolPart): part is SearchToolPart {
-  return part.type === 'tool-searchIndex' || part.type.startsWith('tool-algolia_search_index');
-}
-
 function isMemoryToolPart(part: AIToolPart): part is MemoryToolPart {
   return (
     part.type === 'tool-algolia_ponder' ||

packages/docsearch-react/src/types/AskiAi.ts

@@ -91,3 +91,15 @@ export interface AggregatedToolCallPart {
   type: 'aggregated-tool-call';
   queries: string[];
 }
+
+export type SearchIndexOutputPart = ToolUIPart<{
+  searchIndex: SearchIndexTool;
+}>;
+export type AlgoliaMCPSearchOutputPart = ToolUIPart<
+  {
+    [K in `algolia_search_index_${string}`]: AlgoliaMCPSearchTool;
+  } & {
+    algolia_search_index: AlgoliaMCPSearchTool;
+  }
+>;
+export type SearchOutputPart = AlgoliaMCPSearchOutputPart | SearchIndexOutputPart;

packages/docsearch-react/src/utils/__tests__/ai.test.ts

@@ -1,7 +1,7 @@
 import { describe, it, expect } from 'vitest';
 
 import type { AIMessagePart } from '../../types/AskiAi';
-import { isAIToolPart } from '../ai';
+import { isAIToolPart, isAlgoliaMCPSearchOutputPart } from '../ai';
 
 describe('isAIToolPart', () => {
   it.each([
@@ -39,3 +39,43 @@ describe('isAIToolPart', () => {
     },
   );
 });
+
+describe('isAlgoliaMCPSearchOutputPart', () => {
+  it.each([
+    {
+      part: {
+        type: 'tool-algolia_search_index',
+        toolCallId: 'id-1',
+        state: 'output-available',
+        input: { query: 'foo', index: 'docs' },
+        output: { hits: [] },
+      },
+      expected: true,
+    },
+    {
+      part: {
+        type: 'tool-algolia_search_index_custom',
+        toolCallId: 'id-2',
+        state: 'output-available',
+        input: { query: 'foo', index: 'docs' },
+        output: { hits: [] },
+      },
+      expected: true,
+    },
+    {
+      part: {
+        type: 'tool-algolia_search_indexer',
+        toolCallId: 'id-3',
+        state: 'output-available',
+        input: { query: 'foo', index: 'docs' },
+        output: { hits: [] },
+      },
+      expected: false,
+    },
+  ] satisfies Array<{ part: AIMessagePart; expected: boolean }>)(
+    'returns $expected for $part.type',
+    ({ part, expected }) => {
+      expect(isAlgoliaMCPSearchOutputPart(part)).toBe(expected);
+    },
+  );
+});

packages/docsearch-react/src/utils/__tests__/groupConsecutiveToolResults.test.ts

@@ -0,0 +1,97 @@
+import { describe, it, expect } from 'vitest';
+
+import type { AIMessagePart } from '../../types/AskiAi';
+import { groupConsecutiveToolResults } from '../groupConsecutiveToolResults';
+
+function searchIndexPart(query: string): AIMessagePart {
+  const part: AIMessagePart = {
+    type: 'tool-searchIndex',
+    toolCallId: `searchIndex-${query}`,
+    state: 'output-available',
+    input: { query },
+    output: { query, hits: [] },
+  };
+
+  return part;
+}
+
+function mcpSearchPart(
+  query: string,
+  type: `tool-algolia_search_index${string}` = 'tool-algolia_search_index',
+): AIMessagePart {
+  const part: AIMessagePart = {
+    type,
+    toolCallId: `${type}-${query}`,
+    state: 'output-available',
+    input: { query, index: 'docs' },
+    output: { hits: [] },
+  };
+
+  return part;
+}
+
+function textPart(text: string): AIMessagePart {
+  const part: AIMessagePart = { type: 'text', text };
+
+  return part;
+}
+
+describe('groupConsecutiveToolResults', () => {
+  it('aggregates consecutive algolia_search_index MCP calls', () => {
+    const parts = [mcpSearchPart('foo'), mcpSearchPart('bar')];
+
+    expect(groupConsecutiveToolResults(parts)).toEqual([{ type: 'aggregated-tool-call', queries: ['foo', 'bar'] }]);
+  });
+
+  it('aggregates consecutive algolia_search_index_* MCP calls', () => {
+    const parts = [
+      mcpSearchPart('foo', 'tool-algolia_search_index_custom'),
+      mcpSearchPart('bar', 'tool-algolia_search_index_custom'),
+    ];
+
+    expect(groupConsecutiveToolResults(parts)).toEqual([{ type: 'aggregated-tool-call', queries: ['foo', 'bar'] }]);
+  });
+
+  it('does not aggregate custom tools that only share the algolia_search_index prefix', () => {
+    const parts = [
+      mcpSearchPart('foo', 'tool-algolia_search_indexer'),
+      mcpSearchPart('bar', 'tool-algolia_search_indexer'),
+    ];
+
+    expect(groupConsecutiveToolResults(parts)).toEqual(parts);
+  });
+
+  it('aggregates mixed searchIndex and MCP search calls together', () => {
+    const parts = [searchIndexPart('foo'), mcpSearchPart('bar'), searchIndexPart('baz')];
+
+    expect(groupConsecutiveToolResults(parts)).toEqual([
+      { type: 'aggregated-tool-call', queries: ['foo', 'bar', 'baz'] },
+    ]);
+  });
+
+  it('returns the original part for a single MCP search call', () => {
+    const part = mcpSearchPart('foo');
+
+    expect(groupConsecutiveToolResults([part])).toEqual([part]);
+  });
+
+  it('returns the valid MCP search part when a single valid query follows an empty query', () => {
+    const validPart = mcpSearchPart('foo');
+    const parts = [mcpSearchPart(''), validPart];
+
+    expect(groupConsecutiveToolResults(parts)).toEqual([validPart]);
+  });
+
+  it('ignores empty or whitespace-only MCP queries when aggregating', () => {
+    const parts = [mcpSearchPart('foo'), mcpSearchPart(''), mcpSearchPart('   '), mcpSearchPart('bar')];
+
+    expect(groupConsecutiveToolResults(parts)).toEqual([{ type: 'aggregated-tool-call', queries: ['foo', 'bar'] }]);
+  });
+
+  it('preserves non-search parts and breaks grouping', () => {
+    const text = textPart('hello');
+    const parts = [mcpSearchPart('foo'), text, mcpSearchPart('bar')];
+
+    expect(groupConsecutiveToolResults(parts)).toEqual([mcpSearchPart('foo'), text, mcpSearchPart('bar')]);
+  });
+});

packages/docsearch-react/src/utils/ai.ts

@@ -1,7 +1,17 @@
 import type { TextUIPart } from 'ai';
 
 import type { StoredAskAiState } from '../types';
-import type { AggregatedToolCallPart, AIMessage, AIMessagePart, AIToolPart, ToolCalls } from '../types/AskiAi';
+import type {
+  AggregatedToolCallPart,
+  AIMessage,
+  AIMessagePart,
+  AIToolPart,
+  AlgoliaMCPSearchOutputPart,
+  SearchIndexOutputPart,
+  SearchOutputPart,
+  SearchToolPart,
+  ToolCalls,
+} from '../types/AskiAi';
 
 import { sanitizeUserInput } from './sanitize';
 
@@ -113,3 +123,27 @@ export const EMPTY_TOOLS: Readonly<ToolCalls> = Object.freeze({});
 export function isAIToolPart(part: AggregatedToolCallPart | AIMessagePart): part is AIToolPart {
   return part.type.startsWith('tool-');
 }
+
+export function isSearchToolPart(part: AIToolPart): part is SearchToolPart {
+  return (
+    part.type === 'tool-searchIndex' ||
+    part.type === 'tool-algolia_search_index' ||
+    part.type.startsWith('tool-algolia_search_index_')
+  );
+}
+
+export function isSearchIndexOutputPart(part: AIMessagePart): part is SearchIndexOutputPart {
+  return part.type === 'tool-searchIndex' && part.state === 'output-available';
+}
+
+export function isAlgoliaMCPSearchOutputPart(part: AIMessagePart): part is AlgoliaMCPSearchOutputPart {
+  return (
+    isAIToolPart(part) &&
+    (part.type === 'tool-algolia_search_index' || part.type.startsWith('tool-algolia_search_index_')) &&
+    part.state === 'output-available'
+  );
+}
+
+export function isSearchOutputPart(part: AIMessagePart): part is SearchOutputPart {
+  return isAIToolPart(part) && isSearchToolPart(part) && part.state === 'output-available';
+}

packages/docsearch-react/src/utils/groupConsecutiveToolResults.ts

@@ -1,14 +1,22 @@
-import type { ToolUIPart } from 'ai';
+import type { AIMessagePart, AggregatedToolCallPart, SearchOutputPart } from '../types/AskiAi';
 
-import type { AIMessagePart, SearchIndexTool, AggregatedToolCallPart } from '../types/AskiAi';
+import { isSearchIndexOutputPart, isSearchOutputPart } from './ai';
 
-function isSearchIndexOutputPart(part: AIMessagePart): part is ToolUIPart<{ searchIndex: SearchIndexTool }> {
-  return part.type === 'tool-searchIndex' && part.state === 'output-available';
+/**
+ * Extracts the search query from a search tool result part.
+ * `searchIndex` exposes the query on its output, while the Algolia MCP search
+ * tools expose it on their input.
+ */
+function getSearchQuery(part: SearchOutputPart): string {
+  const query = isSearchIndexOutputPart(part) ? part.output?.query : part.input?.query;
+
+  return (query ?? '').trim();
 }
 
 /**
- * Groups consecutive `searchIndex` tool invocation result parts together.
- * Empty or falsy queries are ignored.
+ * Groups consecutive search tool invocation result parts together. Both the
+ * `searchIndex` tool and the Algolia MCP search tools (`algolia_search_index`
+ * and `algolia_search_index_*`) are aggregated. Empty or falsy queries are ignored.
  */
 export function groupConsecutiveToolResults(parts: AIMessagePart[]): Array<AggregatedToolCallPart | AIMessagePart> {
   const aggregatedParts: Array<AggregatedToolCallPart | AIMessagePart> = [];
@@ -16,18 +24,20 @@ export function groupConsecutiveToolResults(parts: AIMessagePart[]): Array<Aggre
   for (let i = 0; i < parts.length; i++) {
     const part = parts[i];
 
-    if (isSearchIndexOutputPart(part)) {
+    if (isSearchOutputPart(part)) {
       // build list of consecutive result queries
       const queries: string[] = [];
+      let singleQueryPart: SearchOutputPart | undefined;
       let j = i;
       while (j < parts.length) {
         const candidate = parts[j];
-        if (isSearchIndexOutputPart(candidate)) {
-          const q = (candidate.output?.query ?? '').trim();
+        if (isSearchOutputPart(candidate)) {
+          const q = getSearchQuery(candidate);
 
           // eslint-disable-next-line max-depth
           if (q && q.length > 0) {
             queries.push(q);
+            singleQueryPart = candidate;
           }
           j++;
         } else {
@@ -37,9 +47,9 @@ export function groupConsecutiveToolResults(parts: AIMessagePart[]): Array<Aggre
 
       if (queries.length > 1) {
         aggregatedParts.push({ type: 'aggregated-tool-call', queries });
-      } else if (queries.length === 1) {
+      } else if (queries.length === 1 && singleQueryPart) {
         // only one valid query, push the original part so rendering remains unchanged
-        aggregatedParts.push(part);
+        aggregatedParts.push(singleQueryPart);
       }
 
       i = j - 1; // skip processed items