From 67a52ddb1303ceb14a63b1324cb5546ddb3c9054 Mon Sep 17 00:00:00 2001
From: yoyoyohamapi <472285740@qq.com>
Date: Tue, 22 Jul 2025 11:00:35 +0800
Subject: [PATCH] feat(core): extract md & snapshot & attachments from edgeless
 selected

---
 .../ai/components/ai-chat-content/type.ts     |  5 ++
 .../core/src/blocksuite/ai/utils/extract.ts   | 74 ++++++++++++++++++-
 2 files changed, 77 insertions(+), 2 deletions(-)

diff --git a/packages/frontend/core/src/blocksuite/ai/components/ai-chat-content/type.ts b/packages/frontend/core/src/blocksuite/ai/components/ai-chat-content/type.ts
index cd7a456ee7..628c194339 100644
--- a/packages/frontend/core/src/blocksuite/ai/components/ai-chat-content/type.ts
+++ b/packages/frontend/core/src/blocksuite/ai/components/ai-chat-content/type.ts
@@ -1,3 +1,5 @@
+import type { DocSnapshot } from '@blocksuite/affine/store';
+
 import type { AIError } from '../../provider';
 import type { ChatStatus, HistoryMessage } from '../ai-chat-messages';
 
@@ -12,5 +14,8 @@ export type ChatContextValue = {
   markdown: string;
   // images of the selected content or user uploaded
   images: File[];
+  // snapshot of the selected content
+  snapshot: DocSnapshot;
+  attachments: File[];
   abortController: AbortController | null;
 };
diff --git a/packages/frontend/core/src/blocksuite/ai/utils/extract.ts b/packages/frontend/core/src/blocksuite/ai/utils/extract.ts
index bb9e198972..c338c8cf6f 100644
--- a/packages/frontend/core/src/blocksuite/ai/utils/extract.ts
+++ b/packages/frontend/core/src/blocksuite/ai/utils/extract.ts
@@ -1,7 +1,9 @@
+import { WorkspaceImpl } from '@affine/core/modules/workspace/impls/workspace';
 import {
+  AttachmentBlockModel,
   DatabaseBlockModel,
   ImageBlockModel,
-  type NoteBlockModel,
+  NoteBlockModel,
   NoteDisplayMode,
 } from '@blocksuite/affine/model';
 import {
@@ -15,12 +17,18 @@ import {
 } from '@blocksuite/affine/shared/commands';
 import { DocModeProvider } from '@blocksuite/affine/shared/services';
 import {
+  getBlockProps,
   isInsideEdgelessEditor,
   matchModels,
 } from '@blocksuite/affine/shared/utils';
 import { BlockStdScope, type EditorHost } from '@blocksuite/affine/std';
-import type { BlockModel, Store } from '@blocksuite/affine/store';
+import {
+  GfxBlockElementModel,
+  GfxControllerIdentifier,
+} from '@blocksuite/affine/std/gfx';
+import type { BlockModel, DocSnapshot, Store } from '@blocksuite/affine/store';
 import { Slice, toDraftModel } from '@blocksuite/affine/store';
+import { Doc as YDoc } from 'yjs';
 
 import { getStoreManager } from '../../manager/store';
 import type { ChatContextValue } from '../components/ai-chat-content';
@@ -46,6 +54,65 @@ async function extractEdgelessSelected(
   host: EditorHost
 ): Promise<Partial<ChatContextValue> | null> {
   if (!isInsideEdgelessEditor(host)) return null;
+  const gfx = host.std.get(GfxControllerIdentifier);
+  const selectedElements = gfx.selection.selectedElements;
+
+  let snapshot: DocSnapshot | null = null;
+  let markdown = '';
+  const attachments: ChatContextValue['attachments'] = [];
+
+  if (selectedElements.length) {
+    console.log('selectedElements', selectedElements);
+
+    const transformer = host.store.getTransformer();
+    const markdownAdapter = new MarkdownAdapter(
+      transformer,
+      host.store.provider
+    );
+    const collection = new WorkspaceImpl({
+      id: 'AI_EXTRACT',
+      rootDoc: new YDoc({ guid: 'AI_EXTRACT' }),
+    });
+    collection.meta.initialize();
+
+    try {
+      const fragmentDoc = collection.createDoc();
+      const fragment = fragmentDoc.getStore();
+      fragmentDoc.load();
+
+      const rootId = fragment.addBlock('affine:page');
+      const surfaceId = fragment.addBlock('affine:surface', {}, rootId);
+      const noteId = fragment.addBlock('affine:note', {}, rootId);
+      for (const element of selectedElements) {
+        if (element instanceof GfxBlockElementModel) {
+          const props = getBlockProps(element);
+          fragment.addBlock(element.flavour, props, surfaceId);
+        }
+
+        if (element instanceof NoteBlockModel) {
+          for (const child of element.children) {
+            const props = getBlockProps(child);
+            fragment.addBlock(child.flavour, props, noteId);
+          }
+        }
+
+        if (element instanceof AttachmentBlockModel) {
+          const { name, sourceId } = element.props;
+          if (name && sourceId) {
+            const blob = await host.store.blobSync.get(sourceId);
+            if (blob) {
+              attachments.push(new File([blob], name));
+            }
+          }
+        }
+      }
+
+      snapshot = transformer.docToSnapshot(fragment) ?? null;
+      markdown = (await markdownAdapter.fromDoc(fragment))?.file ?? '';
+    } finally {
+      collection.dispose();
+    }
+  }
 
   const canvas = await selectedToCanvas(host);
   if (!canvas) return null;
@@ -57,6 +124,9 @@ async function extractEdgelessSelected(
 
   return {
     images: [new File([blob], 'selected.png')],
+    snapshot: snapshot ?? undefined,
+    markdown: markdown.length ? markdown : undefined,
+    attachments,
   };
 }