feat: refactor copilot module (#14537)

2026-07-19 19:16:29 +08:00 · 2026-03-02 13:57:55 +08:00
parent 60acd81d4b
commit c5d622531c
92 changed files with 5759 additions and 2170 deletions
@@ -43,7 +43,9 @@ Generated by [AVA](https://avajs.dev).
 > Snapshot 5

    Buffer @Uint8Array [
-      66616b65 20696d61 6765
+      89504e47 0d0a1a0a 0000000d 49484452 00000001 00000001 08040000 00b51c0c
+      02000000 0b494441 5478da63 fcff1f00 03030200 efa37c9f 00000000 49454e44
+      ae426082
    ]

 ## should preview link
@@ -12,12 +12,12 @@ Generated by [AVA](https://avajs.dev).
      {
        messages: [
          {
-            content: 'generate text to text',
+            content: 'generate text to text stream',
            role: 'assistant',
          },
        ],
        pinned: false,
-        tokens: 8,
+        tokens: 10,
      },
    ]

@@ -27,12 +27,12 @@ Generated by [AVA](https://avajs.dev).
      {
        messages: [
          {
-            content: 'generate text to text',
+            content: 'generate text to text stream',
            role: 'assistant',
          },
        ],
        pinned: false,
-        tokens: 8,
+        tokens: 10,
      },
    ]

@@ -4,31 +4,31 @@ import type { ExecutionContext, TestFn } from 'ava';
 import ava from 'ava';
 import { z } from 'zod';

-import { ServerFeature, ServerService } from '../core';
-import { AuthService } from '../core/auth';
-import { QuotaModule } from '../core/quota';
-import { Models } from '../models';
-import { CopilotModule } from '../plugins/copilot';
-import { prompts, PromptService } from '../plugins/copilot/prompt';
+import { ServerFeature, ServerService } from '../../core';
+import { AuthService } from '../../core/auth';
+import { QuotaModule } from '../../core/quota';
+import { Models } from '../../models';
+import { CopilotModule } from '../../plugins/copilot';
+import { prompts, PromptService } from '../../plugins/copilot/prompt';
 import {
  CopilotProviderFactory,
  CopilotProviderType,
  StreamObject,
  StreamObjectSchema,
-} from '../plugins/copilot/providers';
-import { TranscriptionResponseSchema } from '../plugins/copilot/transcript/types';
+} from '../../plugins/copilot/providers';
+import { TranscriptionResponseSchema } from '../../plugins/copilot/transcript/types';
 import {
  CopilotChatTextExecutor,
  CopilotWorkflowService,
  GraphExecutorState,
-} from '../plugins/copilot/workflow';
+} from '../../plugins/copilot/workflow';
 import {
  CopilotChatImageExecutor,
  CopilotCheckHtmlExecutor,
  CopilotCheckJsonExecutor,
-} from '../plugins/copilot/workflow/executor';
-import { createTestingModule, TestingModule } from './utils';
-import { TestAssets } from './utils/copilot';
+} from '../../plugins/copilot/workflow/executor';
+import { createTestingModule, TestingModule } from '../utils';
+import { TestAssets } from '../utils/copilot';

 type Tester = {
  auth: AuthService;
@@ -6,25 +6,25 @@ import type { TestFn } from 'ava';
 import ava from 'ava';
 import Sinon from 'sinon';

-import { AppModule } from '../app.module';
-import { JobQueue } from '../base';
-import { ConfigModule } from '../base/config';
-import { AuthService } from '../core/auth';
-import { DocReader } from '../core/doc';
-import { CopilotContextService } from '../plugins/copilot/context';
+import { AppModule } from '../../app.module';
+import { JobQueue } from '../../base';
+import { ConfigModule } from '../../base/config';
+import { AuthService } from '../../core/auth';
+import { DocReader } from '../../core/doc';
+import { CopilotContextService } from '../../plugins/copilot/context';
 import {
  CopilotEmbeddingJob,
  MockEmbeddingClient,
-} from '../plugins/copilot/embedding';
-import { prompts, PromptService } from '../plugins/copilot/prompt';
+} from '../../plugins/copilot/embedding';
+import { prompts, PromptService } from '../../plugins/copilot/prompt';
 import {
  CopilotProviderFactory,
  CopilotProviderType,
  GeminiGenerativeProvider,
  OpenAIProvider,
-} from '../plugins/copilot/providers';
-import { CopilotStorage } from '../plugins/copilot/storage';
-import { MockCopilotProvider } from './mocks';
+} from '../../plugins/copilot/providers';
+import { CopilotStorage } from '../../plugins/copilot/storage';
+import { MockCopilotProvider } from '../mocks';
 import {
  acceptInviteById,
  createTestingApp,
@@ -33,7 +33,7 @@ import {
  smallestPng,
  TestingApp,
  TestUser,
-} from './utils';
+} from '../utils';
 import {
  addContextDoc,
  addContextFile,
@@ -67,7 +67,7 @@ import {
  textToEventStream,
  unsplashSearch,
  updateCopilotSession,
-} from './utils/copilot';
+} from '../utils/copilot';

 const test = ava as TestFn<{
  auth: AuthService;
@@ -513,7 +513,11 @@ test('should be able to chat with api', async t => {
    );
    const messageId = await createCopilotMessage(app, sessionId);
    const ret = await chatWithText(app, sessionId, messageId);
-    t.is(ret, 'generate text to text', 'should be able to chat with text');
+    t.is(
+      ret,
+      'generate text to text stream',
+      'should be able to chat with text'
+    );

    const ret2 = await chatWithTextStream(app, sessionId, messageId);
    t.is(
@@ -657,7 +661,7 @@ test('should be able to retry with api', async t => {
    const histories = await getHistories(app, { workspaceId: id, docId });
    t.deepEqual(
      histories.map(h => h.messages.map(m => m.content)),
-      [['generate text to text', 'generate text to text']],
+      [['generate text to text stream', 'generate text to text stream']],
      'should be able to list history'
    );
  }
@@ -794,7 +798,7 @@ test('should be able to list history', async t => {
    const histories = await getHistories(app, { workspaceId, docId });
    t.deepEqual(
      histories.map(h => h.messages.map(m => m.content)),
-      [['hello', 'generate text to text']],
+      [['hello', 'generate text to text stream']],
      'should be able to list history'
    );
  }
@@ -807,7 +811,7 @@ test('should be able to list history', async t => {
    });
    t.deepEqual(
      histories.map(h => h.messages.map(m => m.content)),
-      [['generate text to text', 'hello']],
+      [['generate text to text stream', 'hello']],
      'should be able to list history'
    );
  }
@@ -858,7 +862,7 @@ test('should reject request that user have not permission', async t => {
    const histories = await getHistories(app, { workspaceId, docId });
    t.deepEqual(
      histories.map(h => h.messages.map(m => m.content)),
-      [['generate text to text']],
+      [['generate text to text stream']],
      'should able to list history'
    );

@@ -8,38 +8,38 @@ import ava from 'ava';
 import { nanoid } from 'nanoid';
 import Sinon from 'sinon';

-import { EventBus, JobQueue } from '../base';
-import { ConfigModule } from '../base/config';
-import { AuthService } from '../core/auth';
-import { QuotaModule } from '../core/quota';
-import { StorageModule, WorkspaceBlobStorage } from '../core/storage';
+import { EventBus, JobQueue } from '../../base';
+import { ConfigModule } from '../../base/config';
+import { AuthService } from '../../core/auth';
+import { QuotaModule } from '../../core/quota';
+import { StorageModule, WorkspaceBlobStorage } from '../../core/storage';
 import {
  ContextCategories,
  CopilotSessionModel,
  WorkspaceModel,
-} from '../models';
-import { CopilotModule } from '../plugins/copilot';
-import { CopilotContextService } from '../plugins/copilot/context';
-import { CopilotCronJobs } from '../plugins/copilot/cron';
+} from '../../models';
+import { CopilotModule } from '../../plugins/copilot';
+import { CopilotContextService } from '../../plugins/copilot/context';
+import { CopilotCronJobs } from '../../plugins/copilot/cron';
 import {
  CopilotEmbeddingJob,
  MockEmbeddingClient,
-} from '../plugins/copilot/embedding';
-import { prompts, PromptService } from '../plugins/copilot/prompt';
+} from '../../plugins/copilot/embedding';
+import { prompts, PromptService } from '../../plugins/copilot/prompt';
 import {
  CopilotProviderFactory,
  CopilotProviderType,
  ModelInputType,
  ModelOutputType,
  OpenAIProvider,
-} from '../plugins/copilot/providers';
+} from '../../plugins/copilot/providers';
 import {
  CitationParser,
  TextStreamParser,
-} from '../plugins/copilot/providers/utils';
-import { ChatSessionService } from '../plugins/copilot/session';
-import { CopilotStorage } from '../plugins/copilot/storage';
-import { CopilotTranscriptionService } from '../plugins/copilot/transcript';
+} from '../../plugins/copilot/providers/utils';
+import { ChatSessionService } from '../../plugins/copilot/session';
+import { CopilotStorage } from '../../plugins/copilot/storage';
+import { CopilotTranscriptionService } from '../../plugins/copilot/transcript';
 import {
  CopilotChatTextExecutor,
  CopilotWorkflowService,
@@ -48,7 +48,7 @@ import {
  WorkflowGraphExecutor,
  type WorkflowNodeData,
  WorkflowNodeType,
-} from '../plugins/copilot/workflow';
+} from '../../plugins/copilot/workflow';
 import {
  CopilotChatImageExecutor,
  CopilotCheckHtmlExecutor,
@@ -56,16 +56,16 @@ import {
  getWorkflowExecutor,
  NodeExecuteState,
  NodeExecutorType,
-} from '../plugins/copilot/workflow/executor';
-import { AutoRegisteredWorkflowExecutor } from '../plugins/copilot/workflow/executor/utils';
-import { WorkflowGraphList } from '../plugins/copilot/workflow/graph';
-import { CopilotWorkspaceService } from '../plugins/copilot/workspace';
-import { PaymentModule } from '../plugins/payment';
-import { SubscriptionService } from '../plugins/payment/service';
-import { SubscriptionStatus } from '../plugins/payment/types';
-import { MockCopilotProvider } from './mocks';
-import { createTestingModule, TestingModule } from './utils';
-import { WorkflowTestCases } from './utils/copilot';
+} from '../../plugins/copilot/workflow/executor';
+import { AutoRegisteredWorkflowExecutor } from '../../plugins/copilot/workflow/executor/utils';
+import { WorkflowGraphList } from '../../plugins/copilot/workflow/graph';
+import { CopilotWorkspaceService } from '../../plugins/copilot/workspace';
+import { PaymentModule } from '../../plugins/payment';
+import { SubscriptionService } from '../../plugins/payment/service';
+import { SubscriptionStatus } from '../../plugins/payment/types';
+import { MockCopilotProvider } from '../mocks';
+import { createTestingModule, TestingModule } from '../utils';
+import { WorkflowTestCases } from '../utils/copilot';

 type Context = {
  auth: AuthService;
@@ -364,6 +364,21 @@ test('should be able to manage chat session', async t => {
    });
    t.is(newSessionId, sessionId, 'should get same session id');
  }
+
+  // should create a fresh session when reuseLatestChat is explicitly disabled
+  {
+    const newSessionId = await session.create({
+      userId,
+      promptName,
+      ...commonParams,
+      reuseLatestChat: false,
+    });
+    t.not(
+      newSessionId,
+      sessionId,
+      'should create new session id when reuseLatestChat is false'
+    );
+  }
 });

 test('should be able to update chat session prompt', async t => {
@@ -881,6 +896,26 @@ test('should be able to get provider', async t => {
  }
 });

+test('should resolve provider by prefixed model id', async t => {
+  const { factory } = t.context;
+
+  const provider = await factory.getProviderByModel('openai-default/test');
+  t.truthy(provider, 'should resolve prefixed model id');
+  t.is(provider?.type, CopilotProviderType.OpenAI);
+
+  const result = await provider?.text({ modelId: 'openai-default/test' }, [
+    { role: 'user', content: 'hello' },
+  ]);
+  t.is(result, 'generate text to text');
+});
+
+test('should fallback to null when prefixed provider id does not exist', async t => {
+  const { factory } = t.context;
+
+  const provider = await factory.getProviderByModel('unknown/test');
+  t.is(provider, null);
+});
+
 // ==================== workflow ====================

 // this test used to preview the final result of the workflow
@@ -2063,25 +2098,23 @@ test('should handle copilot cron jobs correctly', async t => {
 });

 test('should resolve model correctly based on subscription status and prompt config', async t => {
-  const { db, session, subscription } = t.context;
+  const { prompt, session, subscription } = t.context;

  // 1) Seed a prompt that has optionalModels and proModels in config
  const promptName = 'resolve-model-test';
-  await db.aiPrompt.create({
-    data: {
-      name: promptName,
-      model: 'gemini-2.5-flash',
-      messages: {
-        create: [{ idx: 0, role: 'system', content: 'test' }],
-      },
-      config: { proModels: ['gemini-2.5-pro', 'claude-sonnet-4-5@20250929'] },
+  await prompt.set(
+    promptName,
+    'gemini-2.5-flash',
+    [{ role: 'system', content: 'test' }],
+    { proModels: ['gemini-2.5-pro', 'claude-sonnet-4-5@20250929'] },
+    {
      optionalModels: [
        'gemini-2.5-flash',
        'gemini-2.5-pro',
        'claude-sonnet-4-5@20250929',
      ],
-    },
-  });
+    }
+  );

  // 2) Create a chat session with this prompt
  const sessionId = await session.create({
@@ -2106,6 +2139,16 @@ test('should resolve model correctly based on subscription status and prompt con
    const model1 = await s.resolveModel(false, 'gemini-2.5-pro');
    t.snapshot(model1, 'should honor requested pro model');

+    const model1WithPrefix = await s.resolveModel(
+      false,
+      'openai-default/gemini-2.5-pro'
+    );
+    t.is(
+      model1WithPrefix,
+      'openai-default/gemini-2.5-pro',
+      'should honor requested prefixed pro model'
+    );
+
    const model2 = await s.resolveModel(false, 'not-in-optional');
    t.snapshot(model2, 'should fallback to default model');
  }
@@ -2119,6 +2162,16 @@ test('should resolve model correctly based on subscription status and prompt con
      'should fallback to default model when requesting pro model during trialing'
    );

+    const model3WithPrefix = await s.resolveModel(
+      true,
+      'openai-default/gemini-2.5-pro'
+    );
+    t.is(
+      model3WithPrefix,
+      'gemini-2.5-flash',
+      'should fallback to default model when requesting prefixed pro model during trialing'
+    );
+
    const model4 = await s.resolveModel(true, 'gemini-2.5-flash');
    t.snapshot(model4, 'should honor requested non-pro model during trialing');

@@ -2141,6 +2194,16 @@ test('should resolve model correctly based on subscription status and prompt con
    const model7 = await s.resolveModel(true, 'claude-sonnet-4-5@20250929');
    t.snapshot(model7, 'should honor requested pro model during active');

+    const model7WithPrefix = await s.resolveModel(
+      true,
+      'openai-default/claude-sonnet-4-5@20250929'
+    );
+    t.is(
+      model7WithPrefix,
+      'openai-default/claude-sonnet-4-5@20250929',
+      'should honor requested prefixed pro model during active'
+    );
+
    const model8 = await s.resolveModel(true, 'not-in-optional');
    t.snapshot(
      model8,
@@ -0,0 +1,210 @@
+import test from 'ava';
+import { z } from 'zod';
+
+import type { NativeLlmRequest, NativeLlmStreamEvent } from '../../native';
+import {
+  buildNativeRequest,
+  NativeProviderAdapter,
+} from '../../plugins/copilot/providers/native';
+
+const mockDispatch = () =>
+  (async function* (): AsyncIterableIterator<NativeLlmStreamEvent> {
+    yield { type: 'text_delta', text: 'Use [^1] now' };
+    yield { type: 'citation', index: 1, url: 'https://affine.pro' };
+    yield { type: 'done', finish_reason: 'stop' };
+  })();
+
+test('NativeProviderAdapter streamText should append citation footnotes', async t => {
+  const adapter = new NativeProviderAdapter(mockDispatch, {}, 3);
+  const chunks: string[] = [];
+  for await (const chunk of adapter.streamText({
+    model: 'gpt-4.1',
+    stream: true,
+    messages: [{ role: 'user', content: [{ type: 'text', text: 'hi' }] }],
+  })) {
+    chunks.push(chunk);
+  }
+
+  const text = chunks.join('');
+  t.true(text.includes('Use [^1] now'));
+  t.true(
+    text.includes('[^1]: {"type":"url","url":"https%3A%2F%2Faffine.pro"}')
+  );
+});
+
+test('NativeProviderAdapter streamObject should append citation footnotes', async t => {
+  const adapter = new NativeProviderAdapter(mockDispatch, {}, 3);
+  const chunks = [];
+  for await (const chunk of adapter.streamObject({
+    model: 'gpt-4.1',
+    stream: true,
+    messages: [{ role: 'user', content: [{ type: 'text', text: 'hi' }] }],
+  })) {
+    chunks.push(chunk);
+  }
+
+  t.deepEqual(
+    chunks.map(chunk => chunk.type),
+    ['text-delta', 'text-delta']
+  );
+  const text = chunks
+    .filter(chunk => chunk.type === 'text-delta')
+    .map(chunk => chunk.textDelta)
+    .join('');
+  t.true(text.includes('Use [^1] now'));
+  t.true(
+    text.includes('[^1]: {"type":"url","url":"https%3A%2F%2Faffine.pro"}')
+  );
+});
+
+test('NativeProviderAdapter streamObject should append fallback attachment footnotes', async t => {
+  const dispatch = () =>
+    (async function* (): AsyncIterableIterator<NativeLlmStreamEvent> {
+      yield {
+        type: 'tool_result',
+        call_id: 'call_1',
+        name: 'blob_read',
+        arguments: { blob_id: 'blob_1' },
+        output: {
+          blobId: 'blob_1',
+          fileName: 'a.txt',
+          fileType: 'text/plain',
+          content: 'A',
+        },
+      };
+      yield {
+        type: 'tool_result',
+        call_id: 'call_2',
+        name: 'blob_read',
+        arguments: { blob_id: 'blob_2' },
+        output: {
+          blobId: 'blob_2',
+          fileName: 'b.txt',
+          fileType: 'text/plain',
+          content: 'B',
+        },
+      };
+      yield { type: 'text_delta', text: 'Answer from files.' };
+      yield { type: 'done', finish_reason: 'stop' };
+    })();
+
+  const adapter = new NativeProviderAdapter(dispatch, {}, 3);
+  const chunks = [];
+  for await (const chunk of adapter.streamObject({
+    model: 'gpt-4.1',
+    stream: true,
+    messages: [{ role: 'user', content: [{ type: 'text', text: 'hi' }] }],
+  })) {
+    chunks.push(chunk);
+  }
+
+  const text = chunks
+    .filter(chunk => chunk.type === 'text-delta')
+    .map(chunk => chunk.textDelta)
+    .join('');
+  t.true(text.includes('Answer from files.'));
+  t.true(text.includes('[^1][^2]'));
+  t.true(
+    text.includes(
+      '[^1]: {"type":"attachment","blobId":"blob_1","fileName":"a.txt","fileType":"text/plain"}'
+    )
+  );
+  t.true(
+    text.includes(
+      '[^2]: {"type":"attachment","blobId":"blob_2","fileName":"b.txt","fileType":"text/plain"}'
+    )
+  );
+});
+
+test('NativeProviderAdapter streamObject should map tool and text events', async t => {
+  let round = 0;
+  const dispatch = (_request: NativeLlmRequest) =>
+    (async function* (): AsyncIterableIterator<NativeLlmStreamEvent> {
+      round += 1;
+      if (round === 1) {
+        yield {
+          type: 'tool_call',
+          call_id: 'call_1',
+          name: 'doc_read',
+          arguments: { doc_id: 'a1' },
+        };
+        yield { type: 'done', finish_reason: 'tool_calls' };
+        return;
+      }
+      yield { type: 'text_delta', text: 'ok' };
+      yield { type: 'done', finish_reason: 'stop' };
+    })();
+
+  const adapter = new NativeProviderAdapter(
+    dispatch,
+    {
+      doc_read: {
+        inputSchema: z.object({ doc_id: z.string() }),
+        execute: async () => ({ markdown: '# a1' }),
+      },
+    },
+    4
+  );
+
+  const events = [];
+  for await (const event of adapter.streamObject({
+    model: 'gpt-4.1',
+    stream: true,
+    messages: [{ role: 'user', content: [{ type: 'text', text: 'read' }] }],
+  })) {
+    events.push(event);
+  }
+
+  t.deepEqual(
+    events.map(event => event.type),
+    ['tool-call', 'tool-result', 'text-delta']
+  );
+  t.deepEqual(events[0], {
+    type: 'tool-call',
+    toolCallId: 'call_1',
+    toolName: 'doc_read',
+    args: { doc_id: 'a1' },
+  });
+});
+
+test('buildNativeRequest should include rust middleware from profile', async t => {
+  const { request } = await buildNativeRequest({
+    model: 'gpt-4.1',
+    messages: [{ role: 'user', content: 'hello' }],
+    tools: {},
+    middleware: {
+      rust: {
+        request: ['normalize_messages', 'clamp_max_tokens'],
+        stream: ['stream_event_normalize', 'citation_indexing'],
+      },
+      node: {
+        text: ['callout'],
+      },
+    },
+  });
+
+  t.deepEqual(request.middleware, {
+    request: ['normalize_messages', 'clamp_max_tokens'],
+    stream: ['stream_event_normalize', 'citation_indexing'],
+  });
+});
+
+test('NativeProviderAdapter streamText should skip citation footnotes when disabled', async t => {
+  const adapter = new NativeProviderAdapter(mockDispatch, {}, 3, {
+    nodeTextMiddleware: ['callout'],
+  });
+  const chunks: string[] = [];
+  for await (const chunk of adapter.streamText({
+    model: 'gpt-4.1',
+    stream: true,
+    messages: [{ role: 'user', content: [{ type: 'text', text: 'hi' }] }],
+  })) {
+    chunks.push(chunk);
+  }
+
+  const text = chunks.join('');
+  t.true(text.includes('Use [^1] now'));
+  t.false(
+    text.includes('[^1]: {"type":"url","url":"https%3A%2F%2Faffine.pro"}')
+  );
+});
@@ -0,0 +1,56 @@
+import test from 'ava';
+
+import { resolveProviderMiddleware } from '../../plugins/copilot/providers/provider-middleware';
+import { buildProviderRegistry } from '../../plugins/copilot/providers/provider-registry';
+import { CopilotProviderType } from '../../plugins/copilot/providers/types';
+
+test('resolveProviderMiddleware should include anthropic defaults', t => {
+  const middleware = resolveProviderMiddleware(CopilotProviderType.Anthropic);
+
+  t.deepEqual(middleware.rust?.request, [
+    'normalize_messages',
+    'tool_schema_rewrite',
+  ]);
+  t.deepEqual(middleware.rust?.stream, [
+    'stream_event_normalize',
+    'citation_indexing',
+  ]);
+  t.deepEqual(middleware.node?.text, ['citation_footnote', 'callout']);
+});
+
+test('resolveProviderMiddleware should merge defaults and overrides', t => {
+  const middleware = resolveProviderMiddleware(CopilotProviderType.OpenAI, {
+    rust: { request: ['clamp_max_tokens'] },
+    node: { text: ['thinking_format'] },
+  });
+
+  t.deepEqual(middleware.rust?.request, [
+    'normalize_messages',
+    'clamp_max_tokens',
+  ]);
+  t.deepEqual(middleware.node?.text, [
+    'citation_footnote',
+    'callout',
+    'thinking_format',
+  ]);
+});
+
+test('buildProviderRegistry should normalize profile middleware defaults', t => {
+  const registry = buildProviderRegistry({
+    profiles: [
+      {
+        id: 'openai-main',
+        type: CopilotProviderType.OpenAI,
+        config: { apiKey: '1' },
+      },
+    ],
+  });
+
+  const profile = registry.profiles.get('openai-main');
+  t.truthy(profile);
+  t.deepEqual(profile?.middleware.rust?.stream, [
+    'stream_event_normalize',
+    'citation_indexing',
+  ]);
+  t.deepEqual(profile?.middleware.node?.text, ['citation_footnote', 'callout']);
+});
@@ -0,0 +1,99 @@
+import test from 'ava';
+
+import { ProviderMiddlewareConfig } from '../../plugins/copilot/config';
+import { CopilotProvider } from '../../plugins/copilot/providers/provider';
+import {
+  CopilotProviderType,
+  ModelInputType,
+  ModelOutputType,
+} from '../../plugins/copilot/providers/types';
+
+class TestOpenAIProvider extends CopilotProvider<{ apiKey: string }> {
+  readonly type = CopilotProviderType.OpenAI;
+  readonly models = [
+    {
+      id: 'gpt-4.1',
+      capabilities: [
+        {
+          input: [ModelInputType.Text],
+          output: [ModelOutputType.Text],
+          defaultForOutputType: true,
+        },
+      ],
+    },
+  ];
+
+  configured() {
+    return true;
+  }
+
+  async text(_cond: any, _messages: any[], _options?: any) {
+    return '';
+  }
+
+  async *streamText(_cond: any, _messages: any[], _options?: any) {
+    yield '';
+  }
+
+  exposeMetricLabels() {
+    return this.metricLabels('gpt-4.1');
+  }
+
+  exposeMiddleware() {
+    return this.getActiveProviderMiddleware();
+  }
+}
+
+function createProvider(profileMiddleware?: ProviderMiddlewareConfig) {
+  const provider = new TestOpenAIProvider();
+  (provider as any).AFFiNEConfig = {
+    copilot: {
+      providers: {
+        profiles: [
+          {
+            id: 'openai-main',
+            type: CopilotProviderType.OpenAI,
+            config: { apiKey: 'test' },
+            middleware: profileMiddleware,
+          },
+        ],
+        defaults: {},
+        openai: { apiKey: 'legacy' },
+      },
+    },
+  };
+  return provider;
+}
+
+test('metricLabels should include active provider id', t => {
+  const provider = createProvider();
+  const labels = provider.runWithProfile('openai-main', () =>
+    provider.exposeMetricLabels()
+  );
+  t.is(labels.providerId, 'openai-main');
+});
+
+test('getActiveProviderMiddleware should merge defaults with profile override', t => {
+  const provider = createProvider({
+    rust: { request: ['clamp_max_tokens'] },
+    node: { text: ['thinking_format'] },
+  });
+
+  const middleware = provider.runWithProfile('openai-main', () =>
+    provider.exposeMiddleware()
+  );
+
+  t.deepEqual(middleware.rust?.request, [
+    'normalize_messages',
+    'clamp_max_tokens',
+  ]);
+  t.deepEqual(middleware.rust?.stream, [
+    'stream_event_normalize',
+    'citation_indexing',
+  ]);
+  t.deepEqual(middleware.node?.text, [
+    'citation_footnote',
+    'callout',
+    'thinking_format',
+  ]);
+});
@@ -0,0 +1,165 @@
+import test from 'ava';
+
+import {
+  buildProviderRegistry,
+  resolveModel,
+  stripProviderPrefix,
+} from '../../plugins/copilot/providers/provider-registry';
+import {
+  CopilotProviderType,
+  ModelOutputType,
+} from '../../plugins/copilot/providers/types';
+
+test('buildProviderRegistry should keep explicit profile over legacy compatibility profile', t => {
+  const registry = buildProviderRegistry({
+    profiles: [
+      {
+        id: 'openai-default',
+        type: CopilotProviderType.OpenAI,
+        priority: 100,
+        config: { apiKey: 'new' },
+      },
+    ],
+    openai: { apiKey: 'legacy' },
+  });
+
+  const profile = registry.profiles.get('openai-default');
+  t.truthy(profile);
+  t.deepEqual(profile?.config, { apiKey: 'new' });
+});
+
+test('buildProviderRegistry should reject duplicated profile ids', t => {
+  const error = t.throws(() =>
+    buildProviderRegistry({
+      profiles: [
+        {
+          id: 'openai-main',
+          type: CopilotProviderType.OpenAI,
+          config: { apiKey: '1' },
+        },
+        {
+          id: 'openai-main',
+          type: CopilotProviderType.OpenAI,
+          config: { apiKey: '2' },
+        },
+      ],
+    })
+  ) as Error;
+
+  t.truthy(error);
+  t.regex(error.message, /Duplicated copilot provider profile id/);
+});
+
+test('buildProviderRegistry should reject defaults that reference unknown providers', t => {
+  const error = t.throws(() =>
+    buildProviderRegistry({
+      profiles: [
+        {
+          id: 'openai-main',
+          type: CopilotProviderType.OpenAI,
+          config: { apiKey: '1' },
+        },
+      ],
+      defaults: {
+        fallback: 'unknown-provider',
+      },
+    })
+  ) as Error;
+
+  t.truthy(error);
+  t.regex(error.message, /defaults references unknown providerId/);
+});
+
+test('resolveModel should support explicit provider prefix and keep slash models untouched', t => {
+  const registry = buildProviderRegistry({
+    profiles: [
+      {
+        id: 'openai-main',
+        type: CopilotProviderType.OpenAI,
+        config: { apiKey: '1' },
+      },
+      {
+        id: 'fal-main',
+        type: CopilotProviderType.FAL,
+        config: { apiKey: '2' },
+      },
+    ],
+  });
+
+  const prefixed = resolveModel({
+    registry,
+    modelId: 'openai-main/gpt-4.1',
+  });
+  t.deepEqual(prefixed, {
+    rawModelId: 'openai-main/gpt-4.1',
+    modelId: 'gpt-4.1',
+    explicitProviderId: 'openai-main',
+    candidateProviderIds: ['openai-main'],
+  });
+
+  const slashModel = resolveModel({
+    registry,
+    modelId: 'lora/image-to-image',
+  });
+  t.is(slashModel.modelId, 'lora/image-to-image');
+  t.false(slashModel.candidateProviderIds.includes('lora'));
+});
+
+test('resolveModel should follow defaults -> fallback -> order and apply filters', t => {
+  const registry = buildProviderRegistry({
+    profiles: [
+      {
+        id: 'openai-main',
+        type: CopilotProviderType.OpenAI,
+        priority: 10,
+        config: { apiKey: '1' },
+      },
+      {
+        id: 'anthropic-main',
+        type: CopilotProviderType.Anthropic,
+        priority: 5,
+        config: { apiKey: '2' },
+      },
+      {
+        id: 'fal-main',
+        type: CopilotProviderType.FAL,
+        priority: 1,
+        config: { apiKey: '3' },
+      },
+    ],
+    defaults: {
+      [ModelOutputType.Text]: 'anthropic-main',
+      fallback: 'openai-main',
+    },
+  });
+
+  const routed = resolveModel({
+    registry,
+    outputType: ModelOutputType.Text,
+    preferredProviderIds: ['openai-main', 'fal-main'],
+  });
+
+  t.deepEqual(routed.candidateProviderIds, ['openai-main', 'fal-main']);
+});
+
+test('stripProviderPrefix should only strip matched provider prefix', t => {
+  const registry = buildProviderRegistry({
+    profiles: [
+      {
+        id: 'openai-main',
+        type: CopilotProviderType.OpenAI,
+        config: { apiKey: '1' },
+      },
+    ],
+  });
+
+  t.is(
+    stripProviderPrefix(registry, 'openai-main', 'openai-main/gpt-4.1'),
+    'gpt-4.1'
+  );
+  t.is(
+    stripProviderPrefix(registry, 'openai-main', 'another-main/gpt-4.1'),
+    'another-main/gpt-4.1'
+  );
+  t.is(stripProviderPrefix(registry, 'openai-main', 'gpt-4.1'), 'gpt-4.1');
+});
@@ -0,0 +1,134 @@
+import test from 'ava';
+import { z } from 'zod';
+
+import { NativeLlmRequest, NativeLlmStreamEvent } from '../../native';
+import {
+  ToolCallAccumulator,
+  ToolCallLoop,
+  ToolSchemaExtractor,
+} from '../../plugins/copilot/providers/loop';
+
+test('ToolCallAccumulator should merge deltas and complete tool call', t => {
+  const accumulator = new ToolCallAccumulator();
+
+  accumulator.feedDelta({
+    type: 'tool_call_delta',
+    call_id: 'call_1',
+    name: 'doc_read',
+    arguments_delta: '{"doc_id":"',
+  });
+  accumulator.feedDelta({
+    type: 'tool_call_delta',
+    call_id: 'call_1',
+    arguments_delta: 'a1"}',
+  });
+
+  const completed = accumulator.complete({
+    type: 'tool_call',
+    call_id: 'call_1',
+    name: 'doc_read',
+    arguments: { doc_id: 'a1' },
+  });
+
+  t.deepEqual(completed, {
+    id: 'call_1',
+    name: 'doc_read',
+    args: { doc_id: 'a1' },
+    thought: undefined,
+  });
+});
+
+test('ToolSchemaExtractor should convert zod schema to json schema', t => {
+  const toolSet = {
+    doc_read: {
+      description: 'Read doc',
+      inputSchema: z.object({
+        doc_id: z.string(),
+        limit: z.number().optional(),
+      }),
+      execute: async () => ({}),
+    },
+  };
+
+  const extracted = ToolSchemaExtractor.extract(toolSet);
+
+  t.deepEqual(extracted, [
+    {
+      name: 'doc_read',
+      description: 'Read doc',
+      parameters: {
+        type: 'object',
+        properties: {
+          doc_id: { type: 'string' },
+          limit: { type: 'number' },
+        },
+        additionalProperties: false,
+        required: ['doc_id'],
+      },
+    },
+  ]);
+});
+
+test('ToolCallLoop should execute tool call and continue to next round', async t => {
+  const dispatchRequests: NativeLlmRequest[] = [];
+
+  const dispatch = (request: NativeLlmRequest) => {
+    dispatchRequests.push(request);
+    const round = dispatchRequests.length;
+
+    return (async function* (): AsyncIterableIterator<NativeLlmStreamEvent> {
+      if (round === 1) {
+        yield {
+          type: 'tool_call_delta',
+          call_id: 'call_1',
+          name: 'doc_read',
+          arguments_delta: '{"doc_id":"a1"}',
+        };
+        yield {
+          type: 'tool_call',
+          call_id: 'call_1',
+          name: 'doc_read',
+          arguments: { doc_id: 'a1' },
+        };
+        yield { type: 'done', finish_reason: 'tool_calls' };
+        return;
+      }
+
+      yield { type: 'text_delta', text: 'done' };
+      yield { type: 'done', finish_reason: 'stop' };
+    })();
+  };
+
+  let executedArgs: Record<string, unknown> | null = null;
+  const loop = new ToolCallLoop(
+    dispatch,
+    {
+      doc_read: {
+        inputSchema: z.object({ doc_id: z.string() }),
+        execute: async args => {
+          executedArgs = args;
+          return { markdown: '# doc' };
+        },
+      },
+    },
+    4
+  );
+
+  const events: NativeLlmStreamEvent[] = [];
+  for await (const event of loop.run({
+    model: 'gpt-4.1',
+    stream: true,
+    messages: [{ role: 'user', content: [{ type: 'text', text: 'read doc' }] }],
+  })) {
+    events.push(event);
+  }
+
+  t.deepEqual(executedArgs, { doc_id: 'a1' });
+  t.true(
+    dispatchRequests[1]?.messages.some(message => message.role === 'tool')
+  );
+  t.deepEqual(
+    events.map(event => event.type),
+    ['tool_call', 'tool_result', 'text_delta', 'done']
+  );
+});
@@ -0,0 +1,116 @@
+import test from 'ava';
+import { z } from 'zod';
+
+import {
+  chatToGPTMessage,
+  CitationFootnoteFormatter,
+  CitationParser,
+  StreamPatternParser,
+} from '../../plugins/copilot/providers/utils';
+
+test('CitationFootnoteFormatter should format sorted footnotes from citation events', t => {
+  const formatter = new CitationFootnoteFormatter();
+
+  formatter.consume({
+    type: 'citation',
+    index: 2,
+    url: 'https://example.com/b',
+  });
+  formatter.consume({
+    type: 'citation',
+    index: 1,
+    url: 'https://example.com/a',
+  });
+
+  t.is(
+    formatter.end(),
+    [
+      '[^1]: {"type":"url","url":"https%3A%2F%2Fexample.com%2Fa"}',
+      '[^2]: {"type":"url","url":"https%3A%2F%2Fexample.com%2Fb"}',
+    ].join('\n')
+  );
+});
+
+test('CitationFootnoteFormatter should overwrite duplicated index with latest url', t => {
+  const formatter = new CitationFootnoteFormatter();
+
+  formatter.consume({
+    type: 'citation',
+    index: 1,
+    url: 'https://example.com/old',
+  });
+  formatter.consume({
+    type: 'citation',
+    index: 1,
+    url: 'https://example.com/new',
+  });
+
+  t.is(
+    formatter.end(),
+    '[^1]: {"type":"url","url":"https%3A%2F%2Fexample.com%2Fnew"}'
+  );
+});
+
+test('StreamPatternParser should keep state across chunks', t => {
+  const parser = new StreamPatternParser(pattern => {
+    if (pattern.kind === 'wrappedLink') {
+      return `[^${pattern.url}]`;
+    }
+    if (pattern.kind === 'index') {
+      return `[#${pattern.value}]`;
+    }
+    return `[${pattern.text}](${pattern.url})`;
+  });
+
+  const first = parser.write('ref ([AFFiNE](https://affine.pro');
+  const second = parser.write(')) and [2]');
+
+  t.is(first, 'ref ');
+  t.is(second, '[^https://affine.pro] and [#2]');
+  t.is(parser.end(), '');
+});
+
+test('CitationParser should convert wrapped links to numbered footnotes', t => {
+  const parser = new CitationParser();
+
+  const output = parser.parse('Use ([AFFiNE](https://affine.pro)) now');
+  t.is(output, 'Use [^1] now');
+  t.regex(
+    parser.end(),
+    /\[\^1\]: \{"type":"url","url":"https%3A%2F%2Faffine.pro"\}/
+  );
+});
+
+test('chatToGPTMessage should not mutate input and should keep system schema', async t => {
+  const schema = z.object({
+    query: z.string(),
+  });
+  const messages = [
+    {
+      role: 'system' as const,
+      content: 'You are helper',
+      params: { schema },
+    },
+    {
+      role: 'user' as const,
+      content: '',
+      attachments: ['https://example.com/a.png'],
+    },
+  ];
+  const firstRef = messages[0];
+  const secondRef = messages[1];
+  const [system, normalized, parsedSchema] = await chatToGPTMessage(
+    messages,
+    false
+  );
+
+  t.is(system, 'You are helper');
+  t.is(parsedSchema, schema);
+  t.is(messages.length, 2);
+  t.is(messages[0], firstRef);
+  t.is(messages[1], secondRef);
+  t.deepEqual(normalized[0], {
+    role: 'user',
+    content: [{ type: 'text', text: '[no content]' }],
+  });
+});
@@ -0,0 +1,82 @@
+import test from 'ava';
+
+import { NativeStreamAdapter } from '../native';
+
+test('NativeStreamAdapter should support buffered and awaited consumption', async t => {
+  const adapter = new NativeStreamAdapter<number>(undefined);
+
+  adapter.push(1);
+  const first = await adapter.next();
+  t.deepEqual(first, { value: 1, done: false });
+
+  const pending = adapter.next();
+  adapter.push(2);
+  const second = await pending;
+  t.deepEqual(second, { value: 2, done: false });
+
+  adapter.push(null);
+  const done = await adapter.next();
+  t.true(done.done);
+});
+
+test('NativeStreamAdapter return should abort handle and end iteration', async t => {
+  let abortCount = 0;
+  const adapter = new NativeStreamAdapter<number>({
+    abort: () => {
+      abortCount += 1;
+    },
+  });
+
+  const ended = await adapter.return();
+  t.is(abortCount, 1);
+  t.true(ended.done);
+
+  const secondReturn = await adapter.return();
+  t.true(secondReturn.done);
+  t.is(abortCount, 1);
+
+  const next = await adapter.next();
+  t.true(next.done);
+});
+
+test('NativeStreamAdapter should abort when AbortSignal is triggered', async t => {
+  let abortCount = 0;
+  const controller = new AbortController();
+  const adapter = new NativeStreamAdapter<number>(
+    {
+      abort: () => {
+        abortCount += 1;
+      },
+    },
+    controller.signal
+  );
+
+  const pending = adapter.next();
+  controller.abort();
+  const done = await pending;
+  t.true(done.done);
+  t.is(abortCount, 1);
+});
+
+test('NativeStreamAdapter should end immediately for pre-aborted signal', async t => {
+  let abortCount = 0;
+  const controller = new AbortController();
+  controller.abort();
+
+  const adapter = new NativeStreamAdapter<number>(
+    {
+      abort: () => {
+        abortCount += 1;
+      },
+    },
+    controller.signal
+  );
+
+  const next = await adapter.next();
+  t.true(next.done);
+  t.is(abortCount, 1);
+
+  adapter.push(1);
+  const stillDone = await adapter.next();
+  t.true(stillDone.done);
+});
@@ -629,14 +629,35 @@ export async function chatWithText(
  prefix = '',
  retry?: boolean
 ): Promise<string> {
+  const endpoint = prefix || '/stream';
  const query = messageId
    ? `?messageId=${messageId}` + (retry ? '&retry=true' : '')
    : '';
  const res = await app
-    .GET(`/api/copilot/chat/${sessionId}${prefix}${query}`)
+    .GET(`/api/copilot/chat/${sessionId}${endpoint}${query}`)
    .expect(200);

-  return res.text;
+  if (prefix) {
+    return res.text;
+  }
+
+  const events = sse2array(res.text);
+  const errorEvent = events.find(event => event.event === 'error');
+  if (errorEvent?.data) {
+    let message = errorEvent.data;
+    try {
+      const parsed = JSON.parse(errorEvent.data);
+      message = parsed.message || message;
+    } catch {
+      // noop: keep raw error data
+    }
+    throw new Error(message);
+  }
+
+  return events
+    .filter(event => event.event === 'message')
+    .map(event => event.data ?? '')
+    .join('');
 }

 export async function chatWithTextStream(
@@ -38,8 +38,11 @@ test.before(async t => {
  t.context.app = app;
 });

-test.after.always(async t => {
+test.afterEach.always(() => {
  Sinon.restore();
+});
+
+test.after.always(async t => {
  __resetDnsLookupForTests();
  await t.context.app.close();
 });
@@ -80,6 +83,7 @@ const assertAndSnapshotRaw = async (

 test('should proxy image', async t => {
  const assertAndSnapshot = assertAndSnapshotRaw.bind(null, t);
+  const imageUrl = `http://example.com/image-${Date.now()}.png`;

  await assertAndSnapshot(
    '/api/worker/image-proxy',
@@ -105,7 +109,7 @@ test('should proxy image', async t => {

  {
    await assertAndSnapshot(
-      '/api/worker/image-proxy?url=http://example.com/image.png',
+      `/api/worker/image-proxy?url=${imageUrl}`,
      'should return 400 if origin and referer are missing',
      { status: 400, origin: null, referer: null }
    );
@@ -113,14 +117,17 @@ test('should proxy image', async t => {

  {
    await assertAndSnapshot(
-      '/api/worker/image-proxy?url=http://example.com/image.png',
+      `/api/worker/image-proxy?url=${imageUrl}`,
      'should return 400 for invalid origin header',
      { status: 400, origin: 'http://invalid.com' }
    );
  }

  {
-    const fakeBuffer = Buffer.from('fake image');
+    const fakeBuffer = Buffer.from(
+      'iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAQAAAC1HAwCAAAAC0lEQVR42mP8/x8AAwMCAO+jfJ8AAAAASUVORK5CYII=',
+      'base64'
+    );
    const fakeResponse = new Response(fakeBuffer, {
      status: 200,
      headers: {
@@ -130,13 +137,14 @@ test('should proxy image', async t => {
    });

    const fetchSpy = Sinon.stub(global, 'fetch').resolves(fakeResponse);
-
-    await assertAndSnapshot(
-      '/api/worker/image-proxy?url=http://example.com/image.png',
-      'should return image buffer'
-    );
-
-    fetchSpy.restore();
+    try {
+      await assertAndSnapshot(
+        `/api/worker/image-proxy?url=${imageUrl}`,
+        'should return image buffer'
+      );
+    } finally {
+      fetchSpy.restore();
+    }
  }
 });

@@ -200,18 +208,19 @@ test('should preview link', async t => {
    });

    const fetchSpy = Sinon.stub(global, 'fetch').resolves(fakeHTML);
-
-    await assertAndSnapshot(
-      '/api/worker/link-preview',
-      'should process a valid external URL and return link preview data',
-      {
-        status: 200,
-        method: 'POST',
-        body: { url: 'http://external.com/page' },
-      }
-    );
-
-    fetchSpy.restore();
+    try {
+      await assertAndSnapshot(
+        '/api/worker/link-preview',
+        'should process a valid external URL and return link preview data',
+        {
+          status: 200,
+          method: 'POST',
+          body: { url: 'http://external.com/page' },
+        }
+      );
+    } finally {
+      fetchSpy.restore();
+    }
  }

  {
@@ -251,18 +260,19 @@ test('should preview link', async t => {
      });

      const fetchSpy = Sinon.stub(global, 'fetch').resolves(fakeHTML);
-
-      await assertAndSnapshot(
-        '/api/worker/link-preview',
-        'should decode HTML content with charset',
-        {
-          status: 200,
-          method: 'POST',
-          body: { url: `http://example.com/${charset}` },
-        }
-      );
-
-      fetchSpy.restore();
+      try {
+        await assertAndSnapshot(
+          '/api/worker/link-preview',
+          'should decode HTML content with charset',
+          {
+            status: 200,
+            method: 'POST',
+            body: { url: `http://example.com/${charset}` },
+          }
+        );
+      } finally {
+        fetchSpy.restore();
+      }
    }
  }
 });