From 8c3ae99643b054f7fd9119884abf5fa12623b9bf Mon Sep 17 00:00:00 2001
From: haozhenfei <37089575+haozhenfei@users.noreply.github.com>
Date: Wed, 13 Aug 2025 20:55:57 +0800
Subject: [PATCH] fix(frontend): extend image extension to support tos key
 attribute (#723)

---
 .../extensions/image-with-tos-key.ts          | 75 +++++++++++++++++++
 .../editor-actions/upload-image/base.tsx      | 11 ++-
 .../hooks/use-case/use-init-editor.ts         |  4 +-
 3 files changed, 86 insertions(+), 4 deletions(-)
 create mode 100644 frontend/packages/data/knowledge/common/components/src/text-knowledge-editor/extensions/image-with-tos-key.ts

diff --git a/frontend/packages/data/knowledge/common/components/src/text-knowledge-editor/extensions/image-with-tos-key.ts b/frontend/packages/data/knowledge/common/components/src/text-knowledge-editor/extensions/image-with-tos-key.ts
new file mode 100644
index 00000000..d519006c
--- /dev/null
+++ b/frontend/packages/data/knowledge/common/components/src/text-knowledge-editor/extensions/image-with-tos-key.ts
@@ -0,0 +1,75 @@
+/*
+ * Copyright 2025 coze-dev Authors
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+import Image from '@tiptap/extension-image';
+
+// Extend TipTap Image to support a custom data-tos-key HTML attribute
+declare module '@tiptap/core' {
+  interface Commands<ReturnType> {
+    imageWithTosKey: {
+      setImageWithTosKey: (options: {
+        src: string;
+        alt?: string;
+        title?: string;
+        dataTosKey?: string | null;
+      }) => ReturnType;
+    };
+  }
+}
+
+const ImageWithTosKey = Image.extend({
+  addAttributes() {
+    const parentAttributes = (this.parent?.() as Record<string, unknown>) || {};
+    return {
+      ...parentAttributes,
+      dataTosKey: {
+        default: null,
+        parseHTML: (element: HTMLElement) =>
+          element.getAttribute('data-tos-key'),
+        renderHTML: (attributes: { dataTosKey?: string | null }) => {
+          if (!attributes.dataTosKey) {
+            return {};
+          }
+          return { 'data-tos-key': attributes.dataTosKey };
+        },
+      },
+    };
+  },
+
+  addCommands() {
+    return {
+      setImageWithTosKey:
+        (options: {
+          src: string;
+          alt?: string;
+          title?: string;
+          dataTosKey?: string | null;
+        }) =>
+        ({ commands }) =>
+          commands.insertContent({
+            type: this.name,
+            attrs: {
+              src: options.src,
+              alt: options.alt,
+              title: options.title,
+              dataTosKey: options.dataTosKey ?? null,
+            },
+          }),
+    };
+  },
+});
+
+export default ImageWithTosKey;
diff --git a/frontend/packages/data/knowledge/common/components/src/text-knowledge-editor/features/editor-actions/upload-image/base.tsx b/frontend/packages/data/knowledge/common/components/src/text-knowledge-editor/features/editor-actions/upload-image/base.tsx
index ea099d7e..9094e2a9 100644
--- a/frontend/packages/data/knowledge/common/components/src/text-knowledge-editor/features/editor-actions/upload-image/base.tsx
+++ b/frontend/packages/data/knowledge/common/components/src/text-knowledge-editor/features/editor-actions/upload-image/base.tsx
@@ -53,8 +53,15 @@ export const BaseUploadImage = ({
       options: {
         onFinish: (result: { url?: string; tosKey?: string }) => {
           if (result.url && editor) {
-            // Insert pictures into the editor
-            editor.chain().focus().setImage({ src: result.url }).run();
+            // Insert pictures into the editor and set the tosKey
+            editor
+              .chain()
+              .focus()
+              .setImageWithTosKey({
+                src: result.url,
+                dataTosKey: result.tosKey ?? null,
+              })
+              .run();
           }
         },
       },
diff --git a/frontend/packages/data/knowledge/common/components/src/text-knowledge-editor/hooks/use-case/use-init-editor.ts b/frontend/packages/data/knowledge/common/components/src/text-knowledge-editor/hooks/use-case/use-init-editor.ts
index 9c31164b..348e5b0d 100644
--- a/frontend/packages/data/knowledge/common/components/src/text-knowledge-editor/hooks/use-case/use-init-editor.ts
+++ b/frontend/packages/data/knowledge/common/components/src/text-knowledge-editor/hooks/use-case/use-init-editor.ts
@@ -23,11 +23,11 @@ import TableRow from '@tiptap/extension-table-row';
 import TableHeader from '@tiptap/extension-table-header';
 import TableCell from '@tiptap/extension-table-cell';
 import Table from '@tiptap/extension-table';
-import Image from '@tiptap/extension-image';
 
 import { type Chunk } from '@/text-knowledge-editor/types/chunk';
 import { getRenderHtmlContent } from '@/text-knowledge-editor/services/use-case/get-render-editor-content';
 import { getEditorContent } from '@/text-knowledge-editor/services/use-case/get-editor-content';
+import ImageWithTosKey from '@/text-knowledge-editor/extensions/image-with-tos-key';
 
 interface UseDocumentEditorProps {
   chunk: Chunk | null;
@@ -61,7 +61,7 @@ export const useInitEditor = ({
       TableRow,
       TableCell,
       TableHeader,
-      Image.configure({
+      ImageWithTosKey.configure({
         inline: false,
         allowBase64: true,
       }),