From e01ff48f67b0adc4932f2d7a73894dacf3e7d81d Mon Sep 17 00:00:00 2001 From: Paul Jankowski <8bittitan@gmail.com> Date: Wed, 3 Jun 2026 12:48:29 -0400 Subject: [PATCH] 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 | 24 ++--- .../src/components/ToolCall.tsx | 5 +- packages/docsearch-react/src/types/AskiAi.ts | 12 +++ .../src/utils/__tests__/ai.test.ts | 42 +++++++- .../groupConsecutiveToolResults.test.ts | 97 +++++++++++++++++++ packages/docsearch-react/src/utils/ai.ts | 36 ++++++- .../src/utils/groupConsecutiveToolResults.ts | 32 +++--- 7 files changed, 216 insertions(+), 32 deletions(-) create mode 100644 packages/docsearch-react/src/utils/__tests__/groupConsecutiveToolResults.test.ts diff --git a/AGENTS.md b/AGENTS.md index c881fa50..0b131c13 100644 --- a/AGENTS.md +++ b/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 diff --git a/packages/docsearch-react/src/components/ToolCall.tsx b/packages/docsearch-react/src/components/ToolCall.tsx index 4ee082cc..18ffbc28 100644 --- a/packages/docsearch-react/src/components/ToolCall.tsx +++ b/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' || diff --git a/packages/docsearch-react/src/types/AskiAi.ts b/packages/docsearch-react/src/types/AskiAi.ts index 8dbfe18c..8e438f57 100644 --- a/packages/docsearch-react/src/types/AskiAi.ts +++ b/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; diff --git a/packages/docsearch-react/src/utils/__tests__/ai.test.ts b/packages/docsearch-react/src/utils/__tests__/ai.test.ts index ca7c6ebd..30ca778e 100644 --- a/packages/docsearch-react/src/utils/__tests__/ai.test.ts +++ b/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); + }, + ); +}); diff --git a/packages/docsearch-react/src/utils/__tests__/groupConsecutiveToolResults.test.ts b/packages/docsearch-react/src/utils/__tests__/groupConsecutiveToolResults.test.ts new file mode 100644 index 00000000..2e231595 --- /dev/null +++ b/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')]); + }); +}); diff --git a/packages/docsearch-react/src/utils/ai.ts b/packages/docsearch-react/src/utils/ai.ts index e84c88f0..a17a0d8f 100644 --- a/packages/docsearch-react/src/utils/ai.ts +++ b/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 = 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'; +} diff --git a/packages/docsearch-react/src/utils/groupConsecutiveToolResults.ts b/packages/docsearch-react/src/utils/groupConsecutiveToolResults.ts index abaf13f3..9fcc902a 100644 --- a/packages/docsearch-react/src/utils/groupConsecutiveToolResults.ts +++ b/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 { const aggregatedParts: Array = []; @@ -16,18 +24,20 @@ export function groupConsecutiveToolResults(parts: AIMessagePart[]): Array 0) { queries.push(q); + singleQueryPart = candidate; } j++; } else { @@ -37,9 +47,9 @@ export function groupConsecutiveToolResults(parts: AIMessagePart[]): Array 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