chore(server): use jemalloc to reduce RSS

feat(editor): can highlight resolved comment (#13122 )
#### PR Dependency Tree * **PR #13122** 👈 This tree was auto-generated by [Charcoal](https://github.com/danerwilliams/charcoal)  ## Summary by CodeRabbit * **New Features** * Inline comments now visually distinguish between unresolved, resolved, and deleted states. * Only unresolved inline comments are interactive and highlighted in the editor. * **Bug Fixes** * Improved accuracy in fetching and displaying all comments, including resolved ones, during initialization. * **Refactor** * Enhanced handling of comment resolution and deletion to provide clearer differentiation in both behavior and appearance.
2026-07-02 10:10:42 +08:00 · 2025-07-10 11:22:37 +08:00 · 2025-07-10 03:06:05 +00:00 · 2025-07-09 15:56:10 +00:00 · 2025-07-09 22:21:30 +08:00 · 2025-07-09 13:26:06 +00:00
40 changed files with 996 additions and 449 deletions
@@ -7,7 +7,10 @@ COPY ./packages/frontend/apps/mobile/dist /app/static/mobile
 WORKDIR /app

 RUN apt-get update && \
-  apt-get install -y --no-install-recommends openssl && \
+  apt-get install -y --no-install-recommends openssl libjemalloc2 && \
  rm -rf /var/lib/apt/lists/*

+# Enable jemalloc by preloading the library
+ENV LD_PRELOAD=libjemalloc.so.2
+
 CMD ["node", "./dist/main.js"]
@@ -43,10 +43,14 @@ export class InlineCommentManager extends LifeCycleWatcher {

    this._disposables.add(provider.onCommentAdded(this._handleAddComment));
    this._disposables.add(
-      provider.onCommentDeleted(this._handleDeleteAndResolve)
+      provider.onCommentDeleted(id =>
+        this._handleDeleteAndResolve(id, 'delete')
+      )
    );
    this._disposables.add(
-      provider.onCommentResolved(this._handleDeleteAndResolve)
+      provider.onCommentResolved(id =>
+        this._handleDeleteAndResolve(id, 'resolve')
+      )
    );
    this._disposables.add(
      provider.onCommentHighlighted(this._handleHighlightComment)
@@ -64,15 +68,16 @@ export class InlineCommentManager extends LifeCycleWatcher {
    const provider = this._provider;
    if (!provider) return;

-    const commentsInProvider = await provider.getComments('unresolved');
+    const commentsInProvider = await provider.getComments('all');

    const commentsInEditor = this.getCommentsInEditor();

    // remove comments that are in editor but not in provider
    // which means the comment may be removed or resolved in provider side
    difference(commentsInEditor, commentsInProvider).forEach(comment => {
-      this._handleDeleteAndResolve(comment);
-      this.std.get(BlockElementCommentManager).handleDeleteAndResolve(comment);
+      this.std
+        .get(BlockElementCommentManager)
+        .handleDeleteAndResolve(comment, 'delete');
    });
  }

@@ -162,7 +167,10 @@ export class InlineCommentManager extends LifeCycleWatcher {
    });
  };

-  private readonly _handleDeleteAndResolve = (id: CommentId) => {
+  private readonly _handleDeleteAndResolve = (
+    id: CommentId,
+    type: 'delete' | 'resolve'
+  ) => {
    const commentedTexts = findCommentedTexts(this.std.store, id);
    if (commentedTexts.length === 0) return;

@@ -176,7 +184,7 @@ export class InlineCommentManager extends LifeCycleWatcher {
        inlineEditor?.formatText(
          selection.from,
          {
-            [`comment-${id}`]: null,
+            [`comment-${id}`]: type === 'delete' ? null : false,
          },
          {
            withoutTransact: true,
@@ -22,7 +22,7 @@ import { isEqual } from 'lodash-es';
 })
 export class InlineComment extends WithDisposable(ShadowlessElement) {
  static override styles = css`
-    inline-comment {
+    inline-comment.unresolved {
      display: inline-block;
      background-color: ${unsafeCSSVarV2('block/comment/highlightDefault')};
      border-bottom: 2px solid
@@ -41,6 +41,9 @@ export class InlineComment extends WithDisposable(ShadowlessElement) {
  })
  accessor commentIds!: string[];

+  @property({ attribute: false })
+  accessor unresolved = false;
+
  private _index: number = 0;

  @consume({ context: stdContext })
@@ -54,8 +57,10 @@ export class InlineComment extends WithDisposable(ShadowlessElement) {
  }

  private readonly _handleClick = () => {
-    this._provider?.highlightComment(this.commentIds[this._index]);
-    this._index = (this._index + 1) % this.commentIds.length;
+    if (this.unresolved) {
+      this._provider?.highlightComment(this.commentIds[this._index]);
+      this._index = (this._index + 1) % this.commentIds.length;
+    }
  };

  private readonly _handleHighlight = (id: CommentId | null) => {
@@ -89,6 +94,13 @@ export class InlineComment extends WithDisposable(ShadowlessElement) {
        this.classList.remove('highlighted');
      }
    }
+    if (_changedProperties.has('unresolved')) {
+      if (this.unresolved) {
+        this.classList.add('unresolved');
+      } else {
+        this.classList.remove('unresolved');
+      }
+    }
  }

  override render() {
@@ -21,19 +21,25 @@ export const CommentInlineSpecExtension =
    ),
    match: delta => {
      if (!delta.attributes) return false;
-      const comments = Object.entries(delta.attributes).filter(
-        ([key, value]) => isInlineCommendId(key) && value === true
-      );
+      const comments = Object.keys(delta.attributes).filter(isInlineCommendId);
      return comments.length > 0;
    },
-    renderer: ({ delta, children }) =>
-      html`<inline-comment .commentIds=${extractCommentIdFromDelta(delta)}
+    renderer: ({ delta, children }) => {
+      if (!delta.attributes) return html`${nothing}`;
+
+      const unresolved = Object.entries(delta.attributes).some(
+        ([key, value]) => isInlineCommendId(key) && value === true
+      );
+      return html`<inline-comment
+        .unresolved=${unresolved}
+        .commentIds=${extractCommentIdFromDelta(delta)}
        >${when(
          children,
          () => html`${children}`,
          () => nothing
        )}</inline-comment
-      >`,
+      >`;
+    },
    wrapper: true,
  });

@@ -57,10 +57,12 @@ export class BlockElementCommentManager extends LifeCycleWatcher {

    this._disposables.add(provider.onCommentAdded(this._handleAddComment));
    this._disposables.add(
-      provider.onCommentDeleted(this.handleDeleteAndResolve)
+      provider.onCommentDeleted(id => this.handleDeleteAndResolve(id, 'delete'))
    );
    this._disposables.add(
-      provider.onCommentResolved(this.handleDeleteAndResolve)
+      provider.onCommentResolved(id =>
+        this.handleDeleteAndResolve(id, 'resolve')
+      )
    );
    this._disposables.add(
      provider.onCommentHighlighted(this._handleHighlightComment)
@@ -146,18 +148,29 @@ export class BlockElementCommentManager extends LifeCycleWatcher {
    }
  };

-  readonly handleDeleteAndResolve = (id: CommentId) => {
+  readonly handleDeleteAndResolve = (
+    id: CommentId,
+    type: 'delete' | 'resolve'
+  ) => {
    const commentedBlocks = findCommentedBlocks(this.std.store, id);
    this.std.store.withoutTransact(() => {
      commentedBlocks.forEach(block => {
-        delete block.props.comments[id];
+        if (type === 'delete') {
+          delete block.props.comments[id];
+        } else {
+          block.props.comments[id] = false;
+        }
      });
    });

    const commentedElements = findCommentedElements(this.std.store, id);
    this.std.store.withoutTransact(() => {
      commentedElements.forEach(element => {
-        delete element.comments[id];
+        if (type === 'delete') {
+          delete element.comments[id];
+        } else {
+          element.comments[id] = false;
+        }
      });
    });
  };
@@ -369,7 +369,7 @@ The term **“CRDT”** was first introduced by Marc Shapiro, Nuno Preguiça, Ca
              .map(c => JSON.parse(c.citationJson).type)
              .filter(type => ['attachment', 'doc'].includes(type)).length ===
              0,
-          'should not have citation'
+          `should not have citation: ${JSON.stringify(c, null, 2)}`
        );
      });
    },
@@ -73,7 +73,8 @@ e2e('should get comment attachment body', async t => {
    docId,
    key,
    'test.txt',
-    Buffer.from('test')
+    Buffer.from('test'),
+    owner.id
  );

  const res = await app.GET(
@@ -361,7 +361,8 @@ export class CommentResolver {
      docId,
      key,
      attachment.filename ?? key,
-      buffer
+      buffer,
+      me.id
    );
    return this.commentAttachmentStorage.getUrl(workspaceId, docId, key);
  }
@@ -24,11 +24,12 @@ test.after.always(async () => {

 test('should put comment attachment', async t => {
  const workspace = await module.create(Mockers.Workspace);
+  const user = await module.create(Mockers.User);
  const docId = randomUUID();
  const key = randomUUID();
  const blob = Buffer.from('test');

-  await storage.put(workspace.id, docId, key, 'test.txt', blob);
+  await storage.put(workspace.id, docId, key, 'test.txt', blob, user.id);

  const item = await models.commentAttachment.get(workspace.id, docId, key);

@@ -39,15 +40,17 @@ test('should put comment attachment', async t => {
  t.is(item?.mime, 'text/plain');
  t.is(item?.size, blob.length);
  t.is(item?.name, 'test.txt');
+  t.is(item?.createdBy, user.id);
 });

 test('should get comment attachment', async t => {
  const workspace = await module.create(Mockers.Workspace);
+  const user = await module.create(Mockers.User);
  const docId = randomUUID();
  const key = randomUUID();
  const blob = Buffer.from('test');

-  await storage.put(workspace.id, docId, key, 'test.txt', blob);
+  await storage.put(workspace.id, docId, key, 'test.txt', blob, user.id);

  const item = await storage.get(workspace.id, docId, key);

@@ -62,11 +65,12 @@ test('should get comment attachment', async t => {

 test('should get comment attachment with access url', async t => {
  const workspace = await module.create(Mockers.Workspace);
+  const user = await module.create(Mockers.User);
  const docId = randomUUID();
  const key = randomUUID();
  const blob = Buffer.from('test');

-  await storage.put(workspace.id, docId, key, 'test.txt', blob);
+  await storage.put(workspace.id, docId, key, 'test.txt', blob, user.id);

  const url = storage.getUrl(workspace.id, docId, key);

@@ -79,11 +83,12 @@ test('should get comment attachment with access url', async t => {

 test('should delete comment attachment', async t => {
  const workspace = await module.create(Mockers.Workspace);
+  const user = await module.create(Mockers.User);
  const docId = randomUUID();
  const key = randomUUID();
  const blob = Buffer.from('test');

-  await storage.put(workspace.id, docId, key, 'test.txt', blob);
+  await storage.put(workspace.id, docId, key, 'test.txt', blob, user.id);

  await storage.delete(workspace.id, docId, key);

@@ -94,11 +99,12 @@ test('should delete comment attachment', async t => {

 test('should handle comment.attachment.delete event', async t => {
  const workspace = await module.create(Mockers.Workspace);
+  const user = await module.create(Mockers.User);
  const docId = randomUUID();
  const key = randomUUID();
  const blob = Buffer.from('test');

-  await storage.put(workspace.id, docId, key, 'test.txt', blob);
+  await storage.put(workspace.id, docId, key, 'test.txt', blob, user.id);

  await storage.onCommentAttachmentDelete({
    workspaceId: workspace.id,
@@ -113,14 +119,15 @@ test('should handle comment.attachment.delete event', async t => {

 test('should handle workspace.deleted event', async t => {
  const workspace = await module.create(Mockers.Workspace);
+  const user = await module.create(Mockers.User);
  const docId = randomUUID();
  const key1 = randomUUID();
  const key2 = randomUUID();
  const blob1 = Buffer.from('test');
  const blob2 = Buffer.from('test2');

-  await storage.put(workspace.id, docId, key1, 'test.txt', blob1);
-  await storage.put(workspace.id, docId, key2, 'test.txt', blob2);
+  await storage.put(workspace.id, docId, key1, 'test.txt', blob1, user.id);
+  await storage.put(workspace.id, docId, key2, 'test.txt', blob2, user.id);

  const count = module.event.count('comment.attachment.delete');

@@ -59,7 +59,8 @@ export class CommentAttachmentStorage {
    docId: string,
    key: string,
    name: string,
-    blob: Buffer
+    blob: Buffer,
+    userId: string
  ) {
    const meta = autoMetadata(blob);

@@ -75,6 +76,7 @@ export class CommentAttachmentStorage {
      name,
      mime: meta.contentType ?? 'application/octet-stream',
      size: blob.length,
+      createdBy: userId,
    });
  }

@@ -13,6 +13,7 @@ test.after.always(async () => {

 test('should upsert comment attachment', async t => {
  const workspace = await module.create(Mockers.Workspace);
+  const user = await module.create(Mockers.User);

  // add
  const item = await models.commentAttachment.upsert({
@@ -22,6 +23,7 @@ test('should upsert comment attachment', async t => {
    name: 'test-name',
    mime: 'text/plain',
    size: 100,
+    createdBy: user.id,
  });

  t.is(item.workspaceId, workspace.id);
@@ -30,6 +32,7 @@ test('should upsert comment attachment', async t => {
  t.is(item.mime, 'text/plain');
  t.is(item.size, 100);
  t.truthy(item.createdAt);
+  t.is(item.createdBy, user.id);

  // update
  const item2 = await models.commentAttachment.upsert({
@@ -46,6 +49,7 @@ test('should upsert comment attachment', async t => {
  t.is(item2.key, 'test-key');
  t.is(item2.mime, 'text/html');
  t.is(item2.size, 200);
+  t.is(item2.createdBy, user.id);

  // make sure only one blob is created
  const items = await models.commentAttachment.list(workspace.id);
@@ -32,6 +32,7 @@ export class CommentAttachmentModel extends BaseModel {
        name: input.name,
        mime: input.mime,
        size: input.size,
+        createdBy: input.createdBy,
      },
    });
  }
@@ -112,11 +112,14 @@ class ProductionEmbeddingClient extends EmbeddingClient {
    );

    try {
-      return ranks.map((score, chunk) => ({
-        chunk,
-        targetId: this.getTargetId(embeddings[chunk]),
-        score,
-      }));
+      return ranks.map((score, i) => {
+        const chunk = embeddings[i];
+        return {
+          chunk: chunk.chunk,
+          targetId: this.getTargetId(chunk),
+          score: Math.max(score, 1 - (chunk.distance || -Infinity)),
+        };
+      });
    } catch (error) {
      this.logger.error('Failed to parse rerank results', error);
      // silent error, will fallback to default sorting in parent method
@@ -148,7 +151,7 @@ class ProductionEmbeddingClient extends EmbeddingClient {

    const chunks = sortedEmbeddings.reduce(
      (acc, e) => {
-        const targetId = 'docId' in e ? e.docId : 'fileId' in e ? e.fileId : '';
+        const targetId = this.getTargetId(e);
        const key = `${targetId}:${e.chunk}`;
        acc[key] = e;
        return acc;
@@ -179,7 +182,10 @@ class ProductionEmbeddingClient extends EmbeddingClient {
        .filter(Boolean);

      this.logger.verbose(
-        `ReRank completed: ${highConfidenceChunks.length} high-confidence results found`
+        `ReRank completed: ${highConfidenceChunks.length} high-confidence results found, total ${sortedEmbeddings.length} embeddings`,
+        highConfidenceChunks.length !== sortedEmbeddings.length
+          ? JSON.stringify(ranks)
+          : undefined
      );
      return highConfidenceChunks.slice(0, topK);
    } catch (error) {
@@ -338,7 +338,7 @@ Convert a multi-speaker audio recording into a structured JSON format by transcr
  {
    name: 'Rerank results',
    action: 'Rerank results',
-    model: 'gpt-4.1-mini',
+    model: 'gpt-4.1',
    messages: [
      {
        role: 'system',
@@ -1286,7 +1286,7 @@ If there are items in the content that can be used as to-do tasks, please refer
  {
    name: 'Make it real',
    action: 'Make it real',
-    model: 'gpt-4.1-2025-04-14',
+    model: 'claude-sonnet-4@20250514',
    messages: [
      {
        role: 'system',
@@ -1327,7 +1327,7 @@ When sent new wireframes, respond ONLY with the contents of the html file.`,
  {
    name: 'Make it real with text',
    action: 'Make it real with text',
-    model: 'gpt-4.1-2025-04-14',
+    model: 'claude-sonnet-4@20250514',
    messages: [
      {
        role: 'system',
@@ -1677,7 +1677,7 @@ This sentence contains information from the first source[^1]. This sentence refe
 Before starting Tool calling, you need to follow:
 - DO NOT explain what operation you will perform.
 - DO NOT embed a tool call mid-sentence.
- When searching for unknown information or keyword, prioritize searching the user's workspace.
+- When searching for unknown information, personal information or keyword, prioritize searching the user's workspace rather than the web.
 - Depending on the complexity of the question and the information returned by the search tools, you can call different tools multiple times to search.
 </tool-calling-guidelines>

@@ -53,8 +53,11 @@ export class PromptService implements OnApplicationBootstrap {
   * @returns prompt messages
   */
  async get(name: string): Promise<ChatPrompt | null> {
-    const cached = this.cache.get(name);
-    if (cached) return cached;
+    // skip cache in dev mode to ensure the latest prompt is always fetched
+    if (!env.dev) {
+      const cached = this.cache.get(name);
+      if (cached) return cached;
+    }

    const prompt = await this.db.aiPrompt.findUnique({
      where: {
@@ -62,6 +62,7 @@ export abstract class AnthropicProvider<T> extends CopilotProvider<T> {

    try {
      metrics.ai.counter('chat_text_calls').add(1, { model: model.id });
+
      const [system, msgs] = await chatToGPTMessage(messages, true, true);

      const modelInstance = this.instance(model.id);
@@ -88,6 +88,12 @@ export abstract class GeminiProvider<T> extends CopilotProvider<T> {
        system,
        messages: msgs,
        abortSignal: options.signal,
+        providerOptions: {
+          google: this.getGeminiOptions(options, model.id),
+        },
+        tools: await this.getTools(options, model.id),
+        maxSteps: this.MAX_STEPS,
+        experimental_continueSteps: true,
      });

      if (!text) throw new Error('Failed to generate text');
@@ -233,12 +239,16 @@ export abstract class GeminiProvider<T> extends CopilotProvider<T> {
        taskType: 'RETRIEVAL_DOCUMENT',
      });

-      const { embeddings } = await embedMany({
-        model: modelInstance,
-        values: messages,
-      });
+      const embeddings = await Promise.allSettled(
+        messages.map(m =>
+          embedMany({ model: modelInstance, values: [m], maxRetries: 3 })
+        )
+      );

-      return embeddings.filter(v => v && Array.isArray(v));
+      return embeddings
+        .map(e => (e.status === 'fulfilled' ? e.value.embeddings : null))
+        .flat()
+        .filter((v): v is number[] => !!v && Array.isArray(v));
    } catch (e: any) {
      metrics.ai
        .counter('generate_embedding_errors')
@@ -254,16 +264,16 @@ export abstract class GeminiProvider<T> extends CopilotProvider<T> {
  ) {
    const [system, msgs] = await chatToGPTMessage(messages);
    const { fullStream } = streamText({
-      model: this.instance(model.id, {
-        useSearchGrounding: this.useSearchGrounding(options),
-      }),
+      model: this.instance(model.id),
      system,
      messages: msgs,
      abortSignal: options.signal,
-      maxSteps: this.MAX_STEPS,
      providerOptions: {
        google: this.getGeminiOptions(options, model.id),
      },
+      tools: await this.getTools(options, model.id),
+      maxSteps: this.MAX_STEPS,
+      experimental_continueSteps: true,
    });
    return fullStream;
  }
@@ -282,8 +292,4 @@ export abstract class GeminiProvider<T> extends CopilotProvider<T> {
  private isReasoningModel(model: string) {
    return model.startsWith('gemini-2.5');
  }
-
-  private useSearchGrounding(options: CopilotChatOptions) {
-    return options?.tools?.includes('webSearch');
-  }
 }
@@ -274,9 +274,11 @@ export class OpenAIProvider extends CopilotProvider<OpenAIConfig> {
  override getProviderSpecificTools(
    toolName: CopilotChatTools,
    model: string
-  ): [string, Tool] | undefined {
+  ): [string, Tool?] | undefined {
    if (toolName === 'webSearch' && !this.isReasoningModel(model)) {
      return ['web_search_preview', openai.tools.webSearchPreview()];
+    } else if (toolName === 'docEdit') {
+      return ['doc_edit', undefined];
    }
    return;
  }
@@ -126,7 +126,7 @@ export abstract class CopilotProvider<C = any> {
  protected getProviderSpecificTools(
    _toolName: CopilotChatTools,
    _model: string
-  ): [string, Tool] | undefined {
+  ): [string, Tool?] | undefined {
    return;
  }

@@ -143,7 +143,10 @@ export abstract class CopilotProvider<C = any> {
      for (const tool of options.tools) {
        const toolDef = this.getProviderSpecificTools(tool, model);
        if (toolDef) {
-          tools[toolDef[0]] = toolDef[1];
+          // allow provider prevent tool creation
+          if (toolDef[1]) {
+            tools[toolDef[0]] = toolDef[1];
+          }
          continue;
        }
        switch (tool) {
@@ -58,7 +58,7 @@ export const createDocSemanticSearchTool = (
 ) => {
  return tool({
    description:
-      'Retrieve conceptually related passages by performing vector-based semantic similarity search across embedded documents; use this tool only when exact keyword search fails or the user explicitly needs meaning-level matches (e.g., paraphrases, synonyms, broader concepts).',
+      'Retrieve conceptually related passages by performing vector-based semantic similarity search across embedded documents; use this tool only when exact keyword search fails or the user explicitly needs meaning-level matches (e.g., paraphrases, synonyms, broader concepts, recent documents).',
    parameters: z.object({
      query: z
        .string()
@@ -1,6 +1,8 @@
 package app.affine.pro.ai.chat

 import com.affine.pro.graphql.GetCopilotHistoriesQuery
+import com.affine.pro.graphql.fragment.CopilotChatHistory
+import com.affine.pro.graphql.fragment.CopilotChatMessage
 import kotlinx.datetime.Clock
 import kotlinx.datetime.Instant

@@ -51,11 +53,11 @@ data class ChatMessage(
            createAt = Clock.System.now(),
        )

-        fun from(message: GetCopilotHistoriesQuery.Message) = ChatMessage(
+        fun from(message: CopilotChatMessage) = ChatMessage(
            id = message.id,
            role = Role.fromValue(message.role),
            content = message.content,
            createAt = message.createdAt,
        )
    }
-}
+}
@@ -9,7 +9,8 @@ import com.affine.pro.graphql.GetCopilotHistoryIdsQuery
 import com.affine.pro.graphql.GetCopilotSessionsQuery
 import com.affine.pro.graphql.type.CreateChatMessageInput
 import com.affine.pro.graphql.type.CreateChatSessionInput
-import com.affine.pro.graphql.type.QueryChatSessionsInput
+import com.affine.pro.graphql.type.PaginationInput
+import com.affine.pro.graphql.type.QueryChatHistoriesInput
 import com.apollographql.apollo.ApolloClient
 import com.apollographql.apollo.api.Mutation
 import com.apollographql.apollo.api.Optional
@@ -29,12 +30,15 @@ class GraphQLService @Inject constructor() {
        GetCopilotSessionsQuery(
            workspaceId = workspaceId,
            docId = Optional.present(docId),
-            options = Optional.present(QueryChatSessionsInput(action = Optional.present(false)))
+            pagination = PaginationInput(
+              first = Optional.present(100)
+            ),
+            options = Optional.present(QueryChatHistoriesInput(action = Optional.present(false)))
        )
    ).mapCatching { data ->
-        data.currentUser?.copilot?.sessions?.find {
+        data.currentUser?.copilot?.chats?.paginatedCopilotChats?.edges?.map { item -> item.node.copilotChatHistory }?.find {
            it.parentSessionId == null
-        }?.id ?: error(ERROR_NULL_SESSION_ID)
+        }?.sessionId ?: error(ERROR_NULL_SESSION_ID)
    }

    suspend fun createCopilotSession(
@@ -60,12 +64,15 @@ class GraphQLService @Inject constructor() {
    ) = query(
        GetCopilotHistoriesQuery(
            workspaceId = workspaceId,
+            pagination = PaginationInput(
+              first = Optional.present(100)
+            ),
            docId = Optional.present(docId),
        )
    ).mapCatching { data ->
-        data.currentUser?.copilot?.histories?.firstOrNull { history ->
-            history.sessionId == sessionId
-        }?.messages ?: emptyList()
+        data.currentUser?.copilot?.chats?.paginatedCopilotChats?.edges?.map { item -> item.node.copilotChatHistory }?.firstOrNull { history ->
+          history.sessionId == sessionId
+        }?.messages?.map { msg -> msg.copilotChatMessage } ?: emptyList()
    }

    suspend fun getCopilotHistoryIds(
@@ -76,9 +83,12 @@ class GraphQLService @Inject constructor() {
        GetCopilotHistoryIdsQuery(
            workspaceId = workspaceId,
            docId = Optional.present(docId),
+            pagination = PaginationInput(
+              first = Optional.present(100)
+            ),
        )
    ).mapCatching { data ->
-        data.currentUser?.copilot?.histories?.firstOrNull { history ->
+        data.currentUser?.copilot?.chats?.edges?.map { item -> item.node }?.firstOrNull { history ->
            history.sessionId == sessionId
        }?.messages ?: emptyList()
    }
@@ -138,6 +138,8 @@ export class ChatContentStreamObjects extends WithDisposable(
            .std=${this.host?.std}
            .data=${streamObject}
            .width=${this.width}
+            .theme=${this.theme}
+            .notificationService=${this.notificationService}
          ></code-artifact-tool>
        `;
      case 'doc_edit':
@@ -0,0 +1,157 @@
+import { LoadingIcon } from '@blocksuite/affine/components/icons';
+import { SignalWatcher, WithDisposable } from '@blocksuite/affine/global/lit';
+import type { ColorScheme } from '@blocksuite/affine/model';
+import { ShadowlessElement } from '@blocksuite/affine/std';
+import { type NotificationService } from '@blocksuite/affine-shared/services';
+import type { Signal } from '@preact/signals-core';
+import {
+  css,
+  html,
+  nothing,
+  type PropertyValues,
+  type TemplateResult,
+} from 'lit';
+import { property } from 'lit/decorators.js';
+
+import {
+  isPreviewPanelOpen,
+  renderPreviewPanel,
+} from './artifacts-preview-panel';
+
+/**
+ * Base web-component for AI artifact tools.
+ * It encapsulates common reactive properties (data/std/width/…)
+ * and automatically calls `updatePreviewPanel()` when the `data`
+ * property changes while the preview panel is open.
+ */
+export abstract class ArtifactTool<
+  TData extends { type: 'tool-result' | 'tool-call' },
+> extends SignalWatcher(WithDisposable(ShadowlessElement)) {
+  static override styles = css`
+    .artifact-tool-card {
+      cursor: pointer;
+      margin: 8px 0;
+    }
+
+    .artifact-tool-card:hover {
+      opacity: 0.8;
+    }
+  `;
+
+  /** Tool data coming from ChatGPT (tool-call / tool-result). */
+  @property({ attribute: false })
+  accessor data!: TData;
+
+  @property({ attribute: false })
+  accessor width: Signal<number | undefined> | undefined;
+
+  @property({ attribute: false })
+  accessor notificationService!: NotificationService;
+
+  @property({ attribute: false })
+  accessor theme!: Signal<ColorScheme>;
+
+  /* -------------------------- Card meta hooks -------------------------- */
+
+  /**
+   * Sub-class must provide primary information for the card.
+   */
+  protected abstract getCardMeta(): {
+    title: string;
+    /** Page / file icon shown when not loading */
+    icon: TemplateResult | HTMLElement | string | null;
+    /** Whether the spinner should be displayed */
+    loading: boolean;
+    /** Extra css class appended to card root */
+    className?: string;
+  };
+
+  /** Banner shown on the right side of the card (can be undefined). */
+  protected abstract getBanner(
+    theme: ColorScheme
+  ): TemplateResult | HTMLElement | string | null | undefined;
+
+  /**
+   * Provide the main TemplateResult shown in the preview panel.
+   * Called each time the panel opens or the tool data updates.
+   */
+  protected abstract getPreviewContent(): TemplateResult<1>;
+
+  /** Provide the action controls (right-side buttons) for the panel. */
+  protected getPreviewControls(): TemplateResult<1> | undefined {
+    return undefined;
+  }
+
+  /** Open or refresh the preview panel. */
+  private openOrUpdatePreviewPanel() {
+    renderPreviewPanel(
+      this,
+      this.getPreviewContent(),
+      this.getPreviewControls()
+    );
+  }
+
+  protected refreshPreviewPanel() {
+    if (isPreviewPanelOpen(this)) {
+      this.openOrUpdatePreviewPanel();
+    }
+  }
+
+  /** Optionally override to show an error card. Return null if no error. */
+  protected getErrorTemplate(): TemplateResult | null {
+    return null;
+  }
+
+  private readonly onCardClick = (_e: Event) => {
+    this.openOrUpdatePreviewPanel();
+  };
+
+  protected renderCard() {
+    const { title, icon, loading, className } = this.getCardMeta();
+
+    const resolvedIcon = loading
+      ? LoadingIcon({
+          size: '20px',
+        })
+      : icon;
+
+    const banner = this.getBanner(this.theme.value);
+
+    return html`
+      <div
+        class="affine-embed-linked-doc-block artifact-tool-card ${className ??
+        ''} horizontal"
+        @click=${this.onCardClick}
+      >
+        <div class="affine-embed-linked-doc-content">
+          <div class="affine-embed-linked-doc-content-title">
+            <div class="affine-embed-linked-doc-content-title-icon">
+              ${resolvedIcon}
+            </div>
+            <div class="affine-embed-linked-doc-content-title-text">
+              ${title}
+            </div>
+          </div>
+        </div>
+        ${banner
+          ? html`<div class="affine-embed-linked-doc-banner">${banner}</div>`
+          : nothing}
+      </div>
+    `;
+  }
+
+  override render() {
+    const err = this.getErrorTemplate();
+    if (err) {
+      return err;
+    }
+    return this.renderCard();
+  }
+
+  override updated(changed: PropertyValues<this>) {
+    super.updated(changed);
+    if (changed.has('data') && isPreviewPanelOpen(this)) {
+      this.openOrUpdatePreviewPanel();
+    }
+  }
+}
@@ -1,9 +1,14 @@
 import { CodeBlockHighlighter } from '@blocksuite/affine/blocks/code';
-import { toast } from '@blocksuite/affine/components/toast';
 import { SignalWatcher, WithDisposable } from '@blocksuite/affine/global/lit';
+import { ColorScheme } from '@blocksuite/affine/model';
 import { unsafeCSSVar, unsafeCSSVarV2 } from '@blocksuite/affine/shared/theme';
-import { type BlockStdScope, ShadowlessElement } from '@blocksuite/affine/std';
-import { CopyIcon, PageIcon, ToolIcon } from '@blocksuite/icons/lit';
+import { type BlockStdScope } from '@blocksuite/affine/std';
+import {
+  CodeBlockIcon,
+  CopyIcon,
+  PageIcon,
+  ToolIcon,
+} from '@blocksuite/icons/lit';
 import type { Signal } from '@preact/signals-core';
 import { effect, signal } from '@preact/signals-core';
 import { css, html, LitElement, nothing } from 'lit';
@@ -11,7 +16,7 @@ import { property, state } from 'lit/decorators.js';
 import { classMap } from 'lit/directives/class-map.js';
 import { bundledLanguagesInfo, type ThemedToken } from 'shiki';

-import { renderPreviewPanel } from './artifacts-preview-panel';
+import { ArtifactTool } from './artifact-tool';
 import type { ToolError } from './type';

 interface CodeArtifactToolCall {
@@ -103,6 +108,8 @@ export class CodeHighlighter extends SignalWatcher(WithDisposable(LitElement)) {
  override connectedCallback() {
    super.connectedCallback();

+    this.highlighter.mounted();
+
    // recompute highlight when code / language changes
    this.disposables.add(
      effect(() => {
@@ -111,17 +118,25 @@ export class CodeHighlighter extends SignalWatcher(WithDisposable(LitElement)) {
    );
  }

+  override disconnectedCallback() {
+    super.disconnectedCallback();
+    this.highlighter.unmounted();
+  }
+
  private _updateHighlightTokens() {
    let cancelled = false;
    const language = this.language;
    const highlighter = this.highlighter.highlighter$.value;
+
    if (!highlighter) return;

    const updateTokens = () => {
      if (cancelled) return;
-      this.highlightTokens.value = highlighter.codeToTokensBase(this.code, {
-        lang: language,
-        theme: this.highlighter.themeKey,
+      requestIdleCallback(() => {
+        this.highlightTokens.value = highlighter.codeToTokensBase(this.code, {
+          lang: language,
+          theme: this.highlighter.themeKey,
+        });
      });
    };

@@ -199,20 +214,143 @@ export class CodeHighlighter extends SignalWatcher(WithDisposable(LitElement)) {
  }
 }

+const CodeBlockBanner = html`<svg
+  width="204"
+  height="102"
+  viewBox="0 0 204 102"
+  fill="none"
+  xmlns="http://www.w3.org/2000/svg"
+>
+  <g clip-path="url(#clip0_3371_100809)">
+    <g filter="url(#filter0_d_3371_100809)">
+      <rect
+        x="53.5054"
+        width="111.999"
+        height="99.5543"
+        rx="12.4443"
+        transform="rotate(8.37805 53.5054 0)"
+        fill="white"
+      />
+    </g>
+    <path
+      d="M89.7547 40.6581C90.8629 39.8345 92.4285 40.065 93.2522 41.1732C94.0758 42.2813 93.8452 43.847 92.7371 44.6706L79.7618 54.3146L89.4058 67.2899L89.5482 67.5024C90.1977 68.5905 89.9295 70.0161 88.8906 70.7883C87.8516 71.56 86.4104 71.4044 85.5558 70.4689L85.3932 70.2732L74.2581 55.2907C73.4345 54.1826 73.6653 52.617 74.7732 51.7933L89.7547 40.6581ZM114.378 44.2845C115.486 43.4608 117.052 43.6914 117.875 44.7996L129.011 59.7812C129.834 60.8892 129.604 62.4551 128.496 63.2787L113.514 74.4147L113.301 74.5552C112.213 75.2046 110.789 74.9382 110.016 73.8996C109.244 72.8606 109.399 71.4184 110.335 70.5637L110.531 70.4012L123.507 60.7572L113.863 47.7819C113.039 46.6738 113.27 45.1081 114.378 44.2845Z"
+      fill="#F3F3F3"
+    />
+  </g>
+  <defs>
+    <filter
+      id="filter0_d_3371_100809"
+      x="35.6787"
+      y="-3.32129"
+      width="131.951"
+      height="121.453"
+      filterUnits="userSpaceOnUse"
+      color-interpolation-filters="sRGB"
+    >
+      <feFlood flood-opacity="0" result="BackgroundImageFix" />
+      <feColorMatrix
+        in="SourceAlpha"
+        type="matrix"
+        values="0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 127 0"
+        result="hardAlpha"
+      />
+      <feOffset />
+      <feGaussianBlur stdDeviation="2.5" />
+      <feComposite in2="hardAlpha" operator="out" />
+      <feColorMatrix
+        type="matrix"
+        values="0 0 0 0 0.258824 0 0 0 0 0.254902 0 0 0 0 0.286275 0 0 0 0.17 0"
+      />
+      <feBlend
+        mode="normal"
+        in2="BackgroundImageFix"
+        result="effect1_dropShadow_3371_100809"
+      />
+      <feBlend
+        mode="normal"
+        in="SourceGraphic"
+        in2="effect1_dropShadow_3371_100809"
+        result="shape"
+      />
+    </filter>
+    <clipPath id="clip0_3371_100809">
+      <rect width="204" height="102" fill="white" />
+    </clipPath>
+  </defs>
+</svg>`;
+
+const CodeBlockBannerDark = html`<svg
+  width="204"
+  height="102"
+  viewBox="0 0 204 102"
+  fill="none"
+  xmlns="http://www.w3.org/2000/svg"
+>
+  <g clip-path="url(#clip0_3371_101118)">
+    <g filter="url(#filter0_d_3371_101118)">
+      <rect
+        x="53.5055"
+        width="111.999"
+        height="99.5543"
+        rx="12.4443"
+        transform="rotate(8.37805 53.5055 0)"
+        fill="#252525"
+      />
+    </g>
+    <path
+      d="M89.7551 40.6574C90.8631 39.8342 92.429 40.0647 93.2525 41.1725C94.0762 42.2806 93.8455 43.8472 92.7373 44.6709L79.762 54.3149L89.406 67.2902L89.5475 67.5025C90.197 68.5907 89.9298 70.0163 88.8908 70.7886C87.8519 71.5603 86.4106 71.4047 85.5561 70.4692L85.3934 70.2735L74.2574 55.2908C73.4341 54.1829 73.6649 52.6171 74.7725 51.7934L89.7551 40.6574ZM114.378 44.2838C115.486 43.4606 117.052 43.6911 117.876 44.7988L129.011 59.7814C129.834 60.8895 129.604 62.4552 128.496 63.2788L113.514 74.4149L113.301 74.5553C112.213 75.2045 110.789 74.9381 110.016 73.8998C109.244 72.8609 109.398 71.4186 110.334 70.5638L110.532 70.4014L123.507 60.7574L113.863 47.7822C113.039 46.674 113.27 45.1074 114.378 44.2838Z"
+      fill="#565656"
+    />
+  </g>
+  <defs>
+    <filter
+      id="filter0_d_3371_101118"
+      x="35.6787"
+      y="-3.32129"
+      width="131.951"
+      height="121.453"
+      filterUnits="userSpaceOnUse"
+      color-interpolation-filters="sRGB"
+    >
+      <feFlood flood-opacity="0" result="BackgroundImageFix" />
+      <feColorMatrix
+        in="SourceAlpha"
+        type="matrix"
+        values="0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 127 0"
+        result="hardAlpha"
+      />
+      <feOffset />
+      <feGaussianBlur stdDeviation="2.5" />
+      <feComposite in2="hardAlpha" operator="out" />
+      <feColorMatrix
+        type="matrix"
+        values="0 0 0 0 0.258824 0 0 0 0 0.254902 0 0 0 0 0.286275 0 0 0 0.17 0"
+      />
+      <feBlend
+        mode="normal"
+        in2="BackgroundImageFix"
+        result="effect1_dropShadow_3371_101118"
+      />
+      <feBlend
+        mode="normal"
+        in="SourceGraphic"
+        in2="effect1_dropShadow_3371_101118"
+        result="shape"
+      />
+    </filter>
+    <clipPath id="clip0_3371_101118">
+      <rect width="204" height="102" fill="white" />
+    </clipPath>
+  </defs>
+</svg> `;
+
 /**
 * Component to render code artifact tool call/result inside chat.
 */
-export class CodeArtifactTool extends WithDisposable(ShadowlessElement) {
+export class CodeArtifactTool extends ArtifactTool<
+  CodeArtifactToolCall | CodeArtifactToolResult
+> {
  static override styles = css`
-    .code-artifact-result {
-      cursor: pointer;
-      margin: 8px 0;
-    }
-
-    .code-artifact-result:hover {
-      background-color: var(--affine-hover-color);
-    }
-
    .code-artifact-preview {
      padding: 0;
      width: 100%;
@@ -289,165 +427,148 @@ export class CodeArtifactTool extends WithDisposable(ShadowlessElement) {
    }
  `;

-  @property({ attribute: false })
-  accessor data!: CodeArtifactToolCall | CodeArtifactToolResult;
-
-  @property({ attribute: false })
-  accessor width: Signal<number | undefined> | undefined;
-
  @property({ attribute: false })
  accessor std: BlockStdScope | undefined;

  @state()
  private accessor mode: 'preview' | 'code' = 'code';

-  private renderToolCall() {
-    const { args } = this.data as CodeArtifactToolCall;
-    const name = `Generating HTML artifact "${args.title}"`;
-    return html`<tool-call-card
-      .name=${name}
-      .icon=${ToolIcon()}
-    ></tool-call-card>`;
+  /* ---------------- ArtifactTool hooks ---------------- */
+
+  protected getBanner(theme: ColorScheme) {
+    return theme === ColorScheme.Dark ? CodeBlockBannerDark : CodeBlockBanner;
  }

-  private renderToolResult() {
-    if (!this.std) return nothing;
-    if (this.data.type !== 'tool-result') return nothing;
-    const resultData = this.data as CodeArtifactToolResult;
-    const result = resultData.result;
+  protected getCardMeta() {
+    const loading = this.data.type === 'tool-call';
+    return {
+      title: this.data.args.title,
+      icon: CodeBlockIcon({ width: '20', height: '20' }),
+      loading,
+      className: 'code-artifact-result',
+    };
+  }

-    if (result && typeof result === 'object' && 'title' in result) {
-      const { title, html: htmlContent } = result as {
-        title: string;
-        html: string;
-      };
-
-      const onClick = () => {
-        const copyHTML = async () => {
-          if (this.std) {
-            await navigator.clipboard
-              .writeText(htmlContent)
-              .catch(console.error);
-            toast(this.std.host, 'Copied HTML to clipboard');
-          }
-        };
-
-        const downloadHTML = () => {
-          try {
-            const blob = new Blob([htmlContent], { type: 'text/html' });
-            const url = URL.createObjectURL(blob);
-            const a = document.createElement('a');
-            a.href = url;
-            a.download = `${title || 'artifact'}.html`;
-            document.body.append(a);
-            a.click();
-            a.remove();
-            URL.revokeObjectURL(url);
-          } catch (e) {
-            console.error(e);
-          }
-        };
-
-        const setCodeMode = () => {
-          if (this.mode !== 'code') {
-            this.mode = 'code';
-            renderPreview();
-          }
-        };
-
-        const setPreviewMode = () => {
-          if (this.mode !== 'preview') {
-            this.mode = 'preview';
-            renderPreview();
-          }
-        };
-
-        const renderPreview = () => {
-          const controls = html`
-            <div class="code-artifact-toggle-container">
-              <div
-                class=${classMap({
-                  'toggle-button': true,
-                  active: this.mode === 'code',
-                })}
-                @click=${setCodeMode}
-              >
-                Code
-              </div>
-              <div
-                class=${classMap({
-                  'toggle-button': true,
-                  active: this.mode === 'preview',
-                })}
-                @click=${setPreviewMode}
-              >
-                Preview
-              </div>
-            </div>
-            <div style="flex: 1"></div>
-            <button class="code-artifact-control-btn" @click=${downloadHTML}>
-              ${PageIcon({
-                width: '20',
-                height: '20',
-                style: `color: ${unsafeCSSVarV2('icon/primary')}`,
-              })}
-              Download
-            </button>
-            <icon-button @click=${copyHTML} title="Copy HTML">
-              ${CopyIcon({ width: '20', height: '20' })}
-            </icon-button>
-          `;
-          renderPreviewPanel(
-            this,
-            html`<div class="code-artifact-preview">
-              ${this.mode === 'preview'
-                ? html`<html-preview .html=${htmlContent}></html-preview>`
-                : html`<code-highlighter
-                    .std=${this.std}
-                    .code=${htmlContent}
-                    .language=${'html'}
-                    .showLineNumbers=${true}
-                  ></code-highlighter>`}
-            </div>`,
-            controls
-          );
-        };
-
-        renderPreview();
-      };
-
-      return html`
+  protected override getPreviewContent() {
+    if (this.data.type !== 'tool-result' || !this.data.result) {
+      // loading state
+      return html`<div class="code-artifact-preview">
        <div
-          class="affine-embed-linked-doc-block code-artifact-result horizontal"
-          @click=${onClick}
+          style="display:flex;justify-content:center;align-items:center;height:100%"
        >
-          <div class="affine-embed-linked-doc-content">
-            <div class="affine-embed-linked-doc-content-title">
-              <div class="affine-embed-linked-doc-content-title-icon">
-                ${PageIcon({ width: '20', height: '20' })}
-              </div>
-              <div class="affine-embed-linked-doc-content-title-text">
-                ${title}
-              </div>
-            </div>
-          </div>
+          ${CodeBlockIcon({ width: '24', height: '24' })}
        </div>
-      `;
+      </div>`;
    }

-    return html`<tool-call-failed
-      .name=${'Code artifact failed'}
-      .icon=${ToolIcon()}
-    ></tool-call-failed>`;
+    const result = this.data.result;
+    if (typeof result !== 'object' || !('html' in result)) return html``;
+
+    const { html: htmlContent } = result as { html: string };
+
+    return html`<div class="code-artifact-preview">
+      ${this.mode === 'preview'
+        ? html`<html-preview .html=${htmlContent}></html-preview>`
+        : html`<code-highlighter
+            .std=${this.std}
+            .code=${htmlContent}
+            .language=${'html'}
+            .showLineNumbers=${true}
+          ></code-highlighter>`}
+    </div>`;
  }

-  protected override render() {
-    if (this.data.type === 'tool-call') {
-      return this.renderToolCall();
+  protected override getPreviewControls() {
+    if (this.data.type !== 'tool-result' || !this.std || !this.data.result) {
+      return undefined;
    }
-    if (this.data.type === 'tool-result') {
-      return this.renderToolResult();
+
+    const result = this.data.result as { html: string; title: string };
+    const htmlContent = result.html;
+    const title = result.title;
+
+    const copyHTML = async () => {
+      await navigator.clipboard.writeText(htmlContent).catch(console.error);
+      this.notificationService.toast('Copied HTML to clipboard');
+    };
+
+    const downloadHTML = () => {
+      try {
+        const blob = new Blob([htmlContent], { type: 'text/html' });
+        const url = URL.createObjectURL(blob);
+        const a = document.createElement('a');
+        a.href = url;
+        a.download = `${title || 'artifact'}.html`;
+        document.body.append(a);
+        a.click();
+        a.remove();
+        URL.revokeObjectURL(url);
+      } catch (e) {
+        console.error(e);
+      }
+    };
+
+    const setCodeMode = () => {
+      if (this.mode !== 'code') {
+        this.mode = 'code';
+        this.refreshPreviewPanel();
+      }
+    };
+
+    const setPreviewMode = () => {
+      if (this.mode !== 'preview') {
+        this.mode = 'preview';
+        this.refreshPreviewPanel();
+      }
+    };
+
+    return html`
+      <div class="code-artifact-toggle-container">
+        <div
+          class=${classMap({
+            'toggle-button': true,
+            active: this.mode === 'code',
+          })}
+          @click=${setCodeMode}
+        >
+          Code
+        </div>
+        <div
+          class=${classMap({
+            'toggle-button': true,
+            active: this.mode === 'preview',
+          })}
+          @click=${setPreviewMode}
+        >
+          Preview
+        </div>
+      </div>
+      <div style="flex: 1"></div>
+      <button class="code-artifact-control-btn" @click=${downloadHTML}>
+        ${PageIcon({
+          width: '20',
+          height: '20',
+          style: `color: ${unsafeCSSVarV2('icon/primary')}`,
+        })}
+        Download
+      </button>
+      <icon-button @click=${copyHTML} title="Copy HTML">
+        ${CopyIcon({ width: '20', height: '20' })}
+      </icon-button>
+    `;
+  }
+
+  protected override getErrorTemplate() {
+    if (
+      this.data.type === 'tool-result' &&
+      this.data.result &&
+      (this.data.result as any).type === 'error'
+    ) {
+      return html`<tool-call-failed
+        .name=${'Code artifact failed'}
+        .icon=${ToolIcon()}
+      ></tool-call-failed>`;
    }
-    return nothing;
+    return null;
  }
 }
@@ -2,24 +2,18 @@ import { getStoreManager } from '@affine/core/blocksuite/manager/store';
 import { getAFFiNEWorkspaceSchema } from '@affine/core/modules/workspace';
 import { getEmbedLinkedDocIcons } from '@blocksuite/affine/blocks/embed-doc';
 import { LoadingIcon } from '@blocksuite/affine/components/icons';
-import { toast } from '@blocksuite/affine/components/toast';
-import { WithDisposable } from '@blocksuite/affine/global/lit';
 import { RefNodeSlotsProvider } from '@blocksuite/affine/inlines/reference';
 import type { ColorScheme } from '@blocksuite/affine/model';
+import { NotificationProvider } from '@blocksuite/affine/shared/services';
 import { unsafeCSSVarV2 } from '@blocksuite/affine/shared/theme';
-import { type BlockStdScope, ShadowlessElement } from '@blocksuite/affine/std';
 import { MarkdownTransformer } from '@blocksuite/affine/widgets/linked-doc';
-import type { NotificationService } from '@blocksuite/affine-shared/services';
 import { CopyIcon, PageIcon, ToolIcon } from '@blocksuite/icons/lit';
-import { type Signal } from '@preact/signals-core';
-import { css, html, nothing, type PropertyValues } from 'lit';
+import type { BlockStdScope } from '@blocksuite/std';
+import { css, html } from 'lit';
 import { property } from 'lit/decorators.js';

 import { getCustomPageEditorBlockSpecs } from '../text-renderer';
-import {
-  isPreviewPanelOpen,
-  renderPreviewPanel,
-} from './artifacts-preview-panel';
+import { ArtifactTool } from './artifact-tool';
 import type { ToolError } from './type';

 interface DocComposeToolCall {
@@ -47,17 +41,10 @@ interface DocComposeToolResult {
 /**
 * Component to render doc compose tool call/result inside chat.
 */
-export class DocComposeTool extends WithDisposable(ShadowlessElement) {
+export class DocComposeTool extends ArtifactTool<
+  DocComposeToolCall | DocComposeToolResult
+> {
  static override styles = css`
-    .doc-compose-result {
-      cursor: pointer;
-      margin: 8px 0;
-    }
-
-    .doc-compose-result:hover {
-      background-color: var(--affine-hover-color);
-    }
-
    .doc-compose-result-preview {
      padding: 24px;
      height: 100%;
@@ -100,33 +87,60 @@ export class DocComposeTool extends WithDisposable(ShadowlessElement) {
    }
  `;

-  @property({ attribute: false })
-  accessor data!: DocComposeToolCall | DocComposeToolResult;
-
-  @property({ attribute: false })
-  accessor width: Signal<number | undefined> | undefined;
-
  @property({ attribute: false })
  accessor std: BlockStdScope | undefined;

-  @property({ attribute: false })
-  accessor notificationService!: NotificationService;
-
-  @property({ attribute: false })
-  accessor theme!: Signal<ColorScheme>;
-
-  override updated(changedProperties: PropertyValues) {
-    super.updated(changedProperties);
-    if (changedProperties.has('data') && isPreviewPanelOpen(this)) {
-      this.updatePreviewPanel();
-    }
+  protected getBanner(theme: ColorScheme) {
+    const { LinkedDocEmptyBanner } = getEmbedLinkedDocIcons(
+      theme,
+      'page',
+      'horizontal'
+    );
+    return LinkedDocEmptyBanner;
  }

-  private updatePreviewPanel() {
+  protected getCardMeta() {
+    const composing = this.data.type === 'tool-call';
+    return {
+      title: this.data.args.title,
+      icon: PageIcon(),
+      loading: composing,
+      className: 'doc-compose-result',
+    };
+  }
+
+  protected override getPreviewContent() {
+    if (!this.std) return html``;
+    const std = this.std;
+    const resultData = this.data;
+    const title = this.data.args.title;
+    const result = resultData.type === 'tool-result' ? resultData.result : null;
+    const successResult = result && 'markdown' in result ? result : null;
+
+    return html`<div class="doc-compose-result-preview">
+      <div class="doc-compose-result-preview-title">${title}</div>
+      ${successResult
+        ? html`<text-renderer
+            .answer=${successResult.markdown}
+            .host=${std.host}
+            .schema=${std.store.schema}
+            .options=${{
+              customHeading: true,
+              extensions: getCustomPageEditorBlockSpecs(),
+            }}
+          ></text-renderer>`
+        : html`<div class="doc-compose-result-preview-loading">
+            ${LoadingIcon({
+              size: '32px',
+            })}
+          </div>`}
+    </div>`;
+  }
+
+  protected override getPreviewControls() {
    if (!this.std) return;
    const std = this.std;
    const resultData = this.data;
-    const composing = resultData.type === 'tool-call';
    const title = this.data.args.title;
    const result = resultData.type === 'tool-result' ? resultData.result : null;
    const successResult = result && 'markdown' in result ? result : null;
@@ -138,7 +152,7 @@ export class DocComposeTool extends WithDisposable(ShadowlessElement) {
      await navigator.clipboard
        .writeText(successResult.markdown)
        .catch(console.error);
-      toast(std.host, 'Copied markdown to clipboard');
+      this.notificationService.toast('Copied markdown to clipboard');
    };

    const saveAsDoc = async () => {
@@ -147,6 +161,7 @@ export class DocComposeTool extends WithDisposable(ShadowlessElement) {
          return;
        }
        const workspace = std.store.workspace;
+        const notificationService = std.get(NotificationProvider);
        const refNodeSlots = std.getOptional(RefNodeSlotsProvider);
        const docId = await MarkdownTransformer.importMarkdownToDoc({
          collection: workspace,
@@ -156,7 +171,7 @@ export class DocComposeTool extends WithDisposable(ShadowlessElement) {
          extensions: getStoreManager().config.init().value.get('store'),
        });
        if (docId) {
-          const open = await this.notificationService.confirm({
+          const open = await notificationService.confirm({
            title: 'Open the doc you just created',
            message: 'Doc saved successfully! Would you like to open it now?',
            cancelText: 'Cancel',
@@ -170,99 +185,43 @@ export class DocComposeTool extends WithDisposable(ShadowlessElement) {
            });
          }
        } else {
-          toast(std.host, 'Failed to create document');
+          this.notificationService.toast('Failed to create document');
        }
      } catch (e) {
        console.error(e);
-        toast(std.host, 'Failed to create document');
+        this.notificationService.toast('Failed to create document');
      }
    };

-    const controls = html`
-      <button class="doc-compose-result-save-as-doc" @click=${saveAsDoc}>
-        ${PageIcon({
-          width: '20',
-          height: '20',
-          style: `color: ${unsafeCSSVarV2('icon/primary')}`,
-        })}
-        Save as doc
-      </button>
-      <icon-button @click=${copyMarkdown} title="Copy markdown">
-        ${CopyIcon({ width: '20', height: '20' })}
-      </icon-button>
-    `;
-
-    renderPreviewPanel(
-      this,
-      html`<div class="doc-compose-result-preview">
-        <div class="doc-compose-result-preview-title">${title}</div>
-        ${successResult
-          ? html`<text-renderer
-              .answer=${successResult.markdown}
-              .options=${{
-                customHeading: true,
-                extensions: getCustomPageEditorBlockSpecs(),
-              }}
-            ></text-renderer>`
-          : html`<div class="doc-compose-result-preview-loading">
-              ${LoadingIcon({
-                size: '32px',
-              })}
-            </div>`}
-      </div>`,
-      composing ? undefined : controls
-    );
+    return this.data.type === 'tool-call'
+      ? undefined
+      : html`
+          <button class="doc-compose-result-save-as-doc" @click=${saveAsDoc}>
+            ${PageIcon({
+              width: '20',
+              height: '20',
+              style: `color: ${unsafeCSSVarV2('icon/primary')}`,
+            })}
+            Save as doc
+          </button>
+          <icon-button @click=${copyMarkdown} title="Copy markdown">
+            ${CopyIcon({ width: '20', height: '20' })}
+          </icon-button>
+        `;
  }

-  protected override render() {
-    if (!this.std) return nothing;
-    const resultData = this.data;
-    const composing = resultData.type === 'tool-call';
-
-    const title = this.data.args.title;
-
+  protected override getErrorTemplate() {
    if (
-      resultData.type === 'tool-result' &&
-      resultData.result &&
-      'type' in resultData.result &&
-      resultData.result.type === 'error'
+      this.data.type === 'tool-result' &&
+      this.data.result &&
+      'type' in this.data.result &&
+      (this.data.result as any).type === 'error'
    ) {
-      // failed
      return html`<tool-call-failed
        .name=${'Doc compose failed'}
        .icon=${ToolIcon()}
      ></tool-call-failed>`;
    }
-
-    const { LinkedDocEmptyBanner } = getEmbedLinkedDocIcons(
-      this.theme.value,
-      'page',
-      'horizontal'
-    );
-
-    return html`
-      <div
-        class="affine-embed-linked-doc-block doc-compose-result horizontal"
-        @click=${this.updatePreviewPanel}
-      >
-        <div class="affine-embed-linked-doc-content">
-          <div class="affine-embed-linked-doc-content-title">
-            <div class="affine-embed-linked-doc-content-title-icon">
-              ${composing
-                ? LoadingIcon({
-                    size: '20px',
-                  })
-                : PageIcon()}
-            </div>
-            <div class="affine-embed-linked-doc-content-title-text">
-              ${title}
-            </div>
-          </div>
-        </div>
-        <div class="affine-embed-linked-doc-banner">
-          ${LinkedDocEmptyBanner}
-        </div>
-      </div>
-    `;
+    return null;
  }
 }
@@ -1,9 +1,10 @@
-import { IconButton } from '@affine/component';
+import { IconButton, notify } from '@affine/component';
 import { LitDocEditor, type PageEditor } from '@affine/core/blocksuite/editors';
 import { SnapshotHelper } from '@affine/core/modules/comment/services/snapshot-helper';
 import type { CommentAttachment } from '@affine/core/modules/comment/types';
 import { PeekViewService } from '@affine/core/modules/peek-view';
 import { DebugLogger } from '@affine/debug';
+import { getAttachmentFileIconRC } from '@blocksuite/affine/components/icons';
 import { type RichText, selectTextModel } from '@blocksuite/affine/rich-text';
 import { ViewportElementExtension } from '@blocksuite/affine/shared/services';
 import { openFilesWith } from '@blocksuite/affine/shared/utils';
@@ -15,6 +16,7 @@ import {
 } from '@blocksuite/icons/rc';
 import type { TextSelection } from '@blocksuite/std';
 import { useFramework, useService } from '@toeverything/infra';
+import bytes from 'bytes';
 import clsx from 'clsx';
 import {
  forwardRef,
@@ -30,7 +32,7 @@ import { useAsyncCallback } from '../../hooks/affine-async-hooks';
 import { getCommentEditorViewManager } from './specs';
 import * as styles from './style.css';

-const MAX_IMAGE_COUNT = 10;
+const MAX_ATTACHMENT_COUNT = 10;
 const logger = new DebugLogger('CommentEditor');

 const usePatchSpecs = (readonly: boolean) => {
@@ -78,6 +80,16 @@ export interface CommentEditorRef {
  focus: () => void;
 }

+const download = (url: string, name: string) => {
+  const element = document.createElement('a');
+  element.setAttribute('download', name);
+  element.setAttribute('href', url);
+  element.style.display = 'none';
+  document.body.append(element);
+  element.click();
+  element.remove();
+};
+
 // todo: get rid of circular data changes
 const useSnapshotDoc = (
  defaultSnapshotOrDoc: DocSnapshot | Store,
@@ -109,6 +121,75 @@ const useSnapshotDoc = (
  return doc;
 };

+const isImageAttachment = (att: EditorAttachment) => {
+  const type = att.mimeType || att.file?.type || '';
+  if (type) return type.startsWith('image/');
+  return !!att.url && /\.(png|jpe?g|gif|webp|svg)$/i.test(att.url);
+};
+
+const AttachmentPreviewItem: React.FC<{
+  attachment: EditorAttachment;
+  index: number;
+  readonly?: boolean;
+  handleAttachmentClick: (e: React.MouseEvent, index: number) => void;
+  handleAttachmentRemove: (id: string) => void;
+}> = ({
+  attachment,
+  index,
+  readonly,
+  handleAttachmentClick,
+  handleAttachmentRemove,
+}) => {
+  const isImg = isImageAttachment(attachment);
+  const Icon = !isImg
+    ? getAttachmentFileIconRC(
+        attachment.mimeType ||
+          attachment.file?.type ||
+          attachment.filename?.split('.').pop() ||
+          'none'
+      )
+    : undefined;
+
+  return (
+    <div
+      key={attachment.id}
+      className={isImg ? styles.previewBox : styles.filePreviewBox}
+      style={{
+        backgroundImage: isImg
+          ? `url(${attachment.localUrl ?? attachment.url})`
+          : undefined,
+      }}
+      onClick={e => handleAttachmentClick(e, index)}
+    >
+      {!isImg && Icon && <Icon className={styles.fileIcon} />}
+      {!isImg && (
+        <div className={styles.fileInfo}>
+          <span className={styles.fileName}>
+            {attachment.filename || attachment.file?.name || 'File'}
+          </span>
+          <span className={styles.fileSize}>
+            {attachment.size ? bytes(attachment.size) : ''}
+          </span>
+        </div>
+      )}
+
+      {!readonly && (
+        <IconButton
+          size={12}
+          className={styles.attachmentButton}
+          loading={attachment.status === 'uploading'}
+          variant="danger"
+          onClick={e => {
+            e.stopPropagation();
+            handleAttachmentRemove(attachment.id);
+          }}
+          icon={<CloseIcon />}
+        />
+      )}
+    </div>
+  );
+};
+
 export const CommentEditor = forwardRef<CommentEditorRef, CommentEditorProps>(
  function CommentEditor(
    {
@@ -143,25 +224,28 @@ export const CommentEditor = forwardRef<CommentEditorRef, CommentEditorProps>(
      [attachments, onAttachmentsChange]
    );

-    const isImageUploadDisabled = (attachments?.length ?? 0) >= MAX_IMAGE_COUNT;
+    const isUploadDisabled = (attachments?.length ?? 0) >= MAX_ATTACHMENT_COUNT;
    const uploadingAttachments = attachments?.some(
      att => att.status === 'uploading'
    );
    const commitDisabled =
      (empty && (attachments?.length ?? 0) === 0) || uploadingAttachments;

-    const addImages = useAsyncCallback(
+    const addAttachments = useAsyncCallback(
      async (files: File[]) => {
        if (!uploadCommentAttachment) return;
-        const valid = files.filter(f => f.type.startsWith('image/'));
+        const remaining = MAX_ATTACHMENT_COUNT - (attachments?.length ?? 0);
+        const valid = files.slice(0, remaining);
        if (!valid.length) return;
-        logger.info('addImages', { files: valid });
+        logger.info('addAttachments', { files: valid });

        const pendingAttachments: EditorAttachment[] = valid.map(f => ({
          id: nanoid(),
          file: f,
          localUrl: URL.createObjectURL(f),
          status: 'uploading',
+          filename: f.name,
+          mimeType: f.type,
        }));

        setAttachments(prev => [...prev, ...pendingAttachments]);
@@ -189,8 +273,12 @@ export const CommentEditor = forwardRef<CommentEditorRef, CommentEditorProps>(
              };
              return next;
            });
-          } catch (e) {
+          } catch (e: any) {
            logger.error('uploadCommentAttachment failed', { error: e });
+            notify.error({
+              title: 'Failed to upload attachment',
+              message: e.message,
+            });
            pending.localUrl && URL.revokeObjectURL(pending.localUrl);
            setAttachments(prev => {
              const index = prev.findIndex(att => att.id === pending.id);
@@ -202,38 +290,38 @@ export const CommentEditor = forwardRef<CommentEditorRef, CommentEditorProps>(
          }
        }
      },
-      [setAttachments, uploadCommentAttachment]
+      [attachments?.length, setAttachments, uploadCommentAttachment]
    );

-    const handlePasteImage = useCallback(
+    const handlePaste = useCallback(
      (event: React.ClipboardEvent) => {
        const items = event.clipboardData?.items;
        if (!items) return;
        const files: File[] = [];
        for (const index in items) {
          const item = items[index as any];
-          if (item.kind === 'file' && item.type.indexOf('image') >= 0) {
+          if (item.kind === 'file') {
            const blob = item.getAsFile();
            if (blob) files.push(blob);
          }
        }
        if (files.length) {
          event.preventDefault();
-          addImages(files);
+          addAttachments(files);
        }
      },
-      [addImages]
+      [addAttachments]
    );

-    const uploadImageFiles = useAsyncCallback(async () => {
-      if (isImageUploadDisabled) return;
-      const files = await openFilesWith('Images');
+    const openFilePicker = useAsyncCallback(async () => {
+      if (isUploadDisabled) return;
+      const files = await openFilesWith('Any');
      if (files) {
-        addImages(files);
+        addAttachments(files);
      }
-    }, [isImageUploadDisabled, addImages]);
+    }, [isUploadDisabled, addAttachments]);

-    const handleImageRemove = useCallback(
+    const handleAttachmentRemove = useCallback(
      (id: string) => {
        setAttachments(prev => {
          const att = prev.find(att => att.id === id);
@@ -247,10 +335,7 @@ export const CommentEditor = forwardRef<CommentEditorRef, CommentEditorProps>(
    const handleImagePreview = useCallback(
      (index: number) => {
        if (!attachments) return;
-
-        const imageAttachments = attachments.filter(
-          att => att.url || att.localUrl
-        );
+        const imageAttachments = attachments.filter(isImageAttachment);

        if (index >= imageAttachments.length) return;

@@ -291,12 +376,31 @@ export const CommentEditor = forwardRef<CommentEditorRef, CommentEditorProps>(
      [attachments, peekViewService]
    );

-    const handleImageClick = useCallback(
+    const handleAttachmentClick = useCallback(
      (e: React.MouseEvent, index: number) => {
        e.stopPropagation();
-        handleImagePreview(index);
+        if (!attachments) return;
+        const att = attachments[index];
+        if (!att) return;
+        const url = att.url || att.localUrl;
+        if (!url) return;
+        if (isImageAttachment(att)) {
+          // translate attachment index to image index
+          const imageAttachments = attachments.filter(isImageAttachment);
+          const imageIndex = imageAttachments.findIndex(i => i.id === att.id);
+          if (imageIndex >= 0) {
+            handleImagePreview(imageIndex);
+          }
+        } else if (att.url || att.localUrl) {
+          // todo: open attachment preview. for now, just download it
+          download(url, att.filename ?? att.file?.name ?? 'attachment');
+          notify({
+            title: 'Downloading attachment',
+            message: 'The attachment is being downloaded to your computer.',
+          });
+        }
      },
-      [handleImagePreview]
+      [attachments, handleImagePreview]
    );

    // upload attachments and call original onCommit
@@ -433,38 +537,24 @@ export const CommentEditor = forwardRef<CommentEditorRef, CommentEditorProps>(
      <div
        onClick={readonly ? undefined : handleClickEditor}
        onKeyDown={handleKeyDown}
-        onPaste={handlePasteImage}
+        onPaste={handlePaste}
        data-readonly={!!readonly}
        className={clsx(styles.container, 'comment-editor-viewport')}
      >
        {attachments?.length && attachments.length > 0 ? (
          <div
            className={styles.previewRow}
-            data-testid="comment-image-preview"
+            data-testid="comment-attachment-preview"
          >
            {attachments.map((att, index) => (
-              <div
+              <AttachmentPreviewItem
                key={att.id}
-                className={styles.previewBox}
-                style={{
-                  backgroundImage: `url(${att.localUrl ?? att.url})`,
-                }}
-                onClick={e => handleImageClick(e, index)}
-              >
-                {!readonly && (
-                  <IconButton
-                    size={12}
-                    className={styles.attachmentButton}
-                    loading={att.status === 'uploading'}
-                    variant="danger"
-                    onClick={e => {
-                      e.stopPropagation();
-                      handleImageRemove(att.id);
-                    }}
-                    icon={<CloseIcon />}
-                  />
-                )}
-              </div>
+                attachment={att}
+                index={index}
+                readonly={readonly}
+                handleAttachmentClick={handleAttachmentClick}
+                handleAttachmentRemove={handleAttachmentRemove}
+              />
            ))}
          </div>
        ) : null}
@@ -476,8 +566,8 @@ export const CommentEditor = forwardRef<CommentEditorRef, CommentEditorProps>(
          <div className={styles.footer}>
            <IconButton
              icon={<AttachmentIcon />}
-              onClick={uploadImageFiles}
-              aria-disabled={isImageUploadDisabled}
+              onClick={openFilePicker}
+              disabled={isUploadDisabled}
            />
            <button
              onClick={handleCommit}
@@ -106,3 +106,83 @@ export const attachmentButton = style({
    },
  },
 });
+
+// New generic file preview box (non-image attachments)
+export const filePreviewBox = style({
+  position: 'relative',
+  width: 194,
+  height: 62,
+  borderRadius: 4,
+  flex: '0 0 auto',
+  display: 'flex',
+  alignItems: 'center',
+  gap: 8,
+  padding: '0 4px',
+  background: cssVarV2('layer/background/secondary'),
+  border: `1px solid ${cssVarV2('layer/insideBorder/border')}`,
+  cursor: 'pointer',
+  selectors: {
+    '&:hover': {
+      opacity: 0.8,
+    },
+  },
+});
+
+export const fileIcon = style({
+  height: 36,
+  width: 'auto',
+});
+
+export const fileInfo = style({
+  display: 'flex',
+  flexDirection: 'column',
+  gap: 4,
+  overflow: 'hidden',
+});
+
+export const fileName = style({
+  overflow: 'hidden',
+  textOverflow: 'ellipsis',
+  whiteSpace: 'nowrap',
+  fontSize: 14,
+  flex: '1 1 auto',
+  fontWeight: 600,
+});
+
+export const fileSize = style({
+  fontSize: 12,
+  color: cssVarV2('text/secondary'),
+});
+
+export const deleteBtn = style({
+  position: 'absolute',
+  top: -6,
+  right: -6,
+  width: 16,
+  height: 16,
+  borderRadius: '50%',
+  background: cssVarV2('layer/background/primary'),
+  display: 'flex',
+  alignItems: 'center',
+  justifyContent: 'center',
+  border: `1px solid ${cssVarV2('layer/insideBorder/border')}`,
+  selectors: {
+    '&:hover': {
+      background: cssVarV2('layer/background/error'),
+      borderColor: cssVarV2('button/error'),
+    },
+  },
+});
+
+export const spinnerWrapper = style({
+  position: 'absolute',
+  top: 0,
+  left: 0,
+  width: '100%',
+  height: '100%',
+  background: cssVarV2('layer/background/tertiary'),
+  borderRadius: 4,
+  display: 'flex',
+  alignItems: 'center',
+  justifyContent: 'center',
+});
@@ -470,10 +470,6 @@ const CommentItem = ({
  const canDelete =
    (isMyComment && canCreateComment) || (!isMyComment && canDeleteComment);

-  const isCommentInEditor = useLiveData(entity.commentsInEditor$).includes(
-    comment.id
-  );
-
  // invalid comment, should not happen
  if (!comment.content) {
    return null;
@@ -516,12 +512,7 @@ const CommentItem = ({
          onDelete={handleDelete}
        />
      </div>
-      <div
-        data-deleted={!isCommentInEditor}
-        className={styles.previewContainer}
-      >
-        {comment.content?.preview}
-      </div>
+      <div className={styles.previewContainer}>{comment.content?.preview}</div>

      <div className={styles.repliesContainer}>
        {isEditing && editingDoc ? (
@@ -17,7 +17,9 @@ import {
  ViewBody,
  ViewHeader,
  ViewIcon,
+  ViewService,
  ViewTitle,
+  WorkbenchService,
 } from '@affine/core/modules/workbench';
 import { WorkspaceService } from '@affine/core/modules/workspace';
 import { useI18n } from '@affine/i18n';
@@ -212,6 +214,14 @@ export const Component = () => {
      await togglePin();
    };

+    tool.onOpenDoc = (docId: string, sessionId: string) => {
+      const { workbench } = framework.get(WorkbenchService);
+      const viewService = framework.get(ViewService);
+      workbench.open(`/${docId}?sessionId=${sessionId}`, { at: 'active' });
+      workbench.openSidebar();
+      viewService.view.activeSidebarTab('chat');
+    };
+
    // initial props
    if (!chatTool) {
      // mount
@@ -228,6 +238,7 @@ export const Component = () => {
    togglePin,
    workspaceId,
    confirmModal,
+    framework,
  ]);

  const onChatContainerRef = useCallback((node: HTMLDivElement) => {
@@ -351,6 +351,7 @@ export class DocCommentEntity extends Entity<{
      url,
      filename: file.name,
      mimeType: file.type,
+      size: file.size,
    });

    if (isPendingComment) {
@@ -13,6 +13,7 @@ export type CommentAttachment = {
  url?: string; // attachment may not be uploaded yet
  filename?: string;
  mimeType?: string;
+  size?: number; // in bytes
 };

 export interface BaseComment {
@@ -36,6 +36,8 @@ test.describe('AIAction/ImageProcessing', () => {
    await expect(answer.getByTestId('ai-answer-image')).toBeVisible();
    const insert = answer.getByTestId('answer-insert-below');
    await insert.click();
+    await page.reload();
+
    await utils.chatPanel.waitForHistory(page, [
      {
        role: 'action',
@@ -45,28 +45,30 @@ test.describe('AIBasic/Chat', () => {
    // Type and send a message
    await utils.chatPanel.makeChat(
      page,
-      'Introduce AFFiNE to me. Answer in 50 words.'
+      'Introduce AFFiNE to me. Answer in 500 words.'
    );

-    // AI is loading
-    await utils.chatPanel.waitForHistory(page, [
-      {
-        role: 'user',
-        content: 'Introduce AFFiNE to me. Answer in 50 words.',
-      },
-      {
-        role: 'assistant',
-        status: 'loading',
-      },
-    ]);
+    if (!(await page.getByTestId('ai-loading').isVisible())) {
+      // AI is loading
+      await utils.chatPanel.waitForHistory(page, [
+        {
+          role: 'user',
+          content: 'Introduce AFFiNE to me. Answer in 500 words.',
+        },
+        {
+          role: 'assistant',
+          status: 'loading',
+        },
+      ]);

-    await expect(page.getByTestId('ai-loading')).toBeVisible();
+      await expect(page.getByTestId('ai-loading')).toBeVisible();
+    }

    // AI Generating
    await utils.chatPanel.waitForHistory(page, [
      {
        role: 'user',
-        content: 'Introduce AFFiNE to me. Answer in 50 words.',
+        content: 'Introduce AFFiNE to me. Answer in 500 words.',
      },
      {
        role: 'assistant',
@@ -79,7 +81,7 @@ test.describe('AIBasic/Chat', () => {
    await utils.chatPanel.waitForHistory(page, [
      {
        role: 'user',
-        content: 'Introduce AFFiNE to me. Answer in 50 words.',
+        content: 'Introduce AFFiNE to me. Answer in 500 words.',
      },
      {
        role: 'assistant',
@@ -94,14 +96,14 @@ test.describe('AIBasic/Chat', () => {
  }) => {
    await utils.chatPanel.makeChat(
      page,
-      'Introduce AFFiNE to me. Answer in 50 words.'
+      'Introduce AFFiNE to me. Answer in 5000 words.'
    );

    // AI Generating
    await utils.chatPanel.waitForHistory(page, [
      {
        role: 'user',
-        content: 'Introduce AFFiNE to me. Answer in 50 words.',
+        content: 'Introduce AFFiNE to me. Answer in 5000 words.',
      },
      {
        role: 'assistant',
@@ -113,7 +115,7 @@ test.describe('AIBasic/Chat', () => {
    await utils.chatPanel.waitForHistory(page, [
      {
        role: 'user',
-        content: 'Introduce AFFiNE to me. Answer in 50 words.',
+        content: 'Introduce AFFiNE to me. Answer in 5000 words.',
      },
      {
        role: 'assistant',
@@ -183,13 +185,14 @@ test.describe('AIBasic/Chat', () => {
    // Type and send a message
    await utils.chatPanel.makeChat(
      page,
-      'Hello, write a poem about the moon. Answer in 50 words.'
+      'Hello, give a introduction about the moon. Answer in 500 words.'
    );

    await utils.chatPanel.waitForHistory(page, [
      {
        role: 'user',
-        content: 'Hello, write a poem about the moon. Answer in 50 words.',
+        content:
+          'Hello, give a introduction about the moon. Answer in 500 words.',
      },
      {
        role: 'assistant',
@@ -2,9 +2,14 @@ import { expect } from '@playwright/test';

 import { test } from '../base/base-test';

+test.describe.configure({ mode: 'serial' });
+
 test.describe('AIChatWith/Attachments', () => {
  test.beforeEach(async ({ loggedInPage: page, utils }) => {
-    await utils.testUtils.setupTestEnvironment(page);
+    await utils.testUtils.setupTestEnvironment(
+      page,
+      'claude-sonnet-4@20250514'
+    );
    await utils.chatPanel.openChatPanel(page);
  });

@@ -48,8 +53,10 @@ test.describe('AIChatWith/Attachments', () => {
    loggedInPage: page,
    utils,
  }) => {
-    const textContent1 = 'AttachmentEEee is a cute cat';
-    const textContent2 = 'AttachmentFFff is a cute dog';
+    const randomStr1 = Math.random().toString(36).substring(2, 6);
+    const randomStr2 = Math.random().toString(36).substring(2, 6);
+    const textContent1 = `Attachment${randomStr1} is a cute cat`;
+    const textContent2 = `Attachment${randomStr2} is a cute dog`;
    const buffer1 = Buffer.from(textContent1);
    const buffer2 = Buffer.from(textContent2);

@@ -67,13 +74,13 @@ test.describe('AIChatWith/Attachments', () => {
          buffer: buffer2,
        },
      ],
-      'What is AttachmentEEee? What is AttachmentFFff?'
+      `What is Attachment${randomStr1}? What is Attachment${randomStr2}?`
    );

    await utils.chatPanel.waitForHistory(page, [
      {
        role: 'user',
-        content: 'What is AttachmentEEee? What is AttachmentFFff?',
+        content: `What is Attachment${randomStr1}? What is Attachment${randomStr2}?`,
      },
      {
        role: 'assistant',
@@ -84,8 +91,8 @@ test.describe('AIChatWith/Attachments', () => {
    await expect(async () => {
      const { content, message } =
        await utils.chatPanel.getLatestAssistantMessage(page);
-      expect(content).toMatch(/AttachmentEEee/);
-      expect(content).toMatch(/AttachmentFFff/);
+      expect(content).toMatch(new RegExp(`Attachment${randomStr1}`));
+      expect(content).toMatch(new RegExp(`Attachment${randomStr2}`));
      expect(await message.locator('affine-footnote-node').count()).toBe(2);
    }).toPass({ timeout: 20000 });
  });
@@ -6,7 +6,10 @@ test.describe.configure({ mode: 'serial' });

 test.describe('AIChatWith/Collections', () => {
  test.beforeEach(async ({ loggedInPage: page, utils }) => {
-    await utils.testUtils.setupTestEnvironment(page);
+    await utils.testUtils.setupTestEnvironment(
+      page,
+      'claude-sonnet-4@20250514'
+    );
    await utils.chatPanel.openChatPanel(page);
    await utils.editor.clearAllCollections(page);

@@ -1,12 +1,16 @@
 import { createLocalWorkspace } from '@affine-test/kit/utils/workspace';
-import { faker } from '@faker-js/faker';
 import { expect } from '@playwright/test';

 import { test } from '../base/base-test';

+test.describe.configure({ mode: 'serial' });
+
 test.describe('AISettings/Embedding', () => {
  test.beforeEach(async ({ loggedInPage: page, utils }) => {
-    await utils.testUtils.setupTestEnvironment(page);
+    await utils.testUtils.setupTestEnvironment(
+      page,
+      'claude-sonnet-4@20250514'
+    );
    await utils.chatPanel.openChatPanel(page);
  });

@@ -246,7 +250,7 @@ test.describe('AISettings/Embedding', () => {
    await createLocalWorkspace({ name: 'test' }, page, false, 'affine-cloud');
    await utils.settings.openSettingsPanel(page);
    await utils.settings.enableWorkspaceEmbedding(page);
-    const person = faker.person.fullName();
+    const person = 'test123';

    const hobby1 = Buffer.from(`${person} love climbing`);
    const hobby2 = Buffer.from(`${person} love skating`);
@@ -201,7 +201,7 @@ export class ChatPanelUtils {
  public static async chatWithDoc(page: Page, docName: string) {
    const withButton = page.getByTestId('chat-panel-with-button');
    await withButton.hover();
-    await withButton.click();
+    await withButton.click({ delay: 200 });
    const withMenu = page.getByTestId('ai-add-popover');
    await withMenu.waitFor({ state: 'visible' });
    await withMenu.getByText(docName).click();
@@ -221,7 +221,7 @@ export class ChatPanelUtils {
      const fileChooserPromise = page.waitForEvent('filechooser');
      const withButton = page.getByTestId('chat-panel-with-button');
      await withButton.hover();
-      await withButton.click();
+      await withButton.click({ delay: 200 });
      const withMenu = page.getByTestId('ai-add-popover');
      await withMenu.waitFor({ state: 'visible' });
      await withMenu.getByTestId('ai-chat-with-files').click();
@@ -282,7 +282,7 @@ export class ChatPanelUtils {
    for (const tag of tags) {
      const withButton = page.getByTestId('chat-panel-with-button');
      await withButton.hover();
-      await withButton.click();
+      await withButton.click({ delay: 200 });
      const withMenu = page.getByTestId('ai-add-popover');
      await withMenu.waitFor({ state: 'visible' });
      await withMenu.getByTestId('ai-chat-with-tags').click();
@@ -299,7 +299,7 @@ export class ChatPanelUtils {
    for (const collection of collections) {
      const withButton = page.getByTestId('chat-panel-with-button');
      await withButton.hover();
-      await withButton.click();
+      await withButton.click({ delay: 200 });
      const withMenu = page.getByTestId('ai-add-popover');
      await withMenu.waitFor({ state: 'visible' });
      await withMenu.getByTestId('ai-chat-with-collections').click();
@@ -1,5 +1,8 @@
 import { skipOnboarding } from '@affine-test/kit/playwright';
-import { createRandomAIUser } from '@affine-test/kit/utils/cloud';
+import {
+  createRandomAIUser,
+  switchDefaultChatModel,
+} from '@affine-test/kit/utils/cloud';
 import { openHomePage, setCoreUrl } from '@affine-test/kit/utils/load-page';
 import {
  clickNewPageButton,
@@ -58,7 +61,12 @@ export class TestUtils {
    await waitForEditorLoad(page);
  }

-  public async setupTestEnvironment(page: Page) {
+  public async setupTestEnvironment(
+    page: Page,
+    defaultModel = 'gemini-2.5-flash'
+  ) {
+    await switchDefaultChatModel(defaultModel);
+
    await skipOnboarding(page.context());
    await openHomePage(page);
    await this.createNewPage(page);
@@ -152,6 +152,22 @@ export async function createRandomUser(): Promise<{
  } as any;
 }

+export async function switchDefaultChatModel(model: string) {
+  await runPrisma(async client => {
+    const promptId = await client.aiPrompt
+      .findFirst({
+        where: { name: 'Chat With AFFiNE AI' },
+        select: { id: true },
+      })
+      .then(f => f!.id);
+
+    await client.aiPrompt.update({
+      where: { id: promptId },
+      data: { model },
+    });
+  });
+}
+
 export async function createRandomAIUser(): Promise<{
  name: string;
  email: string;
Author	SHA1	Message	Date
fengmk2	b1d7011047	chore(server): use jemalloc to reduce RSS	2025-07-10 11:22:37 +08:00
L-Sun	1fe07410c0	feat(editor): can highlight resolved comment (#13122 ) #### PR Dependency Tree * PR #13122 👈 This tree was auto-generated by [Charcoal](https://github.com/danerwilliams/charcoal) <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit * New Features * Inline comments now visually distinguish between unresolved, resolved, and deleted states. * Only unresolved inline comments are interactive and highlighted in the editor. * Bug Fixes * Improved accuracy in fetching and displaying all comments, including resolved ones, during initialization. * Refactor * Enhanced handling of comment resolution and deletion to provide clearer differentiation in both behavior and appearance. <!-- end of auto-generated comment: release notes by coderabbit.ai -->	2025-07-10 03:06:05 +00:00
DarkSky	0f3066f7d0	fix(server): batch size in gemini embedding (#13120 ) <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit * Bug Fixes * Improved reliability of embedding generation for multiple messages, allowing partial results even if some embeddings fail. * Enhanced error handling to ensure only valid embeddings are returned. <!-- end of auto-generated comment: release notes by coderabbit.ai -->	2025-07-09 15:56:10 +00:00
DarkSky	c4c11da976	feat(server): use faster model in ci test (#13038 ) fix AI-329	2025-07-09 22:21:30 +08:00
Peng Xiao	38537bf310	fix(core): code block artifact styles (#13116 ) fix AI-314 #### PR Dependency Tree * PR #13116 👈 This tree was auto-generated by [Charcoal](https://github.com/danerwilliams/charcoal) <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit * New Features * Improved theme support for AI artifact tools, with banners and UI adapting to light or dark mode. * Enhanced notification handling for user actions like copying or saving content. * Refactor * Streamlined the structure of AI artifact tools for better maintainability and a more consistent user experience. * Unified and modernized preview and control panels for code and document compose tools. * Updated component integrations to consistently pass theme and notification services. * Style * Updated hover effects and visual feedback for artifact tool cards. <!-- end of auto-generated comment: release notes by coderabbit.ai --> #### PR Dependency Tree * PR #13116 👈 This tree was auto-generated by [Charcoal](https://github.com/danerwilliams/charcoal)	2025-07-09 13:26:06 +00:00
Wu Yue	1f87cd8752	feat(core): add onOpenDoc handler for AFFiNE Intelligence page (#13118 ) Close [AI-240](https://linear.app/affine-design/issue/AI-240) <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit * New Features * Enabled opening specific documents directly from the chat toolbar, automatically displaying the document in the workbench and focusing the chat tab in the sidebar. <!-- end of auto-generated comment: release notes by coderabbit.ai -->	2025-07-09 12:55:31 +00:00
EYHN	f54cb5c296	fix(android): fix android build error (#13117 ) <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit * New Features * Improved chat session and history retrieval with support for paginated results in the chat interface. * Bug Fixes * Enhanced reliability when loading large numbers of chat messages and histories. * Refactor * Updated chat data handling to align with the latest backend schema and pagination model. <!-- end of auto-generated comment: release notes by coderabbit.ai -->	2025-07-09 12:52:05 +00:00
fengmk2	45c016af8b	fix(server): add user id to comment-attachment model (#13113 ) close AF-2723 #### PR Dependency Tree * PR #13113 👈 This tree was auto-generated by [Charcoal](https://github.com/danerwilliams/charcoal) <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit * New Features * Comment attachments now track and display the user who uploaded them. * Tests * Updated tests to verify that the uploader’s information is correctly stored and retrieved with comment attachments. <!-- end of auto-generated comment: release notes by coderabbit.ai -->	2025-07-09 12:16:09 +00:00
Peng Xiao	d4c905600b	feat(core): support normal attachments (#13112 ) fix AF-2722 ![image](https://github.com/user-attachments/assets/376a0119-ae8e-4cb4-a31c-2eb6bb56c868) #### PR Dependency Tree * PR #13112 👈 This tree was auto-generated by [Charcoal](https://github.com/danerwilliams/charcoal) <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit * New Features * Expanded comment editor attachment support to include any file type, not just images. * Added file preview and download functionality for non-image attachments. * Introduced notifications for attachment upload failures and downloads. * Added a new AI artifact tool component for enhanced AI tool integrations. * Style * Added new styles for generic file previews, including icons, file info, and delete button. * Bug Fixes * Improved error handling and user feedback for attachment uploads. * Refactor * Unified attachment UI rendering and handling for both images and other file types. * Chores * Removed obsolete editor state attribute from comment preview sidebar. <!-- end of auto-generated comment: release notes by coderabbit.ai --> #### PR Dependency Tree * PR #13112 👈 This tree was auto-generated by [Charcoal](https://github.com/danerwilliams/charcoal)	2025-07-09 11:22:04 +00:00
Peng Xiao	f839e5c136	fix(core): should use sonnet 4 for make it real (#13106 ) #### PR Dependency Tree * PR #13106 👈 This tree was auto-generated by [Charcoal](https://github.com/danerwilliams/charcoal) <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit * New Features * Improved code highlighting performance and resource management for AI-generated code artifacts, resulting in smoother user experience and more efficient updates. * Chores * Updated underlying AI model for "Make it real" features, which may affect AI-generated outputs. <!-- end of auto-generated comment: release notes by coderabbit.ai -->	2025-07-09 11:10:53 +00:00