fix: Agent Studio MCP search tool

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': {
+      const queries = getSearchToolQueries(part);
+
       return (
-        <ToolState
-          shimmer={true}
-          icon={
-            <LoadingIcon className="DocSearch-AskAiScreen-SmallerLoadingIcon" />
-          }
-          variant="Call"
-        >
-          <span>
-            {preToolCallText} {`"${part.input.query || ''}" ...`}
-          </span>
-        </ToolState>
+        <>
+          {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>
+            </ToolState>
+          ))}
+        </>
       );
+    }
     case 'output-available': {
-      const query =
-        part.type === 'tool-searchIndex' ? part.output.query : part.input.query;
+      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;
-              </span>
-            ) : (
-              <span className="DocSearch-AskAiScreen-MessageContent-Tool-Query">
-                {' '}
-                &quot;{query || ''}&quot;
+        <>
+          {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:

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', () => {

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 {

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']);
+  });
+});