diff --git a/packages/frontend/core/src/blocksuite/ai/components/ai-chat-chips/chat-panel-chips.ts b/packages/frontend/core/src/blocksuite/ai/components/ai-chat-chips/chat-panel-chips.ts
index 9361320769..9ce2e1ac47 100644
--- a/packages/frontend/core/src/blocksuite/ai/components/ai-chat-chips/chat-panel-chips.ts
+++ b/packages/frontend/core/src/blocksuite/ai/components/ai-chat-chips/chat-panel-chips.ts
@@ -448,11 +448,10 @@ export class ChatPanelChips extends SignalWatcher(
       if (!contextId || !AIProvider.context) {
         throw new Error('Context not found');
       }
-      const collection = this._collections.value.find(
-        collection => collection.id === chip.collectionId
-      );
       // TODO: server side docIds calculation
-      const docIds = collection?.allowList ?? [];
+      const docIds = this.docDisplayConfig.getCollectionPageIds(
+        chip.collectionId
+      );
       await AIProvider.context.addContextCollection({
         contextId,
         collectionId: chip.collectionId,
diff --git a/packages/frontend/core/src/blocksuite/ai/components/ai-chat-chips/type.ts b/packages/frontend/core/src/blocksuite/ai/components/ai-chat-chips/type.ts
index bf906f298a..d75bcd0cab 100644
--- a/packages/frontend/core/src/blocksuite/ai/components/ai-chat-chips/type.ts
+++ b/packages/frontend/core/src/blocksuite/ai/components/ai-chat-chips/type.ts
@@ -74,6 +74,7 @@ export interface DocDisplayConfig {
     signal: Signal<Collection[]>;
     cleanup: () => void;
   };
+  getCollectionPageIds: (collectionId: string) => string[];
 }
 
 export interface SearchMenuConfig {
diff --git a/packages/frontend/core/src/blocksuite/ai/components/ai-chat-input/ai-chat-input.ts b/packages/frontend/core/src/blocksuite/ai/components/ai-chat-input/ai-chat-input.ts
index 56a5e9ca89..a5fcc52332 100644
--- a/packages/frontend/core/src/blocksuite/ai/components/ai-chat-input/ai-chat-input.ts
+++ b/packages/frontend/core/src/blocksuite/ai/components/ai-chat-input/ai-chat-input.ts
@@ -23,7 +23,12 @@ import type {
   DocDisplayConfig,
   FileChip,
 } from '../ai-chat-chips/type';
-import { isDocChip, isFileChip } from '../ai-chat-chips/utils';
+import {
+  isCollectionChip,
+  isDocChip,
+  isFileChip,
+  isTagChip,
+} from '../ai-chat-chips/utils';
 import type { ChatMessage } from '../ai-chat-messages';
 import type { AIChatInputContext, AINetworkSearchConfig } from './type';
 
@@ -521,7 +526,8 @@ export class AIChatInput extends SignalWatcher(WithDisposable(LitElement)) {
       await this._preUpdateMessages(userInput, attachments);
 
       const sessionId = await this.createSessionId();
-      const contexts = await this._getMatchedContexts(userInput);
+      let contexts = await this._getMatchedContexts(userInput);
+      contexts = this._filterContexts(contexts);
       if (abortController.signal.aborted) {
         return;
       }
@@ -684,6 +690,42 @@ export class AIChatInput extends SignalWatcher(WithDisposable(LitElement)) {
       files: Array.from(fileContexts.values()),
     };
   }
+
+  // TODO: remove this function after workspace embedding is ready
+  private _filterContexts(contexts: {
+    docs: BlockSuitePresets.AIDocContextOption[];
+    files: BlockSuitePresets.AIFileContextOption[];
+  }) {
+    const docIds = this.chips.reduce((acc, chip) => {
+      if (isDocChip(chip)) {
+        acc.push(chip.docId);
+      }
+      if (isTagChip(chip)) {
+        const docIds = this.docDisplayConfig.getTagPageIds(chip.tagId);
+        acc.push(...docIds);
+      }
+      if (isCollectionChip(chip)) {
+        const docIds = this.docDisplayConfig.getCollectionPageIds(
+          chip.collectionId
+        );
+        acc.push(...docIds);
+      }
+      return acc;
+    }, [] as string[]);
+
+    const fileIds = this.chips.reduce((acc, chip) => {
+      if (isFileChip(chip) && chip.blobId) {
+        acc.push(chip.blobId);
+      }
+      return acc;
+    }, [] as string[]);
+
+    const { docs, files } = contexts;
+    return {
+      docs: docs.filter(doc => docIds.includes(doc.docId)),
+      files: files.filter(file => fileIds.includes(file.blobId)),
+    };
+  }
 }
 
 declare global {
diff --git a/packages/frontend/core/src/components/hooks/affine/use-ai-chat-config.ts b/packages/frontend/core/src/components/hooks/affine/use-ai-chat-config.ts
index fce345ed05..a91eda7463 100644
--- a/packages/frontend/core/src/components/hooks/affine/use-ai-chat-config.ts
+++ b/packages/frontend/core/src/components/hooks/affine/use-ai-chat-config.ts
@@ -77,6 +77,11 @@ export function useAIChatConfig() {
       const collections$ = collectionService.collections$;
       return createSignalFromObservable(collections$, []);
     },
+    getCollectionPageIds: (collectionId: string) => {
+      const collection$ = collectionService.collection$(collectionId);
+      // TODO: lack of documents that meet the collection rules
+      return collection$?.value?.allowList ?? [];
+    },
   };
 
   const searchMenuConfig = {