diff --git a/tests/ui-smoke.spec.ts b/tests/ui-smoke.spec.ts
index f0b94b865..a44ba8ef3 100644
--- a/tests/ui-smoke.spec.ts
+++ b/tests/ui-smoke.spec.ts
@@ -29,6 +29,10 @@ function visibleQuestionInput(page: Page) {
   return page.locator('[aria-label="Search indexed guidelines by question or keyword"]:visible').first();
 }
 
+async function isVisibleWithoutThrow(locator: Locator) {
+  return locator.isVisible().catch(() => false);
+}
+
 async function fillVisibleQuestionInput(page: Page, value: string) {
   const questionInput = visibleQuestionInput(page);
   const submitAnswer = page.getByRole("button", { name: "Generate source-backed answer" });
@@ -47,13 +51,31 @@ async function fillVisibleQuestionInput(page: Page, value: string) {
 }
 
 async function switchToDocumentSearchMode(page: Page) {
-  const documentsMode = page.getByRole("button", { name: "Switch to document search mode" });
-  await expect(documentsMode).toBeVisible();
-  await expect(documentsMode).toBeEnabled();
+  const legacyDocumentsMode = page.getByRole("button", { name: "Switch to document search mode" });
+  if (await isVisibleWithoutThrow(legacyDocumentsMode)) {
+    await expect(legacyDocumentsMode).toBeEnabled();
+    await expect(async () => {
+      await legacyDocumentsMode.click();
+      await expect(legacyDocumentsMode).toHaveAttribute("aria-pressed", "true", { timeout: 1_000 });
+    }).toPass({ timeout: 8_000 });
+    return;
+  }
+
+  const appModeMenu = page.getByRole("button", { name: /^Current app mode:/ });
+  if (!(await isVisibleWithoutThrow(appModeMenu))) {
+    throw new Error(
+      "Could not switch to document search mode: neither the legacy mode toggle nor the app mode menu is visible.",
+    );
+  }
+  await expect(appModeMenu).toBeEnabled();
 
   await expect(async () => {
+    await appModeMenu.click();
+    await expect(appModeMenu).toHaveAttribute("aria-expanded", "true", { timeout: 2_000 });
+    const documentsMode = page.getByRole("menuitemradio", { name: /^Documents\b/ });
+    await expect(documentsMode).toBeVisible({ timeout: 3_000 });
     await documentsMode.click();
-    await expect(documentsMode).toHaveAttribute("aria-pressed", "true", { timeout: 1_000 });
+    await expect(appModeMenu).toHaveAccessibleName("Current app mode: Documents", { timeout: 2_000 });
   }).toPass({ timeout: 8_000 });
 }