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

Close [AI-244](https://linear.app/affine-design/issue/AI-244)

<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->
## 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.
<!-- end of auto-generated comment: release notes by coderabbit.ai -->

packages/frontend/core/src/blocksuite/ai/chat-panel/index.ts

@@ -1,3 +1,4 @@
+import type { AIDraftService } from '@affine/core/modules/ai-button';
 import type { WorkspaceDialogService } from '@affine/core/modules/dialogs';
 import type { FeatureFlagService } from '@affine/core/modules/feature-flag';
 import type { AppThemeService } from '@affine/core/modules/theme';
@@ -115,6 +116,9 @@ export class ChatPanel extends SignalWatcher(
   @property({ attribute: false })
   accessor notificationService!: NotificationService;
 
+  @property({ attribute: false })
+  accessor aiDraftService!: AIDraftService;
+
   @state()
   accessor session: CopilotChatHistoryFragment | null | undefined;
 
@@ -408,6 +412,7 @@ export class ChatPanel extends SignalWatcher(
           .affineWorkspaceDialogService=${this.affineWorkspaceDialogService}
           .affineThemeService=${this.affineThemeService}
           .notificationService=${this.notificationService}
+          .aiDraftService=${this.aiDraftService}
           .onEmbeddingProgressChange=${this.onEmbeddingProgressChange}
           .onContextChange=${this.onContextChange}
           .width=${this.sidebarWidth}

packages/frontend/core/src/blocksuite/ai/components/ai-chat-composer/ai-chat-composer.ts

@@ -1,5 +1,6 @@
 import './ai-chat-composer-tip';
 
+import type { AIDraftService } from '@affine/core/modules/ai-button';
 import type { WorkspaceDialogService } from '@affine/core/modules/dialogs';
 import type {
   ContextEmbedStatus,
@@ -116,6 +117,9 @@ export class AIChatComposer extends SignalWatcher(
   @property({ attribute: false })
   accessor notificationService!: NotificationService;
 
+  @property({ attribute: false })
+  accessor aiDraftService!: AIDraftService;
+
   @state()
   accessor chips: ChatChip[] = [];
 
@@ -161,6 +165,7 @@ export class AIChatComposer extends SignalWatcher(
         .reasoningConfig=${this.reasoningConfig}
         .docDisplayConfig=${this.docDisplayConfig}
         .searchMenuConfig=${this.searchMenuConfig}
+        .aiDraftService=${this.aiDraftService}
         .portalContainer=${this.portalContainer}
         .onChatSuccess=${this.onChatSuccess}
         .trackOptions=${this.trackOptions}

packages/frontend/core/src/blocksuite/ai/components/ai-chat-content/ai-chat-content.ts

@@ -1,3 +1,5 @@
+import type { AIDraftService } from '@affine/core/modules/ai-button';
+import type { AIDraftState } from '@affine/core/modules/ai-button/services/ai-draft';
 import type { WorkspaceDialogService } from '@affine/core/modules/dialogs';
 import type { FeatureFlagService } from '@affine/core/modules/feature-flag';
 import type { AppThemeService } from '@affine/core/modules/theme';
@@ -15,6 +17,7 @@ import { property, state } from 'lit/decorators.js';
 import { classMap } from 'lit/directives/class-map.js';
 import { createRef, type Ref, ref } from 'lit/directives/ref.js';
 import { styleMap } from 'lit/directives/style-map.js';
+import { pick } from 'lodash-es';
 
 import { HISTORY_IMAGE_ACTIONS } from '../../chat-panel/const';
 import { type AIChatParams, AIProvider } from '../../provider/ai-provider';
@@ -149,6 +152,9 @@ export class AIChatContent extends SignalWatcher(
   @property({ attribute: false })
   accessor notificationService!: NotificationService;
 
+  @property({ attribute: false })
+  accessor aiDraftService!: AIDraftService;
+
   @property({ attribute: false })
   accessor onEmbeddingProgressChange:
     | ((count: Record<ContextEmbedStatus, number>) => void)
@@ -263,6 +269,19 @@ export class AIChatContent extends SignalWatcher(
   private readonly updateContext = (context: Partial<ChatContextValue>) => {
     this.chatContextValue = { ...this.chatContextValue, ...context };
     this.onContextChange?.(context);
+    this.updateDraft(context).catch(console.error);
+  };
+
+  private readonly updateDraft = async (context: Partial<ChatContextValue>) => {
+    const draft: Partial<AIDraftState> = pick(context, [
+      'quote',
+      'images',
+      'markdown',
+    ]);
+    if (!Object.keys(draft).length) {
+      return;
+    }
+    await this.aiDraftService.setDraft(draft);
   };
 
   private readonly initChatContent = async () => {
@@ -322,8 +341,19 @@ export class AIChatContent extends SignalWatcher(
 
   override connectedCallback() {
     super.connectedCallback();
+
     this.initChatContent().catch(console.error);
 
+    this.aiDraftService
+      .getDraft()
+      .then(draft => {
+        this.chatContextValue = {
+          ...this.chatContextValue,
+          ...draft,
+        };
+      })
+      .catch(console.error);
+
     this._disposables.add(
       AIProvider.slots.actions.subscribe(({ event }) => {
         const { status } = this.chatContextValue;
@@ -403,6 +433,7 @@ export class AIChatContent extends SignalWatcher(
         .searchMenuConfig=${this.searchMenuConfig}
         .affineWorkspaceDialogService=${this.affineWorkspaceDialogService}
         .notificationService=${this.notificationService}
+        .aiDraftService=${this.aiDraftService}
         .trackOptions=${{
           where: 'chat-panel',
           control: 'chat-send',

packages/frontend/core/src/blocksuite/ai/components/ai-chat-input/ai-chat-input.ts

@@ -1,10 +1,11 @@
+import type { AIDraftService } from '@affine/core/modules/ai-button';
 import type { CopilotChatHistoryFragment } from '@affine/graphql';
 import { SignalWatcher, WithDisposable } from '@blocksuite/affine/global/lit';
 import { unsafeCSSVar, unsafeCSSVarV2 } from '@blocksuite/affine/shared/theme';
 import type { EditorHost } from '@blocksuite/affine/std';
 import { ShadowlessElement } from '@blocksuite/affine/std';
 import { ArrowUpBigIcon, CloseIcon } from '@blocksuite/icons/lit';
-import { css, html, nothing } from 'lit';
+import { css, html, nothing, type PropertyValues } from 'lit';
 import { property, query, state } from 'lit/decorators.js';
 import { repeat } from 'lit/directives/repeat.js';
 import { styleMap } from 'lit/directives/style-map.js';
@@ -351,6 +352,9 @@ export class AIChatInput extends SignalWatcher(
   @property({ attribute: false })
   accessor searchMenuConfig!: SearchMenuConfig;
 
+  @property({ attribute: false })
+  accessor aiDraftService!: AIDraftService;
+
   @property({ attribute: false })
   accessor isRootSession: boolean = true;
 
@@ -379,6 +383,7 @@ export class AIChatInput extends SignalWatcher(
 
   override connectedCallback() {
     super.connectedCallback();
+
     this._disposables.add(
       AIProvider.slots.requestSendWithChat.subscribe(
         (params: AISendParams | null) => {
@@ -399,6 +404,17 @@ export class AIChatInput extends SignalWatcher(
     );
   }
 
+  protected override firstUpdated(changedProperties: PropertyValues): void {
+    super.firstUpdated(changedProperties);
+    this.aiDraftService
+      .getDraft()
+      .then(draft => {
+        this.textarea.value = draft.input;
+        this.isInputEmpty = !this.textarea.value.trim();
+      })
+      .catch(console.error);
+  }
+
   protected override render() {
     const { images, status } = this.chatContextValue;
     const hasImages = images.length > 0;
@@ -506,9 +522,11 @@ export class AIChatInput extends SignalWatcher(
     }
   };
 
-  private readonly _handleInput = () => {
+  private readonly _handleInput = async () => {
     const { textarea } = this;
-    this.isInputEmpty = !textarea.value.trim();
+    const value = textarea.value.trim();
+    this.isInputEmpty = !value;
+
     textarea.style.height = 'auto';
     textarea.style.height = textarea.scrollHeight + 'px';
     let imagesHeight = this.imagePreviewGrid?.scrollHeight ?? 0;
@@ -517,6 +535,10 @@ export class AIChatInput extends SignalWatcher(
       textarea.style.height = '148px';
       textarea.style.overflowY = 'scroll';
     }
+
+    await this.aiDraftService.setDraft({
+      input: value,
+    });
   };
 
   private readonly _handleKeyDown = async (evt: KeyboardEvent) => {
@@ -572,6 +594,9 @@ export class AIChatInput extends SignalWatcher(
     this.textarea.style.height = 'unset';
 
     await this.send(value);
+    await this.aiDraftService.setDraft({
+      input: '',
+    });
   };
 
   private readonly _handleModelChange = (modelId: string) => {