fix: Agent Studio MCP search tool (#2927)

* fix: Agent Studio MCP search tool * Add changeset
2026-07-26 21:14:05 -04:00 · 2026-07-26 21:14:05 -04:00 · f8e06784be
commit f8e06784be
parent 82716c0888
9 changed files with 324 additions and 108 deletions
--- a/.changeset/sharp-donkeys-search.md
+++ b/.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
--- a/packages/docsearch-react/src/AggregatedSearchBlock.tsx
+++ b/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);
                }
--- a/packages/docsearch-react/src/components/ToolCall.tsx
+++ b/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':
-      return (
-        <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;
+    case 'input-available': {
+      const queries = getSearchToolQueries(part);

      return (
-        <ToolState 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(query || '');
-                  }
-                }}
-                onClick={() => onSearchQueryClick(query || '')}
-              >
-                {' '}
-                &quot;{query || ''}&quot;
+        <>
+          {queries.map((q, index) => (
+            <ToolState
+              // oxlint-disable-next-line react/no-array-index-key
+              key={`${part.toolCallId}-call-${index}`}
+              shimmer={true}
+              icon={
+                <LoadingIcon className="DocSearch-AskAiScreen-SmallerLoadingIcon" />
+              }
+              variant="Call"
+            >
+              <span>
+                {preToolCallText} {`"${q}" ...`}
              </span>
-            ) : (
-              <span className="DocSearch-AskAiScreen-MessageContent-Tool-Query">
-                {' '}
-                &quot;{query || ''}&quot;
+            </ToolState>
+          ))}
+        </>
+      );
+    }
+    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>
-            )}
-          </span>
-        </ToolState>
+            </ToolState>
+          ))}
+        </>
      );
    }
    default:
--- a/packages/docsearch-react/src/components/tests/ToolCall.test.tsx
+++ b/packages/docsearch-react/src/components/tests/ToolCall.test.tsx
@ -48,12 +48,15 @@ describe('ToolCall', () => {
          toolCallId: 'id-3',
          state: 'output-available',
          input: {
-            index: 'docs',
-            query: 'test',
-            number_of_results: 10,
-            facet_filters: null,
+            clickAnalytics: false,
+            originalQuery: 'testing',
+            queries: [
+              {
+                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', () => {
--- a/packages/docsearch-react/src/types/AskiAi.ts
+++ b/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: {
-    query: string;
-    index: string;
-    number_of_results?: number;
-    facet_filters?: string[];
-  };
-  output: {
-    hits?: any[];
-    nbHits?: number;
-    queryId?: string;
-  };
+  input: MCPSearchToolInputV1 | MCPSearchToolInputV2;
+  output:
+    | {
+        hits?: unknown[];
+        nbHits?: number;
+        queryId?: string;
+      }
+    | undefined;
 }

 export interface MemoryTool {
--- a/packages/docsearch-react/src/utils/tests/ai.test.ts
+++ b/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']);
+  });
+});
--- a/packages/docsearch-react/src/utils/tests/groupConsecutiveToolResults.test.ts
+++ b/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(
-  query: string,
-  type: `tool-algolia_search_index${string}` = 'tool-algolia_search_index'
-): AIMessagePart {
-  const part: AIMessagePart = {
-    type,
-    toolCallId: `${type}-${query}`,
+function messagePart(query: string): AIMessagePart {
+  return {
+    type: 'tool-algolia_search_indexer',
+    toolCallId: 'tool-123',
+    input: {},
    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('bar', 'tool-algolia_search_index_custom'),
+      mcpSearchPart(['foo'], '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 = [
-      mcpSearchPart('foo', 'tool-algolia_search_indexer'),
-      mcpSearchPart('bar', 'tool-algolia_search_indexer'),
-    ];
+    const parts = [messagePart('foo'), messagePart('bar')];

    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 parts = [mcpSearchPart(''), validPart];
+    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'),
+      mcpSearchPart(['foo']),
+      mcpSearchPart(['']),
+      mcpSearchPart(['   ']),
+      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]);
+  });
 });
--- a/packages/docsearch-react/src/utils/ai.ts
+++ b/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 [];
+}
--- a/packages/docsearch-react/src/utils/groupConsecutiveToolResults.ts
+++ b/packages/docsearch-react/src/utils/groupConsecutiveToolResults.ts
@ -4,20 +4,7 @@ import type {
  SearchOutputPart,
 } from '../types/AskiAi';

-import { isSearchIndexOutputPart, 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();
-}
+import { getSearchToolQueries, isSearchOutputPart } from './ai';

 /**
 * 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) {
-            queries.push(q);
+          if (queriesForPart.length > 0) {
            singleQueryPart = candidate;
          }
          j++;