feat(core): add ai draft service (#13252)

Close [AI-244](https://linear.app/affine-design/issue/AI-244)  ## Summary by CodeRabbit * **New Features** * Added AI chat draft persistence, allowing your chat input, quotes, markdown, and images to be automatically saved and restored across sessions. * Drafts are now synchronized across chat components, so you won’t lose your progress if you navigate away or refresh the page. * **Improvements** * Enhanced chat experience with seamless restoration of previously entered content and attachments.
2026-07-16 17:46:18 +08:00 · 2025-07-17 17:42:01 +08:00
parent 4018b3aeca
commit 0770b109cb
9 changed files with 244 additions and 4 deletions
@@ -1,12 +1,15 @@
 export { AIButtonProvider } from './provider/ai-button';
 export { AIButtonService } from './services/ai-button';
+export { AIDraftService } from './services/ai-draft';

 import type { Framework } from '@toeverything/infra';

 import { FeatureFlagService } from '../feature-flag';
-import { GlobalStateService } from '../storage';
+import { CacheStorage, GlobalStateService } from '../storage';
+import { WorkspaceScope } from '../workspace';
 import { AIButtonProvider } from './provider/ai-button';
 import { AIButtonService } from './services/ai-button';
+import { AIDraftService } from './services/ai-draft';
 import { AINetworkSearchService } from './services/network-search';
 import { AIPlaygroundService } from './services/playground';
 import { AIReasoningService } from './services/reasoning';
@@ -31,3 +34,9 @@ export function configureAIReasoningModule(framework: Framework) {
 export function configureAIPlaygroundModule(framework: Framework) {
  framework.service(AIPlaygroundService, [FeatureFlagService]);
 }
+
+export function configureAIDraftModule(framework: Framework) {
+  framework
+    .scope(WorkspaceScope)
+    .service(AIDraftService, [GlobalStateService, CacheStorage]);
+}
@@ -0,0 +1,159 @@
+import { Service } from '@toeverything/infra';
+
+import type { CacheStorage, GlobalStateService } from '../../storage';
+
+const AI_DRAFTS_KEY = 'AIDrafts';
+const AI_DRAFT_FILES_PREFIX = 'AIDraftFile:';
+
+export interface CacheFile {
+  name: string;
+  size: number;
+  type: string;
+  cacheKey: string;
+}
+
+export interface AIDraftState {
+  input: string;
+  quote: string;
+  markdown: string;
+  images: File[];
+}
+
+export interface AIDraftGlobal {
+  input: string;
+  quote: string;
+  markdown: string;
+  images: CacheFile[];
+}
+
+const DEFAULT_VALUE = {
+  input: '',
+  quote: '',
+  markdown: '',
+  images: [],
+};
+
+export class AIDraftService extends Service {
+  private state: AIDraftState | null = null;
+
+  constructor(
+    private readonly globalStateService: GlobalStateService,
+    private readonly cacheStorage: CacheStorage
+  ) {
+    super();
+  }
+
+  setDraft = async (data: Partial<AIDraftState>) => {
+    const state = await this.getState();
+    const newState = {
+      ...state,
+      ...data,
+    };
+    this.state = newState;
+    await this.saveDraft(newState);
+  };
+
+  getDraft = async () => {
+    const state = await this.getState();
+    return state;
+  };
+
+  private readonly saveDraft = async (state: AIDraftState) => {
+    const draft =
+      this.globalStateService.globalState.get<AIDraftGlobal>(AI_DRAFTS_KEY) ||
+      DEFAULT_VALUE;
+
+    const addedImages = state.images.filter(image => {
+      return !draft.images.some(cacheImage => {
+        return cacheImage.cacheKey === this.getCacheKey(image);
+      });
+    });
+    const removedImages = draft.images.filter(cacheImage => {
+      return !state.images.some(image => {
+        return cacheImage.cacheKey === this.getCacheKey(image);
+      });
+    });
+
+    const cacheKeys = removedImages.map(image => image.cacheKey);
+    await this.removeFilesFromCache(cacheKeys);
+    await this.addFilesToCache(addedImages);
+
+    this.globalStateService.globalState.set<AIDraftGlobal>(AI_DRAFTS_KEY, {
+      input: state.input,
+      quote: state.quote,
+      markdown: state.markdown,
+      images: state.images.map(image => {
+        return {
+          name: image.name,
+          size: image.size,
+          type: image.type,
+          cacheKey: this.getCacheKey(image),
+        };
+      }),
+    });
+  };
+
+  private readonly initState = async () => {
+    if (this.state) {
+      return;
+    }
+    const draft =
+      this.globalStateService.globalState.get<AIDraftGlobal>(AI_DRAFTS_KEY);
+    if (draft) {
+      const images = await this.restoreFilesFromData(draft.images);
+      this.state = {
+        input: draft.input,
+        quote: draft.quote,
+        markdown: draft.markdown,
+        images,
+      };
+    } else {
+      this.state = DEFAULT_VALUE;
+    }
+  };
+
+  private readonly getState = async () => {
+    await this.initState();
+    return this.state as AIDraftState;
+  };
+
+  private readonly getCacheKey = (file: File) => {
+    return AI_DRAFT_FILES_PREFIX + file.name + file.size;
+  };
+
+  private readonly addFilesToCache = async (files: File[]) => {
+    for (const file of files) {
+      const arrayBuffer = await file.arrayBuffer();
+      const cacheKey = this.getCacheKey(file);
+      await this.cacheStorage.set(cacheKey, arrayBuffer);
+    }
+  };
+
+  private readonly removeFilesFromCache = async (cacheKeys: string[]) => {
+    for (const cacheKey of cacheKeys) {
+      await this.cacheStorage.del(cacheKey);
+    }
+  };
+
+  private readonly restoreFilesFromData = async (
+    cacheFiles: CacheFile[]
+  ): Promise<File[]> => {
+    const files: File[] = [];
+    for (const cacheFile of cacheFiles) {
+      try {
+        const arrayBuffer = await this.cacheStorage.get<ArrayBuffer>(
+          cacheFile.cacheKey
+        );
+        if (arrayBuffer) {
+          const file = new File([arrayBuffer], cacheFile.name, {
+            type: cacheFile.type,
+          });
+          files.push(file);
+        }
+      } catch (error) {
+        console.warn(`Failed to restore file ${cacheFile.name}:`, error);
+      }
+    }
+    return files;
+  };
+}
@@ -3,6 +3,7 @@ import { type Framework } from '@toeverything/infra';

 import {
  configureAIButtonModule,
+  configureAIDraftModule,
  configureAINetworkSearchModule,
  configureAIPlaygroundModule,
  configureAIReasoningModule,
@@ -110,6 +111,7 @@ export function configureCommonModules(framework: Framework) {
  configureAIReasoningModule(framework);
  configureAIPlaygroundModule(framework);
  configureAIButtonModule(framework);
+  configureAIDraftModule(framework);
  configureTemplateDocModule(framework);
  configureBlobManagementModule(framework);
  configureMediaModule(framework);