fix: Agent Studio MCP search tool (#2927)

* fix: Agent Studio MCP search tool

* Add changeset

.changeset/sharp-donkeys-search.md

@@ -0,0 +1,20 @@
+---
+"@docsearch/react": patch
+---
+
+fix(askai): support Agent Studio's batched MCP search tool
+
+The Algolia MCP search tool can now issue multiple queries in a single tool
+call (`queries[]`), which previously rendered as an empty query and broke
+consecutive tool-call aggregation.
+
+- `AlgoliaMCPSearchTool["input"]` accepts both the legacy single-query shape
+  and the new batched `queries[]` shape, and tolerates a missing `output`
+- New `getSearchToolQueries()` helper normalizes query extraction across
+  `searchIndex`, batched MCP, and legacy MCP tool parts
+- `ToolCall` renders one tool state per query for both `input-available` and
+  `output-available` states
+- `groupConsecutiveToolResults` reuses the shared helper so batched calls
+  aggregate correctly
+- Fix `AggregatedSearchBlock` keyboard handler comparing `e.key === 'enter'`,
+  which never matched and made query chips unusable via keyboard

packages/docsearch-react/src/AggregatedSearchBlock.tsx

@@ -65,7 +65,7 @@ export function AggregatedSearchBlock({
               tabIndex={0}
               className="DocSearch-AskAiScreen-MessageContent-Tool-Query"
               onKeyDown={(e) => {
-                if (e.key === 'enter' || e.key === ' ') {
+                if (e.key === 'Enter' || e.key === ' ') {
                   e.preventDefault();
                   onSearchQueryClick(q);
                 }

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

@@ -9,7 +9,7 @@ import type {
   ToolCalls,
   ToolDefinition,
 } from '../types/AskiAi';
-import { isSearchToolPart } from '../utils/ai';
+import { getSearchToolQueries, isSearchToolPart } from '../utils/ai';
 
 import { ToolState } from './ui/ToolState';
 
@@ -60,52 +60,69 @@ function SearchTool({
           <span>{searchingText}</span>
         </ToolState>
       );
-    case 'input-available':
+    case 'input-available': {
-      return (
+      const queries = getSearchToolQueries(part);
-        <ToolState
-          shimmer={true}
-          icon={
-            <LoadingIcon className="DocSearch-AskAiScreen-SmallerLoadingIcon" />
-          }
-          variant="Call"
-        >
-          <span>
-            {preToolCallText} {`"${part.input.query || ''}" ...`}
-          </span>
-        </ToolState>
-      );
-    case 'output-available': {
-      const query =
-        part.type === 'tool-searchIndex' ? part.output.query : part.input.query;
 
       return (
-        <ToolState icon={<SearchIcon />} variant="Result">
+        <>
-          <span>
+          {queries.map((q, index) => (
-            {toolCallResultText}{' '}
+            <ToolState
-            {onSearchQueryClick ? (
+              // oxlint-disable-next-line react/no-array-index-key
-              <span
+              key={`${part.toolCallId}-call-${index}`}
-                role="button"
+              shimmer={true}
-                tabIndex={0}
+              icon={
-                className="DocSearch-AskAiScreen-MessageContent-Tool-Query"
+                <LoadingIcon className="DocSearch-AskAiScreen-SmallerLoadingIcon" />
-                onKeyDown={(e) => {
+              }
-                  if (e.key === 'Enter' || e.key === ' ') {
+              variant="Call"
-                    e.preventDefault();
+            >
-                    onSearchQueryClick(query || '');
+              <span>
-                  }
+                {preToolCallText} {`"${q}" ...`}
-                }}
-                onClick={() => onSearchQueryClick(query || '')}
-              >
-                {' '}
-                &quot;{query || ''}&quot;
               </span>
-            ) : (
+            </ToolState>
-              <span className="DocSearch-AskAiScreen-MessageContent-Tool-Query">
+          ))}
-                {' '}
+        </>
-                &quot;{query || ''}&quot;
+      );
+    }
+    case 'output-available': {
+      const queries = getSearchToolQueries(part);
+
+      return (
+        <>
+          {queries.map((q, index) => (
+            <ToolState
+              // oxlint-disable-next-line react/no-array-index-key
+              key={`${part.toolCallId}-result-${index}`}
+              icon={<SearchIcon />}
+              variant="Result"
+            >
+              <span>
+                {toolCallResultText}{' '}
+                {onSearchQueryClick ? (
+                  <span
+                    role="button"
+                    tabIndex={0}
+                    className="DocSearch-AskAiScreen-MessageContent-Tool-Query"
+                    onKeyDown={(e) => {
+                      if (e.key === 'Enter' || e.key === ' ') {
+                        e.preventDefault();
+                        onSearchQueryClick(q);
+                      }
+                    }}
+                    onClick={() => onSearchQueryClick(q)}
+                  >
+                    {' '}
+                    &quot;{q}&quot;
+                  </span>
+                ) : (
+                  <span className="DocSearch-AskAiScreen-MessageContent-Tool-Query">
+                    {' '}
+                    &quot;{q}&quot;
+                  </span>
+                )}
               </span>
-            )}
+            </ToolState>
-          </span>
+          ))}
-        </ToolState>
+        </>
       );
     }
     default:

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

@@ -48,12 +48,15 @@ describe('ToolCall', () => {
           toolCallId: 'id-3',
           state: 'output-available',
           input: {
-            index: 'docs',
+            clickAnalytics: false,
-            query: 'test',
+            originalQuery: 'testing',
-            number_of_results: 10,
+            queries: [
-            facet_filters: null,
+              {
+                query: 'test',
+              },
+            ],
           },
-          output: { hits: [{}, {}] },
+          output: { hits: [], nbHits: 7 },
         },
       },
       {
@@ -62,7 +65,15 @@ describe('ToolCall', () => {
           type: 'tool-algolia_search_index_custom',
           toolCallId: 'id-4',
           state: 'output-available',
-          input: { query: 'test' },
+          input: {
+            clickAnalytics: false,
+            originalQuery: 'testing',
+            queries: [
+              {
+                query: 'test',
+              },
+            ],
+          },
           output: { hits: [{}] },
         },
       },
@@ -82,6 +93,34 @@ describe('ToolCall', () => {
         expect(within(container).getByText(/"test"/)).toBeInTheDocument();
       }
     );
+
+    it('renders multiple MCP search tool queries', () => {
+      const part = {
+        type: 'tool-algolia_search_index_test',
+        toolCallId: 'multiple-queries',
+        state: 'output-available',
+        input: {
+          clickAnalytics: false,
+          originalQuery: 'testing',
+          queries: [
+            {
+              query: 'first',
+            },
+            {
+              query: 'second',
+            },
+          ],
+        },
+        output: {
+          hits: [],
+        },
+      } satisfies AIToolPart;
+
+      render(<ToolCall part={part} translations={TRANSLATIONS} tools={{}} />);
+
+      expect(screen.getByText(/"first"/)).toBeInTheDocument();
+      expect(screen.getByText(/"second"/)).toBeInTheDocument();
+    });
   });
 
   describe('memory tools', () => {

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

@@ -22,18 +22,33 @@ export interface SearchIndexTool {
   };
 }
 
+interface MCPSearchToolQuery {
+  query: string;
+  [key: string]: unknown;
+}
+
+interface MCPSearchToolInputV1 {
+  query: string;
+  index: string;
+  number_of_results?: number;
+  facet_filters?: string[];
+}
+
+interface MCPSearchToolInputV2 {
+  queries: MCPSearchToolQuery[];
+  clickAnalytics: boolean;
+  originalQuery: string;
+}
+
 export interface AlgoliaMCPSearchTool {
-  input: {
+  input: MCPSearchToolInputV1 | MCPSearchToolInputV2;
-    query: string;
+  output:
-    index: string;
+    | {
-    number_of_results?: number;
+        hits?: unknown[];
-    facet_filters?: string[];
+        nbHits?: number;
-  };
+        queryId?: string;
-  output: {
+      }
-    hits?: any[];
+    | undefined;
-    nbHits?: number;
-    queryId?: string;
-  };
 }
 
 export interface MemoryTool {

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

@@ -1,8 +1,13 @@
 import { describe, it, expect } from 'vitest';
 
-import type { AIMessage, AIMessagePart } from '../../types/AskiAi';
+import type {
+  AIMessage,
+  AIMessagePart,
+  SearchToolPart,
+} from '../../types/AskiAi';
 import {
   getAgentPromptSuggestions,
+  getSearchToolQueries,
   isAIToolPart,
   isAlgoliaMCPSearchOutputPart,
   sanitizeMessagesForRequest,
@@ -185,3 +190,60 @@ describe('getAgentPromptSuggestions', () => {
     ).toEqual(['First suggestion']);
   });
 });
+
+describe('getSearchToolQueries', () => {
+  it('returns input query for tool-searchIndex', () => {
+    const queries = getSearchToolQueries({
+      toolCallId: 'testing-123',
+      type: 'tool-searchIndex',
+      state: 'input-available',
+      input: {
+        query: 'testing',
+      },
+      output: undefined,
+    });
+
+    expect(queries).toEqual(['testing']);
+  });
+
+  it('returns queries for MCP search tool', () => {
+    const queries = getSearchToolQueries({
+      type: 'tool-algolia_search_index_testing',
+      toolCallId: 'testing-456',
+      state: 'input-available',
+      input: {
+        clickAnalytics: false,
+        originalQuery: 'testing',
+        queries: [
+          {
+            query: 'first',
+          },
+          {
+            query: '',
+          },
+          {
+            query: 'second',
+          },
+        ],
+      },
+      output: undefined,
+    });
+
+    expect(queries).toEqual(['first', 'second']);
+  });
+
+  it('extracts query from stored MCP tool call with v1 input', () => {
+    const part: SearchToolPart = {
+      type: 'tool-algolia_search_index',
+      toolCallId: 'legacy-id',
+      state: 'output-available',
+      input: {
+        query: '  foo  ',
+        index: 'docs',
+      },
+      output: { hits: [] },
+    };
+
+    expect(getSearchToolQueries(part)).toEqual(['foo']);
+  });
+});

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

@@ -1,6 +1,10 @@
 import { describe, it, expect } from 'vitest';
 
-import type { AIMessagePart } from '../../types/AskiAi';
+import type {
+  AIMessagePart,
+  AlgoliaMCPSearchOutputPart,
+  SearchToolPart,
+} from '../../types/AskiAi';
 import { groupConsecutiveToolResults } from '../groupConsecutiveToolResults';
 
 function searchIndexPart(query: string): AIMessagePart {
@@ -15,19 +19,49 @@ function searchIndexPart(query: string): AIMessagePart {
   return part;
 }
 
-function mcpSearchPart(
+function messagePart(query: string): AIMessagePart {
-  query: string,
+  return {
-  type: `tool-algolia_search_index${string}` = 'tool-algolia_search_index'
+    type: 'tool-algolia_search_indexer',
-): AIMessagePart {
+    toolCallId: 'tool-123',
-  const part: AIMessagePart = {
+    input: {},
-    type,
-    toolCallId: `${type}-${query}`,
     state: 'output-available',
-    input: { query, index: 'docs' },
+    output: {
+      query,
+    },
+  };
+}
+
+function mcpSearchPart(
+  queries: string[],
+  type:
+    | 'tool-algolia_search_index'
+    | `tool-algolia_search_index_${string}` = 'tool-algolia_search_index'
+): AlgoliaMCPSearchOutputPart {
+  return {
+    type,
+    toolCallId: `${type}-123`,
+    state: 'output-available',
+    input: {
+      clickAnalytics: false,
+      originalQuery: 'What is composable api?',
+      queries: queries.map((q) => ({ query: q })),
+    },
     output: { hits: [] },
   };
+}
 
-  return part;
+function v1McpSearchPart(query: string): SearchToolPart {
+  return {
+    type: 'tool-algolia_search_index_testing',
+    state: 'output-available',
+    toolCallId: 'legacy-tool-123',
+    input: {
+      query,
+      number_of_results: 2,
+      index: 'test_index',
+    },
+    output: undefined,
+  };
 }
 
 function textPart(text: string): AIMessagePart {
@@ -38,7 +72,7 @@ function textPart(text: string): AIMessagePart {
 
 describe('groupConsecutiveToolResults', () => {
   it('aggregates consecutive algolia_search_index MCP calls', () => {
-    const parts = [mcpSearchPart('foo'), mcpSearchPart('bar')];
+    const parts = [mcpSearchPart(['foo']), mcpSearchPart(['bar'])];
 
     expect(groupConsecutiveToolResults(parts)).toEqual([
       { type: 'aggregated-tool-call', queries: ['foo', 'bar'] },
@@ -47,8 +81,8 @@ describe('groupConsecutiveToolResults', () => {
 
   it('aggregates consecutive algolia_search_index_* MCP calls', () => {
     const parts = [
-      mcpSearchPart('foo', 'tool-algolia_search_index_custom'),
+      mcpSearchPart(['foo'], 'tool-algolia_search_index_custom'),
-      mcpSearchPart('bar', 'tool-algolia_search_index_custom'),
+      mcpSearchPart(['bar'], 'tool-algolia_search_index_custom'),
     ];
 
     expect(groupConsecutiveToolResults(parts)).toEqual([
@@ -57,10 +91,7 @@ describe('groupConsecutiveToolResults', () => {
   });
 
   it('does not aggregate custom tools that only share the algolia_search_index prefix', () => {
-    const parts = [
+    const parts = [messagePart('foo'), messagePart('bar')];
-      mcpSearchPart('foo', 'tool-algolia_search_indexer'),
-      mcpSearchPart('bar', 'tool-algolia_search_indexer'),
-    ];
 
     expect(groupConsecutiveToolResults(parts)).toEqual(parts);
   });
@@ -68,7 +99,7 @@ describe('groupConsecutiveToolResults', () => {
   it('aggregates mixed searchIndex and MCP search calls together', () => {
     const parts = [
       searchIndexPart('foo'),
-      mcpSearchPart('bar'),
+      mcpSearchPart(['bar']),
       searchIndexPart('baz'),
     ];
 
@@ -78,39 +109,60 @@ describe('groupConsecutiveToolResults', () => {
   });
 
   it('returns the original part for a single MCP search call', () => {
-    const part = mcpSearchPart('foo');
+    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 validPart = mcpSearchPart(['foo']);
-    const parts = [mcpSearchPart(''), validPart];
+    const parts = [mcpSearchPart(['']), validPart];
 
     expect(groupConsecutiveToolResults(parts)).toEqual([validPart]);
   });
 
   it('ignores empty or whitespace-only MCP queries when aggregating', () => {
     const parts = [
-      mcpSearchPart('foo'),
+      mcpSearchPart(['foo']),
-      mcpSearchPart(''),
+      mcpSearchPart(['']),
-      mcpSearchPart('   '),
+      mcpSearchPart(['   ']),
-      mcpSearchPart('bar'),
+      mcpSearchPart(['bar']),
+      mcpSearchPart([' baz  ']),
     ];
 
     expect(groupConsecutiveToolResults(parts)).toEqual([
-      { type: 'aggregated-tool-call', queries: ['foo', 'bar'] },
+      { type: 'aggregated-tool-call', queries: ['foo', 'bar', 'baz'] },
     ]);
   });
 
   it('preserves non-search parts and breaks grouping', () => {
     const text = textPart('hello');
-    const parts = [mcpSearchPart('foo'), text, mcpSearchPart('bar')];
+    const parts = [mcpSearchPart(['foo']), text, mcpSearchPart(['bar'])];
 
     expect(groupConsecutiveToolResults(parts)).toEqual([
-      mcpSearchPart('foo'),
+      mcpSearchPart(['foo']),
       text,
-      mcpSearchPart('bar'),
+      mcpSearchPart(['bar']),
     ]);
   });
+
+  it('ignores MCP part with blank queries', () => {
+    const part = mcpSearchPart([]);
+
+    expect(groupConsecutiveToolResults([part])).toEqual([]);
+  });
+
+  it('aggregates multiple queries from a MCP call', () => {
+    const part = mcpSearchPart(['foo', 'bar']);
+
+    expect(groupConsecutiveToolResults([part])).toEqual([
+      { type: 'aggregated-tool-call', queries: ['foo', 'bar'] },
+    ]);
+  });
+
+  it('keeps legacy MCP results in grouped output', () => {
+    const part = v1McpSearchPart('foo');
+
+    expect(groupConsecutiveToolResults([part])).toEqual([part]);
+  });
 });

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

@@ -204,3 +204,27 @@ export function getAgentPromptSuggestions(parts: AIMessagePart[]): string[] {
 
   return suggestionsPart.data.suggestions;
 }
+
+export function getSearchToolQueries(part: SearchToolPart): string[] {
+  if (part.state !== 'input-available' && part.state !== 'output-available') {
+    return [];
+  }
+
+  if (part.type === 'tool-searchIndex') {
+    const query = (part.output?.query ?? part.input?.query ?? '').trim();
+    return query ? [query] : [];
+  }
+
+  if ('queries' in part.input && Array.isArray(part.input.queries)) {
+    return part.input.queries.map(({ query }) => query.trim()).filter(Boolean);
+  }
+
+  // There could be older stored MCP search tool calls,
+  // we should parse it's input properly
+  if ('query' in part.input && typeof part.input.query === 'string') {
+    const query = part.input.query.trim();
+    return query ? [query] : [];
+  }
+
+  return [];
+}

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

@@ -4,20 +4,7 @@ import type {
   SearchOutputPart,
 } from '../types/AskiAi';
 
-import { isSearchIndexOutputPart, isSearchOutputPart } from './ai';
+import { getSearchToolQueries, isSearchOutputPart } from './ai';
-
-/**
- * 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 search tool invocation result parts together. Both the
@@ -41,11 +28,11 @@ export function groupConsecutiveToolResults(
       while (j < parts.length) {
         const candidate = parts[j];
         if (isSearchOutputPart(candidate)) {
-          const q = getSearchQuery(candidate);
+          const queriesForPart = getSearchToolQueries(candidate);
+          queries.push(...queriesForPart);
 
           // eslint-disable-next-line max-depth
-          if (q && q.length > 0) {
+          if (queriesForPart.length > 0) {
-            queries.push(q);
             singleQueryPart = candidate;
           }
           j++;