From f8e06784be52527edf6d29ef10e7af8f7fe88981 Mon Sep 17 00:00:00 2001 From: Paul Jankowski <8bittitan@gmail.com> Date: Sun, 26 Jul 2026 21:14:05 -0400 Subject: [PATCH] fix: Agent Studio MCP search tool (#2927) * fix: Agent Studio MCP search tool * Add changeset --- .changeset/sharp-donkeys-search.md | 20 ++++ .../src/AggregatedSearchBlock.tsx | 2 +- .../src/components/ToolCall.tsx | 103 +++++++++------- .../components/__tests__/ToolCall.test.tsx | 51 +++++++- packages/docsearch-react/src/types/AskiAi.ts | 37 ++++-- .../src/utils/__tests__/ai.test.ts | 64 +++++++++- .../groupConsecutiveToolResults.test.ts | 110 +++++++++++++----- packages/docsearch-react/src/utils/ai.ts | 24 ++++ .../src/utils/groupConsecutiveToolResults.ts | 21 +--- 9 files changed, 324 insertions(+), 108 deletions(-) create mode 100644 .changeset/sharp-donkeys-search.md diff --git a/.changeset/sharp-donkeys-search.md b/.changeset/sharp-donkeys-search.md new file mode 100644 index 00000000..3284a142 --- /dev/null +++ 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 diff --git a/packages/docsearch-react/src/AggregatedSearchBlock.tsx b/packages/docsearch-react/src/AggregatedSearchBlock.tsx index 59e0c4f2..30bdda00 100644 --- 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); } diff --git a/packages/docsearch-react/src/components/ToolCall.tsx b/packages/docsearch-react/src/components/ToolCall.tsx index 9882d0af..e824a93a 100644 --- 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({ {searchingText} ); - case 'input-available': - return ( - - } - variant="Call" - > - - {preToolCallText} {`"${part.input.query || ''}" ...`} - - - ); - case 'output-available': { - const query = - part.type === 'tool-searchIndex' ? part.output.query : part.input.query; + case 'input-available': { + const queries = getSearchToolQueries(part); return ( - } variant="Result"> - - {toolCallResultText}{' '} - {onSearchQueryClick ? ( - { - if (e.key === 'Enter' || e.key === ' ') { - e.preventDefault(); - onSearchQueryClick(query || ''); - } - }} - onClick={() => onSearchQueryClick(query || '')} - > - {' '} - "{query || ''}" + <> + {queries.map((q, index) => ( + + } + variant="Call" + > + + {preToolCallText} {`"${q}" ...`} - ) : ( - - {' '} - "{query || ''}" + + ))} + + ); + } + case 'output-available': { + const queries = getSearchToolQueries(part); + + return ( + <> + {queries.map((q, index) => ( + } + variant="Result" + > + + {toolCallResultText}{' '} + {onSearchQueryClick ? ( + { + if (e.key === 'Enter' || e.key === ' ') { + e.preventDefault(); + onSearchQueryClick(q); + } + }} + onClick={() => onSearchQueryClick(q)} + > + {' '} + "{q}" + + ) : ( + + {' '} + "{q}" + + )} - )} - - + + ))} + ); } default: diff --git a/packages/docsearch-react/src/components/__tests__/ToolCall.test.tsx b/packages/docsearch-react/src/components/__tests__/ToolCall.test.tsx index d7fa5c8d..0d50f7d9 100644 --- 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(); + + expect(screen.getByText(/"first"/)).toBeInTheDocument(); + expect(screen.getByText(/"second"/)).toBeInTheDocument(); + }); }); describe('memory tools', () => { diff --git a/packages/docsearch-react/src/types/AskiAi.ts b/packages/docsearch-react/src/types/AskiAi.ts index bf03b595..d68bb733 100644 --- 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 { diff --git a/packages/docsearch-react/src/utils/__tests__/ai.test.ts b/packages/docsearch-react/src/utils/__tests__/ai.test.ts index bc81687d..0593a4f3 100644 --- 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']); + }); +}); diff --git a/packages/docsearch-react/src/utils/__tests__/groupConsecutiveToolResults.test.ts b/packages/docsearch-react/src/utils/__tests__/groupConsecutiveToolResults.test.ts index 696ac598..06f4ad1c 100644 --- 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]); + }); }); diff --git a/packages/docsearch-react/src/utils/ai.ts b/packages/docsearch-react/src/utils/ai.ts index 5698bb65..a0f9427f 100644 --- 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 []; +} diff --git a/packages/docsearch-react/src/utils/groupConsecutiveToolResults.ts b/packages/docsearch-react/src/utils/groupConsecutiveToolResults.ts index e0700227..70eff3a3 100644 --- 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++;