feat(askai): add Agent Studio memory support (#2888)

packages/docsearch-react/src/AskAiScreen.tsx

@@ -3,7 +3,7 @@ import React, { type JSX, useMemo, useState, useEffect } from 'react';
 
 import { AggregatedSearchBlock } from './AggregatedSearchBlock';
 import type { AskAiScreenStateProps } from './AskAiScreenState';
-import { ToolCall } from './components/ui/ToolCall';
+import { ToolCall, type ToolCallTranslations } from './components/ToolCall';
 import { AlertIcon, LoadingIcon } from './icons';
 import { MemoizedMarkdown } from './MemoizedMarkdown';
 import type { StoredSearchPlugin } from './stored-searches';
@@ -12,56 +12,88 @@ import { type AIMessage, type ToolCalls } from './types/AskiAi';
 import { extractLinksFromMessage, getMessageContent, isThreadDepthError, isAIToolPart } from './utils/ai';
 import { groupConsecutiveToolResults } from './utils/groupConsecutiveToolResults';
 
-export type AskAiScreenTranslations = Partial<{
-  // Misc texts
-  disclaimerText: string;
-  relatedSourcesText: string;
-  thinkingText: string;
-  copyButtonText: string;
-  copyButtonCopiedText: string;
-  // Feedback buttons
-  copyButtonTitle: string;
-  likeButtonTitle: string;
-  dislikeButtonTitle: string;
-  thanksForFeedbackText: string;
-  // Tool call texts
-  preToolCallText: string;
-  duringToolCallText: string;
-  afterToolCallText: string;
-  /**
-   * Build the full jsx element for the aggregated search block.
-   * If provided, completely overrides the default english renderer.
-   */
-  aggregatedToolCallNode?: (queries: string[], onSearchQueryClick: (query: string) => void) => React.ReactNode;
+export type AskAiScreenTranslations = Partial<
+  // Inherit the shared tool-call translations, but expose the search-related
+  // keys under AskAiScreen's own public names (see mapping below).
+  Omit<ToolCallTranslations, 'searchingText' | 'toolCallResultText'> & {
+    // Misc texts
+    disclaimerText: string;
+    relatedSourcesText: string;
+    thinkingText: string;
+    copyButtonText: string;
+    copyButtonCopiedText: string;
+    // Feedback buttons
+    copyButtonTitle: string;
+    likeButtonTitle: string;
+    dislikeButtonTitle: string;
+    thanksForFeedbackText: string;
+    // Tool call texts
+    /**
+     * Text shown while assistant is performing search tool call.
+     * Maps to `ToolCallTranslations.searchingText`.
+     */
+    duringToolCallText: string;
+    /**
+     * Text shown while assistant is finished performing tool call.
+     * Maps to `ToolCallTranslations.toolCallResultText`.
+     */
+    afterToolCallText: string;
+    /**
+     * Build the full jsx element for the aggregated search block.
+     * If provided, completely overrides the default english renderer.
+     */
+    aggregatedToolCallNode?: (queries: string[], onSearchQueryClick: (query: string) => void) => React.ReactNode;
+    /**
+     * Generate the list connective parts only (backwards compatibility).
+     * Receives full list of queries and should return translation parts for before/after/separators.
+     * Example: (qs) => ({ before: 'searched for ', separator: ', ', lastSeparator: ' and ', after: '' }).
+     */
+    aggregatedToolCallText?: (queries: string[]) => {
+      before?: string;
+      separator?: string;
+      lastSeparator?: string;
+      after?: string;
+    };
+    /**
+     * Message that's shown when user has stopped the streaming of a message.
+     */
+    stoppedStreamingText: string;
+    /**
+     * Error title shown if there is an error while chatting.
+     */
+    errorTitleText: string;
+    /**
+     * Message shown when thread depth limit is exceeded (AI-217 error).
+     */
+    threadDepthExceededMessage: string;
+    /**
+     * Button text for starting a new conversation after thread depth error.
+     */
+    startNewConversationButtonText: string;
+  }
+>;
 
-  /**
-   * Generate the list connective parts only (backwards compatibility).
-   * Receives full list of queries and should return translation parts for before/after/separators.
-   * Example: (qs) => ({ before: 'searched for ', separator: ', ', lastSeparator: ' and ', after: '' }).
-   */
-  aggregatedToolCallText?: (queries: string[]) => {
-    before?: string;
-    separator?: string;
-    lastSeparator?: string;
-    after?: string;
+/**
+ * Maps AskAiScreen's public translation keys to the shared `ToolCallTranslations`
+ * shape consumed by the `ToolCall` component, applying default English values.
+ */
+function toToolCallTranslations(translations: AskAiScreenTranslations): ToolCallTranslations {
+  const {
+    preToolCallText = 'Searching...',
+    duringToolCallText = 'Searching...',
+    afterToolCallText = 'Searched for',
+    savedMemoryToolResultText = 'Saved to memory',
+    memoryToolResultText = 'Used memory to enhance results',
+  } = translations;
+
+  return {
+    preToolCallText,
+    searchingText: duringToolCallText,
+    toolCallResultText: afterToolCallText,
+    savedMemoryToolResultText,
+    memoryToolResultText,
   };
-  /**
-   * Message that's shown when user has stopped the streaming of a message.
-   */
-  stoppedStreamingText: string;
-  /**
-   * Error title shown if there is an error while chatting.
-   */
-  errorTitleText: string;
-  /**
-   * Message shown when thread depth limit is exceeded (AI-217 error).
-   */
-  threadDepthExceededMessage: string;
-  /**
-   * Button text for starting a new conversation after thread depth error.
-   */
-  startNewConversationButtonText: string;
-}>;
+}
 
 type AskAiScreenProps = Omit<AskAiScreenStateProps<InternalDocSearchHit>, 'translations'> & {
   messages: AIMessage[];
@@ -71,6 +103,7 @@ type AskAiScreenProps = Omit<AskAiScreenStateProps<InternalDocSearchHit>, 'trans
   translations?: AskAiScreenTranslations;
   onNewConversation: () => void;
   agentStudio?: boolean;
+  memoryEnabled?: boolean;
 };
 
 interface AskAiScreenHeaderProps {
@@ -98,6 +131,7 @@ interface AskAiExchangeCardProps {
   conversations: StoredSearchPlugin<StoredAskAiState>;
   onFeedback?: (messageId: string, thumbs: 0 | 1) => Promise<void>;
   agentStudio?: boolean;
+  memoryEnabled?: boolean;
 }
 
 function AskAiExchangeCard({
@@ -111,16 +145,13 @@ function AskAiExchangeCard({
   conversations,
   onFeedback,
   agentStudio,
+  memoryEnabled,
 }: AskAiExchangeCardProps): JSX.Element {
   const { userMessage, assistantMessage } = exchange;
 
-  const {
-    stoppedStreamingText = 'You stopped this response',
-    errorTitleText = 'Chat error',
-    preToolCallText = 'Searching...',
-    afterToolCallText = 'Searched for',
-    duringToolCallText = 'Searching...',
-  } = translations;
+  const { stoppedStreamingText = 'You stopped this response', errorTitleText = 'Chat error' } = translations;
+
+  const toolCallTranslations = useMemo(() => toToolCallTranslations(translations), [translations]);
 
   const isThreadDepth = isThreadDepthError(askAiError);
 
@@ -191,13 +222,10 @@ function AskAiExchangeCard({
                 return (
                   <ToolCall
                     key={index}
-                    translations={{
-                      preToolCallText,
-                      searchingText: duringToolCallText,
-                      toolCallResultText: afterToolCallText,
-                    }}
+                    translations={toolCallTranslations}
                     part={part}
                     tools={tools}
+                    memoryEnabled={memoryEnabled}
                     onSearchQueryClick={onSearchQueryClick}
                   />
                 );
@@ -377,7 +405,7 @@ export function AskAiScreen({ translations = {}, ...props }: AskAiScreenProps):
     startNewConversationButtonText = 'Start a new conversation',
   } = translations;
 
-  const { messages, tools, askAiError, status, agentStudio } = props;
+  const { messages, tools, askAiError, status, agentStudio, memoryEnabled } = props;
 
   // Check if there's a thread depth error
   const hasThreadDepthError = useMemo(() => {
@@ -454,6 +482,7 @@ export function AskAiScreen({ translations = {}, ...props }: AskAiScreenProps):
                 tools={tools}
                 conversations={props.conversations}
                 agentStudio={agentStudio}
+                memoryEnabled={memoryEnabled}
                 onSearchQueryClick={handleSearchQueryClick}
                 onFeedback={props.onFeedback}
               />

packages/docsearch-react/src/AskAiScreenState.tsx

@@ -58,6 +58,7 @@ export interface AskAiScreenStateProps<TItem extends BaseItem>
   selectSuggestedQuestion: (question: SuggestedQuestionHit) => void;
   onNewConversation: () => void;
   agentStudio?: boolean;
+  memoryEnabled?: boolean;
 }
 
 export const AskAiScreenState = React.memo(
@@ -86,6 +87,7 @@ export const AskAiScreenState = React.memo(
           askAiError={props.askAiError}
           translations={translations?.askAiScreen}
           agentStudio={props.agentStudio}
+          memoryEnabled={props.memoryEnabled}
         />
       );
     }

packages/docsearch-react/src/DocSearch.tsx

@@ -234,6 +234,21 @@ export interface DocSearchProps {
   keyboardShortcuts?: DocSearchModalShortcuts;
 }
 
+export interface Memory {
+  /**
+   * Determines whether or not to display the memory based tool calls.
+   *
+   * @default false
+   */
+  enabled?: boolean;
+  /**
+   * The JWT used by the agent to know which user's memory to read.
+   *
+   * @see https://www.algolia.com/doc/guides/algolia-ai/agent-studio/how-to/user-authentication
+   */
+  userToken?: string;
+}
+
 export interface DocSearchAIProps extends DocSearchProps {
   /**
    * Configuration or assistant id to enable ask ai mode. Pass a string assistant id or a full config object.
@@ -254,6 +269,10 @@ export interface DocSearchAIProps extends DocSearchProps {
    * render but will not affect correctness.
    **/
   tools?: ToolCalls;
+  /**
+   * Configuration for the Agent Studio memory feature.
+   */
+  memory?: Memory;
 }
 
 function DocSearchComponent(props: DocSearchProps, ref: React.ForwardedRef<DocSearchRef>): JSX.Element {

packages/docsearch-react/src/DocSearchAskAiModal.tsx

@@ -117,6 +117,7 @@ export function DocSearchAskAiModal({
     suggestedQuestionsEnabled: askAiConfig?.suggestedQuestions,
   });
   const agentStudio = askAiConfig?.agentStudio ?? false;
+  const memoryEnabled = props.memory?.enabled ?? false;
 
   const indexes = normalizeDocSearchIndexes({
     indexName,
@@ -143,6 +144,7 @@ export function DocSearchAskAiModal({
       useStagingEnv: askAiUseStagingEnv,
       agentStudio,
       tools,
+      memory: props.memory,
     });
 
   const prevStatus = React.useRef(status);
@@ -464,6 +466,7 @@ export function DocSearchAskAiModal({
           suggestedQuestions={suggestedQuestions}
           selectSuggestedQuestion={selectSuggestedQuestion}
           agentStudio={agentStudio}
+          memoryEnabled={memoryEnabled}
           onAskAiToggle={onAskAiToggle}
           onNewConversation={handleNewConversation}
           onItemClick={(item, event) => {

packages/docsearch-react/src/Sidepanel.tsx

@@ -4,7 +4,7 @@ import type { JSX } from 'react';
 import React from 'react';
 import { createPortal } from 'react-dom';
 
-import type { AgentStudioSearchParameters, AskAiSearchParameters } from './DocSearch';
+import type { AgentStudioSearchParameters, AskAiSearchParameters, Memory } from './DocSearch';
 import type { SidepanelButtonProps, SidepanelProps as SidepanelPanelProps } from './Sidepanel/index';
 import { SidepanelButton, Sidepanel } from './Sidepanel/index';
 import type { ToolCalls } from './types/AskiAi';
@@ -93,6 +93,10 @@ export type DocSearchSidepanelProps = DocSearchCallbacks & {
    * render but will not affect correctness.
    **/
   tools?: ToolCalls;
+  /**
+   * Configuration for the Agent Studio memory feature.
+   */
+  memory?: Memory;
 };
 
 type SidepanelProps = DocSearchSidepanelProps & SidepanelSearchParameters;

packages/docsearch-react/src/Sidepanel/ConversationScreen.tsx

@@ -3,7 +3,7 @@ import type { JSX } from 'react';
 import React, { memo, useMemo } from 'react';
 
 import { AskAiSourcesPanel, type Exchange } from '../AskAiScreen';
-import { ToolCall, type ToolCallTranslations } from '../components/ui/ToolCall';
+import { ToolCall, type ToolCallTranslations } from '../components/ToolCall';
 import { AlertIcon, LoadingIcon } from '../icons';
 import { MemoizedMarkdown } from '../MemoizedMarkdown';
 import type { StoredSearchPlugin } from '../stored-searches';
@@ -72,6 +72,7 @@ export type ConversationScreenProps = {
   handleFeedback?: (messageId: string, thumbs: 0 | 1) => Promise<void>;
   streamError?: Error;
   agentStudio?: boolean;
+  memoryEnabled?: boolean;
   tools?: ToolCalls;
 };
 
@@ -84,12 +85,24 @@ type ConversationnExchangeProps = {
   onFeedback?: ConversationScreenProps['handleFeedback'];
   streamError?: ConversationScreenProps['streamError'];
   agentStudio?: boolean;
+  memoryEnabled?: boolean;
   tools: ToolCalls;
 };
 
 const ConversationExchange = React.forwardRef<HTMLDivElement, ConversationnExchangeProps>(
   (
-    { exchange, translations = {}, isLastExchange, conversations, onFeedback, status, streamError, agentStudio, tools },
+    {
+      exchange,
+      translations = {},
+      isLastExchange,
+      conversations,
+      onFeedback,
+      status,
+      streamError,
+      agentStudio,
+      memoryEnabled,
+      tools,
+    },
     conversationRef,
   ): JSX.Element => {
     const { userMessage, assistantMessage } = exchange;
@@ -105,6 +118,8 @@ const ConversationExchange = React.forwardRef<HTMLDivElement, ConversationnExcha
       copyButtonText = 'Copy',
       copyButtonCopiedText = 'Copied!',
       errorTitleText = 'Chat error',
+      savedMemoryToolResultText = 'Saved to memory',
+      memoryToolResultText = 'Used memory to enhance results',
     } = translations;
 
     const assistantContent = useMemo(() => getMessageContent(assistantMessage), [assistantMessage]);
@@ -172,8 +187,11 @@ const ConversationExchange = React.forwardRef<HTMLDivElement, ConversationnExcha
                         preToolCallText,
                         searchingText,
                         toolCallResultText,
+                        savedMemoryToolResultText,
+                        memoryToolResultText,
                       }}
                       tools={tools}
+                      memoryEnabled={memoryEnabled}
                     />
                   );
                 }

packages/docsearch-react/src/Sidepanel/Sidepanel.tsx

@@ -153,6 +153,7 @@ function SidepanelInner(
     useStagingEnv = false,
     agentStudio = false,
     tools = EMPTY_TOOLS,
+    memory,
   }: Props,
   ref: React.ForwardedRef<SidepanelRef>,
 ): JSX.Element {
@@ -197,6 +198,7 @@ function SidepanelInner(
     useStagingEnv,
     agentStudio,
     tools,
+    memory,
   });
 
   const suggestedQuestions = useSuggestedQuestions({
@@ -400,6 +402,7 @@ function SidepanelInner(
               translations={translations.conversationScreen}
               streamError={askAiError}
               agentStudio={agentStudio}
+              memoryEnabled={memory?.enabled ?? false}
               tools={tools}
             />
           )}

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

@@ -1,15 +1,17 @@
-import { render, screen } from '@testing-library/react';
+import { render, screen, within } from '@testing-library/react';
 import '@testing-library/jest-dom/vitest';
 import React from 'react';
 import { describe, it, expect, vi } from 'vitest';
 
-import { ToolCall, type ToolCallTranslations } from '../components/ui/ToolCall';
-import type { AIToolPart, ToolCalls } from '../types/AskiAi';
+import { ToolCall, type ToolCallTranslations } from '../components/ToolCall';
+import type { AIToolPart, MemoryToolPart, ToolCalls } from '../types/AskiAi';
 
 const TRANSLATIONS: ToolCallTranslations = {
   preToolCallText: 'Searching for',
   searchingText: 'Searching...',
   toolCallResultText: 'Searched',
+  savedMemoryToolResultText: 'Saved to memory',
+  memoryToolResultText: 'Used memory to enhance results',
 };
 
 describe('ToolCall', () => {
@@ -75,6 +77,172 @@ describe('ToolCall', () => {
     });
   });
 
+  describe('memory tools', () => {
+    it.each([
+      {
+        description: 'tool-algolia_ponder',
+        part: {
+          type: 'tool-algolia_ponder',
+          toolCallId: 'memory-ponder',
+          state: 'output-available',
+          input: { value: 'remember this' },
+          output: { stored: true },
+        },
+        expectedMessage: 'Used memory to enhance results',
+      },
+      {
+        description: 'tool-algolia_memorize',
+        part: {
+          type: 'tool-algolia_memorize',
+          toolCallId: 'memory-memorize',
+          state: 'output-available',
+          input: { value: 'remember this' },
+          output: { stored: true },
+        },
+        expectedMessage: 'Saved to memory',
+      },
+      {
+        description: 'tool-algolia_memory_search',
+        part: {
+          type: 'tool-algolia_memory_search',
+          toolCallId: 'memory-search',
+          state: 'output-available',
+          input: { value: 'remember this' },
+          output: { stored: true },
+        },
+        expectedMessage: 'Used memory to enhance results',
+      },
+    ] satisfies Array<{
+      description: string;
+      part: MemoryToolPart;
+      expectedMessage: string;
+    }>)(
+      'renders the memory result text for $description when output is available and memory is enabled',
+      ({ part, expectedMessage }) => {
+        const { container } = render(
+          <ToolCall part={part} translations={TRANSLATIONS} tools={{}} memoryEnabled={true} />,
+        );
+
+        expect(within(container).getByText(expectedMessage)).toBeInTheDocument();
+      },
+    );
+
+    it.each([
+      {
+        description: 'input-streaming',
+        part: {
+          type: 'tool-algolia_ponder',
+          toolCallId: 'memory-input-streaming',
+          state: 'input-streaming',
+          input: { value: 'remember this' },
+        },
+        expectedMessage: 'Used memory to enhance results',
+      },
+      {
+        description: 'input-available',
+        part: {
+          type: 'tool-algolia_memorize',
+          toolCallId: 'memory-input-available',
+          state: 'input-available',
+          input: { value: 'remember this' },
+        },
+        expectedMessage: 'Saved to memory',
+      },
+      {
+        description: 'output-available',
+        part: {
+          type: 'tool-algolia_memory_search',
+          toolCallId: 'memory-output-available',
+          state: 'output-available',
+          input: { value: 'remember this' },
+          output: { stored: true },
+        },
+        expectedMessage: 'Used memory to enhance results',
+      },
+    ] satisfies Array<{
+      description: string;
+      part: MemoryToolPart;
+      expectedMessage: string;
+    }>)('renders the memory result text in $description state when memory is enabled', ({ part, expectedMessage }) => {
+      const { container } = render(
+        <ToolCall part={part} translations={TRANSLATIONS} tools={{}} memoryEnabled={true} />,
+      );
+
+      expect(within(container).getByText(expectedMessage)).toBeInTheDocument();
+    });
+
+    it('renders nothing for memory tool output errors when memory is enabled', () => {
+      const part: MemoryToolPart = {
+        type: 'tool-algolia_memorize',
+        toolCallId: 'memory-error',
+        state: 'output-error',
+        input: { value: 'remember this' },
+        errorText: 'Memory tool failed',
+      };
+
+      const { container } = render(
+        <ToolCall part={part} translations={TRANSLATIONS} tools={{}} memoryEnabled={true} />,
+      );
+
+      expect(container).toBeEmptyDOMElement();
+    });
+
+    it.each([
+      {
+        description: 'memoryEnabled is false',
+        part: {
+          type: 'tool-algolia_ponder',
+          toolCallId: 'memory-disabled',
+          state: 'output-available',
+          input: { value: 'remember this' },
+          output: { stored: true },
+        },
+        memoryEnabled: false,
+      },
+      {
+        description: 'memoryEnabled is omitted (defaults to false)',
+        part: {
+          type: 'tool-algolia_memory_search',
+          toolCallId: 'memory-default-off',
+          state: 'output-available',
+          input: { value: 'remember this' },
+          output: { stored: true },
+        },
+        memoryEnabled: undefined,
+      },
+    ] satisfies Array<{
+      description: string;
+      part: MemoryToolPart;
+      memoryEnabled: boolean | undefined;
+    }>)('renders nothing for memory tools when $description', ({ part, memoryEnabled }) => {
+      const { container } = render(
+        <ToolCall part={part} translations={TRANSLATIONS} tools={{}} memoryEnabled={memoryEnabled} />,
+      );
+
+      expect(container).toBeEmptyDOMElement();
+    });
+
+    it('prioritizes a matching custom tool over the memory tool rendering', () => {
+      const part: MemoryToolPart = {
+        type: 'tool-algolia_ponder',
+        toolCallId: 'memory-custom',
+        state: 'output-available',
+        input: { value: 'remember this' },
+        output: { stored: true },
+      };
+      const tools: ToolCalls = {
+        algolia_ponder: {
+          render: () => 'Custom memory render',
+        },
+      };
+
+      const { container } = render(<ToolCall part={part} translations={TRANSLATIONS} tools={tools} />);
+
+      expect(within(container).getByText('Custom memory render')).toBeInTheDocument();
+      expect(within(container).queryByText('Used memory to enhance results')).not.toBeInTheDocument();
+    });
+  });
+
   describe('custom tools', () => {
     it('renders the custom loading text while output is unavailable', () => {
       const part: AIToolPart = {

packages/docsearch-react/src/__tests__/useAskAi.test.tsx

@@ -11,6 +11,7 @@ type ToolCall = {
 
 type ChatOptions = {
   onToolCall: (params: { toolCall: ToolCall }) => unknown;
+  transport: { options: { headers?: Record<string, string> } };
 };
 
 type CustomOnToolCallParams = ToolCall & {
@@ -48,6 +49,14 @@ describe('useAskAi', () => {
     return chatOptions.onToolCall;
   }
 
+  function getTransportHeaders(): Record<string, string> {
+    if (!chatOptions) {
+      throw new Error('useChat was not initialized');
+    }
+
+    return chatOptions.transport.options.headers ?? {};
+  }
+
   beforeEach(() => {
     vi.clearAllMocks();
 
@@ -149,4 +158,37 @@ describe('useAskAi', () => {
     expect(onToolCall).toHaveBeenCalledTimes(1);
     expect(result).toBeUndefined();
   });
+
+  it('sends the secure user token header when memory.userToken is provided in Agent Studio mode', () => {
+    renderHook(() =>
+      useAskAi({
+        agentStudio: true,
+        apiKey: 'api-key',
+        appId: 'app-id',
+        assistantId: 'assistant-id',
+        indexName: 'index-name',
+        tools: {},
+        memory: { userToken: 'secure-user-token' },
+      }),
+    );
+
+    expect(getTransportHeaders()).toMatchObject({
+      'x-algolia-secure-user-token': 'secure-user-token',
+    });
+  });
+
+  it('omits the secure user token header when no memory token is provided', () => {
+    renderHook(() =>
+      useAskAi({
+        agentStudio: true,
+        apiKey: 'api-key',
+        appId: 'app-id',
+        assistantId: 'assistant-id',
+        indexName: 'index-name',
+        tools: {},
+      }),
+    );
+
+    expect(getTransportHeaders()).not.toHaveProperty('x-algolia-secure-user-token');
+  });
 });

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

@@ -1,8 +1,10 @@
 import type { JSX } from 'react';
 import React, { useMemo } from 'react';
 
-import { LoadingIcon, SearchIcon, ToolIcon } from '../../icons';
-import type { AIToolPart, SearchToolPart, ToolCalls, ToolDefinition } from '../../types/AskiAi';
+import { LoadingIcon, MemoryIcon, SearchIcon, ToolIcon } from '../icons';
+import type { AIToolPart, MemoryToolPart, SearchToolPart, ToolCalls, ToolDefinition } from '../types/AskiAi';
+
+import { ToolState } from './ui/ToolState';
 
 export type ToolCallTranslations = {
   /**
@@ -17,6 +19,14 @@ export type ToolCallTranslations = {
    * Text shown while assistant is finished performing tool call.
    */
   toolCallResultText: string;
+  /**
+   * Text shown when the agent saved related information to memory.
+   */
+  savedMemoryToolResultText: string;
+  /**
+   * Text shown when the agent used the memory tool to enhance results.
+   */
+  memoryToolResultText: string;
 };
 
 interface ToolCallProps {
@@ -24,24 +34,7 @@ interface ToolCallProps {
   translations: ToolCallTranslations;
   tools: ToolCalls;
   onSearchQueryClick?: (query: string) => void;
-}
-
-interface ToolStateProps {
-  variant: 'Call' | 'PartialCall' | 'Result';
-  icon: React.ReactNode;
-  shimmer?: boolean;
-  children: React.ReactNode;
-}
-
-function ToolState({ icon, shimmer = false, children, variant }: ToolStateProps) {
-  const className = `DocSearch-AskAiScreen-MessageContent-Tool Tool--${variant}${shimmer ? ' shimmer' : ''}`;
-
-  return (
-    <div className={className}>
-      {icon}
-      {children}
-    </div>
-  );
+  memoryEnabled?: boolean;
 }
 
 interface SearchToolProps {
@@ -148,21 +141,59 @@ function CustomTool({ tool, part }: CustomToolProps) {
   );
 }
 
+function MemoryTool({ part, translations }: { part: MemoryToolPart; translations: ToolCallTranslations }) {
+  const { savedMemoryToolResultText, memoryToolResultText } = translations;
+
+  if (part.state === 'output-error') return null;
+
+  if (part.type === 'tool-algolia_memorize') {
+    return (
+      <ToolState variant="Result" icon={<MemoryIcon />}>
+        {savedMemoryToolResultText}
+      </ToolState>
+    );
+  }
+
+  return (
+    <ToolState variant="Result" icon={<MemoryIcon />}>
+      {memoryToolResultText}
+    </ToolState>
+  );
+}
+
 function isSearchToolPart(part: AIToolPart): part is SearchToolPart {
   return part.type === 'tool-searchIndex' || part.type.startsWith('tool-algolia_search_index');
 }
 
-export function ToolCall({ part, translations, tools, onSearchQueryClick }: ToolCallProps): JSX.Element | null {
-  const normalizedToolName = part.type.replace('tool-', '');
-  const dynamicTool = tools[normalizedToolName];
-
-  if (dynamicTool) {
-    return <CustomTool tool={dynamicTool} part={part} />;
-  }
-
-  if (!isSearchToolPart(part)) {
-    return null;
-  }
-
-  return <SearchTool part={part} translations={translations} onSearchQueryClick={onSearchQueryClick} />;
+function isMemoryToolPart(part: AIToolPart): part is MemoryToolPart {
+  return (
+    part.type === 'tool-algolia_ponder' ||
+    part.type === 'tool-algolia_memorize' ||
+    part.type === 'tool-algolia_memory_search'
+  );
+}
+
+export function ToolCall({
+  part,
+  translations,
+  tools,
+  onSearchQueryClick,
+  memoryEnabled = false,
+}: ToolCallProps): JSX.Element | null {
+  const normalizedToolName = part.type.replace('tool-', '');
+  const customTool = tools[normalizedToolName];
+
+  if (customTool) {
+    return <CustomTool tool={customTool} part={part} />;
+  }
+
+  if (memoryEnabled && isMemoryToolPart(part)) {
+    return <MemoryTool part={part} translations={translations} />;
+  }
+
+  if (isSearchToolPart(part)) {
+    return <SearchTool part={part} translations={translations} onSearchQueryClick={onSearchQueryClick} />;
+  }
+
+  return null;
 }

packages/docsearch-react/src/components/ui/ToolState.tsx

@@ -0,0 +1,19 @@
+import React from 'react';
+
+export interface ToolStateProps {
+  variant: 'Call' | 'PartialCall' | 'Result';
+  icon: React.ReactNode;
+  shimmer?: boolean;
+  children: React.ReactNode;
+}
+
+export function ToolState({ icon, shimmer = false, children, variant }: ToolStateProps): React.JSX.Element {
+  const className = `DocSearch-AskAiScreen-MessageContent-Tool Tool--${variant}${shimmer ? ' shimmer' : ''}`;
+
+  return (
+    <div className={className}>
+      {icon}
+      {children}
+    </div>
+  );
+}

packages/docsearch-react/src/icons/MemoryIcon.tsx

@@ -0,0 +1,20 @@
+import React, { type JSX } from 'react';
+
+export function MemoryIcon(): JSX.Element {
+  return (
+    <svg
+      xmlns="http://www.w3.org/2000/svg"
+      width="24"
+      height="24"
+      viewBox="0 0 24 24"
+      fill="none"
+      stroke="currentColor"
+      strokeWidth="2"
+      strokeLinecap="round"
+      strokeLinejoin="round"
+    >
+      <path d="M2 3h6a4 4 0 0 1 4 4v14a3 3 0 0 0-3-3H2z" />
+      <path d="M22 3h-6a4 4 0 0 0-4 4v14a3 3 0 0 1 3-3h7z" />
+    </svg>
+  );
+}

packages/docsearch-react/src/icons/index.ts

@@ -19,3 +19,4 @@ export * from './ExpandIcon';
 export * from './FolderIcon';
 export * from './EditIcon';
 export * from './ToolIcon';
+export * from './MemoryIcon';

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

@@ -32,6 +32,11 @@ export interface AlgoliaMCPSearchTool {
   };
 }
 
+export interface MemoryTool {
+  input: unknown;
+  output: unknown;
+}
+
 export type ToolDefinition = {
   /**
    * Use the tool's input and output to build a string output for the tool result.
@@ -57,6 +62,12 @@ export type ToolDefinition = {
 
 export type ToolCalls = Record<string, ToolDefinition>;
 
+type AgentStudioMemoryTools = {
+  algolia_ponder: MemoryTool;
+  algolia_memorize: MemoryTool;
+  algolia_memory_search: MemoryTool;
+};
+
 type SearchTools = {
   [K in `algolia_search_index_${string}`]: AlgoliaMCPSearchTool;
 } & {
@@ -66,13 +77,14 @@ type SearchTools = {
 type CustomTools = {
   [K in string as K extends keyof SearchTools | `algolia_search_index_${string}` ? never : K]: CustomTool;
 };
-type Tools = CustomTools & SearchTools;
+type Tools = AgentStudioMemoryTools & CustomTools & SearchTools;
 
 export type AIMessage = UIMessage<{ stopped?: boolean }, UIDataTypes, Tools>;
 
 export type AIMessagePart = UIMessagePart<UIDataTypes, Tools>;
 
 export type SearchToolPart = ToolUIPart<SearchTools>;
+export type MemoryToolPart = ToolUIPart<AgentStudioMemoryTools>;
 export type AIToolPart = ToolUIPart<Tools>;
 
 export interface AggregatedToolCallPart {

packages/docsearch-react/src/useAskAi.ts

@@ -18,7 +18,7 @@ import { createStoredConversations } from './stored-searches';
 import { type AIMessage, type ToolCalls } from './types/AskiAi';
 import { EMPTY_TOOLS } from './utils/ai';
 
-import type { AgentStudioSearchParameters, AskAiSearchParameters, StoredAskAiState } from '.';
+import type { AgentStudioSearchParameters, AskAiSearchParameters, Memory, StoredAskAiState } from '.';
 
 type UseChat = UseChatHelpers<AIMessage>;
 
@@ -31,6 +31,7 @@ type UseAskAiParams = {
   searchParameters?: AskAiSearchParameters;
   tools: ToolCalls;
   agentStudio: boolean;
+  memory?: Memory;
 } & (
   | {
       agentStudio: false;
@@ -59,6 +60,7 @@ type UseAskAi = (params: UseAskAiParams) => UseAskAiReturn;
 
 type AgentStudioTransportParams = Pick<UseAskAiParams, 'apiKey' | 'appId' | 'assistantId'> & {
   searchParameters?: AgentStudioSearchParameters;
+  userToken?: string;
 };
 
 const getAgentStudioTransport = ({
@@ -66,12 +68,14 @@ const getAgentStudioTransport = ({
   apiKey,
   assistantId,
   searchParameters,
+  userToken,
 }: AgentStudioTransportParams): DefaultChatTransport<AIMessage> => {
   return new DefaultChatTransport({
     api: `${agentStudioBaseUrl(appId)}/agents/${assistantId}/completions?stream=true&compatibilityMode=ai-sdk-5`,
     headers: {
       'x-algolia-application-id': appId,
       'x-algolia-api-key': apiKey,
+      ...(userToken ? { 'x-algolia-secure-user-token': userToken } : {}),
     },
     body: searchParameters ? { algolia: { searchParameters } } : {},
   });
@@ -123,6 +127,7 @@ export const useAskAi: UseAskAi = ({
   tools = EMPTY_TOOLS,
   agentStudio,
   searchParameters,
+  memory,
 }) => {
   const abortControllerRef = useRef(new AbortController());
 
@@ -134,6 +139,7 @@ export const useAskAi: UseAskAi = ({
             appId,
             assistantId: assistantId ?? '',
             searchParameters,
+            userToken: memory?.userToken,
           })
         : getAskAiTransport({
             assistantId: assistantId ?? '',
@@ -144,7 +150,7 @@ export const useAskAi: UseAskAi = ({
             abortController: abortControllerRef.current,
             useStagingEnv,
           }),
-    [apiKey, appId, assistantId, indexName, useStagingEnv, agentStudio, searchParameters],
+    [apiKey, appId, assistantId, indexName, useStagingEnv, agentStudio, searchParameters, memory?.userToken],
   );
 
   // Sync ref during render so the stable `handleToolCall` (registered once