From ccc210f88a994ff5c08ab1191f3e1eee944d17a8 Mon Sep 17 00:00:00 2001
From: fundon <fundon@pindash.io>
Date: Thu, 20 Mar 2025 02:08:19 +0000
Subject: [PATCH] refactor(core): edgeless toolbar ai action config extension
 (#10884)

---
 .../src/edgeless/configs/toolbar/connector.ts | 11 +--
 .../blocksuite/ai/entries/edgeless/index.ts   | 71 +++++++++++++++++++
 .../ai/extensions/ai-edgeless-root.ts         |  6 ++
 3 files changed, 80 insertions(+), 8 deletions(-)

diff --git a/blocksuite/affine/blocks/block-root/src/edgeless/configs/toolbar/connector.ts b/blocksuite/affine/blocks/block-root/src/edgeless/configs/toolbar/connector.ts
index 568b6c009..88abbcebf 100644
--- a/blocksuite/affine/blocks/block-root/src/edgeless/configs/toolbar/connector.ts
+++ b/blocksuite/affine/blocks/block-root/src/edgeless/configs/toolbar/connector.ts
@@ -51,12 +51,11 @@ import {
 import { html } from 'lit';
 import { styleMap } from 'lit/directives/style-map.js';
 
-import { EdgelessRootBlockComponent } from '../..';
 import { mountConnectorLabelEditor } from '../../utils/text';
 import { LINE_STYLE_LIST } from './consts';
 import { createTextActions } from './text-common';
 import type { MenuItem } from './types';
-import { renderMenu } from './utils';
+import { getEdgelessWith, renderMenu } from './utils';
 
 const FRONT_ENDPOINT_STYLE_LIST = [
   {
@@ -336,12 +335,8 @@ export const builtinConnectorToolbarConfig = {
         const rootModel = ctx.store.root;
         if (!rootModel) return;
 
-        // TODO(@fundon): it should be simple
-        const edgeless = ctx.view.getBlock(rootModel.id);
-        if (!ctx.matchBlock(edgeless, EdgelessRootBlockComponent)) {
-          console.error('edgeless view is not found.');
-          return;
-        }
+        const edgeless = getEdgelessWith(ctx);
+        if (!edgeless) return;
 
         mountConnectorLabelEditor(model, edgeless);
       },
diff --git a/packages/frontend/core/src/blocksuite/ai/entries/edgeless/index.ts b/packages/frontend/core/src/blocksuite/ai/entries/edgeless/index.ts
index 8f0182d2a..79a3ca0e3 100644
--- a/packages/frontend/core/src/blocksuite/ai/entries/edgeless/index.ts
+++ b/packages/frontend/core/src/blocksuite/ai/entries/edgeless/index.ts
@@ -4,6 +4,10 @@ import type {
 } from '@blocksuite/affine/blocks/root';
 import { noop } from '@blocksuite/affine/global/utils';
 import type { DocMode } from '@blocksuite/affine/model';
+import {
+  ActionPlacement,
+  type ToolbarModuleConfig,
+} from '@blocksuite/affine/shared/services';
 import { html } from 'lit';
 
 import type { AIItemGroupConfig } from '../../components/ai-item/types';
@@ -78,3 +82,70 @@ export function setupEdgelessElementToolbarAIEntry(
     },
   });
 }
+
+export function edgelessToolbarAIEntryConfig(): ToolbarModuleConfig {
+  return {
+    actions: [
+      {
+        placement: ActionPlacement.Start,
+        id: 'A.ai',
+        score: -1,
+        when(ctx) {
+          const models = ctx.getSurfaceModels();
+          return models.length > 0 && !models.some(model => model.isLocked());
+        },
+        content: ({ host, chain }) => {
+          const filteredGroups = edgelessAIGroups.reduce<AIItemGroupConfig[]>(
+            (pre, group) => {
+              const filtered = group.items.filter(item =>
+                item.showWhen?.(chain, 'edgeless' as DocMode, host)
+              );
+
+              if (filtered.length > 0) pre.push({ ...group, items: filtered });
+
+              return pre;
+            },
+            []
+          );
+
+          if (filteredGroups.every(group => group.items.length === 0))
+            return null;
+
+          const handler = () => {
+            const aiPanel = getAIPanelWidget(host);
+            if (aiPanel.config) {
+              aiPanel.config.generateAnswer = ({ finish, input }) => {
+                finish('success');
+                aiPanel.hide();
+                extractSelectedContent(host)
+                  .then(context => {
+                    AIProvider.slots.requestSendWithChat.next({
+                      input,
+                      context,
+                      host,
+                    });
+                  })
+                  .catch(console.error);
+              };
+              aiPanel.config.inputCallback = text => {
+                const copilotWidget = getEdgelessCopilotWidget(host);
+                const panel = copilotWidget.shadowRoot?.querySelector(
+                  'edgeless-copilot-panel'
+                );
+                if (panel instanceof HTMLElement) {
+                  panel.style.visibility = text ? 'hidden' : 'visible';
+                }
+              };
+            }
+          };
+
+          return html`<edgeless-copilot-toolbar-entry
+            .host=${host}
+            .groups=${edgelessAIGroups}
+            .onClick=${handler}
+          ></edgeless-copilot-toolbar-entry>`;
+        },
+      },
+    ],
+  };
+}
diff --git a/packages/frontend/core/src/blocksuite/ai/extensions/ai-edgeless-root.ts b/packages/frontend/core/src/blocksuite/ai/extensions/ai-edgeless-root.ts
index 300038d81..ca6862544 100644
--- a/packages/frontend/core/src/blocksuite/ai/extensions/ai-edgeless-root.ts
+++ b/packages/frontend/core/src/blocksuite/ai/extensions/ai-edgeless-root.ts
@@ -13,6 +13,7 @@ import type { FrameworkProvider } from '@toeverything/infra';
 import { buildAIPanelConfig } from '../ai-panel';
 import { toolbarAIEntryConfig } from '../entries';
 import {
+  edgelessToolbarAIEntryConfig,
   setupEdgelessCopilot,
   setupEdgelessElementToolbarAIEntry,
 } from '../entries/edgeless/index';
@@ -37,10 +38,15 @@ export function createAIEdgelessRootBlockSpec(
     aiPanelWidget,
     edgelessCopilotWidget,
     getAIEdgelessRootWatcher(framework),
+    // In note
     ToolbarModuleExtension({
       id: BlockFlavourIdentifier('custom:affine:note'),
       config: toolbarAIEntryConfig(),
     }),
+    ToolbarModuleExtension({
+      id: BlockFlavourIdentifier('custom:affine:surface:*'),
+      config: edgelessToolbarAIEntryConfig(),
+    }),
     AiSlashMenuConfigExtension(),
   ];
 }