Swizzle and use agentStudio for ask AI

docusaurus.config.js

@@ -9,6 +9,7 @@ const codeTheme = themes.dracula
 const productsDropdown = fs.readFileSync('./src/components/NavDropdown/Products.html', 'utf-8')
 const baseUrl = process.env.DEST || '/'
 const siteUrl = 'https://docs.metamask.io'
+const algoliaAssistantId = process.env.ALGOLIA_ASSISTANT_ID
 
 // Options for the `llms-html-injector` plugin (which wraps
 // `docusaurus-plugin-llms`). Centralized in a standalone CommonJS module so
@@ -649,9 +650,11 @@ const config = {
             buttonAriaLabel: 'Search or Ask AI',
           },
         },
-        askAi: {
-          assistantId: 'REak1eiP5wfp',
-        },
+        ...(algoliaAssistantId && {
+          askAi: {
+            assistantId: algoliaAssistantId,
+          },
+        }),
         // Disable the standalone `/search/` results page. The Algolia DocSearch
         // modal still works; the dedicated page was being indexed as an orphan
         // (Ahrefs orphan report, 2026-05-25).

src/globals.d.ts

@@ -1,3 +1,9 @@
+// Docusaurus declares `@docusaurus/*` and `@theme/*` here. A full `tsc` run reaches this package
+// transitively through whichever source file happens to import `@theme/*`, but the `tsc-files`
+// pre-commit hook builds a program from the staged files alone, so nothing guarantees that import
+// is present. Referencing it directly makes the declarations unconditional in both.
+/// <reference types="@docusaurus/module-type-aliases" />
+
 declare module '*.svg' {
   import { FC, SVGProps } from 'react'
   const content: FC<SVGProps<SVGElement>>

src/lib/algolia-ask-ai.ts

@@ -0,0 +1,177 @@
+import type { ThemeConfigAlgolia } from '@docusaurus/theme-search-algolia'
+
+export type AskAiConfig = Exclude<ThemeConfigAlgolia['askAi'], string | undefined>
+
+/**
+ * DocSearch only routes Ask AI through Agent Studio when this flag is passed, and it must be
+ * passed on every surface that renders Ask AI: the search modal (`src/theme/SearchBar`) and the
+ * sidepanel (`src/theme/Root`). It cannot live in `themeConfig.algolia.askAi` because Docusaurus
+ * validates that object against a Joi schema that rejects unknown keys.
+ */
+export const ASK_AI_AGENT_STUDIO = true
+
+const AGENT_STUDIO_COMPLETIONS = /\/agent-studio\/\d+\/agents\/[^/]+\/completions/
+
+/** DocSearch renders the modal in a portal, so these are matched from the document. */
+const SEARCH_INPUT_SELECTOR = '.DocSearch-Modal .DocSearch-Input'
+const HIGHLIGHTED_ASK_AI_OPTION_SELECTOR = '[id^="docsearch-askAI-item"][aria-selected="true"]'
+const ASK_AI_OPTION_QUERY_SELECTOR = '.DocSearch-Hit-AskAIButton-title-query'
+
+/** Comfortably longer than a search round-trip, so the wait ends in a catch-up rather than a
+ * timeout under normal conditions. */
+const OPTION_CATCH_UP_TIMEOUT_MS = 2000
+const OPTION_CATCH_UP_POLL_MS = 50
+
+interface AskAiMessage {
+  role?: string
+  parts?: Array<{ type?: string; text?: string }>
+}
+
+interface AskAiRequestBody {
+  messages?: AskAiMessage[]
+}
+
+function parseRequestBody(body: BodyInit | null | undefined): AskAiRequestBody | null {
+  if (typeof body !== 'string') {
+    return null
+  }
+
+  try {
+    return JSON.parse(body) as AskAiRequestBody
+  } catch {
+    return null
+  }
+}
+
+/**
+ * DocSearch builds its chat with the AI SDK `lastAssistantMessageIsCompleteWithToolCalls` auto-send
+ * predicate so that client-side tools can hand their results back to the model. Agent Studio runs
+ * its search tool server-side but never marks the resulting parts `providerExecuted`, so whenever
+ * the agent emits its tool calls and its answer text in a single step, the predicate misfires and
+ * re-posts the already finished conversation. Agent Studio rejects that with a 422 and DocSearch
+ * renders it as a "Chat error" next to an otherwise correct answer.
+ *
+ * A conversation ending in an assistant message is always one of these re-posts, because every
+ * genuine request ends with the question the user just asked.
+ */
+function isAutoResubmit({ messages }: AskAiRequestBody): boolean {
+  return Array.isArray(messages) && messages[messages.length - 1]?.role === 'assistant'
+}
+
+function normalizeQuery(value: string): string {
+  return value.replace(/\s+/g, ' ').trim()
+}
+
+/** The question the highlighted "Ask AI" option would ask, or `null` when it isn't highlighted. */
+function highlightedAskAiQuery(): string | null {
+  const option = document.querySelector(HIGHLIGHTED_ASK_AI_OPTION_SELECTOR)
+  return option?.querySelector(ASK_AI_OPTION_QUERY_SELECTOR)?.textContent ?? null
+}
+
+/**
+ * Holds Enter back while the "Ask AI" option is still showing an older query than the input, then
+ * replays it once the option catches up, so DocSearch builds the conversation from the question the
+ * user actually typed. Fixing it here rather than on the request keeps the transcript honest: the
+ * message DocSearch renders is the one it sends.
+ *
+ * Replaying is abandoned if the user keeps typing or leaves, and falls through to the original
+ * behaviour on timeout so Enter can never be swallowed outright.
+ */
+function installStaleQuestionGuard(): () => void {
+  let replaying = false
+
+  const submit = (input: HTMLInputElement) => {
+    replaying = true
+    input.dispatchEvent(new KeyboardEvent('keydown', { key: 'Enter', bubbles: true }))
+    replaying = false
+  }
+
+  const onKeyDown = (event: KeyboardEvent) => {
+    const input = event.target
+    if (
+      replaying ||
+      event.key !== 'Enter' ||
+      event.shiftKey ||
+      event.metaKey ||
+      event.ctrlKey ||
+      event.altKey ||
+      !(input instanceof HTMLInputElement) ||
+      !input.matches(SEARCH_INPUT_SELECTOR)
+    ) {
+      return
+    }
+
+    const optionQuery = highlightedAskAiQuery()
+    const intendedQuery = input.value
+    if (optionQuery === null || normalizeQuery(optionQuery) === normalizeQuery(intendedQuery)) {
+      return
+    }
+
+    event.preventDefault()
+    event.stopImmediatePropagation()
+
+    const deadline = Date.now() + OPTION_CATCH_UP_TIMEOUT_MS
+    const poll = window.setInterval(() => {
+      const caughtUp =
+        normalizeQuery(highlightedAskAiQuery() ?? '') === normalizeQuery(intendedQuery)
+      const abandoned = !input.isConnected || input.value !== intendedQuery
+
+      if (abandoned) {
+        window.clearInterval(poll)
+        return
+      }
+
+      if (caughtUp || Date.now() > deadline) {
+        window.clearInterval(poll)
+        submit(input)
+      }
+    }, OPTION_CATCH_UP_POLL_MS)
+  }
+
+  document.addEventListener('keydown', onKeyDown, true)
+
+  return () => {
+    document.removeEventListener('keydown', onKeyDown, true)
+  }
+}
+
+/** Suppresses the spurious re-post that Agent Studio answers with a 422. */
+function installResubmitGuard(): () => void {
+  const originalFetch = window.fetch
+
+  window.fetch = function patchedFetch(input, init) {
+    const url = typeof input === 'string' ? input : input instanceof URL ? input.href : input.url
+    const body = AGENT_STUDIO_COMPLETIONS.test(url) ? parseRequestBody(init?.body) : null
+
+    if (!body) {
+      return originalFetch.call(this, input, init)
+    }
+
+    // The re-post has to fail as an `AbortError`: the AI SDK returns early from that case, which
+    // both leaves the chat in the `ready` state and skips the auto-send it would otherwise
+    // re-evaluate after every request. Resolving it, even with an empty stream, re-satisfies the
+    // predicate and spins forever. This is only safe because Ask AI registers no client-side tools;
+    // if it ever does, the auto-send becomes load bearing and the guard has to go.
+    if (isAutoResubmit(body)) {
+      return Promise.reject(new DOMException('Ask AI auto-resubmit suppressed', 'AbortError'))
+    }
+
+    return originalFetch.call(this, input, init)
+  }
+
+  return () => {
+    window.fetch = originalFetch
+  }
+}
+
+/**
+ * Works around DocSearch v4 defects that surface against Agent Studio, none of which have a
+ * supported override.
+ */
+export function installAskAiWorkarounds(): () => void {
+  const cleanups = [installStaleQuestionGuard(), installResubmitGuard()]
+
+  return () => {
+    cleanups.forEach(cleanup => cleanup())
+  }
+}

src/theme/Root.tsx

@@ -20,6 +20,11 @@ import AuthModal, {
   WALLET_LINK_TYPE,
 } from '@site/src/components/AuthLogin/AuthModal'
 import { DocSearchSidepanel } from '@docsearch/react/sidepanel'
+import {
+  ASK_AI_AGENT_STUDIO,
+  installAskAiWorkarounds,
+  type AskAiConfig,
+} from '@site/src/lib/algolia-ask-ai'
 import '@docsearch/css/dist/sidepanel.css'
 
 interface Project {
@@ -67,10 +72,7 @@ interface AlgoliaThemeConfig {
   appId: string
   apiKey: string
   indexName: string
-  assistantId?: string
-  askAi?: {
-    assistantId: string
-  }
+  askAi?: AskAiConfig
 }
 
 export const MetamaskProviderContext = createContext<IMetamaskProviderContext>({
@@ -301,17 +303,26 @@ export default function Root({ children }: { children: ReactElement }) {
   const { siteConfig } = useDocusaurusContext()
   const isBrowser = useIsBrowser()
   const algolia = siteConfig?.themeConfig?.algolia as AlgoliaThemeConfig | undefined
+  const assistantId = algolia?.askAi?.assistantId
+
+  useEffect(() => {
+    if (!isBrowser || !assistantId) {
+      return undefined
+    }
+    return installAskAiWorkarounds()
+  }, [isBrowser, assistantId])
 
   return (
     <LoginProvider>
       <AlertProvider template={AlertTemplate} {...options}>
         {children}
-        {isBrowser && (algolia?.assistantId || algolia?.askAi?.assistantId) ? (
+        {isBrowser && assistantId ? (
           <DocSearchSidepanel
             appId={algolia.appId}
             apiKey={algolia.apiKey}
-            assistantId={algolia.assistantId || algolia.askAi?.assistantId}
+            assistantId={assistantId}
             indexName={algolia.indexName}
+            agentStudio={ASK_AI_AGENT_STUDIO}
             panel={{
               translations: {
                 newConversationScreen: {

src/theme/SearchBar/index.tsx

@@ -0,0 +1,27 @@
+import React, { useMemo, type ComponentType, type ReactNode } from 'react'
+import SearchBar from '@theme-original/SearchBar'
+import useDocusaurusContext from '@docusaurus/useDocusaurusContext'
+import { ASK_AI_AGENT_STUDIO, type AskAiConfig } from '@site/src/lib/algolia-ask-ai'
+
+type AskAiProps = { askAi?: AskAiConfig & { agentStudio?: boolean } }
+
+/**
+ * The theme types `SearchBar` as taking no props, but its implementation spreads whatever it
+ * receives over `themeConfig.algolia` so that props win, which is the only way to hand DocSearch a
+ * key that Docusaurus's Joi schema rejects in the config.
+ */
+const SearchBarWithAskAi = SearchBar as ComponentType<AskAiProps>
+
+export default function SearchBarWrapper(): ReactNode {
+  const { siteConfig } = useDocusaurusContext()
+  const { askAi } = siteConfig.themeConfig.algolia as { askAi?: AskAiConfig }
+
+  // `useAlgoliaAskAi` memoizes on the identity of this object, and DocSearch derives the Ask AI
+  // chat config from that memo. Rebuilding it on every render would defeat both.
+  const askAiWithAgentStudio = useMemo(
+    () => (askAi ? { ...askAi, agentStudio: ASK_AI_AGENT_STUDIO } : undefined),
+    [askAi]
+  )
+
+  return <SearchBarWithAskAi askAi={askAiWithAgentStudio} />
+}