From d1bc3e429dd3a1c48f8ba06743525f7d7e98fbed Mon Sep 17 00:00:00 2001
From: Vincent Driessen <me@nvie.com>
Date: Tue, 5 Aug 2025 12:23:10 +0200
Subject: [PATCH] Fixes knowledge leaking bug in AiChat (#2570)

---
 CHANGELOG.md                                  |   5 +
 .../app/dual-chat/page.tsx                    | 111 +++++++
 .../app/knowledge/page.tsx                    |  17 +-
 e2e/next-ai-kitchen-sink/app/page.tsx         |   3 +
 .../test/knowledge-isolation.test.ts          | 287 ++++++++++++++++++
 .../test/knowledge.test.ts                    |  71 ++---
 .../test/simple-chat.test.ts                  |  48 ++-
 .../test/tool-calling.test.ts                 |  14 +-
 packages/liveblocks-core/src/ai.ts            |  15 +-
 .../src/components/AiChat.tsx                 |  26 +-
 .../components/internal/AiChatComposer.tsx    |  11 +
 11 files changed, 523 insertions(+), 85 deletions(-)
 create mode 100644 e2e/next-ai-kitchen-sink/app/dual-chat/page.tsx
 create mode 100644 e2e/next-ai-kitchen-sink/test/knowledge-isolation.test.ts

diff --git a/CHANGELOG.md b/CHANGELOG.md
index e7d1a737498..f355b664573 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -1,5 +1,10 @@
 ## vNEXT (not yet published)
 
+### `@liveblocks/react-ui`
+
+- Knowledge passed as a prop to `AiChat` no longer leaks that knowledge to other
+  instances of `AiChat` that are currently mounted on screen.
+
 ## v3.2.1
 
 ### `@liveblocks/react-ui`
diff --git a/e2e/next-ai-kitchen-sink/app/dual-chat/page.tsx b/e2e/next-ai-kitchen-sink/app/dual-chat/page.tsx
new file mode 100644
index 00000000000..23ceedbb18d
--- /dev/null
+++ b/e2e/next-ai-kitchen-sink/app/dual-chat/page.tsx
@@ -0,0 +1,111 @@
+"use client";
+
+import { LiveblocksProvider, RegisterAiKnowledge } from "@liveblocks/react";
+import { AiChat } from "@liveblocks/react-ui";
+import { useState } from "react";
+
+function ChatWithLocalKnowledge() {
+  const [localKnowledge, setLocalKnowledge] = useState("Spaghetti Carbonara");
+
+  return (
+    <div className="w-1/2 h-full border-r border-gray-300 p-4">
+      <div className="bg-blue-50 border border-blue-200 p-4 rounded-lg mb-4">
+        <label className="block text-sm font-medium mb-2 text-blue-700">
+          My favorite pasta dish:
+        </label>
+        <input
+          type="text"
+          value={localKnowledge}
+          onChange={(e) => setLocalKnowledge(e.target.value)}
+          className="w-full p-2 border border-blue-300 rounded bg-blue-50"
+          data-testid="chat-a-knowledge-input"
+        />
+        <p className="text-xs text-blue-600 mt-3">
+          This knowledge is passed via the <code>knowledge</code> prop to chat A
+          only.
+        </p>
+      </div>
+
+      <h2 className="text-xl font-bold mb-4">Chat A</h2>
+      <div className="h-80" data-testid="chat-a">
+        <AiChat
+          chatId="chat-a"
+          className="h-full border border-gray-200 rounded"
+          knowledge={[
+            { description: "My favorite pasta dish", value: localKnowledge },
+          ]}
+        />
+      </div>
+    </div>
+  );
+}
+
+function ChatWithoutLocalKnowledge() {
+  return (
+    <div className="w-1/2 h-full p-4">
+      <div className="bg-gray-50 border border-gray-200 p-4 rounded-lg mb-4">
+        <h3 className="text-sm font-semibold mb-2 text-gray-800 h-12">
+          No local knowledge
+        </h3>
+        <p className="text-sm text-gray-600">
+          This chat has no additional local knowledge passed via props. It can
+          only access the global knowledge.
+        </p>
+      </div>
+
+      <h2 className="text-xl font-bold mb-4">Chat B</h2>
+      <div className="h-80" data-testid="chat-b">
+        <AiChat
+          chatId="chat-b"
+          className="h-full border border-gray-200 rounded"
+        />
+      </div>
+    </div>
+  );
+}
+
+export default function DualChatPage() {
+  const [globalKnowledge, setGlobalKnowledge] = useState("Tiramisu");
+
+  return (
+    <LiveblocksProvider
+      authEndpoint="/api/auth/liveblocks"
+      // @ts-expect-error
+      baseUrl={process.env.NEXT_PUBLIC_LIVEBLOCKS_BASE_URL}
+    >
+      {/* Global knowledge that should be accessible to ALL chat instances */}
+      <RegisterAiKnowledge
+        description="My favorite dessert"
+        value={globalKnowledge}
+      />
+
+      <main className="h-screen w-full flex flex-col">
+        <div className="bg-yellow-50 border-b border-yellow-200 p-4">
+          <h1 className="text-2xl font-bold mb-2">Knowledge Isolation Test</h1>
+          <div className="mb-4">
+            <label className="block text-sm font-medium mb-2 text-yellow-800">
+              My favorite dessert:
+            </label>
+            <input
+              type="text"
+              value={globalKnowledge}
+              onChange={(e) => setGlobalKnowledge(e.target.value)}
+              className="w-full p-2 border border-yellow-300 rounded bg-yellow-50"
+              data-testid="global-knowledge-input"
+            />
+            <p className="text-sm text-yellow-700 mt-1">
+              This knowledge is registered via{" "}
+              <code>&lt;RegisterAiKnowledge /&gt;</code> and should be
+              accessible to both chat A and B.
+            </p>
+          </div>
+        </div>
+
+        <div className="flex-1 flex">
+          <ChatWithLocalKnowledge />
+          <ChatWithoutLocalKnowledge />
+        </div>
+      </main>
+    </LiveblocksProvider>
+  );
+}
diff --git a/e2e/next-ai-kitchen-sink/app/knowledge/page.tsx b/e2e/next-ai-kitchen-sink/app/knowledge/page.tsx
index fc95726597d..a22b80c7d78 100644
--- a/e2e/next-ai-kitchen-sink/app/knowledge/page.tsx
+++ b/e2e/next-ai-kitchen-sink/app/knowledge/page.tsx
@@ -28,6 +28,7 @@ function DarkModeToggle() {
     <div className="flex flex-col mx-auto px-4 max-w-4xl py-8 border-b-1">
       <label>
         <input
+          data-testid="expose-dark-mode-checkbox"
           type="checkbox"
           checked={exposed}
           onChange={() => toggleExposeTool()}
@@ -71,6 +72,7 @@ function DarkModeToggle() {
 
       <label>
         <input
+          data-testid="dark-mode-toggle"
           type="checkbox"
           checked={mode === "dark"}
           onChange={() => toggleDarkMode()}
@@ -94,7 +96,12 @@ function MyNickName() {
       ) : null}
       <div className="flex flex-col mx-auto px-4 max-w-4xl py-8 border-b-1">
         <label>
-          <input type="checkbox" checked={enabled} onChange={() => toggle()} />
+          <input
+            data-testid="share-nickname-checkbox"
+            type="checkbox"
+            checked={enabled}
+            onChange={() => toggle()}
+          />
           <span className="ml-2">Share my nickname</span>
         </label>
       </div>
@@ -230,6 +237,7 @@ export default function Page() {
         <div className="flex flex-col px-4 max-w-4xl py-8 border-b-1">
           <div className="flex gap-4 mx-auto">
             <button
+              data-testid="tab-todo-app"
               className={
                 selectedTab === 1
                   ? "cursor-pointer font-bold"
@@ -240,6 +248,7 @@ export default function Page() {
               Todo app
             </button>
             <button
+              data-testid="tab-another-app"
               className={
                 selectedTab === 2
                   ? "cursor-pointer font-bold"
@@ -250,6 +259,7 @@ export default function Page() {
               Another app
             </button>
             <button
+              data-testid="tab-both"
               className={
                 selectedTab === 3
                   ? "cursor-pointer font-bold"
@@ -272,7 +282,10 @@ export default function Page() {
 
         <div className="fixed bottom-8 right-8">
           <Popover.Root>
-            <Popover.Trigger className="inline-flex items-center justify-center rounded-full p-4 shadow-[0_0_0_1px_#0000000a,0_2px_6px_#0000000f,0_8px_26px_#00000014] hover:shadow-[0_0_0_1px_#00000014,0_2px_6px_#00000014,0_8px_26px_#00000014] dark:shadow-[0_0_0_1px_#ffffff0f] dark:hover:shadow-[0_0_0_1px_#ffffff14,0_2px_6px_#ffffff14,0_8px_26px_#ffffff14]">
+            <Popover.Trigger
+              data-testid="ai-chat-trigger"
+              className="inline-flex items-center justify-center rounded-full p-4 shadow-[0_0_0_1px_#0000000a,0_2px_6px_#0000000f,0_8px_26px_#00000014] hover:shadow-[0_0_0_1px_#00000014,0_2px_6px_#00000014,0_8px_26px_#00000014] dark:shadow-[0_0_0_1px_#ffffff0f] dark:hover:shadow-[0_0_0_1px_#ffffff14,0_2px_6px_#ffffff14,0_8px_26px_#ffffff14]"
+            >
               <svg
                 xmlns="http://www.w3.org/2000/svg"
                 width={32}
diff --git a/e2e/next-ai-kitchen-sink/app/page.tsx b/e2e/next-ai-kitchen-sink/app/page.tsx
index 3cfd400d91d..07afdac83e3 100644
--- a/e2e/next-ai-kitchen-sink/app/page.tsx
+++ b/e2e/next-ai-kitchen-sink/app/page.tsx
@@ -20,6 +20,9 @@ export default function Home() {
         <Link href="/knowledge" className="underline">
           Registering client-side knowledge
         </Link>
+        <Link href="/dual-chat" className="underline">
+          Dual chat - Knowledge isolation & sharing
+        </Link>
         <Link href="/styles" className="underline">
           Default styles playground
         </Link>
diff --git a/e2e/next-ai-kitchen-sink/test/knowledge-isolation.test.ts b/e2e/next-ai-kitchen-sink/test/knowledge-isolation.test.ts
new file mode 100644
index 00000000000..786319aadc0
--- /dev/null
+++ b/e2e/next-ai-kitchen-sink/test/knowledge-isolation.test.ts
@@ -0,0 +1,287 @@
+import { test, expect } from "@playwright/test";
+import type { Page } from "@playwright/test";
+
+async function setupChatA(page: Page) {
+  const textInput = page.locator(
+    '[data-testid="chat-a"] .lb-ai-chat-composer-editor'
+  );
+  const sendButton = page.locator(
+    '[data-testid="chat-a"] .lb-ai-chat-composer-action'
+  );
+
+  // If there's an ongoing operation (abort button is enabled), click it first to clear state
+  if ((await sendButton.getAttribute("data-variant")) === "secondary") {
+    await sendButton.click();
+    // Wait for it to return to send state
+    await expect(sendButton).toHaveAttribute("data-variant", "primary");
+  }
+
+  // Clear any existing text
+  await textInput.clear();
+
+  return { textInput, sendButton };
+}
+
+async function setupChatB(page: Page) {
+  const textInput = page.locator(
+    '[data-testid="chat-b"] .lb-ai-chat-composer-editor'
+  );
+  const sendButton = page.locator(
+    '[data-testid="chat-b"] .lb-ai-chat-composer-action'
+  );
+
+  // If there's an ongoing operation (abort button is enabled), click it first to clear state
+  if ((await sendButton.getAttribute("data-variant")) === "secondary") {
+    await sendButton.click();
+    // Wait for it to return to send state
+    await expect(sendButton).toHaveAttribute("data-variant", "primary");
+  }
+
+  // Clear any existing text
+  await textInput.clear();
+
+  return { textInput, sendButton };
+}
+
+async function sendMessageToChatA(page: Page, message: string) {
+  const { textInput, sendButton } = await setupChatA(page);
+
+  await textInput.fill(message);
+  await sendButton.click({ timeout: 5000 });
+
+  // Verify the message was sent (appears in chat A)
+  await expect(
+    page.locator('[data-testid="chat-a"]').locator(`text=${message}`)
+  ).toBeVisible();
+
+  // Wait for generation to finish by monitoring button state change
+  // Button should become "secondary" (abort) during generation, then back to "primary" (send)
+  await expect(sendButton).toHaveAttribute("data-variant", "secondary", {
+    timeout: 10000,
+  });
+  await expect(sendButton).toHaveAttribute("data-variant", "primary", {
+    timeout: 60000,
+  });
+}
+
+async function sendMessageToChatB(page: Page, message: string) {
+  const { textInput, sendButton } = await setupChatB(page);
+
+  await textInput.fill(message);
+  await sendButton.click({ timeout: 5000 });
+
+  // Verify the message was sent (appears in chat B)
+  await expect(
+    page.locator('[data-testid="chat-b"]').locator(`text=${message}`)
+  ).toBeVisible();
+
+  // Wait for generation to finish by monitoring button state change
+  // Button should become "secondary" (abort) during generation, then back to "primary" (send)
+  await expect(sendButton).toHaveAttribute("data-variant", "secondary", {
+    timeout: 10000,
+  });
+  await expect(sendButton).toHaveAttribute("data-variant", "primary", {
+    timeout: 60000,
+  });
+}
+
+test.describe("AiChat Knowledge Isolation", () => {
+  test.beforeEach(async ({ page }) => {
+    // Clean up existing chats
+    await page.goto("/chats");
+    await expect(page.locator("h1")).toHaveText("List of all chats");
+
+    // Delete chat-a if it exists
+    const chatALink = page.locator('a[href="/chats/chat-a"]');
+    if (await chatALink.isVisible()) {
+      const deleteButton = chatALink
+        .locator("..")
+        .locator('button:has-text("Delete")');
+      await deleteButton.click();
+      await expect(chatALink).not.toBeVisible();
+    }
+
+    // Delete chat-b if it exists
+    const chatBLink = page.locator('a[href="/chats/chat-b"]');
+    if (await chatBLink.isVisible()) {
+      const deleteButton = chatBLink
+        .locator("..")
+        .locator('button:has-text("Delete")');
+      await deleteButton.click();
+      await expect(chatBLink).not.toBeVisible();
+    }
+  });
+
+  test("should NOT share knowledge between separate AiChat instances", async ({
+    page,
+  }) => {
+    // Go to the dual chat page
+    await page.goto("/dual-chat");
+
+    // Wait for the page title and both chats to be visible
+    await expect(
+      page.locator('h1:has-text("Knowledge Isolation Test")')
+    ).toBeVisible();
+    await expect(page.locator('h2:has-text("Chat A")')).toBeVisible();
+    await expect(page.locator('h2:has-text("Chat B")')).toBeVisible();
+
+    // Verify the default pasta knowledge is set
+    const knowledgeInput = page.locator(
+      '[data-testid="chat-a-knowledge-input"]'
+    );
+    await expect(knowledgeInput).toHaveValue("Spaghetti Carbonara");
+
+    // Wait for the AiChat components to be fully loaded
+    await expect(
+      page.locator('[data-testid="chat-a"] .lb-ai-chat-composer')
+    ).toBeVisible({
+      timeout: 10000,
+    });
+    await expect(
+      page.locator('[data-testid="chat-b"] .lb-ai-chat-composer')
+    ).toBeVisible({
+      timeout: 10000,
+    });
+
+    // Ask both chats about favorite pasta - only one should know it
+    await sendMessageToChatA(page, "What's my favorite pasta?");
+    await sendMessageToChatB(page, "What's my favorite pasta?");
+
+    // Wait for Chat A's response containing the pasta name
+    await expect(
+      page.locator('[data-testid="chat-a"]').locator("text=Spaghetti Carbonara")
+    ).toBeVisible({ timeout: 10000 });
+
+    // Get Chat B's assistant response text and verify it doesn't include "Spaghetti Carbonara"
+    const chatBAssistantMessage = page
+      .locator('[data-testid="chat-b"] .lb-ai-chat-assistant-message')
+      .last();
+    const chatBText = await chatBAssistantMessage.textContent();
+    expect(chatBText).not.toContain("Spaghetti Carbonara");
+  });
+
+  test("should isolate different knowledge values between chats", async ({
+    page,
+  }) => {
+    await page.goto("/dual-chat");
+
+    // Wait for the page title to be visible
+    await expect(
+      page.locator('h1:has-text("Knowledge Isolation Test")')
+    ).toBeVisible();
+
+    // Wait for components to load
+    await expect(
+      page.locator('[data-testid="chat-a"] .lb-ai-chat-composer')
+    ).toBeVisible({
+      timeout: 10000,
+    });
+    await expect(
+      page.locator('[data-testid="chat-b"] .lb-ai-chat-composer')
+    ).toBeVisible({
+      timeout: 10000,
+    });
+
+    // Change the pasta knowledge for Chat A
+    const knowledgeInput = page.locator(
+      '[data-testid="chat-a-knowledge-input"]'
+    );
+    await knowledgeInput.clear();
+    await knowledgeInput.fill("Penne Arrabiata");
+
+    // Ask Chat A about favorite pasta
+    await sendMessageToChatA(page, "What's my favorite pasta?");
+
+    // Chat A should know the pasta is Penne Arrabiata
+    await expect(
+      page.locator('[data-testid="chat-a"]').locator("text=Penne Arrabiata")
+    ).toBeVisible({ timeout: 30000 });
+
+    // Ask Chat B about favorite pasta
+    await sendMessageToChatB(page, "What's my favorite pasta?");
+
+    // Get Chat B's assistant response text and verify it doesn't include "Penne Arrabiata"
+    // This test should FAIL with the current bug
+    const chatBAssistantMessage = page
+      .locator('[data-testid="chat-b"] .lb-ai-chat-assistant-message')
+      .last();
+    const chatBText = await chatBAssistantMessage.textContent();
+    expect(chatBText).not.toContain("Penne Arrabiata");
+  });
+
+  test("should SHARE knowledge when using RegisterAiKnowledge globally", async ({
+    page,
+  }) => {
+    // Go to the dual chat page (now contains both local and global knowledge testing)
+    await page.goto("/dual-chat");
+
+    // Wait for the page title and both chats to be visible
+    await expect(
+      page.locator('h1:has-text("Knowledge Isolation Test")')
+    ).toBeVisible();
+    await expect(page.locator('h2:has-text("Chat A")')).toBeVisible();
+    await expect(page.locator('h2:has-text("Chat B")')).toBeVisible();
+
+    // Verify the default global knowledge is set
+    const globalKnowledgeInput = page.locator(
+      '[data-testid="global-knowledge-input"]'
+    );
+    await expect(globalKnowledgeInput).toHaveValue("Tiramisu");
+
+    // Wait for the AiChat components to be fully loaded
+    await expect(
+      page.locator('[data-testid="chat-a"] .lb-ai-chat-composer')
+    ).toBeVisible({
+      timeout: 10000,
+    });
+    await expect(
+      page.locator('[data-testid="chat-b"] .lb-ai-chat-composer')
+    ).toBeVisible({
+      timeout: 10000,
+    });
+
+    // Ask Chat A about the global knowledge - it should have access to both local and global
+    await sendMessageToChatA(page, "What dessert do I like?");
+
+    // Wait for Chat A's response containing the global knowledge
+    await expect(
+      page
+        .locator('[data-testid="chat-a"] .lb-ai-chat-assistant-message')
+        .last()
+        .locator("text=Tiramisu")
+    ).toBeVisible({ timeout: 30000 });
+
+    // Ask Chat B about the same global knowledge - it should ALSO have access
+    await sendMessageToChatB(page, "What dessert do I like?");
+
+    // Chat B should ALSO have access to the global knowledge
+    await expect(
+      page
+        .locator('[data-testid="chat-b"] .lb-ai-chat-assistant-message')
+        .last()
+        .locator("text=Tiramisu")
+    ).toBeVisible({ timeout: 30000 });
+
+    // Change the global knowledge
+    await globalKnowledgeInput.clear();
+    await globalKnowledgeInput.fill("Gelato");
+
+    // Ask both chats again - they should both know the new global knowledge
+    await sendMessageToChatA(page, "What dessert do I like now?");
+    await sendMessageToChatB(page, "What dessert do I like now?");
+
+    // Both chats should have access to the updated global knowledge
+    await expect(
+      page
+        .locator('[data-testid="chat-a"] .lb-ai-chat-assistant-message')
+        .last()
+        .locator("text=Gelato")
+    ).toBeVisible({ timeout: 30000 });
+    await expect(
+      page
+        .locator('[data-testid="chat-b"] .lb-ai-chat-assistant-message')
+        .last()
+        .locator("text=Gelato")
+    ).toBeVisible({ timeout: 30000 });
+  });
+});
diff --git a/e2e/next-ai-kitchen-sink/test/knowledge.test.ts b/e2e/next-ai-kitchen-sink/test/knowledge.test.ts
index 963e8df7f1c..002d741aa90 100644
--- a/e2e/next-ai-kitchen-sink/test/knowledge.test.ts
+++ b/e2e/next-ai-kitchen-sink/test/knowledge.test.ts
@@ -2,8 +2,8 @@ import { test, expect } from "@playwright/test";
 import type { Page } from "@playwright/test";
 
 async function setupAiChat(page: Page) {
-  // The knowledge page has the chat trigger in the bottom right corner
-  const triggerButton = page.locator(".fixed.bottom-8.right-8 button");
+  // The knowledge page has the chat trigger button with test ID
+  const triggerButton = page.getByTestId("ai-chat-trigger");
 
   // Check if popover is already open
   const isOpen = (await triggerButton.getAttribute("data-state")) === "open";
@@ -20,11 +20,22 @@ async function setupAiChat(page: Page) {
   const textInput = page.locator(".lb-ai-chat-composer-editor");
   const sendButton = page.locator(".lb-ai-chat-composer-action");
 
+  // Wait for the send button to be present and ready
+  await expect(sendButton).toBeVisible({ timeout: 10000 });
+
+  // Wait for the composer form to not be disabled
+  const composerForm = page.locator(".lb-ai-chat-composer-form");
+  await expect(composerForm).not.toHaveAttribute("disabled", {
+    timeout: 10000,
+  });
+
   // If there's an ongoing operation (abort button is enabled), click it first to clear state
-  if ((await sendButton.getAttribute("data-variant")) === "secondary") {
+  if ((await sendButton.getAttribute("aria-label")) === "Abort response") {
     await sendButton.click();
     // Wait for it to return to send state
-    await expect(sendButton).toHaveAttribute("data-variant", "primary");
+    await expect(sendButton).toHaveAttribute("aria-label", "Send", {
+      timeout: 15000,
+    });
   }
 
   // Clear any existing text
@@ -37,10 +48,15 @@ async function sendAiMessage(page: Page, message: string) {
   const { textInput, sendButton } = await setupAiChat(page);
 
   await textInput.fill(message);
+
+  // Wait for the button to be enabled after filling the text
+  await expect(sendButton).toBeEnabled({ timeout: 10000 });
+
   await sendButton.click({ timeout: 5000 });
 
   // Verify the message was sent (appears in the chat)
-  await expect(page.locator(`text=${message}`)).toBeVisible();
+  // Use a user message selector to be more specific and avoid duplicates
+  await expect(page.locator(".lb-ai-chat-user-message").last()).toContainText(message);
 }
 
 test.describe("Knowledge Registration", () => {
@@ -207,48 +223,37 @@ test.describe("Knowledge Registration", () => {
     await page.goto("/knowledge");
 
     // Start on Todo app tab
-    await expect(
-      page.locator('button.font-bold:has-text("Todo app")')
-    ).toBeVisible();
+    await expect(page.getByTestId("tab-todo-app")).toHaveClass(/font-bold/);
 
     // Ask about current view
     await sendAiMessage(page, "What view am I currently in?");
-    await expect(
-      page.locator("text=Todo").or(page.locator("text=todo"))
-    ).toBeVisible({ timeout: 30000 });
+    await expect(page.locator(".lb-ai-chat-messages")).toContainText(/todo/i, {
+      timeout: 30000,
+    });
 
     // Switch to "Another app" tab
-    await page.click('button:has-text("Another app")');
-    await expect(
-      page.locator('button.font-bold:has-text("Another app")')
-    ).toBeVisible();
+    await page.getByTestId("tab-another-app").click();
+    await expect(page.getByTestId("tab-another-app")).toHaveClass(/font-bold/);
     await expect(page.locator("text=Another part of the app")).toBeVisible();
 
     // Ask about current view again - should now know it's "Another app"
     await sendAiMessage(page, "What is my current view now?");
     // Look for "Another" in the AI chat response, not in the UI
-    await expect(
-      page
-        .locator(".lb-ai-chat-messages")
-        .locator("text=Another")
-        .or(page.locator(".lb-ai-chat-messages").locator("text=another"))
-    ).toBeVisible({ timeout: 30000 });
+    await expect(page.locator(".lb-ai-chat-messages")).toContainText(
+      /another/i,
+      { timeout: 30000 }
+    );
 
     // Switch to "Both" tab
-    await page.click('button:has-text("Both")');
-    await expect(
-      page.locator('button.font-bold:has-text("Both")')
-    ).toBeVisible();
+    await page.getByTestId("tab-both").click();
+    await expect(page.getByTestId("tab-both")).toHaveClass(/font-bold/);
 
     // Ask about current view - should know it's "Both apps"
     await sendAiMessage(page, "What view am I in now?");
     // Look for "Both" in the AI chat response, not in the UI
-    await expect(
-      page
-        .locator(".lb-ai-chat-messages")
-        .locator("text=Both")
-        .or(page.locator(".lb-ai-chat-messages").locator("text=both"))
-    ).toBeVisible({ timeout: 30000 });
+    await expect(page.locator(".lb-ai-chat-messages")).toContainText(/both/i, {
+      timeout: 30000,
+    });
   });
 
   test("should handle knowledge being disabled", async ({ page }) => {
@@ -266,9 +271,7 @@ test.describe("Knowledge Registration", () => {
     await page.goto("/knowledge");
 
     // Find and disable dark mode knowledge exposure
-    const exposeCheckbox = page.locator(
-      'label:has-text("Expose dark mode as knowledge & tool") input[type="checkbox"]'
-    );
+    const exposeCheckbox = page.getByTestId("expose-dark-mode-checkbox");
     await exposeCheckbox.uncheck();
     await expect(exposeCheckbox).not.toBeChecked();
 
diff --git a/e2e/next-ai-kitchen-sink/test/simple-chat.test.ts b/e2e/next-ai-kitchen-sink/test/simple-chat.test.ts
index 501e8f8bca8..c54d739bfda 100644
--- a/e2e/next-ai-kitchen-sink/test/simple-chat.test.ts
+++ b/e2e/next-ai-kitchen-sink/test/simple-chat.test.ts
@@ -14,11 +14,13 @@ async function setupSimpleChat(page: Page) {
   // If there's an ongoing operation (abort button is enabled), click it first to clear state
   if (
     (await sendButton.isEnabled()) &&
-    (await sendButton.getAttribute("data-variant")) === "secondary"
+    (await sendButton.getAttribute("aria-label")) === "Abort response"
   ) {
     await sendButton.click();
     // Wait for it to return to send state
-    await expect(sendButton).toHaveAttribute("data-variant", "primary");
+    await expect(sendButton).toHaveAttribute("aria-label", "Send", {
+      timeout: 15000,
+    });
   }
 
   // Clear any existing text
@@ -28,6 +30,22 @@ async function setupSimpleChat(page: Page) {
 }
 
 test.describe("Simple Chat", () => {
+  test.beforeEach(async ({ page }) => {
+    // Clean up - delete the "ai-chat" chat to start fresh
+    await page.goto("/chats");
+    await expect(page.locator("h1")).toHaveText("List of all chats");
+
+    // Look for the "ai-chat" chat and delete it if it exists
+    const aiChatLink = page.locator('a[href="/chats/ai-chat"]');
+    if (await aiChatLink.isVisible()) {
+      const deleteButton = aiChatLink
+        .locator("..")
+        .locator('button:has-text("Delete")');
+      await deleteButton.click();
+      await expect(aiChatLink).not.toBeVisible();
+    }
+  });
+
   test("should handle ping-pong interaction", async ({ page }) => {
     const { textInput, sendButton } = await setupSimpleChat(page);
 
@@ -36,12 +54,12 @@ test.describe("Simple Chat", () => {
     await sendButton.click();
 
     // Ensure the send button turns into an abort button (should now show StopIcon)
-    // The button should change to have StopIcon and different data-variant
-    await expect(sendButton).toHaveAttribute("data-variant", "secondary");
+    // The button should change to show "Abort response" aria-label
+    await expect(sendButton).toHaveAttribute("aria-label", "Abort response");
 
-    // Wait for the send button to become enabled again (back to primary variant)
-    await expect(sendButton).toHaveAttribute("data-variant", "primary", {
-      timeout: 5000, // Give it up to 5 seconds for the AI response
+    // Wait for the send button to become enabled again (back to send state)
+    await expect(sendButton).toHaveAttribute("aria-label", "Send", {
+      timeout: 15000, // Give it up to 15 seconds for the AI response
     });
 
     // Check that a response message containing "pong" is received
@@ -49,11 +67,11 @@ test.describe("Simple Chat", () => {
     const assistantMessage = page
       .locator(".lb-ai-chat-assistant-message")
       .last();
-    await expect(assistantMessage).toBeVisible({ timeout: 5000 });
+    await expect(assistantMessage).toBeVisible({ timeout: 15000 });
 
     // Check if it contains "pong"
     await expect(assistantMessage).toContainText("pong", {
-      timeout: 5000,
+      timeout: 15000,
     });
   });
 
@@ -68,21 +86,19 @@ test.describe("Simple Chat", () => {
     );
 
     // The button should be enabled once we have text
-    await expect(sendButton).toBeEnabled({ timeout: 5000 });
+    await expect(sendButton).toBeEnabled({ timeout: 15000 });
     await sendButton.click();
 
-    // Verify the button changes to abort state (secondary variant)
-    await expect(sendButton).toHaveAttribute("data-variant", "secondary");
+    // Verify the button changes to abort state
     await expect(sendButton).toHaveAttribute("aria-label", "Abort response");
 
     // Click the abort button while the AI is generating
     await sendButton.click();
 
-    // Verify the button goes back to send state (primary variant)
-    await expect(sendButton).toHaveAttribute("data-variant", "primary", {
-      timeout: 5000,
+    // Verify the button goes back to send state
+    await expect(sendButton).toHaveAttribute("aria-label", "Send", {
+      timeout: 15000,
     });
-    await expect(sendButton).toHaveAttribute("aria-label", "Send");
 
     // Verify the user message exists
     const userMessage = page.locator(".lb-ai-chat-user-message").last();
diff --git a/e2e/next-ai-kitchen-sink/test/tool-calling.test.ts b/e2e/next-ai-kitchen-sink/test/tool-calling.test.ts
index 8aba9ff119a..0f0a48bba97 100644
--- a/e2e/next-ai-kitchen-sink/test/tool-calling.test.ts
+++ b/e2e/next-ai-kitchen-sink/test/tool-calling.test.ts
@@ -12,10 +12,12 @@ async function setupAiChat(page: Page) {
   const sendButton = page.locator(".lb-ai-chat-composer-action");
 
   // If there's an ongoing operation (abort button is enabled), click it first to clear state
-  if ((await sendButton.getAttribute("data-variant")) === "secondary") {
+  if ((await sendButton.getAttribute("aria-label")) === "Abort response") {
     await sendButton.click();
     // Wait for it to return to send state
-    await expect(sendButton).toHaveAttribute("data-variant", "primary");
+    await expect(sendButton).toHaveAttribute("aria-label", "Send", {
+      timeout: 15000,
+    });
   }
 
   // Clear any existing text
@@ -32,7 +34,7 @@ async function sendAiMessage(page: Page, message: string) {
   await sendButton.click();
 
   // Verify the message was sent (appears in the chat)
-  await expect(page.locator(`text=${message}`)).toBeVisible();
+  await expect(page.locator(".lb-ai-chat-user-message").last()).toContainText(message);
 }
 
 test.describe("Tool Calling", () => {
@@ -220,8 +222,8 @@ test.describe("Tool Calling", () => {
     ).toBeVisible();
 
     // Verify the AI shows the request was denied
-    await expect(
-      page.locator("text=denied").or(page.locator("text=cancelled"))
-    ).toBeVisible({ timeout: 10000 });
+    await expect(page.locator(".lb-ai-chat-messages")).toContainText(/denied|cancelled/i, {
+      timeout: 10000
+    });
   });
 });
diff --git a/packages/liveblocks-core/src/ai.ts b/packages/liveblocks-core/src/ai.ts
index 685483010b5..fec3a42525d 100644
--- a/packages/liveblocks-core/src/ai.ts
+++ b/packages/liveblocks-core/src/ai.ts
@@ -84,10 +84,7 @@ export type AiToolTypePack<
   R: R;
 };
 
-export type AskUserMessageInChatOptions = Omit<
-  AiGenerationOptions,
-  "tools" | "knowledge"
->;
+export type AskUserMessageInChatOptions = Omit<AiGenerationOptions, "tools">;
 
 export type SetToolResultOptions = Omit<
   AiGenerationOptions,
@@ -1301,7 +1298,9 @@ export function createAi(config: AiConfig): Ai {
         targetMessageId: MessageId,
         options?: AskUserMessageInChatOptions
       ): Promise<AskInChatResponse> => {
-        const knowledge = context.knowledge.get();
+        const globalKnowledge = context.knowledge.get();
+        const requestKnowledge = options?.knowledge || [];
+        const combinedKnowledge = [...globalKnowledge, ...requestKnowledge];
         const tools = context.toolsStore.getToolDescriptions(chatId);
 
         const resp: AskInChatResponse = await sendClientMsgWithResponse({
@@ -1314,9 +1313,9 @@ export function createAi(config: AiConfig): Ai {
             stream: options?.stream,
             timeout: options?.timeout,
 
-            // Knowledge and tools aren't coming from the options, but retrieved
-            // from the global context
-            knowledge: knowledge.length > 0 ? knowledge : undefined,
+            // Combine global knowledge with request-specific knowledge
+            knowledge:
+              combinedKnowledge.length > 0 ? combinedKnowledge : undefined,
             tools: tools.length > 0 ? tools : undefined,
           },
         });
diff --git a/packages/liveblocks-react-ui/src/components/AiChat.tsx b/packages/liveblocks-react-ui/src/components/AiChat.tsx
index ab60b8f8c7a..17a7349bda4 100644
--- a/packages/liveblocks-react-ui/src/components/AiChat.tsx
+++ b/packages/liveblocks-react-ui/src/components/AiChat.tsx
@@ -4,11 +4,7 @@ import type {
   CopilotId,
   MessageId,
 } from "@liveblocks/core";
-import {
-  RegisterAiKnowledge,
-  RegisterAiTool,
-  useAiChatMessages,
-} from "@liveblocks/react";
+import { RegisterAiTool, useAiChatMessages } from "@liveblocks/react";
 import { useLatest } from "@liveblocks/react/_private";
 import {
   type ComponentProps,
@@ -92,8 +88,10 @@ export interface AiChatProps extends ComponentProps<"div"> {
   copilotId?: string;
 
   /**
-   * The contextual knowledge to include in the chat. May be used by the assistant when generating responses.
-   * Any knowledge you provide via this prop will be added to any already globally registered knowledge via <RegisterAiKnowledge />.
+   * The contextual knowledge to include in the chat. May be used by the
+   * assistant when generating responses. In addition to the knowledge passed
+   * in via this prop, the AiChat instance will also have access to any
+   * globally registered knowledge via <RegisterAiKnowledge />.
    */
   knowledge?: AiKnowledgeSource[];
 
@@ -418,7 +416,7 @@ export const AiChat = forwardRef<HTMLDivElement, AiChatProps>(
       copilotId,
       autoFocus,
       overrides,
-      knowledge,
+      knowledge: localKnowledge,
       tools = {},
       layout = "inset",
       components,
@@ -491,17 +489,6 @@ export const AiChat = forwardRef<HTMLDivElement, AiChatProps>(
           className
         )}
       >
-        {knowledge
-          ? knowledge.map((source, index) => (
-              <RegisterAiKnowledge
-                key={index}
-                description={source.description}
-                value={source.value}
-                // knowledgeKey={source.knowledgeKey}
-              />
-            ))
-          : null}
-
         {Object.entries(tools).map(([name, tool]) => (
           <RegisterAiTool key={name} chatId={chatId} name={name} tool={tool} />
         ))}
@@ -573,6 +560,7 @@ export const AiChat = forwardRef<HTMLDivElement, AiChatProps>(
             copilotId={copilotId as CopilotId}
             overrides={overrides}
             autoFocus={autoFocus}
+            knowledge={localKnowledge}
             onUserMessageCreate={({ id }) => setLastSentMessageId(id)}
             className={
               layout === "inset"
diff --git a/packages/liveblocks-react-ui/src/components/internal/AiChatComposer.tsx b/packages/liveblocks-react-ui/src/components/internal/AiChatComposer.tsx
index 73f024155e7..210d535dc58 100644
--- a/packages/liveblocks-react-ui/src/components/internal/AiChatComposer.tsx
+++ b/packages/liveblocks-react-ui/src/components/internal/AiChatComposer.tsx
@@ -1,5 +1,6 @@
 import type {
   AiChatMessage,
+  AiKnowledgeSource,
   CopilotId,
   MessageId,
   WithNavigation,
@@ -79,6 +80,13 @@ export interface AiChatComposerProps extends ComponentProps<"form"> {
    * The id of the copilot to use to send the message.
    */
   copilotId?: CopilotId;
+  /**
+   * The contextual knowledge to include in the chat. May be used by the
+   * assistant when generating responses. In addition to the knowledge passed
+   * in via this prop, the AiChat instance will also have access to any
+   * globally registered knowledge via <RegisterAiKnowledge />.
+   */
+  knowledge?: AiKnowledgeSource[];
   /**
    * @internal
    */
@@ -101,6 +109,7 @@ export const AiChatComposer = forwardRef<HTMLFormElement, AiChatComposerProps>(
       chatId,
       branchId,
       copilotId,
+      knowledge: localKnowledge,
       stream = true,
       onUserMessageCreate,
       ...props
@@ -168,6 +177,7 @@ export const AiChatComposer = forwardRef<HTMLFormElement, AiChatComposerProps>(
           {
             stream,
             copilotId,
+            knowledge: localKnowledge,
           }
         );
       },
@@ -180,6 +190,7 @@ export const AiChatComposer = forwardRef<HTMLFormElement, AiChatComposerProps>(
         abortableMessageId,
         stream,
         copilotId,
+        localKnowledge,
       ]
     );