feat(server): improve context management (#15448)

#### PR Dependency Tree * **PR #15448** 👈 This tree was auto-generated by [Charcoal](https://github.com/danerwilliams/charcoal)  ## Summary by CodeRabbit * **New Features** * Added workspace artifact upload, browsing, removal, deduplication, and library ownership support. * Copilot now supports scoped document and artifact search, canvas reading, live editor context, and frontend tools. * Added scope and focus selectors with source-resolution receipts in chat. * Added embedding health, progress, synchronization, and retrieval capabilities. * Added BYOK policy visibility, provider restrictions, endpoint dialect selection, and validation. * Added delegated editor interactions and userdata document authorization. * **Bug Fixes** * Improved attachment handling, cancellation, access control, retrieval fallbacks, workspace synchronization, and configuration validation.
2026-08-11 22:18:54 +08:00 · 2026-08-10 09:27:58 +08:00
parent 42322d13fe
commit ee899a267b
311 changed files with 20468 additions and 14806 deletions
@@ -384,8 +384,7 @@ test.describe('AIBasic/Chat', () => {
    await utils.chatPanel.waitForHistory(page, [
      {
        role: 'user',
-        content:
-          'AFFiNE is an open source all in one workspace.\nTranslate to chinese.',
+        content: 'Translate to chinese.',
      },
      {
        role: 'assistant',
@@ -470,8 +469,7 @@ test.describe('AIBasic/Chat', () => {
    await utils.chatPanel.waitForHistory(page, [
      {
        role: 'user',
-        content:
-          'AFFiNE is an open source all in one workspace.\nTranslate to chinese.',
+        content: 'Translate to chinese.',
      },
      {
        role: 'assistant',
@@ -6,7 +6,7 @@ test.describe.configure({ mode: 'serial' });

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

@@ -26,13 +26,13 @@ test.describe('AIChatWith/Attachments', () => {
          buffer: buffer,
        },
      ],
-      'What is AttachmentEEee?'
+      'What is AttachmentEEee? Answer with a citation.'
    );

    await utils.chatPanel.waitForHistory(page, [
      {
        role: 'user',
-        content: 'What is AttachmentEEee?',
+        content: 'What is AttachmentEEee? Answer with a citation.',
      },
      {
        role: 'assistant',
@@ -41,8 +41,29 @@ test.describe('AIChatWith/Attachments', () => {
    ]);

    await expect(async () => {
-      const { content } = await utils.chatPanel.getLatestAssistantMessage(page);
+      const { content, message } =
+        await utils.chatPanel.getLatestAssistantMessage(page);
      expect(content).toMatch(/EEee/);
+      const references = await message
+        .locator('affine-footnote-node')
+        .evaluateAll(nodes =>
+          nodes.map(
+            node =>
+              (
+                node as HTMLElement & {
+                  footnote?: { reference?: Record<string, unknown> };
+                }
+              ).footnote?.reference
+          )
+        );
+      expect(references).toEqual(
+        expect.arrayContaining([
+          expect.objectContaining({
+            type: 'attachment',
+            artifactId: expect.any(String),
+          }),
+        ])
+      );
    }).toPass({ timeout: 10000 });
  });

@@ -6,7 +6,7 @@ test.describe.configure({ mode: 'serial' });

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

@@ -31,14 +31,31 @@ test.describe('AIChatWith/Collections', () => {
    );

    await utils.chatPanel.chatWithCollections(page, ['Collection 1']);
+    await utils.settings.openSettingsPanel(page);
+    await utils.settings.disableWorkspaceEmbedding(page);
+    await utils.settings.closeSettingsPanel(page);
    await utils.chatPanel.makeChat(
      page,
-      `What is Collection${randomStr}(Use English)`
+      `What is Collection${randomStr}? Use English and answer with a citation.`
    );
+    await expect(
+      page.getByText(
+        'One or more selected sources are unavailable. Check the sources or workspace AI indexing settings, then retry.'
+      )
+    ).toBeVisible();
+    expect(
+      (await utils.chatPanel.collectHistory(page)).filter(
+        message => message.role === 'user'
+      )
+    ).toHaveLength(1);
+    await utils.settings.openSettingsPanel(page);
+    await utils.settings.enableWorkspaceEmbedding(page);
+    await utils.settings.closeSettingsPanel(page);
+    await page.getByTestId('ai-error-action-button').click();
    await utils.chatPanel.waitForHistory(page, [
      {
        role: 'user',
-        content: `What is Collection${randomStr}(Use English)`,
+        content: `What is Collection${randomStr}? Use English and answer with a citation.`,
      },
      {
        role: 'assistant',
@@ -50,7 +67,9 @@ test.describe('AIChatWith/Collections', () => {
      const { content, message } =
        await utils.chatPanel.getLatestAssistantMessage(page);
      expect(content).toMatch(new RegExp(`Collection${randomStr}.*dog`));
-      expect(await message.locator('affine-footnote-node').count()).toBe(1);
+      expect(
+        await message.locator('affine-footnote-node').count()
+      ).toBeGreaterThanOrEqual(1);
    }).toPass();
  });

@@ -79,12 +98,12 @@ test.describe('AIChatWith/Collections', () => {
    ]);
    await utils.chatPanel.makeChat(
      page,
-      `What is Collection${randomStr1}? What is Collection${randomStr2}?(Use English)`
+      `What is Collection${randomStr1}? What is Collection${randomStr2}? Use English and cite both sources.`
    );
    await utils.chatPanel.waitForHistory(page, [
      {
        role: 'user',
-        content: `What is Collection${randomStr1}? What is Collection${randomStr2}?(Use English)`,
+        content: `What is Collection${randomStr1}? What is Collection${randomStr2}? Use English and cite both sources.`,
      },
      {
        role: 'assistant',
@@ -97,7 +116,9 @@ test.describe('AIChatWith/Collections', () => {
        await utils.chatPanel.getLatestAssistantMessage(page);
      expect(content).toMatch(new RegExp(`Collection${randomStr1}.*cat`));
      expect(content).toMatch(new RegExp(`Collection${randomStr2}.*dog`));
-      expect(await message.locator('affine-footnote-node').count()).toBe(2);
+      expect(
+        await message.locator('affine-footnote-node').count()
+      ).toBeGreaterThanOrEqual(2);
    }).toPass();
  });
 });
@@ -22,6 +22,7 @@ test.describe('AIChatWith/Doc', () => {
    await page.keyboard.insertText('Test Doc');
    await page.keyboard.press('Enter');
    await page.keyboard.insertText('DocEEee is a cute cat');
+    await utils.editor.waitForCurrentDocSynced(page);

    await utils.chatPanel.chatWithDoc(page, 'Test Doc');

@@ -53,6 +54,7 @@ test.describe('AIChatWith/Doc', () => {
    await page.keyboard.insertText('Test Doc1');
    await page.keyboard.press('Enter');
    await page.keyboard.insertText('DocEEee is a cute cat');
+    await utils.editor.waitForCurrentDocSynced(page);

    // Initialize the doc 2
    await clickNewPageButton(page);
@@ -61,6 +63,7 @@ test.describe('AIChatWith/Doc', () => {
    await page.keyboard.insertText('Test Doc2');
    await page.keyboard.press('Enter');
    await page.keyboard.insertText('DocFFff is a cute dog');
+    await utils.editor.waitForCurrentDocSynced(page);

    await utils.chatPanel.chatWithDoc(page, 'Test Doc1');
    await utils.chatPanel.chatWithDoc(page, 'Test Doc2');
@@ -7,7 +7,7 @@ test.describe.configure({ mode: 'serial' });

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

@@ -158,28 +158,8 @@ test.describe('AISettings/Embedding', () => {
      },
    ];

-    const client = await page.context().newCDPSession(page);
-    await client.send('Network.enable');
-    await client.send('Network.emulateNetworkConditions', {
-      offline: false,
-      latency: 1000,
-      downloadThroughput: (50 * 1024) / 8,
-      uploadThroughput: (50 * 1024) / 8,
-      connectionType: 'cellular3g',
-    });
-
    await utils.settings.uploadWorkspaceEmbedding(page, attachments);

-    await client.send('Network.emulateNetworkConditions', {
-      offline: false,
-      latency: 0,
-      downloadThroughput: -1,
-      uploadThroughput: -1,
-    });
-
-    await utils.settings.disableWorkspaceEmbedding(page);
-    await utils.settings.enableWorkspaceEmbedding(page);
-
    await utils.settings.waitForFileEmbeddingReadiness(page, 2);

    await utils.settings.closeSettingsPanel(page);
@@ -236,7 +216,7 @@ test.describe('AISettings/Embedding', () => {
    await page.context().setOffline(false);
  });

-  test('should support hybrid search for both globally uploaded attachments and those uploaded in the current session', async ({
+  test('should isolate selected attachments from workspace artifact search', async ({
    loggedInPage: page,
    utils,
  }) => {
@@ -287,8 +267,8 @@ test.describe('AISettings/Embedding', () => {
    await expect(async () => {
      const { message } = await utils.chatPanel.getLatestAssistantMessage(page);
      const fullText = await message.innerText();
-      expect(fullText).toMatch(/climbing/i);
      expect(fullText).toMatch(/skating/i);
+      expect(fullText).not.toMatch(/climbing/i);
      expect(
        await message.locator('affine-footnote-node').count()
      ).toBeGreaterThanOrEqual(1);
@@ -305,7 +285,7 @@ test.describe('AISettings/Embedding', () => {
    const attachments = Array.from({ length: 11 }, (_, i) => ({
      name: `document${i + 1}.txt`,
      mimeType: 'text/plain',
-      buffer: Buffer.from('attachment content'),
+      buffer: Buffer.from(`attachment content ${i + 1}`),
    }));

    await utils.settings.uploadWorkspaceEmbedding(page, attachments);
@@ -18,6 +18,37 @@ import type {
 import type { GfxModel } from '@blocksuite/std/gfx';
 import { type Page } from '@playwright/test';
 export class EditorUtils {
+  public static async waitForCurrentDocSynced(page: Page) {
+    await page.evaluate(async () => {
+      const [, , workspaceId, docId] = location.pathname.split('/');
+      const workspace = (
+        window as typeof window & {
+          currentWorkspace?: {
+            engine: {
+              doc: {
+                waitForSynced(docId: string, abort: AbortSignal): Promise<void>;
+              };
+            };
+          };
+        }
+      ).currentWorkspace;
+      if (!workspaceId || !docId || !workspace) {
+        throw new Error('Current document is unavailable');
+      }
+      const abort = AbortSignal.timeout(60_000);
+      await Promise.all([
+        workspace.engine.doc.waitForSynced(docId, abort),
+        workspace.engine.doc.waitForSynced('db$docProperties', abort),
+      ]);
+      const response = await fetch(
+        `/api/workspaces/${workspaceId}/docs/${docId}`
+      );
+      if (!response.ok) {
+        throw new Error('Current document could not be persisted');
+      }
+    });
+  }
+
  public static async focusToEditor(page: Page) {
    const title = getBlockSuiteEditorTitle(page);
    await title.focus();
@@ -351,8 +382,7 @@ export class EditorUtils {
        await page.keyboard.press('Enter');
      }
    }
-    // sleep 1 sec to wait the doc sync
-    await page.waitForTimeout(1000);
+    await this.waitForCurrentDocSynced(page);
  }

  public static async createTagAndDoc(
@@ -1,4 +1,3 @@
-import { cleanupWorkspace, runPrisma } from '@affine-test/kit/utils/cloud';
 import { expect, type Page } from '@playwright/test';

 const WORKSPACE_EMBEDDING_SWITCH_TEST_ID = 'workspace-embedding-setting-switch';
@@ -95,7 +94,7 @@ export class SettingsPanelUtils {

      await page
        .getByTestId('workspace-embedding-setting-attachment-uploading-item')
-        .waitFor({ state: 'hidden' });
+        .waitFor({ state: 'hidden', timeout: 30_000 });
    }
  }

@@ -226,7 +225,6 @@ export class SettingsPanelUtils {
    timeout: number,
    status = 'synced'
  ) {
-    await cleanupWorkspace(page.url().split('/').slice(-2)[0] || '');
    await expect(async () => {
      await this.openSettingsPanel(page);
      const title = page.getByTestId('embedding-progress-title');
@@ -266,17 +264,17 @@ export class SettingsPanelUtils {
        'workspace-embedding-setting-attachment-item'
      );
      await expect(attachmentItems).toHaveCount(expectedFileCount);
+      await expect(
+        attachmentList.getByTestId(
+          'workspace-embedding-setting-attachment-uploading-item'
+        )
+      ).toHaveCount(0);
    }).toPass({ timeout });

-    const workspaceId = page.url().split('/').slice(-2)[0] || '';
-    await cleanupWorkspace(workspaceId);
    await expect(async () => {
-      const embeddedFileCount = await runPrisma(client =>
-        client.aiWorkspaceFiles.count({
-          where: { workspaceId, embeddings: { some: {} } },
-        })
-      );
-      expect(embeddedFileCount).toBe(expectedFileCount);
+      await expect(
+        page.getByTestId('workspace-embedding-setting-attachment-ready-item')
+      ).toHaveCount(expectedFileCount);
    }).toPass({ timeout });
  }
 }
@@ -58,16 +58,8 @@ export class TestUtils {
    await waitForEditorLoad(page);
  }

-  public async setupTestEnvironment(page: Page, defaultModel?: string) {
-    const selectedModel = defaultModel ?? 'gpt-5.6-luna';
-
+  public async setupTestEnvironment(page: Page) {
    await skipOnboarding(page.context());
-    await page.context().addInitScript(model => {
-      window.localStorage.setItem(
-        'global-state:AIModelId',
-        JSON.stringify(model)
-      );
-    }, selectedModel);
    await openHomePage(page);
    await this.createNewPage(page);
  }
@@ -1,9 +1,15 @@
+import { fileURLToPath } from 'node:url';
+
 import { testResultDir } from '@affine-test/kit/playwright';
 import type {
  PlaywrightTestConfig,
  PlaywrightWorkerOptions,
 } from '@playwright/test';

+const runtimeConfigPath = fileURLToPath(
+  new URL('./runtime-config.json', import.meta.url)
+);
+
 const config: PlaywrightTestConfig = {
  testDir: './e2e',
  fullyParallel: true,
@@ -50,6 +56,7 @@ const config: PlaywrightTestConfig = {
          'postgresql://affine:affine@localhost:5432/affine',
        NODE_ENV: 'test',
        AFFINE_ENV: process.env.AFFINE_ENV ?? 'dev',
+        AFFINE_BACKEND_RUNTIME_CONFIG_PATH: runtimeConfigPath,
        DEBUG: 'affine:*',
        FORCE_COLOR: 'true',
        DEBUG_COLORS: 'true',
@@ -0,0 +1,21 @@
+{
+  "storages": {
+    "blob.storage": {
+      "provider": "assetpack",
+      "bucket": "blobs",
+      "config": {
+        "path": "/tmp/affine-cloud-copilot-storage"
+      }
+    }
+  },
+  "copilot": {
+    "enabled": true,
+    "storage": {
+      "provider": "assetpack",
+      "bucket": "copilot",
+      "config": {
+        "path": "/tmp/affine-cloud-copilot-storage"
+      }
+    }
+  }
+}
@@ -308,6 +308,7 @@ export async function enableCloudWorkspace(page: Page) {
  await waitForAllPagesLoad(page);
  await dismissBlockingModal(page);
  await clickNewPageButton(page);
+  await waitForWorkspaceSynced(page);
 }

 export async function enableCloudWorkspaceFromShareButton(page: Page) {
@@ -324,6 +325,32 @@ export async function enableCloudWorkspaceFromShareButton(page: Page) {
  await waitForEditorLoad(page);
  await dismissBlockingModal(page);
  await clickNewPageButton(page);
+  await waitForWorkspaceSynced(page);
+}
+
+async function waitForWorkspaceSynced(page: Page) {
+  await page.evaluate(async () => {
+    const workspaceId = location.pathname.split('/')[2];
+    const workspace = (
+      window as typeof window & {
+        currentWorkspace?: {
+          engine: {
+            doc: {
+              waitForSynced(docId: string, abort: AbortSignal): Promise<void>;
+            };
+          };
+        };
+      }
+    ).currentWorkspace;
+    if (!workspaceId || !workspace) {
+      throw new Error('Cloud workspace is unavailable');
+    }
+    const abort = AbortSignal.timeout(60_000);
+    await Promise.all([
+      workspace.engine.doc.waitForSynced(workspaceId, abort),
+      workspace.engine.doc.waitForSynced('db$docProperties', abort),
+    ]);
+  });
 }

 export async function enableShare(page: Page) {