feat: refactor doc write in native (#14272)

packages/frontend/core/src/blocksuite/ai/components/ai-message-content/stream-objects.ts

@@ -131,6 +131,16 @@ export class ChatContentStreamObjects extends WithDisposable(
           .data=${streamObject}
           .width=${this.width}
         ></doc-read-result>`;
+      case 'doc_create':
+      case 'doc_update':
+      case 'doc_update_meta':
+        return html`<doc-write-tool
+          .data=${streamObject}
+          .width=${this.width}
+          .peekViewService=${this.peekViewService}
+          .docDisplayService=${this.docDisplayService}
+          .onOpenDoc=${this.onOpenDoc}
+        ></doc-write-tool>`;
       case 'section_edit':
         return html`
           <section-edit-tool
@@ -223,6 +233,16 @@ export class ChatContentStreamObjects extends WithDisposable(
           .peekViewService=${this.peekViewService}
           .onOpenDoc=${this.onOpenDoc}
         ></doc-read-result>`;
+      case 'doc_create':
+      case 'doc_update':
+      case 'doc_update_meta':
+        return html`<doc-write-tool
+          .data=${streamObject}
+          .width=${this.width}
+          .peekViewService=${this.peekViewService}
+          .docDisplayService=${this.docDisplayService}
+          .onOpenDoc=${this.onOpenDoc}
+        ></doc-write-tool>`;
       case 'section_edit':
         return html`
           <section-edit-tool

packages/frontend/core/src/blocksuite/ai/components/ai-tools/doc-write.ts

@@ -0,0 +1,194 @@
+import type { PeekViewService } from '@affine/core/modules/peek-view';
+import { WithDisposable } from '@blocksuite/global/lit';
+import { PageIcon, PenIcon } from '@blocksuite/icons/lit';
+import { ShadowlessElement } from '@blocksuite/std';
+import type { Signal } from '@preact/signals-core';
+import { html, nothing } from 'lit';
+import { property } from 'lit/decorators.js';
+
+import type { DocDisplayConfig } from '../ai-chat-chips';
+import type { ToolError } from './type';
+
+type DocWriteToolName = 'doc_create' | 'doc_update' | 'doc_update_meta';
+
+type DocWriteToolArgs = {
+  doc_id?: string;
+  title?: string;
+  content?: string;
+};
+
+interface DocWriteToolCall {
+  type: 'tool-call';
+  toolCallId: string;
+  toolName: DocWriteToolName;
+  args: DocWriteToolArgs;
+}
+
+interface DocWriteToolResult {
+  type: 'tool-result';
+  toolCallId: string;
+  toolName: DocWriteToolName;
+  args: DocWriteToolArgs;
+  result:
+    | {
+        success?: boolean;
+        docId?: string;
+        message?: string;
+      }
+    | ToolError
+    | null;
+}
+
+const isToolError = (result: unknown): result is ToolError =>
+  !!result &&
+  typeof result === 'object' &&
+  'type' in result &&
+  (result as ToolError).type === 'error';
+
+export class DocWriteTool extends WithDisposable(ShadowlessElement) {
+  @property({ attribute: false })
+  accessor data!: DocWriteToolCall | DocWriteToolResult;
+
+  @property({ attribute: false })
+  accessor width: Signal<number | undefined> | undefined;
+
+  @property({ attribute: false })
+  accessor peekViewService!: PeekViewService;
+
+  @property({ attribute: false })
+  accessor docDisplayService!: DocDisplayConfig;
+
+  @property({ attribute: false })
+  accessor onOpenDoc!: (docId: string, sessionId?: string) => void;
+
+  private getDocId() {
+    const { data } = this;
+    if (
+      data.type === 'tool-result' &&
+      data.result &&
+      !isToolError(data.result)
+    ) {
+      const docId =
+        typeof data.result.docId === 'string' ? data.result.docId : undefined;
+      if (docId) return docId;
+    }
+    const docId = data.args.doc_id;
+    return typeof docId === 'string' && docId.trim() ? docId : undefined;
+  }
+
+  private getDocTitle(docId?: string) {
+    const { data } = this;
+    if (data.toolName === 'doc_create' || data.toolName === 'doc_update_meta') {
+      const title = data.args.title;
+      if (title) return title;
+    }
+    if (docId && this.docDisplayService) {
+      const title = this.docDisplayService.getTitle(docId);
+      if (title) return title;
+    }
+    return undefined;
+  }
+
+  private getToolIcon() {
+    return this.data.toolName === 'doc_create' ? PageIcon() : PenIcon();
+  }
+
+  private getCallLabel(title?: string) {
+    switch (this.data.toolName) {
+      case 'doc_create':
+        return title ? `Creating "${title}"` : 'Creating document';
+      case 'doc_update':
+        return title ? `Updating "${title}"` : 'Updating document';
+      case 'doc_update_meta':
+        return title ? `Renaming to "${title}"` : 'Updating document title';
+      default:
+        return 'Updating document';
+    }
+  }
+
+  private getResultLabel(title?: string) {
+    switch (this.data.toolName) {
+      case 'doc_create':
+        return title ? `Created "${title}"` : 'Document created';
+      case 'doc_update':
+        return title ? `Updated "${title}"` : 'Document updated';
+      case 'doc_update_meta':
+        return title ? `Renamed "${title}"` : 'Document title updated';
+      default:
+        return 'Document updated';
+    }
+  }
+
+  private openDoc(docId?: string) {
+    if (!docId) return;
+    if (this.peekViewService) {
+      this.peekViewService.peekView
+        .open({ type: 'doc', docRef: { docId } })
+        .catch(console.error);
+      return;
+    }
+    this.onOpenDoc?.(docId);
+  }
+
+  renderToolCall() {
+    const docId = this.getDocId();
+    const title = this.getDocTitle(docId);
+    return html`<tool-call-card
+      .name=${this.getCallLabel(title)}
+      .icon=${this.getToolIcon()}
+      .width=${this.width}
+    ></tool-call-card>`;
+  }
+
+  renderToolResult() {
+    if (this.data.type !== 'tool-result') {
+      return nothing;
+    }
+
+    const result = this.data.result;
+    if (!result || isToolError(result)) {
+      const name = isToolError(result) ? result.name : 'Document action failed';
+      return html`<tool-call-failed
+        .name=${name}
+        .icon=${this.getToolIcon()}
+      ></tool-call-failed>`;
+    }
+
+    const docId = this.getDocId();
+    const title = this.getDocTitle(docId) ?? 'Document';
+    const parts: string[] = [];
+    if (result.message) parts.push(result.message);
+    if (docId) parts.push(`Doc ID: ${docId}`);
+    const content = parts.length ? parts.join('\n') : undefined;
+
+    return html`<tool-result-card
+      .name=${this.getResultLabel(title)}
+      .icon=${this.getToolIcon()}
+      .width=${this.width}
+      .results=${[
+        {
+          title,
+          icon: PageIcon(),
+          content,
+          onClick: () => this.openDoc(docId),
+        },
+      ]}
+    ></tool-result-card>`;
+  }
+
+  protected override render() {
+    if (this.data.type === 'tool-call') {
+      return this.renderToolCall();
+    }
+    if (this.data.type === 'tool-result') {
+      return this.renderToolResult();
+    }
+    return nothing;
+  }
+}
+
+declare global {
+  interface HTMLElementTagNameMap {
+    'doc-write-tool': DocWriteTool;
+  }
+}

packages/frontend/core/src/blocksuite/ai/effects.ts

@@ -64,6 +64,7 @@ import { DocEditTool } from './components/ai-tools/doc-edit';
 import { DocKeywordSearchResult } from './components/ai-tools/doc-keyword-search-result';
 import { DocReadResult } from './components/ai-tools/doc-read-result';
 import { DocSemanticSearchResult } from './components/ai-tools/doc-semantic-search-result';
+import { DocWriteTool } from './components/ai-tools/doc-write';
 import { SectionEditTool } from './components/ai-tools/section-edit';
 import { ToolCallCard } from './components/ai-tools/tool-call-card';
 import { ToolFailedCard } from './components/ai-tools/tool-failed-card';
@@ -222,6 +223,7 @@ export function registerAIEffects() {
   customElements.define('doc-semantic-search-result', DocSemanticSearchResult);
   customElements.define('doc-keyword-search-result', DocKeywordSearchResult);
   customElements.define('doc-read-result', DocReadResult);
+  customElements.define('doc-write-tool', DocWriteTool);
   customElements.define('web-crawl-tool', WebCrawlTool);
   customElements.define('web-search-tool', WebSearchTool);
   customElements.define('section-edit-tool', SectionEditTool);

packages/frontend/core/src/blocksuite/ai/provider/setup-provider.tsx

@@ -102,6 +102,7 @@ export function setupAIProvider(
         selectedSnapshot: contexts?.selectedSnapshot,
         selectedMarkdown: contexts?.selectedMarkdown,
         html: contexts?.html,
+        ...(options.docId ? { currentDocId: options.docId } : {}),
       },
       endpoint: Endpoint.StreamObject,
     });