feat(core): improve byok editing (#15427)

fix #14287
fix #15359
fix #15424

<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->
## Summary by CodeRabbit

* **New Features**
* Redesigned workspace AI provider settings with connection testing,
storage options, model selection, capability management, ordering, and
custom endpoints.
* AI chat model choices now adapt to the selected workspace and
conversation route.
  * Added support for image-based AI requests.
* **Bug Fixes**
  * Improved handling of unavailable or outdated model selections.
* App configuration updates now reject overlapping paths and load
deterministically.
* **Tests**
* Expanded coverage for provider models, AI chat scoping, image
requests, and configuration validation.
<!-- end of auto-generated comment: release notes by coderabbit.ai -->

packages/frontend/core/src/blocksuite/ai/actions/types.ts

@@ -367,7 +367,6 @@ declare global {
     // TODO(@Peng): should be refactored to get rid of implement details (like messages, action, role, etc.)
     interface AIHistory {
       sessionId: string;
-      tokens: number;
       action: string | null;
       createdAt: string;
       messages: {

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

@@ -2,9 +2,9 @@ import './ai-chat-composer-tip';
 
 import type {
   AIDraftService,
+  AIModelService,
   AIToolsConfigService,
 } from '@affine/core/modules/ai-button';
-import type { AIModelService } from '@affine/core/modules/ai-button/services/models';
 import type {
   ServerService,
   SubscriptionService,
@@ -128,15 +128,15 @@ export class AIChatComposer extends SignalWatcher(
   @property({ attribute: false })
   accessor aiToolsConfigService!: AIToolsConfigService;
 
+  @property({ attribute: false })
+  accessor aiModelService!: AIModelService;
+
   @property({ attribute: false })
   accessor affineFeatureFlagService!: FeatureFlagService;
 
   @property({ attribute: false })
   accessor subscriptionService!: SubscriptionService;
 
-  @property({ attribute: false })
-  accessor aiModelService!: AIModelService;
-
   @property({ attribute: false })
   accessor onAISubscribe!: () => Promise<void>;
 
@@ -183,9 +183,9 @@ export class AIChatComposer extends SignalWatcher(
         .affineFeatureFlagService=${this.affineFeatureFlagService}
         .aiDraftService=${this.aiDraftService}
         .aiToolsConfigService=${this.aiToolsConfigService}
+        .aiModelService=${this.aiModelService}
         .notificationService=${this.notificationService}
         .subscriptionService=${this.subscriptionService}
-        .aiModelService=${this.aiModelService}
         .onAISubscribe=${this.onAISubscribe}
         .portalContainer=${this.portalContainer}
         .onChatSuccess=${this.onChatSuccess}

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

@@ -1,9 +1,9 @@
 import type {
   AIDraftService,
+  AIModelService,
   AIToolsConfigService,
 } from '@affine/core/modules/ai-button';
 import type { AIDraftState } from '@affine/core/modules/ai-button/services/ai-draft';
-import type { AIModelService } from '@affine/core/modules/ai-button/services/models';
 import type {
   ServerService,
   SubscriptionService,
@@ -404,8 +404,8 @@ export class AIChatContent extends SignalWatcher(
         .notificationService=${this.notificationService}
         .aiDraftService=${this.aiDraftService}
         .aiToolsConfigService=${this.aiToolsConfigService}
-        .subscriptionService=${this.subscriptionService}
         .aiModelService=${this.aiModelService}
+        .subscriptionService=${this.subscriptionService}
         .onAISubscribe=${this.onAISubscribe}
         .trackOptions=${{
           where: 'chat-panel',

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

@@ -1,8 +1,8 @@
 import type {
   AIDraftService,
+  AIModelService,
   AIToolsConfigService,
 } from '@affine/core/modules/ai-button';
-import type { AIModelService } from '@affine/core/modules/ai-button/services/models';
 import type {
   ServerService,
   SubscriptionService,
@@ -399,6 +399,9 @@ export class AIChatInput extends SignalWatcher(
   @property({ attribute: false })
   accessor aiToolsConfigService!: AIToolsConfigService;
 
+  @property({ attribute: false })
+  accessor aiModelService!: AIModelService;
+
   @property({ attribute: false })
   accessor affineFeatureFlagService!: FeatureFlagService;
 
@@ -408,9 +411,6 @@ export class AIChatInput extends SignalWatcher(
   @property({ attribute: false })
   accessor subscriptionService!: SubscriptionService;
 
-  @property({ attribute: false })
-  accessor aiModelService!: AIModelService;
-
   @property({ attribute: false })
   accessor onAISubscribe!: () => Promise<void>;
 
@@ -483,6 +483,12 @@ export class AIChatInput extends SignalWatcher(
     window.addEventListener('dragend', this._resetDragState);
   }
 
+  protected override updated(changedProperties: PropertyValues<this>) {
+    if (changedProperties.has('workspaceId')) {
+      this.aiModelService.setScope(this.workspaceId, 'Chat With AFFiNE AI');
+    }
+  }
+
   protected override firstUpdated(changedProperties: PropertyValues): void {
     super.firstUpdated(changedProperties);
     if (this.aiDraftService) {
@@ -631,9 +637,9 @@ export class AIChatInput extends SignalWatcher(
           .onExtendedThinkingChange=${this._toggleReasoning}
           .serverService=${this.serverService}
           .toolsConfigService=${this.aiToolsConfigService}
+          .aiModelService=${this.aiModelService}
           .notificationService=${this.notificationService}
           .subscriptionService=${this.subscriptionService}
-          .aiModelService=${this.aiModelService}
           .onAISubscribe=${this.onAISubscribe}
         ></chat-input-preference>
         ${status === 'transmitting' || status === 'loading'
@@ -858,7 +864,7 @@ export class AIChatInput extends SignalWatcher(
       control: this.trackOptions?.control,
       reasoning: this._isReasoningActive,
       toolsConfig: this.aiToolsConfigService.config.value,
-      modelId: this.aiModelService.modelId.value,
+      routeTargetId: this.aiModelService.modelId.value,
       userInfo: {
         userId: userInfo?.id,
         userName: userInfo?.name,

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

@@ -1,14 +1,5 @@
 import type { AIToolsConfigService } from '@affine/core/modules/ai-button';
 import type { AIModelService } from '@affine/core/modules/ai-button/services/models';
-import type {
-  ServerService,
-  SubscriptionService,
-} from '@affine/core/modules/cloud';
-import {
-  type CopilotChatHistoryFragment,
-  ServerDeploymentType,
-  SubscriptionStatus,
-} from '@affine/graphql';
 import {
   menu,
   popMenu,
@@ -70,10 +61,10 @@ export class ChatInputPreference extends SignalWatcher(
       min-width: 220px;
     }
     .ai-active-model-name {
-      font-size: 14px;
-      color: ${unsafeCSSVarV2('text/secondary')};
-      line-height: 22px;
       margin-left: 40px;
+      color: ${unsafeCSSVarV2('text/secondary')};
+      font-size: 14px;
+      line-height: 22px;
     }
     .ai-model-prefix {
       width: 20px;
@@ -82,21 +73,21 @@ export class ChatInputPreference extends SignalWatcher(
     .ai-model-prefix svg {
       color: ${unsafeCSSVarV2('icon/activated')};
     }
+    .ai-model-postfix {
+      width: 20px;
+      height: 20px;
+    }
     .ai-model-postfix svg:hover {
       color: ${unsafeCSSVarV2('icon/activated')};
     }
     .ai-model-version {
-      font-size: 12px;
-      color: ${unsafeCSSVarV2('text/tertiary')};
-      line-height: 20px;
       margin-right: 40px;
+      color: ${unsafeCSSVarV2('text/tertiary')};
+      font-size: 12px;
+      line-height: 20px;
     }
   `;
 
-  @property({ attribute: false })
-  accessor session!: CopilotChatHistoryFragment | null | undefined;
-  // --------- model props end ---------
-
   // --------- extended thinking props start ---------
   @property({ attribute: false })
   accessor extendedThinking: boolean = false;
@@ -107,90 +98,91 @@ export class ChatInputPreference extends SignalWatcher(
     | undefined;
   // --------- extended thinking props end ---------
 
-  @property({ attribute: false })
-  accessor serverService!: ServerService;
-
   @property({ attribute: false })
   accessor toolsConfigService!: AIToolsConfigService;
 
-  @property({ attribute: false })
-  accessor notificationService!: NotificationService;
-
-  @property({ attribute: false })
-  accessor subscriptionService!: SubscriptionService;
-
   @property({ attribute: false })
   accessor aiModelService!: AIModelService;
 
+  @property({ attribute: false })
+  accessor notificationService!: NotificationService;
+
   @property({ attribute: false })
   accessor onAISubscribe!: () => Promise<void>;
 
-  model = computed(() => {
-    const modelId = this.aiModelService.modelId.value;
-    const activeModel = this.aiModelService.models.value.find(
-      model => model.id === modelId
-    );
-    const defaultModel = this.aiModelService.models.value.find(
-      model => model.isDefault
-    );
-    return activeModel || defaultModel;
-  });
+  private readonly model = computed(() =>
+    this.aiModelService.models.value.find(
+      model => model.id === this.aiModelService.modelId.value
+    )
+  );
 
   openPreference(e: Event) {
     const element = e.currentTarget;
     if (!(element instanceof HTMLElement)) return;
-    const modelItems = [];
+    const preferenceItems = [];
     const searchItems = [];
 
-    // model switch
-    modelItems.push(
-      menu.subMenu({
-        name: 'Model',
-        prefix: AiOutlineIcon(),
-        middleware: modelSubMenuMiddleware,
-        postfix: html`
-          <span class="ai-active-model-name"> ${this.model.value?.name} </span>
-        `,
-        options: {
-          items: this.aiModelService.models.value.map(model => {
-            const isSelected = model.id === this.model.value?.id;
-            const isSelfHosted =
-              this.serverService.server.config$.value?.type ===
-              ServerDeploymentType.Selfhosted;
-            const status =
-              this.subscriptionService.subscription.ai$.value?.status;
-            const isSubscribed = status === SubscriptionStatus.Active;
-            return menu.action({
-              name: model.category,
-              info: html`
-                <span class="ai-model-version">${model.version}</span>
-              `,
-              prefix: html`
-                <div class="ai-model-prefix">
-                  ${isSelected ? DoneIcon() : undefined}
-                </div>
-              `,
-              postfix: html`
-                <div class="ai-model-postfix" @click=${this.onAISubscribe}>
-                  ${model.isPro && !isSubscribed ? LockIcon() : undefined}
-                </div>
-              `,
-              select: () => {
-                if (model.isPro && !isSelfHosted && !isSubscribed) {
-                  this.notificationService.toast(
-                    `Pro models require an AFFiNE AI subscription.`
-                  );
-                  return;
-                }
-                this.aiModelService.setModel(model.id);
-              },
-            });
-          }),
-        },
-      })
-    );
+    if (this.aiModelService.models.value.length) {
+      preferenceItems.push(
+        menu.subMenu({
+          name: 'Model',
+          prefix: AiOutlineIcon(),
+          middleware: modelSubMenuMiddleware,
+          postfix: html`
+            <span class="ai-active-model-name">
+              ${this.model.value?.name ?? 'Auto'}
+            </span>
+          `,
+          options: {
+            items: [
+              menu.action({
+                name: 'Auto',
+                prefix: html`
+                  <div class="ai-model-prefix">
+                    ${this.aiModelService.modelId.value
+                      ? undefined
+                      : DoneIcon()}
+                  </div>
+                `,
+                select: () => this.aiModelService.resetModel(),
+              }),
+              ...this.aiModelService.models.value.map(model =>
+                menu.action({
+                  name: model.category,
+                  info: html`
+                    <span class="ai-model-version">${model.version}</span>
+                  `,
+                  prefix: html`
+                    <div class="ai-model-prefix">
+                      ${model.id === this.aiModelService.modelId.value
+                        ? DoneIcon()
+                        : undefined}
+                    </div>
+                  `,
+                  postfix: html`
+                    <div class="ai-model-postfix">
+                      ${model.available ? undefined : LockIcon()}
+                    </div>
+                  `,
+                  select: () => {
+                    if (!model.available) {
+                      this.notificationService.toast(
+                        'This model requires an AFFiNE AI subscription.'
+                      );
+                      this.onAISubscribe().catch(console.error);
+                      return;
+                    }
+                    this.aiModelService.setModel(model.id);
+                  },
+                })
+              ),
+            ],
+          },
+        })
+      );
+    }
 
-    modelItems.push(
+    preferenceItems.push(
       menu.toggleSwitch({
         name: 'Extended Thinking',
         prefix: ThinkingIcon(),
@@ -220,7 +212,7 @@ export class ChatInputPreference extends SignalWatcher(
       options: {
         items: [
           menu.group({
-            items: [...modelItems],
+            items: [...preferenceItems],
           }),
           menu.group({
             items: [...searchItems],
@@ -238,7 +230,7 @@ export class ChatInputPreference extends SignalWatcher(
       class="chat-input-preference-trigger"
     >
       <span class="chat-input-preference-trigger-label">
-        ${this.model.value?.category}
+        ${this.model.value?.category ?? 'Auto'}
       </span>
       <span class="chat-input-preference-trigger-icon">
         ${ArrowDownSmallIcon()}

packages/frontend/core/src/blocksuite/ai/components/playground/chat.ts

@@ -1,5 +1,7 @@
-import type { AIToolsConfigService } from '@affine/core/modules/ai-button';
-import type { AIModelService } from '@affine/core/modules/ai-button/services/models';
+import type {
+  AIModelService,
+  AIToolsConfigService,
+} from '@affine/core/modules/ai-button';
 import type {
   ServerService,
   SubscriptionService,
@@ -185,10 +187,10 @@ export class PlaygroundChat extends SignalWatcher(
   accessor aiToolsConfigService!: AIToolsConfigService;
 
   @property({ attribute: false })
-  accessor subscriptionService!: SubscriptionService;
+  accessor aiModelService!: AIModelService;
 
   @property({ attribute: false })
-  accessor aiModelService!: AIModelService;
+  accessor subscriptionService!: SubscriptionService;
 
   @property({ attribute: false })
   accessor onAISubscribe: (() => Promise<void>) | undefined;

packages/frontend/core/src/blocksuite/ai/components/playground/content.ts

@@ -1,5 +1,7 @@
-import type { AIToolsConfigService } from '@affine/core/modules/ai-button';
-import type { AIModelService } from '@affine/core/modules/ai-button/services/models';
+import type {
+  AIModelService,
+  AIToolsConfigService,
+} from '@affine/core/modules/ai-button';
 import type {
   ServerService,
   SubscriptionService,
@@ -103,15 +105,15 @@ export class PlaygroundContent extends SignalWatcher(
   @property({ attribute: false })
   accessor aiToolsConfigService!: AIToolsConfigService;
 
+  @property({ attribute: false })
+  accessor aiModelService!: AIModelService;
+
   @property({ attribute: false })
   accessor affineWorkspaceDialogService!: WorkspaceDialogService;
 
   @property({ attribute: false })
   accessor subscriptionService!: SubscriptionService;
 
-  @property({ attribute: false })
-  accessor aiModelService!: AIModelService;
-
   @state()
   accessor sessions: CopilotChatHistoryFragment[] = [];
 
@@ -380,10 +382,10 @@ export class PlaygroundContent extends SignalWatcher(
                 .affineThemeService=${this.affineThemeService}
                 .notificationService=${this.notificationService}
                 .aiToolsConfigService=${this.aiToolsConfigService}
+                .aiModelService=${this.aiModelService}
                 .affineWorkspaceDialogService=${this
                   .affineWorkspaceDialogService}
                 .subscriptionService=${this.subscriptionService}
-                .aiModelService=${this.aiModelService}
                 .addChat=${this.addChat}
               ></playground-chat>
             </div>

packages/frontend/core/src/blocksuite/ai/peek-view/chat-block-peek-view.ts

@@ -1,8 +1,8 @@
 import type {
   AIDraftService,
+  AIModelService,
   AIToolsConfigService,
 } from '@affine/core/modules/ai-button';
-import type { AIModelService } from '@affine/core/modules/ai-button/services/models';
 import type {
   ServerService,
   SubscriptionService,
@@ -584,6 +584,7 @@ export class AIChatBlockPeekView extends LitElement {
         .affineWorkspaceDialogService=${this.affineWorkspaceDialogService}
         .notificationService=${notificationService}
         .aiToolsConfigService=${this.aiToolsConfigService}
+        .aiModelService=${this.aiModelService}
         .affineFeatureFlagService=${this.affineFeatureFlagService}
         .onChatSuccess=${this._onChatSuccess}
         .trackOptions=${{
@@ -594,7 +595,6 @@ export class AIChatBlockPeekView extends LitElement {
         .reasoningConfig=${this.reasoningConfig}
         .serverService=${this.serverService}
         .subscriptionService=${this.subscriptionService}
-        .aiModelService=${this.aiModelService}
         .onAISubscribe=${this.onAISubscribe}
       ></ai-chat-composer>
     </div> `;
@@ -681,8 +681,8 @@ export const AIChatBlockPeekViewTemplate = (
   affineWorkspaceDialogService: WorkspaceDialogService,
   aiDraftService: AIDraftService,
   aiToolsConfigService: AIToolsConfigService,
-  subscriptionService: SubscriptionService,
   aiModelService: AIModelService,
+  subscriptionService: SubscriptionService,
   onAISubscribe: (() => Promise<void>) | undefined
 ) => {
   return html`<ai-chat-block-peek-view
@@ -696,8 +696,8 @@ export const AIChatBlockPeekViewTemplate = (
     .affineWorkspaceDialogService=${affineWorkspaceDialogService}
     .aiDraftService=${aiDraftService}
     .aiToolsConfigService=${aiToolsConfigService}
-    .subscriptionService=${subscriptionService}
     .aiModelService=${aiModelService}
+    .subscriptionService=${subscriptionService}
     .onAISubscribe=${onAISubscribe}
   ></ai-chat-block-peek-view>`;
 };

packages/frontend/core/src/blocksuite/ai/runtime/chat/actions.ts

@@ -18,7 +18,7 @@ export type AIChatSendOptions = {
   control?: BlockSuitePresets.TrackerControl;
   reasoning?: boolean;
   toolsConfig?: unknown;
-  modelId?: string;
+  routeTargetId?: string;
   userInfo?: {
     userId?: string;
     userName?: string;
@@ -45,7 +45,7 @@ export type AIChatAction =
   | { type: 'clearError' }
   | { type: 'setComposerText'; text: string }
   | { type: 'setReasoning'; reasoning: boolean }
-  | { type: 'setModel'; modelId?: string }
+  | { type: 'setRouteTarget'; routeTargetId?: string }
   | { type: 'addAttachment'; attachment: string | Blob | File }
   | { type: 'removeAttachment'; index: number }
   | { type: 'addContextItem'; item: AIChatContextItem }

packages/frontend/core/src/blocksuite/ai/runtime/chat/runtime.spec.ts

@@ -34,8 +34,6 @@ function session(
     parentSessionId: null,
     promptName: 'Chat With AFFiNE AI',
     action: null,
-    optionalModels: null,
-    tokens: 0,
     ...overrides,
   } as CopilotChatHistoryFragment;
 }

packages/frontend/core/src/blocksuite/ai/runtime/chat/runtime.ts

@@ -152,8 +152,8 @@ export class AIChatRuntime {
       case 'setReasoning':
         this.updateComposer({ reasoning: action.reasoning });
         return;
-      case 'setModel':
-        this.updateComposer({ modelId: action.modelId });
+      case 'setRouteTarget':
+        this.updateComposer({ routeTargetId: action.routeTargetId });
         return;
       case 'addAttachment':
         this.updateComposer({
@@ -388,7 +388,8 @@ export class AIChatRuntime {
         contextId: this.snapshot.composer.context.contextId,
         reasoning: options.reasoning ?? this.snapshot.composer.reasoning,
         toolsConfig: options.toolsConfig ?? this.snapshot.composer.toolsConfig,
-        modelId: options.modelId ?? this.snapshot.composer.modelId,
+        routeTargetId:
+          options.routeTargetId ?? this.snapshot.composer.routeTargetId,
         isRootSession: options.isRootSession,
         where: options.where,
         control: options.control,

packages/frontend/core/src/blocksuite/ai/runtime/chat/state.ts

@@ -136,7 +136,7 @@ export type AIChatComposerState = {
   context: AIChatContextState;
   reasoning: boolean;
   toolsConfig?: AIToolsConfig;
-  modelId?: string;
+  routeTargetId?: string;
 };
 
 export type AIChatNavigationRequest = {

packages/frontend/core/src/blocksuite/ai/runtime/request/byok-local-lease.ts

@@ -67,9 +67,8 @@ export async function createWorkspaceByokLocalLease(
               provider: gqlProvider,
               name: provider.name,
               description: provider.description ?? null,
-              apiKey: provider.apiKey,
-              endpoint: provider.endpoint ?? null,
-              sortOrder: provider.sortOrder ?? 0,
+              credential: provider.credential,
+              definition: provider.definition,
               enabled: provider.enabled ?? true,
             },
           ]

packages/frontend/core/src/blocksuite/ai/runtime/request/copilot-client.ts

@@ -467,7 +467,9 @@ export class CopilotClient {
       sessionId,
       messageId,
       reasoning,
+      profileId,
       modelId,
+      routeTargetId,
       toolsConfig,
       actionId,
       actionVersion,
@@ -478,7 +480,9 @@ export class CopilotClient {
       sessionId: string;
       messageId?: string;
       reasoning?: boolean;
+      profileId?: string;
       modelId?: string;
+      routeTargetId?: string;
       toolsConfig?: AIToolsConfig;
       actionId?: string;
       actionVersion?: string;
@@ -495,7 +499,9 @@ export class CopilotClient {
     const queryString = this.paramsToQueryString({
       messageId,
       reasoning,
+      profileId,
       modelId,
+      routeTargetId,
       toolsConfig,
       actionId,
       actionVersion,

packages/frontend/core/src/blocksuite/ai/runtime/request/message-transport.ts

@@ -24,7 +24,9 @@ export type TextToTextOptions = {
   runId?: string;
   isRootSession?: boolean;
   reasoning?: boolean;
+  profileId?: string;
   modelId?: string;
+  routeTargetId?: string;
   toolsConfig?: AIToolsConfig;
 };
 
@@ -127,7 +129,9 @@ export function textToText({
   actionVersion,
   runId,
   reasoning,
+  profileId,
   modelId,
+  routeTargetId,
   toolsConfig,
 }: TextToTextOptions) {
   let messageId: string | undefined;
@@ -161,7 +165,9 @@ export function textToText({
             sessionId,
             messageId,
             reasoning,
+            profileId,
             modelId,
+            routeTargetId,
             toolsConfig,
             actionId,
             actionVersion,
@@ -229,7 +235,9 @@ export function textToText({
           sessionId,
           messageId,
           reasoning,
+          profileId,
           modelId,
+          routeTargetId,
           toolsConfig,
           actionId,
           actionVersion,

packages/frontend/core/src/blocksuite/ai/runtime/request/service.spec.ts

@@ -23,9 +23,13 @@ const electronApis = vi.hoisted(() => ({
           Array<{
             provider: string;
             name: string;
-            apiKey: string;
+            credential: string;
+            definition: {
+              version: number;
+              endpoint: { kind: string; url?: string | null };
+              models: unknown[];
+            };
             description?: string | null;
-            endpoint?: string | null;
             sortOrder?: number | null;
             enabled?: boolean | null;
           }>
@@ -117,7 +121,12 @@ describe('runtime request transport BYOK local lease handling', () => {
         {
           provider: 'openai',
           name: 'OpenAI',
-          apiKey: 'sk-local',
+          credential: 'sk-local',
+          definition: {
+            version: 1,
+            endpoint: { kind: 'provider_default' },
+            models: [{ modelId: 'model-1', capabilities: [] }],
+          },
         },
       ]),
     };

packages/frontend/core/src/desktop/dialogs/setting/workspace-setting/byok/add-key-modal.tsx

@@ -1,12 +1,12 @@
-import { Button, Modal, notify } from '@affine/component';
+import { Button, Input, Modal, notify } from '@affine/component';
 import {
-  ByokKeyStorage,
   ByokProvider,
-  testWorkspaceByokConfigMutation as testByokMutation,
-  upsertWorkspaceByokConfigMutation as upsertByokMutation,
+  createWorkspaceByokProfileMutation,
+  probeWorkspaceByokDraftMutation,
+  replaceWorkspaceByokProfileMutation,
 } from '@affine/graphql';
 import { useI18n } from '@affine/i18n';
-import { useCallback, useEffect, useState } from 'react';
+import { useCallback, useEffect, useMemo, useRef, useState } from 'react';
 
 import { logByokError } from './errors';
 import * as styles from './index.css';
@@ -18,13 +18,16 @@ import {
   shouldShowEndpoint,
   storageLabel,
 } from './metadata';
-import type {
-  ByokKey,
-  ByokSettings,
-  ByokStorage,
-  ByokTestResult,
-  GqlFn,
-} from './types';
+import { ModelSelector } from './model-selector';
+import {
+  catalogModels,
+  defaultModels,
+  type ModelDeclaration,
+  modelUseCases,
+  probeChecks,
+} from './model-utils';
+import type { ByokDefinition, ByokKey, ByokSettings, GqlFn } from './types';
+import { ByokStorage } from './types';
 
 export const AddKeyModal = ({
   workspaceId,
@@ -59,101 +62,133 @@ export const AddKeyModal = ({
   const [provider, setProvider] = useState<ByokProvider>(ByokProvider.openai);
   const [name, setName] = useState('');
   const [description, setDescription] = useState('');
-  const [storage, setStorage] = useState<ByokStorage>(ByokKeyStorage.server);
+  const [profileEnabled, setProfileEnabled] = useState(true);
+  const [storage, setStorage] = useState<ByokStorage>(ByokStorage.server);
   const [apiKey, setApiKey] = useState('');
+  const [customEndpoint, setCustomEndpoint] = useState(false);
   const [endpoint, setEndpoint] = useState('');
-  const [testResult, setTestResult] = useState<ByokTestResult | null>(null);
-  const [testing, setTesting] = useState(false);
-  const canTestStoredConfig =
-    storage === ByokKeyStorage.server &&
-    editingKey?.storage === ByokKeyStorage.server &&
-    editingKey.provider === provider;
-  const canTest = !!apiKey || canTestStoredConfig;
+  const [models, setModels] = useState<ModelDeclaration[]>([]);
+  const [testStatus, setTestStatus] = useState<'passed' | 'failed' | null>(
+    null
+  );
+  const [includeImageProbe, setIncludeImageProbe] = useState(false);
+  const [busy, setBusy] = useState(false);
+  const busyRef = useRef(false);
+  const localStorageUnavailable = !localStorageSupported || !canAddLocalKey;
+  const localStorageDisabled = !!editingKey || localStorageUnavailable;
+  const showCustomEndpoint = shouldShowEndpoint(
+    isSelfHosted,
+    settings.customEndpointSupported
+  );
+
   const endpointHint = endpointHintKey(
     settings.customEndpointSupported,
     settings.privateEndpointSupported
   );
+  const providerCatalog = useMemo(
+    () => catalogModels(settings, provider),
+    [provider, settings]
+  );
 
   useEffect(() => {
-    if (!open) {
-      return;
-    }
-    setProvider(editingKey?.provider ?? ByokProvider.openai);
-    setName(editingKey?.name ?? '');
+    if (!open) return;
+    const nextProvider = editingKey?.provider ?? ByokProvider.openai;
+    setProvider(nextProvider);
+    setName(editingKey?.name ?? providerLabels[nextProvider]);
     setDescription(editingKey?.description ?? '');
+    setProfileEnabled(editingKey?.enabled ?? true);
     setStorage(
       editingKey?.storage ??
-        (canAddServerKey ? ByokKeyStorage.server : ByokKeyStorage.local)
+        (canAddServerKey ? ByokStorage.server : ByokStorage.local)
     );
     setApiKey('');
-    setEndpoint(editingKey?.endpoint ?? '');
-    setTestResult(null);
-  }, [canAddServerKey, editingKey, open]);
+    setEndpoint(editingKey?.definition.endpoint.url ?? '');
+    setCustomEndpoint(editingKey?.definition.endpoint.kind === 'custom');
+    setModels(
+      editingKey?.definition.models ?? defaultModels(settings, nextProvider)
+    );
+    setTestStatus(null);
+    setIncludeImageProbe(false);
+  }, [canAddServerKey, editingKey, open, settings]);
 
-  const testKey = useCallback(async () => {
-    if (!gql) {
-      return;
-    }
-    setTesting(true);
-    try {
-      const result = await gql({
-        query: testByokMutation,
-        variables: {
-          input: {
-            workspaceId,
-            provider,
-            storage,
-            apiKey: apiKey || null,
-            endpoint: endpoint || null,
-            configId: canTestStoredConfig ? editingKey.id : null,
-          },
+  const definition = useMemo<ByokDefinition>(
+    () => ({
+      version: editingKey?.definition.version ?? 1,
+      endpoint: customEndpoint
+        ? { kind: 'custom', url: endpoint }
+        : { kind: 'provider_default', url: null },
+      models,
+    }),
+    [customEndpoint, editingKey?.definition.version, endpoint, models]
+  );
+
+  const invalidateTest = () => setTestStatus(null);
+  const runProbe = useCallback(async () => {
+    if (!gql) return false;
+    const canReuseServerCredential =
+      editingKey?.storage === ByokStorage.server && !apiKey;
+    const checks = probeChecks(models, includeImageProbe);
+    const result = await gql({
+      query: probeWorkspaceByokDraftMutation,
+      variables: {
+        input: {
+          workspaceId,
+          provider,
+          credential: apiKey || null,
+          profileId: canReuseServerCredential ? editingKey.id : null,
+          expectedRevision: canReuseServerCredential
+            ? (editingKey.revision ?? null)
+            : null,
+          definition,
+          checks,
         },
-      });
-      const nextResult = result.testWorkspaceByokConfig as
-        | ByokTestResult
-        | undefined;
-      setTestResult(nextResult ?? null);
-      if (nextResult && !nextResult.ok) {
-        notify.error({
-          title: byokT(t, 'notify.test-failed.title'),
-          message: nextResult.message,
-        });
-      }
-    } finally {
-      setTesting(false);
-    }
+      },
+    });
+    const probe = result.probeWorkspaceByokDraft;
+    const verifiedChecks = new Set(
+      probe.models.flatMap(model =>
+        model.checks
+          .filter(check => check.status.kind === 'verified')
+          .map(check => `${model.modelId}\0${check.operation}`)
+      )
+    );
+    const passed =
+      checks.length > 0 &&
+      probe.connection.kind === 'verified' &&
+      checks.every(check =>
+        verifiedChecks.has(`${check.modelId}\0${check.operation}`)
+      );
+    setTestStatus(passed ? 'passed' : 'failed');
+    return passed;
   }, [
     apiKey,
-    canTestStoredConfig,
+    definition,
     editingKey,
-    endpoint,
     gql,
+    includeImageProbe,
+    models,
     provider,
-    storage,
-    t,
     workspaceId,
   ]);
 
-  const save = useCallback(async () => {
-    if (!testResult?.ok || !gql) {
-      return;
-    }
-    if (storage === ByokKeyStorage.local) {
+  const persist = useCallback(async () => {
+    if (!gql) return;
+    if (storage === ByokStorage.local) {
       const saved = await upsertLocalKey(workspaceId, {
         id:
-          editingKey?.storage === ByokKeyStorage.local
+          editingKey?.storage === ByokStorage.local
             ? editingKey.id
             : crypto.randomUUID(),
         provider,
         name,
         description,
-        apiKey,
-        endpoint: endpoint || null,
+        credential: apiKey,
+        definition,
         sortOrder:
-          editingKey?.storage === ByokKeyStorage.local
+          editingKey?.storage === ByokStorage.local
             ? editingKey.sortOrder
             : localKeys.length,
-        enabled: true,
+        enabled: profileEnabled,
       });
       if (!saved) {
         notify.error({
@@ -163,189 +198,359 @@ export const AddKeyModal = ({
         return;
       }
       setLocalKeys(await readLocalKeys(workspaceId));
-    } else {
+    } else if (editingKey?.storage === ByokStorage.server) {
+      if (editingKey.revision === undefined) {
+        notify.error({
+          title: byokT(t, 'notify.reload-required.title'),
+          message: byokT(t, 'notify.reload-required.message'),
+        });
+        return;
+      }
       await gql({
-        query: upsertByokMutation,
+        query: replaceWorkspaceByokProfileMutation,
+        variables: {
+          input: {
+            workspaceId,
+            profileId: editingKey.id,
+            expectedRevision: editingKey.revision,
+            name,
+            description: description || null,
+            credential: apiKey || null,
+            definition,
+            enabled: profileEnabled,
+          },
+        },
+      });
+      await onSaved();
+    } else {
+      await gql({
+        query: createWorkspaceByokProfileMutation,
         variables: {
           input: {
             workspaceId,
-            id:
-              editingKey?.storage === ByokKeyStorage.server
-                ? editingKey.id
-                : null,
             provider,
             name,
-            description,
-            storage,
-            apiKey: apiKey || null,
-            endpoint: endpoint || null,
-            enabled: true,
+            description: description || null,
+            credential: apiKey,
+            definition,
+            enabled: profileEnabled,
           },
         },
       });
       await onSaved();
     }
     onOpenChange(false);
-    setApiKey('');
-    setTestResult(null);
   }, [
     apiKey,
+    definition,
     description,
     editingKey,
-    endpoint,
     gql,
-    localKeys,
+    localKeys.length,
     name,
     onOpenChange,
     onSaved,
     provider,
+    profileEnabled,
     setLocalKeys,
     storage,
     t,
-    testResult?.ok,
     workspaceId,
   ]);
 
+  const connect = useCallback(async () => {
+    if (busyRef.current) return;
+    busyRef.current = true;
+    setBusy(true);
+    try {
+      const passed = testStatus === 'passed' || (await runProbe());
+      if (!passed) {
+        notify.error({
+          title: byokT(t, 'notify.test-failed.title'),
+          message: byokT(t, 'notify.operation-failed.message'),
+        });
+        return;
+      }
+      await persist();
+    } finally {
+      busyRef.current = false;
+      setBusy(false);
+    }
+  }, [persist, runProbe, t, testStatus]);
+
+  const testConnection = useCallback(async () => {
+    if (busyRef.current) return;
+    busyRef.current = true;
+    setBusy(true);
+    try {
+      await runProbe();
+    } finally {
+      busyRef.current = false;
+      setBusy(false);
+    }
+  }, [runProbe]);
+
+  const hasCredential = !!apiKey || editingKey?.storage === ByokStorage.server;
+  const valid =
+    !!name.trim() &&
+    hasCredential &&
+    models.length > 0 &&
+    models.every(model => model.modelId.trim() && model.capabilities.length) &&
+    new Set(models.map(model => model.modelId.trim())).size === models.length &&
+    (!customEndpoint || !!endpoint.trim());
+
   return (
     <Modal
-      width={520}
+      width={640}
       open={open}
       onOpenChange={onOpenChange}
-      title={
-        editingKey ? byokT(t, 'modal.edit-title') : byokT(t, 'modal.add-title')
-      }
-      description={byokT(t, 'modal.description')}
+      title={byokT(
+        t,
+        editingKey ? 'modal.manage-title' : 'modal.connect-title'
+      )}
+      description={byokT(t, 'modal.connect-description')}
     >
       <div className={styles.form}>
-        <label className={styles.field}>
-          <span className={styles.label}>{byokT(t, 'field.provider')}</span>
-          <select
-            className={styles.input}
-            value={provider}
-            onChange={event => {
-              setProvider(event.target.value as ByokProvider);
-              setTestResult(null);
-            }}
-          >
-            {settings.allowedProviders.map(provider => (
-              <option key={provider} value={provider}>
-                {providerLabels[provider]}
-              </option>
-            ))}
-          </select>
-        </label>
-        <label className={styles.field}>
-          <span className={styles.label}>{byokT(t, 'field.key-name')}</span>
-          <input
-            className={styles.input}
-            value={name}
-            onChange={event => setName(event.target.value)}
-            placeholder={byokT(t, 'placeholder.key-name')}
-          />
-        </label>
-        <label className={styles.field}>
-          <span className={styles.label}>{byokT(t, 'field.description')}</span>
-          <input
-            className={styles.input}
-            value={description}
-            onChange={event => setDescription(event.target.value)}
-            placeholder={byokT(t, 'placeholder.description')}
-          />
-        </label>
-        <label className={styles.field}>
-          <span className={styles.label}>{byokT(t, 'field.storage')}</span>
-          <select
-            className={styles.input}
-            value={storage}
-            disabled={!!editingKey}
-            onChange={event => {
-              setStorage(event.target.value as ByokStorage);
-              setTestResult(null);
-            }}
-          >
-            <option value={ByokKeyStorage.server} disabled={!canAddServerKey}>
-              {storageLabel(t, ByokKeyStorage.server)}
-            </option>
-            <option
-              value={ByokKeyStorage.local}
-              disabled={!localStorageSupported || !canAddLocalKey}
-            >
-              {canAddLocalKey
-                ? byokT(t, 'storage.local-this-device')
-                : byokT(t, 'storage.local-desktop-only')}
-            </option>
-          </select>
-        </label>
-        <label className={styles.field}>
-          <span className={styles.label}>{byokT(t, 'field.api-key')}</span>
-          <input
-            className={styles.input}
-            value={apiKey}
-            onChange={event => {
-              setApiKey(event.target.value);
-              setTestResult(null);
-            }}
-            type="password"
-          />
-        </label>{' '}
-        {shouldShowEndpoint(isSelfHosted, settings.customEndpointSupported) ? (
-          <label className={styles.endpointField}>
-            <span className={styles.label}>{byokT(t, 'field.endpoint')}</span>
-            <input
+        <div className={styles.formSection}>
+          <div className={styles.sectionTitle}>
+            {byokT(t, 'section.connection')}
+          </div>
+          <label className={styles.field}>
+            <span className={styles.label}>{byokT(t, 'field.provider')}</span>
+            <select
               className={styles.input}
-              value={endpoint}
-              disabled={!settings.customEndpointSupported}
+              value={provider}
+              disabled={!!editingKey}
               onChange={event => {
-                setEndpoint(event.target.value);
-                setTestResult(null);
+                const next = event.target.value as ByokProvider;
+                setProvider(next);
+                setName(providerLabels[next]);
+                setModels(defaultModels(settings, next));
+                invalidateTest();
               }}
-              placeholder="https://api.example.com/v1"
+            >
+              {settings.allowedProviders.map(item => (
+                <option key={item} value={item}>
+                  {providerLabels[item]}
+                </option>
+              ))}
+            </select>
+          </label>
+          <div className={styles.storageOptions}>
+            <label
+              className={styles.storageOption}
+              data-disabled={!canAddServerKey}
+            >
+              <input
+                className={styles.storageRadio}
+                type="radio"
+                name="byok-storage"
+                checked={storage === ByokStorage.server}
+                disabled={!!editingKey || !canAddServerKey}
+                onChange={() => setStorage(ByokStorage.server)}
+              />
+              <span className={styles.storageCopy}>
+                <strong>{storageLabel(t, ByokStorage.server)}</strong>
+                <small className={styles.storageDescription}>
+                  {byokT(t, 'storage.server.description')}
+                </small>
+              </span>
+            </label>
+            <label
+              className={styles.storageOption}
+              data-disabled={localStorageUnavailable}
+            >
+              <input
+                className={styles.storageRadio}
+                type="radio"
+                name="byok-storage"
+                checked={storage === ByokStorage.local}
+                disabled={localStorageDisabled}
+                onChange={() => setStorage(ByokStorage.local)}
+              />
+              <span className={styles.storageCopy}>
+                <strong>{storageLabel(t, ByokStorage.local)}</strong>
+                <small className={styles.storageDescription}>
+                  {!BUILD_CONFIG.isElectron
+                    ? byokT(t, 'storage.local.desktop-only')
+                    : !localStorageSupported
+                      ? byokT(t, 'storage.local.unavailable')
+                      : byokT(t, 'storage.local.description')}
+                </small>
+              </span>
+            </label>
+          </div>
+          <label className={styles.field}>
+            <span className={styles.label}>{byokT(t, 'field.api-key')}</span>
+            <Input
+              size="large"
+              value={apiKey}
+              onChange={value => {
+                setApiKey(value);
+                invalidateTest();
+              }}
+              type="password"
+              placeholder={
+                editingKey?.storage === ByokStorage.server
+                  ? byokT(t, 'placeholder.keep-current-key')
+                  : ''
+              }
             />
-            {endpointHint ? (
-              <span className={styles.fieldHint}>{byokT(t, endpointHint)}</span>
+            {storage === ByokStorage.local ? (
+              <span className={styles.fieldHint}>
+                {byokT(t, 'storage.local.test-disclosure')}
+              </span>
             ) : null}
           </label>
+          {showCustomEndpoint ? (
+            <>
+              <label className={styles.checkboxRow}>
+                <input
+                  type="checkbox"
+                  checked={customEndpoint}
+                  disabled={!settings.customEndpointSupported}
+                  onChange={event => {
+                    setCustomEndpoint(event.target.checked);
+                    if (event.target.checked && !editingKey) setModels([]);
+                    if (!event.target.checked)
+                      setModels(defaultModels(settings, provider));
+                    invalidateTest();
+                  }}
+                />
+                {byokT(t, 'endpoint.use-custom')}
+              </label>
+              {!settings.customEndpointSupported && endpointHint ? (
+                <span className={styles.fieldHint}>
+                  {byokT(t, endpointHint)}
+                </span>
+              ) : null}
+              {customEndpoint ? (
+                <label className={styles.endpointField}>
+                  <span className={styles.label}>
+                    {byokT(t, 'field.endpoint')}
+                  </span>
+                  <Input
+                    size="large"
+                    value={endpoint}
+                    onChange={value => {
+                      setEndpoint(value);
+                      invalidateTest();
+                    }}
+                    placeholder="https://api.example.com/v1"
+                  />
+                  {endpointHint ? (
+                    <span className={styles.fieldHint}>
+                      {byokT(t, endpointHint)}
+                    </span>
+                  ) : null}
+                </label>
+              ) : null}
+            </>
+          ) : null}
+        </div>
+
+        <div className={styles.formSection}>
+          <div className={styles.sectionHeading}>
+            <div>
+              <div className={styles.sectionTitle}>
+                {byokT(t, 'section.models')}
+              </div>
+              <div className={styles.description}>
+                {byokT(t, 'models.description.selected')}
+              </div>
+            </div>
+          </div>
+          <ModelSelector
+            customEndpoint={customEndpoint}
+            catalog={providerCatalog}
+            models={models}
+            validation={editingKey?.validation}
+            onChange={models => {
+              setModels(models);
+              invalidateTest();
+            }}
+          />
+        </div>
+
+        <details className={styles.advanced}>
+          <summary>{byokT(t, 'section.advanced')}</summary>
+          <div className={styles.advancedFields}>
+            <label className={styles.checkboxRow}>
+              <input
+                type="checkbox"
+                checked={profileEnabled}
+                onChange={event => setProfileEnabled(event.target.checked)}
+              />
+              {byokT(t, 'field.provider-enabled')}
+            </label>
+            <label className={styles.field}>
+              <span className={styles.label}>{byokT(t, 'field.key-name')}</span>
+              <Input size="large" value={name} onChange={setName} />
+            </label>
+            <label className={styles.field}>
+              <span className={styles.label}>
+                {byokT(t, 'field.description')}
+              </span>
+              <Input
+                size="large"
+                value={description}
+                onChange={setDescription}
+              />
+            </label>
+          </div>
+        </details>
+
+        {models.some(model => modelUseCases(model).includes('image')) ? (
+          <label className={styles.checkboxRow}>
+            <input
+              type="checkbox"
+              checked={includeImageProbe}
+              onChange={event => {
+                setIncludeImageProbe(event.target.checked);
+                invalidateTest();
+              }}
+            />
+            {byokT(t, 'probe.include-image')}
+          </label>
         ) : null}
+
         <div className={styles.modalActions}>
           <span
             className={`${styles.testStatus} ${
-              testResult?.ok
+              testStatus === 'passed'
                 ? styles.success
-                : testResult && !testResult.ok
+                : testStatus === 'failed'
                   ? styles.error
                   : ''
             }`}
           >
-            {testResult?.ok
-              ? byokT(t, 'status.key-verified')
-              : testResult
-                ? byokT(t, 'status.key-test-failed')
+            {testStatus === 'passed'
+              ? byokT(t, 'probe.verified')
+              : testStatus === 'failed'
+                ? byokT(t, 'probe.failed')
                 : ''}
           </span>
           <Button
             variant="secondary"
-            disabled={!canTest || testing}
+            disabled={!valid || busy}
             onClick={() => {
-              testKey().catch(error => {
-                logByokError('Failed to test BYOK key', error);
-                notify.error({
-                  title: byokT(t, 'notify.test-failed.title'),
-                  message: byokT(t, 'notify.operation-failed.message'),
-                });
+              testConnection().catch(error => {
+                logByokError('Failed to test BYOK provider', error);
+                setTestStatus('failed');
               });
             }}
           >
-            {byokT(t, 'action.test-key')}
+            {byokT(t, 'action.test-connection')}
           </Button>
           <Button variant="secondary" onClick={() => onOpenChange(false)}>
             {byokT(t, 'action.cancel')}
           </Button>
           <Button
             variant="primary"
-            disabled={!testResult?.ok || !name}
+            disabled={!valid || busy}
             onClick={() => {
-              save().catch(error => {
-                logByokError('Failed to save BYOK key', error);
+              connect().catch(error => {
+                logByokError('Failed to save BYOK provider', error);
                 notify.error({
                   title: byokT(t, 'notify.save-failed.title'),
                   message: byokT(t, 'notify.operation-failed.message'),
@@ -353,7 +558,14 @@ export const AddKeyModal = ({
               });
             }}
           >
-            {byokT(t, 'action.save-key')}
+            {byokT(
+              t,
+              busy
+                ? 'action.connecting'
+                : editingKey
+                  ? 'action.save-changes'
+                  : 'action.connect'
+            )}
           </Button>
         </div>
       </div>

packages/frontend/core/src/desktop/dialogs/setting/workspace-setting/byok/coverage.tsx

@@ -9,8 +9,8 @@ import {
 import type { ReactNode } from 'react';
 
 import * as styles from './index.css';
-import { byokT, capabilityRows, warningDescription } from './metadata';
-import type { ByokKey, ByokSettings } from './types';
+import { byokT, capabilityRows } from './metadata';
+import type { ByokKey } from './types';
 
 function coverageIcon(
   icon: (typeof capabilityRows)[number]['icon']
@@ -47,13 +47,7 @@ function isRowCovered(row: (typeof capabilityRows)[number], keys: ByokKey[]) {
   });
 }
 
-export const CoveragePanel = ({
-  keys,
-  settings,
-}: {
-  keys: ByokKey[];
-  settings: ByokSettings;
-}) => {
+export const CoveragePanel = ({ keys }: { keys: ByokKey[] }) => {
   const t = useI18n();
 
   return (
@@ -63,9 +57,6 @@ export const CoveragePanel = ({
       </div>
       <div className={styles.rows}>
         {capabilityRows.map(row => {
-          const warning = settings.warnings.find(
-            w => w.featureKind === row.featureKind
-          );
           const covered = isRowCovered(row, keys);
           return (
             <div
@@ -86,7 +77,7 @@ export const CoveragePanel = ({
               <div className={styles.rowMain}>
                 <div className={styles.rowTitle}>{byokT(t, row.titleKey)}</div>
                 <div className={styles.rowDescription}>
-                  {warningDescription(t, warning) ?? byokT(t, row.fallbackKey)}
+                  {byokT(t, row.fallbackKey)}
                 </div>
               </div>
             </div>

packages/frontend/core/src/desktop/dialogs/setting/workspace-setting/byok/index.css.ts

@@ -86,7 +86,7 @@ export const capabilityIcon = style({
 
 export const capabilityIconActive = style({
   color: cssVarV2('button/primary'),
-  background: '#f0f7ff',
+  background: cssVarV2('chip/label/blue'),
 });
 
 export const capabilityIconSvg = style({
@@ -176,14 +176,16 @@ export const locked = style({
 export const form = style({
   display: 'flex',
   flexDirection: 'column',
-  gap: 12,
+  gap: 16,
+  maxHeight: 'min(720px, calc(100vh - 180px))',
+  overflowY: 'auto',
+  paddingRight: 2,
 });
 
 export const field = style({
   display: 'flex',
   flexDirection: 'column',
   gap: 4,
-  height: '3em',
 });
 
 export const endpointField = style([field, { height: 'auto' }]);
@@ -201,6 +203,8 @@ export const fieldHint = style({
 
 export const input = style({
   height: 32,
+  minHeight: 32,
+  maxHeight: 32,
   width: '100%',
   boxSizing: 'border-box',
   borderRadius: 8,
@@ -228,6 +232,319 @@ export const modalActions = style({
   justifyContent: 'flex-end',
   gap: 8,
   marginTop: 8,
+  position: 'sticky',
+  bottom: 0,
+  paddingTop: 12,
+  background: cssVarV2('layer/background/primary'),
+});
+
+export const formSection = style({
+  display: 'flex',
+  flexDirection: 'column',
+  gap: 12,
+  padding: 14,
+  border: `1px solid ${cssVarV2('layer/insideBorder/border')}`,
+  borderRadius: 10,
+});
+
+export const sectionHeading = style({
+  display: 'flex',
+  alignItems: 'flex-start',
+  justifyContent: 'space-between',
+  gap: 12,
+});
+
+export const sectionTitle = style({
+  fontSize: cssVar('fontSm'),
+  fontWeight: 600,
+  color: cssVarV2('text/primary'),
+});
+
+export const storageOptions = style({
+  display: 'grid',
+  gridTemplateColumns: 'repeat(2, minmax(0, 1fr))',
+  gap: 8,
+});
+
+export const storageOption = style({
+  position: 'relative',
+  display: 'flex',
+  alignItems: 'flex-start',
+  minHeight: 76,
+  boxSizing: 'border-box',
+  padding: 12,
+  border: `1px solid ${cssVarV2('layer/insideBorder/border')}`,
+  borderRadius: 8,
+  color: cssVarV2('text/primary'),
+  fontSize: cssVar('fontSm'),
+  cursor: 'pointer',
+  selectors: {
+    '&:has(input:checked)': {
+      borderColor: cssVarV2('button/primary'),
+      background: cssVarV2('layer/background/secondary'),
+    },
+    '&:has(input:focus-visible)': {
+      boxShadow: '0px 0px 0px 2px rgba(30, 150, 235, 0.30)',
+    },
+    '&[data-disabled="true"]': {
+      cursor: 'not-allowed',
+      color: cssVarV2('text/disable'),
+      background: cssVarV2('layer/background/secondary'),
+    },
+  },
+});
+
+export const storageRadio = style({
+  position: 'absolute',
+  width: 1,
+  height: 1,
+  margin: 0,
+  opacity: 0,
+  pointerEvents: 'none',
+});
+
+export const storageCopy = style({
+  display: 'flex',
+  minWidth: 0,
+  flexDirection: 'column',
+  gap: 2,
+  lineHeight: '20px',
+});
+
+export const storageDescription = style({
+  color: cssVarV2('text/secondary'),
+  fontSize: cssVar('fontXs'),
+  selectors: {
+    [`${storageOption}[data-disabled="true"] &`]: {
+      color: cssVarV2('text/disable'),
+    },
+  },
+});
+
+export const checkboxRow = style({
+  display: 'flex',
+  alignItems: 'center',
+  gap: 6,
+  fontSize: cssVar('fontXs'),
+  color: cssVarV2('text/primary'),
+});
+
+export const modelToolbar = style({
+  display: 'flex',
+  alignItems: 'center',
+  justifyContent: 'space-between',
+  gap: 12,
+});
+
+export const selectedModels = style({
+  display: 'flex',
+  flexDirection: 'column',
+  gap: 6,
+  margin: 0,
+  padding: 0,
+  listStyle: 'none',
+});
+
+export const selectedModel = style({
+  display: 'grid',
+  gridTemplateColumns: '20px minmax(0, 1fr) auto auto auto',
+  alignItems: 'center',
+  gap: 10,
+  minHeight: 72,
+  padding: '10px 12px',
+  borderRadius: 8,
+  background: cssVarV2('layer/background/secondary'),
+  color: cssVarV2('text/primary'),
+  fontSize: cssVar('fontSm'),
+});
+
+export const selectedModelDisabled = style({
+  opacity: 0.58,
+});
+
+export const modelDragHandle = style({
+  display: 'flex',
+  alignItems: 'center',
+  justifyContent: 'center',
+  color: cssVarV2('text/secondary'),
+  cursor: 'grab',
+});
+
+export const modelCopy = style({
+  display: 'flex',
+  minWidth: 0,
+  flexDirection: 'column',
+  gap: 2,
+});
+
+export const modelStatus = style({
+  color: cssVarV2('text/secondary'),
+  fontSize: cssVar('fontXs'),
+  whiteSpace: 'nowrap',
+});
+
+export const recommended = style({
+  padding: '2px 6px',
+  borderRadius: 999,
+  color: cssVarV2('button/primary'),
+  background: cssVarV2('chip/label/blue'),
+  fontSize: 11,
+  fontWeight: 400,
+  lineHeight: '16px',
+});
+
+export const modelEmpty = style({
+  padding: '24px 16px',
+  borderRadius: 8,
+  textAlign: 'center',
+  color: cssVarV2('text/secondary'),
+  background: cssVarV2('layer/background/secondary'),
+  fontSize: cssVar('fontXs'),
+});
+
+export const modelModalBody = style({
+  display: 'flex',
+  flexDirection: 'column',
+  gap: 12,
+  maxHeight: 'min(440px, calc(100dvh - 220px))',
+  overflowY: 'auto',
+});
+
+export const modelModalDescription = style({
+  margin: '0 0 12px',
+  color: cssVarV2('text/secondary'),
+  fontSize: cssVar('fontSm'),
+  lineHeight: '20px',
+});
+
+export const modelSearch = style({
+  flexShrink: 0,
+});
+
+export const modelFieldLabel = style({
+  color: cssVarV2('text/secondary'),
+  fontSize: cssVar('fontSm'),
+  fontWeight: 500,
+  lineHeight: '20px',
+});
+
+export const catalogChoices = style({
+  display: 'flex',
+  flexDirection: 'column',
+  gap: 4,
+  paddingBottom: 8,
+});
+
+export const catalogChoice = style({
+  display: 'grid',
+  gridTemplateColumns: '16px minmax(0, 1fr)',
+  alignItems: 'center',
+  columnGap: 10,
+  minHeight: 56,
+  boxSizing: 'border-box',
+  padding: '8px 10px',
+  border: '1px solid transparent',
+  borderRadius: 8,
+  background: cssVarV2('layer/background/secondary'),
+  color: cssVarV2('text/primary'),
+  fontSize: cssVar('fontSm'),
+  cursor: 'pointer',
+  selectors: {
+    '&[data-selected="true"]': {
+      borderColor: cssVarV2('button/primary'),
+      background: cssVarV2('chip/label/blue'),
+    },
+  },
+});
+
+export const modelCheckbox = style({
+  flex: '0 0 auto',
+  fontSize: 16,
+});
+
+export const catalogModelCopy = style({
+  display: 'flex',
+  minWidth: 0,
+  flexDirection: 'column',
+  gap: 2,
+});
+
+export const catalogModelTitle = style({
+  display: 'flex',
+  minWidth: 0,
+  alignItems: 'center',
+  gap: 6,
+  fontSize: cssVar('fontSm'),
+  lineHeight: '20px',
+});
+
+export const catalogModelMeta = style({
+  overflow: 'hidden',
+  color: cssVarV2('text/secondary'),
+  fontSize: cssVar('fontXs'),
+  lineHeight: '16px',
+  textOverflow: 'ellipsis',
+  whiteSpace: 'nowrap',
+});
+
+export const modelCapabilities = style({
+  minWidth: 0,
+  margin: 0,
+  padding: 0,
+  border: 0,
+});
+
+export const useCaseGrid = style({
+  display: 'grid',
+  gridTemplateColumns: 'repeat(2, minmax(0, 1fr))',
+  columnGap: 16,
+  rowGap: 4,
+  marginTop: 8,
+  '@media': {
+    '(max-width: 600px)': {
+      gridTemplateColumns: '1fr',
+    },
+  },
+});
+
+export const modelUseCase = style({
+  width: '100%',
+  minHeight: 28,
+  gap: 8,
+  color: cssVarV2('text/primary'),
+  fontSize: 16,
+  lineHeight: '20px',
+});
+
+export const modelUseCaseLabel = style({
+  fontSize: cssVar('fontSm'),
+});
+
+export const modelModalActions = style({
+  display: 'flex',
+  justifyContent: 'flex-end',
+  gap: 8,
+  marginTop: 12,
+  paddingTop: 12,
+  borderTop: `1px solid ${cssVarV2('layer/insideBorder/border')}`,
+});
+
+export const advanced = style({
+  color: cssVarV2('text/secondary'),
+  fontSize: cssVar('fontSm'),
+});
+
+export const advancedFields = style({
+  display: 'flex',
+  flexDirection: 'column',
+  gap: 10,
+  marginTop: 10,
+});
+
+export const inputStack = style({
+  display: 'flex',
+  flexDirection: 'column',
+  gap: 4,
 });
 
 export const testStatus = style({

packages/frontend/core/src/desktop/dialogs/setting/workspace-setting/byok/index.tsx

@@ -6,10 +6,10 @@ import {
 import { WorkspaceServerService } from '@affine/core/modules/cloud';
 import { WorkspaceService } from '@affine/core/modules/workspace';
 import {
-  ByokKeyStorage,
-  clearWorkspaceByokConfigsMutation as clearByokMutation,
-  deleteWorkspaceByokConfigMutation as deleteByokMutation,
+  deleteWorkspaceByokProfileMutation,
   type GraphQLQuery,
+  probeWorkspaceByokProfileMutation,
+  reorderWorkspaceByokProfilesMutation,
   ServerDeploymentType,
   workspaceByokSettingsQuery as byokSettingsQuery,
 } from '@affine/graphql';
@@ -29,27 +29,12 @@ import {
   readLocalKeys,
   reorderLocalKeys,
 } from './local-storage';
-import { byokT } from './metadata';
-import type {
-  ByokKey,
-  ByokSettings,
-  ByokStorage,
-  ByokUsagePoint,
-  GqlFn,
-} from './types';
+import { byokT, capabilitiesFor } from './metadata';
+import { probeChecks } from './model-utils';
+import type { ByokKey, ByokSettings, ByokUsagePoint, GqlFn } from './types';
+import { ByokStorage } from './types';
 import { UsagePanel } from './usage';
 
-const reorderByokMutation = {
-  id: 'reorderWorkspaceByokConfigsMutation',
-  op: 'reorderWorkspaceByokConfigs',
-  query: `mutation reorderWorkspaceByokConfigs($input: ReorderWorkspaceByokConfigsInput!) {
-    reorderWorkspaceByokConfigs(input: $input) {
-      id
-      sortOrder
-    }
-  }`,
-} satisfies GraphQLQuery;
-
 export const WorkspaceByokSetting = () => {
   const t = useI18n();
   const workspace = useService(WorkspaceService).workspace;
@@ -59,6 +44,7 @@ export const WorkspaceByokSetting = () => {
   const [localKeys, setLocalKeys] = useState<ByokKey[]>([]);
   const [modalOpen, setModalOpen] = useState(false);
   const [editingKey, setEditingKey] = useState<ByokKey | null>(null);
+  const [testingKeyId, setTestingKeyId] = useState<string | null>(null);
   const [draggingKey, setDraggingKey] = useState<{
     id: string;
     storage: ByokStorage;
@@ -83,8 +69,27 @@ export const WorkspaceByokSetting = () => {
       localByokStorageSupported(),
       readLocalKeys(workspace.id),
     ]);
+    const serverKeys = data.workspace.byokSettings.profiles.map(profile => {
+      const key: ByokKey = {
+        id: profile.profileId,
+        provider: profile.provider,
+        name: profile.name,
+        description: profile.description,
+        storage: ByokStorage.server,
+        configured: true,
+        enabled: profile.enabled,
+        sortOrder: profile.sortOrder,
+        revision: profile.revision,
+        definition: profile.definition,
+        capabilities: [],
+        validation: profile.validation,
+      };
+      key.capabilities = capabilitiesFor(key);
+      return key;
+    });
     setSettings({
       ...data.workspace.byokSettings,
+      keys: serverKeys,
       localStorageSupported:
         data.workspace.byokSettings.localEntitled && localStorageSupported,
     });
@@ -105,7 +110,7 @@ export const WorkspaceByokSetting = () => {
   const keys = useMemo(() => {
     return [...localKeys, ...(settings?.keys ?? [])].toSorted((a, b) => {
       if (a.storage !== b.storage) {
-        return a.storage === ByokKeyStorage.local ? -1 : 1;
+        return a.storage === ByokStorage.local ? -1 : 1;
       }
       return a.sortOrder - b.sortOrder;
     });
@@ -123,26 +128,35 @@ export const WorkspaceByokSetting = () => {
     if (!settings) {
       return;
     }
-    if (!workspaceServer.server && settings.serverEntitled) {
-      return;
-    }
+    const deletions: Promise<unknown>[] = [];
     if (settings.serverEntitled && workspaceServer.server) {
       const gql = workspaceServer.server.gql as GqlFn;
-      await gql({
-        query: clearByokMutation,
-        variables: { workspaceId: workspace.id },
-      });
+      deletions.push(
+        ...settings.keys.map(key =>
+          gql({
+            query: deleteWorkspaceByokProfileMutation,
+            variables: { workspaceId: workspace.id, profileId: key.id },
+          })
+        )
+      );
     }
     if (settings.localStorageSupported) {
-      await clearLocalKeys(workspace.id);
+      deletions.push(clearLocalKeys(workspace.id));
     }
-    setLocalKeys([]);
+    const results = await Promise.allSettled(deletions);
     await load();
+    if (
+      results.some(
+        result => result.status === 'rejected' || result.value === false
+      )
+    ) {
+      throw new Error('Some BYOK profiles could not be deleted');
+    }
   }, [load, settings, workspace.id, workspaceServer.server]);
 
   const deleteKey = useCallback(
     async (key: ByokKey) => {
-      if (key.storage === ByokKeyStorage.local) {
+      if (key.storage === ByokStorage.local) {
         await deleteLocalKey(workspace.id, key.id);
         setLocalKeys(await readLocalKeys(workspace.id));
         return;
@@ -154,14 +168,39 @@ export const WorkspaceByokSetting = () => {
           }) => Promise<unknown>)
         | undefined;
       await gql?.({
-        query: deleteByokMutation,
-        variables: { workspaceId: workspace.id, id: key.id },
+        query: deleteWorkspaceByokProfileMutation,
+        variables: { workspaceId: workspace.id, profileId: key.id },
       });
       await load();
     },
     [load, workspace.id, workspaceServer.server]
   );
 
+  const testKey = useCallback(
+    async (key: ByokKey) => {
+      if (key.storage !== ByokStorage.server || !workspaceServer.server) {
+        return;
+      }
+      setTestingKeyId(key.id);
+      try {
+        await (workspaceServer.server.gql as GqlFn)({
+          query: probeWorkspaceByokProfileMutation,
+          variables: {
+            input: {
+              workspaceId: workspace.id,
+              profileId: key.id,
+              checks: probeChecks(key.definition.models, false),
+            },
+          },
+        });
+        await load();
+      } finally {
+        setTestingKeyId(null);
+      }
+    },
+    [load, workspace.id, workspaceServer.server]
+  );
+
   const reorderKey = useCallback(
     async (targetKey: ByokKey) => {
       if (!draggingKey || draggingKey.id === targetKey.id) {
@@ -187,38 +226,44 @@ export const WorkspaceByokSetting = () => {
       nextBucket.splice(toIndex, 0, moved);
       const nextBucketIds = nextBucket.map(key => key.id);
 
-      if (targetKey.storage === ByokKeyStorage.local) {
+      if (targetKey.storage === ByokStorage.local) {
         setLocalKeys(await reorderLocalKeys(workspace.id, nextBucketIds));
-        return;
-      }
-
-      const gql = workspaceServer.server?.gql as
-        | ((input: {
-            query: GraphQLQuery;
-            variables?: Record<string, unknown>;
-          }) => Promise<unknown>)
-        | undefined;
-      await gql?.({
-        query: reorderByokMutation,
-        variables: {
-          input: {
-            workspaceId: workspace.id,
-            storage: ByokKeyStorage.server,
-            ids: nextBucketIds,
+      } else if (workspaceServer.server) {
+        if (nextBucket.some(key => key.revision === undefined)) {
+          notify.error({
+            title: byokT(t, 'notify.reload-required.title'),
+            message: byokT(t, 'notify.reload-required.message'),
+          });
+          await load();
+          return;
+        }
+        await (workspaceServer.server.gql as GqlFn)({
+          query: reorderWorkspaceByokProfilesMutation,
+          variables: {
+            input: {
+              workspaceId: workspace.id,
+              profiles: nextBucket.flatMap(key =>
+                key.revision === undefined
+                  ? []
+                  : [
+                      {
+                        profileId: key.id,
+                        expectedRevision: key.revision,
+                      },
+                    ]
+              ),
+            },
           },
-        },
-      });
-      await load();
+        });
+        await load();
+      }
     },
     [draggingKey, keys, load, t, workspace.id, workspaceServer.server]
   );
 
   if (!settings) {
     return (
-      <SettingHeader
-        title={byokT(t, 'title-beta')}
-        subtitle={byokT(t, 'loading')}
-      />
+      <SettingHeader title={byokT(t, 'title')} subtitle={byokT(t, 'loading')} />
     );
   }
 
@@ -226,7 +271,7 @@ export const WorkspaceByokSetting = () => {
     return (
       <>
         <SettingHeader
-          title={byokT(t, 'title-beta')}
+          title={byokT(t, 'title')}
           subtitle={byokT(t, 'subtitle')}
         />
         <SettingWrapper>
@@ -237,13 +282,7 @@ export const WorkspaceByokSetting = () => {
                 {byokT(t, 'locked.description')}
               </div>
             </div>
-            <div className={styles.tags}>
-              {settings.entitlementRequired.map(plan => (
-                <span className={styles.tag} key={plan}>
-                  {plan}
-                </span>
-              ))}
-            </div>
+            <div className={styles.tags}></div>
           </div>
         </SettingWrapper>
       </>
@@ -252,21 +291,9 @@ export const WorkspaceByokSetting = () => {
 
   return (
     <>
-      <SettingHeader
-        title={byokT(t, 'title-beta')}
-        subtitle={byokT(t, 'header')}
-      />
+      <SettingHeader title={byokT(t, 'title')} subtitle={byokT(t, 'header')} />
       <SettingWrapper>
         <div className={styles.stack}>
-          {settings.hasAiPlan ? (
-            <div className={styles.notice}>
-              <div className={styles.title}>{byokT(t, 'notice.title')}</div>
-              <div className={styles.description}>
-                {byokT(t, 'notice.description')}
-              </div>
-            </div>
-          ) : null}
-
           <div className={styles.panel} data-testid="workspace-byok-keys">
             <div className={styles.panelHeader}>
               <div>
@@ -289,6 +316,7 @@ export const WorkspaceByokSetting = () => {
             {keys.length ? (
               <KeyList
                 keys={keys}
+                testingKeyId={testingKeyId}
                 onEdit={key => {
                   setEditingKey(key);
                   setModalOpen(true);
@@ -302,6 +330,15 @@ export const WorkspaceByokSetting = () => {
                     });
                   });
                 }}
+                onTest={key => {
+                  testKey(key).catch(error => {
+                    logByokError('Failed to test BYOK provider', error);
+                    notify.error({
+                      title: byokT(t, 'notify.test-failed.title'),
+                      message: byokT(t, 'notify.operation-failed.message'),
+                    });
+                  });
+                }}
                 onDragStart={key => {
                   setDraggingKey({ id: key.id, storage: key.storage });
                 }}
@@ -326,7 +363,7 @@ export const WorkspaceByokSetting = () => {
             )}
           </div>
 
-          <CoveragePanel keys={keys} settings={settings} />
+          <CoveragePanel keys={keys} />
 
           <UsagePanel
             keys={keys}

packages/frontend/core/src/desktop/dialogs/setting/workspace-setting/byok/key-list.tsx

@@ -1,4 +1,4 @@
-import { DragHandle, IconButton } from '@affine/component';
+import { Button, DragHandle, IconButton } from '@affine/component';
 import { useI18n } from '@affine/i18n';
 import { DeleteIcon, EditIcon } from '@blocksuite/icons/rc';
 import type { DragEvent } from 'react';
@@ -11,19 +11,23 @@ import {
   rowDescription,
   storageLabel,
 } from './metadata';
-import type { ByokKey } from './types';
+import { type ByokKey, ByokStorage } from './types';
 
 export const KeyList = ({
   keys,
+  testingKeyId,
   onEdit,
   onDelete,
+  onTest,
   onDragStart,
   onDragEnd,
   onDrop,
 }: {
   keys: ByokKey[];
+  testingKeyId: string | null;
   onEdit: (key: ByokKey) => void;
   onDelete: (key: ByokKey) => void;
+  onTest: (key: ByokKey) => void;
   onDragStart: (key: ByokKey) => void;
   onDragEnd: () => void;
   onDrop: (key: ByokKey) => void;
@@ -72,6 +76,18 @@ export const KeyList = ({
             </div>
           </div>
           <div className={styles.rowActions}>
+            {key.storage === ByokStorage.server ? (
+              <Button
+                variant="plain"
+                disabled={testingKeyId !== null}
+                onClick={() => onTest(key)}
+              >
+                {byokT(
+                  t,
+                  testingKeyId === key.id ? 'action.testing' : 'action.test'
+                )}
+              </Button>
+            ) : null}
             <IconButton
               size="20"
               title={byokT(t, 'action.edit')}

packages/frontend/core/src/desktop/dialogs/setting/workspace-setting/byok/local-storage.ts

@@ -1,8 +1,12 @@
 import { apis } from '@affine/electron-api';
-import { ByokKeyStorage, ByokKeyTestStatus } from '@affine/graphql';
 
 import { capabilitiesFor } from './metadata';
-import type { ByokKey, LocalByokKeyInput, LocalByokPublicKey } from './types';
+import {
+  type ByokKey,
+  ByokStorage,
+  type LocalByokKeyInput,
+  type LocalByokPublicKey,
+} from './types';
 
 function byokStorageApi() {
   return BUILD_CONFIG.isElectron ? apis?.byokStorage : undefined;
@@ -26,14 +30,12 @@ function toLocalByokKey(key: LocalByokPublicKey): ByokKey {
     provider: key.provider,
     name: key.name,
     description: key.description ?? null,
-    storage: ByokKeyStorage.local,
+    storage: ByokStorage.local,
     configured: key.configured ?? true,
     enabled: key.enabled ?? true,
-    endpoint: key.endpoint ?? null,
-    endpointEditable: key.endpointEditable ?? false,
     sortOrder: key.sortOrder ?? 0,
-    capabilities: capabilitiesFor(key.provider, ByokKeyStorage.local),
-    testStatus: key.testStatus ?? ByokKeyTestStatus.passed,
+    definition: key.definition,
+    capabilities: capabilitiesFor(key),
   };
 }
 

packages/frontend/core/src/desktop/dialogs/setting/workspace-setting/byok/metadata.spec.ts

@@ -1,34 +0,0 @@
-import { describe, expect, test } from 'vitest';
-
-import { endpointHintKey, shouldShowEndpoint } from './metadata';
-
-describe('endpointHintKey', () => {
-  test.each([
-    [false, false, 'endpoint.custom-disabled'],
-    [true, false, 'endpoint.private-disabled'],
-    [true, true, null],
-  ] as const)(
-    'maps custom=%s private=%s to %s',
-    (customEndpointSupported, privateEndpointSupported, expected) => {
-      expect(
-        endpointHintKey(customEndpointSupported, privateEndpointSupported)
-      ).toBe(expected);
-    }
-  );
-});
-
-describe('shouldShowEndpoint', () => {
-  test.each([
-    [true, false, true],
-    [true, true, true],
-    [false, true, true],
-    [false, false, false],
-  ])(
-    'self-hosted %s with custom endpoint support %s returns %s',
-    (isSelfHosted, customEndpointSupported, expected) => {
-      expect(shouldShowEndpoint(isSelfHosted, customEndpointSupported)).toBe(
-        expected
-      );
-    }
-  );
-});

packages/frontend/core/src/desktop/dialogs/setting/workspace-setting/byok/metadata.ts

@@ -1,7 +1,7 @@
-import { ByokKeyStorage, ByokProvider } from '@affine/graphql';
+import { ByokProvider } from '@affine/graphql';
 import type { I18nInstance } from '@affine/i18n';
 
-import type { ByokKey, ByokStorage } from './types';
+import { type ByokKey, ByokStorage } from './types';
 
 export function byokT(
   t: I18nInstance,
@@ -19,7 +19,7 @@ export const providerLabels: Record<ByokProvider, string> = {
 };
 
 export function storageLabel(t: I18nInstance, storage: ByokStorage) {
-  return storage === ByokKeyStorage.local
+  return storage === ByokStorage.local
     ? byokT(t, 'storage.local')
     : byokT(t, 'storage.server');
 }
@@ -44,26 +44,20 @@ export function shouldShowEndpoint(
   return isSelfHosted || customEndpointSupported;
 }
 
-export function capabilitiesFor(provider: ByokProvider, storage: ByokStorage) {
-  switch (provider) {
-    case ByokProvider.openai:
-      return ['Text', 'Image input', 'Actions', 'Image generate'];
-    case ByokProvider.anthropic:
-      return ['Text', 'Image input'];
-    case ByokProvider.gemini:
-      return storage === ByokKeyStorage.server
-        ? [
-            'Text',
-            'Image input',
-            'Actions',
-            'Image generate',
-            'Transcript',
-            'Indexing',
-          ]
-        : ['Text', 'Image input', 'Actions', 'Image generate'];
-    case ByokProvider.fal:
-      return ['Image generate'];
+export function capabilitiesFor(key: Pick<ByokKey, 'definition'>) {
+  const capabilities = key.definition.models.flatMap(model =>
+    model.enabled ? model.capabilities : []
+  );
+  const labels = new Set<string>();
+  for (const capability of capabilities) {
+    if (capability.output.includes('text')) labels.add('Text');
+    if (capability.input.includes('image')) labels.add('Image input');
+    if (capability.output.includes('image')) labels.add('Image generate');
+    if (capability.features.includes('tools')) labels.add('Actions');
+    if (capability.input.includes('audio')) labels.add('Transcript');
+    if (capability.output.includes('embedding')) labels.add('Indexing');
   }
+  return [...labels];
 }
 
 export function capabilityLabel(t: I18nInstance, capability: string) {
@@ -121,7 +115,7 @@ export const capabilityRows = [
     icon: 'transcript',
     providers: [ByokProvider.gemini],
     coverageCapabilities: ['Transcript'],
-    storage: ByokKeyStorage.server,
+    storage: ByokStorage.server,
   },
   {
     titleKey: 'feature.workspace-indexing.title',
@@ -130,7 +124,7 @@ export const capabilityRows = [
     icon: 'indexing',
     providers: [ByokProvider.gemini],
     coverageCapabilities: ['Indexing'],
-    storage: ByokKeyStorage.server,
+    storage: ByokStorage.server,
   },
 ] as const;
 
@@ -145,16 +139,13 @@ function formatDate(value?: string | null) {
 }
 
 export function rowDescription(t: I18nInstance, key: ByokKey) {
-  const failed = formatDate(key.lastErrorAt);
-  const used = formatDate(key.lastUsedAt);
-  const today = formatDate(new Date().toISOString());
-  const activity = failed
-    ? byokT(t, 'row.activity.failed', { date: failed })
-    : used
-      ? used === today
-        ? byokT(t, 'row.activity.used-today')
-        : byokT(t, 'row.activity.used', { date: used })
-      : byokT(t, 'row.activity.unused');
+  const tested = formatDate(key.validation?.connection.testedAt);
+  const activity =
+    key.validation?.connection.kind === 'failed'
+      ? byokT(t, 'row.activity.failed', { date: tested ?? '' })
+      : key.validation?.connection.kind === 'verified'
+        ? byokT(t, 'status.key-verified')
+        : byokT(t, 'row.activity.unused');
 
   return [storageLabel(t, key.storage), activity, key.description]
     .filter(Boolean)

packages/frontend/core/src/desktop/dialogs/setting/workspace-setting/byok/model-editor-modal.tsx

@@ -0,0 +1,256 @@
+import { Button, Checkbox, Input, Modal } from '@affine/component';
+import { useI18n } from '@affine/i18n';
+import { useEffect, useMemo, useState } from 'react';
+
+import * as styles from './index.css';
+import { byokT } from './metadata';
+import {
+  capabilitiesForUseCases,
+  type catalogModels,
+  type ModelDeclaration,
+  modelUseCases,
+  type UseCase,
+  useCases,
+} from './model-utils';
+
+export const ModelEditorModal = ({
+  open,
+  customEndpoint,
+  catalog,
+  models,
+  editingModel,
+  onOpenChange,
+  onSubmit,
+}: {
+  open: boolean;
+  customEndpoint: boolean;
+  catalog: ReturnType<typeof catalogModels>;
+  models: ModelDeclaration[];
+  editingModel: ModelDeclaration | null;
+  onOpenChange: (open: boolean) => void;
+  onSubmit: (models: ModelDeclaration[]) => void;
+}) => {
+  const t = useI18n();
+  const [search, setSearch] = useState('');
+  const [selectedIds, setSelectedIds] = useState<string[]>([]);
+  const [modelId, setModelId] = useState('');
+  const [selectedUseCases, setSelectedUseCases] = useState<UseCase[]>(['chat']);
+
+  useEffect(() => {
+    if (!open) return;
+    setSearch('');
+    setSelectedIds([]);
+    setModelId(editingModel?.modelId ?? '');
+    setSelectedUseCases(editingModel ? modelUseCases(editingModel) : ['chat']);
+  }, [editingModel, open]);
+
+  const availableCatalog = useMemo(() => {
+    return catalog
+      .filter(item => !models.some(model => model.modelId === item.modelId))
+      .sort(
+        (left, right) => Number(right.recommended) - Number(left.recommended)
+      );
+  }, [catalog, models]);
+  const available = useMemo(() => {
+    const query = search.trim().toLocaleLowerCase();
+    return availableCatalog.filter(
+      item =>
+        !query ||
+        item.displayName.toLocaleLowerCase().includes(query) ||
+        item.modelId.toLocaleLowerCase().includes(query)
+    );
+  }, [availableCatalog, search]);
+
+  const submit = () => {
+    if (customEndpoint) {
+      onSubmit([
+        {
+          modelId: modelId.trim(),
+          enabled: editingModel?.enabled ?? true,
+          capabilities: capabilitiesForUseCases(editingModel, selectedUseCases),
+        },
+      ]);
+    } else {
+      onSubmit(
+        selectedIds.flatMap(id => {
+          const model = catalog.find(item => item.modelId === id);
+          return model
+            ? [
+                {
+                  modelId: model.modelId,
+                  enabled: true,
+                  capabilities: model.capabilities,
+                },
+              ]
+            : [];
+        })
+      );
+    }
+    onOpenChange(false);
+  };
+
+  const normalizedModelId = modelId.trim();
+  const duplicateModelId = models.some(
+    model => model !== editingModel && model.modelId === normalizedModelId
+  );
+  const valid = customEndpoint
+    ? !!normalizedModelId && !duplicateModelId && selectedUseCases.length > 0
+    : selectedIds.length > 0;
+
+  return (
+    <Modal
+      width={customEndpoint ? 520 : 560}
+      open={open}
+      onOpenChange={onOpenChange}
+      descriptionClassName={styles.modelModalDescription}
+      title={byokT(
+        t,
+        editingModel
+          ? 'modal.edit-model-title'
+          : customEndpoint
+            ? 'modal.add-custom-model-title'
+            : 'modal.add-model-title'
+      )}
+      description={byokT(
+        t,
+        customEndpoint
+          ? 'modal.custom-model-description'
+          : 'modal.catalog-model-description'
+      )}
+    >
+      {customEndpoint ? (
+        <div className={styles.modelModalBody}>
+          <label className={styles.field}>
+            <span className={styles.modelFieldLabel}>
+              {byokT(t, 'field.model-id')}
+            </span>
+            <Input
+              size="large"
+              value={modelId}
+              onChange={setModelId}
+              placeholder={byokT(t, 'placeholder.model-id')}
+            />
+            {duplicateModelId ? (
+              <span className={styles.error}>
+                {byokT(t, 'model.duplicate-id')}
+              </span>
+            ) : null}
+          </label>
+          <fieldset className={styles.modelCapabilities}>
+            <legend className={styles.modelFieldLabel}>
+              {byokT(t, 'model.use-this-for')}
+            </legend>
+            <div className={styles.useCaseGrid}>
+              {useCases.map(useCase => (
+                <Checkbox
+                  className={styles.modelUseCase}
+                  labelClassName={styles.modelUseCaseLabel}
+                  key={useCase.id}
+                  name={`byok-model-use-${useCase.id}`}
+                  label={byokT(t, useCase.labelKey)}
+                  checked={selectedUseCases.includes(useCase.id)}
+                  onChange={(_, checked) =>
+                    setSelectedUseCases(
+                      checked
+                        ? [...selectedUseCases, useCase.id]
+                        : selectedUseCases.filter(item => item !== useCase.id)
+                    )
+                  }
+                />
+              ))}
+            </div>
+          </fieldset>
+        </div>
+      ) : (
+        <div className={styles.modelModalBody}>
+          {availableCatalog.length > 6 ? (
+            <Input
+              className={styles.modelSearch}
+              size="large"
+              value={search}
+              onChange={setSearch}
+              placeholder={byokT(t, 'placeholder.search-models')}
+            />
+          ) : null}
+          <div className={styles.catalogChoices}>
+            {available.length ? (
+              available.map(model => {
+                const modelUses = modelUseCases({
+                  modelId: model.modelId,
+                  enabled: true,
+                  capabilities: model.capabilities,
+                });
+                const visibleUses = modelUses.slice(0, 3).flatMap(useCase => {
+                  const item = useCases.find(item => item.id === useCase);
+                  return item ? [byokT(t, item.labelKey)] : [];
+                });
+                if (modelUses.length > 3) {
+                  visibleUses.push(`+${modelUses.length - 3}`);
+                }
+                return (
+                  <label
+                    className={styles.catalogChoice}
+                    data-selected={selectedIds.includes(model.modelId)}
+                    key={model.modelId}
+                  >
+                    <Checkbox
+                      className={styles.modelCheckbox}
+                      aria-label={model.displayName}
+                      checked={selectedIds.includes(model.modelId)}
+                      onChange={(_, checked) =>
+                        setSelectedIds(
+                          checked
+                            ? [...selectedIds, model.modelId]
+                            : selectedIds.filter(id => id !== model.modelId)
+                        )
+                      }
+                    />
+                    <span className={styles.catalogModelCopy}>
+                      <span className={styles.catalogModelTitle}>
+                        <strong>{model.displayName}</strong>
+                        {model.recommended ? (
+                          <span className={styles.recommended}>
+                            {byokT(t, 'model.recommended')}
+                          </span>
+                        ) : null}
+                      </span>
+                      <span
+                        className={styles.catalogModelMeta}
+                        title={[model.modelId, ...visibleUses].join(' · ')}
+                      >
+                        {[model.modelId, ...visibleUses].join(' · ')}
+                      </span>
+                    </span>
+                  </label>
+                );
+              })
+            ) : (
+              <div className={styles.modelEmpty}>
+                {byokT(
+                  t,
+                  search ? 'models.no-search-results' : 'models.all-added'
+                )}
+              </div>
+            )}
+          </div>
+        </div>
+      )}
+      <div className={styles.modelModalActions}>
+        <Button variant="secondary" onClick={() => onOpenChange(false)}>
+          {byokT(t, 'action.cancel')}
+        </Button>
+        <Button variant="primary" disabled={!valid} onClick={submit}>
+          {byokT(
+            t,
+            editingModel
+              ? 'action.save-model'
+              : customEndpoint
+                ? 'action.add-model'
+                : 'action.add-selected-models',
+            { count: selectedIds.length }
+          )}
+        </Button>
+      </div>
+    </Modal>
+  );
+};

packages/frontend/core/src/desktop/dialogs/setting/workspace-setting/byok/model-selector.tsx

@@ -0,0 +1,227 @@
+import {
+  Button,
+  DragHandle,
+  IconButton,
+  Menu,
+  MenuItem,
+  Switch,
+} from '@affine/component';
+import { useI18n } from '@affine/i18n';
+import { MoreHorizontalIcon } from '@blocksuite/icons/rc';
+import { useState } from 'react';
+
+import * as styles from './index.css';
+import { byokT } from './metadata';
+import { ModelEditorModal } from './model-editor-modal';
+import {
+  type catalogModels,
+  type ModelDeclaration,
+  modelUseCases,
+  useCases,
+} from './model-utils';
+import type { ByokKey } from './types';
+
+export const ModelSelector = ({
+  customEndpoint,
+  catalog,
+  models,
+  validation,
+  onChange,
+}: {
+  customEndpoint: boolean;
+  catalog: ReturnType<typeof catalogModels>;
+  models: ModelDeclaration[];
+  validation?: ByokKey['validation'];
+  onChange: (models: ModelDeclaration[]) => void;
+}) => {
+  const t = useI18n();
+  const [editorOpen, setEditorOpen] = useState(false);
+  const [editingIndex, setEditingIndex] = useState<number | null>(null);
+  const [draggingIndex, setDraggingIndex] = useState<number | null>(null);
+
+  const update = (index: number, model: ModelDeclaration) => {
+    onChange(models.map((current, i) => (i === index ? model : current)));
+  };
+  const move = (index: number, offset: number) => {
+    const target = index + offset;
+    if (target < 0 || target >= models.length) return;
+    const next = [...models];
+    [next[index], next[target]] = [next[target], next[index]];
+    onChange(next);
+  };
+  const drop = (targetIndex: number) => {
+    if (draggingIndex === null || draggingIndex === targetIndex) return;
+    const next = [...models];
+    const [dragged] = next.splice(draggingIndex, 1);
+    next.splice(targetIndex, 0, dragged);
+    onChange(next);
+    setDraggingIndex(null);
+  };
+  const evidence = (modelId: string) => {
+    const checks = validation?.models.find(
+      model => model.modelId === modelId
+    )?.checks;
+    if (!checks?.length) return byokT(t, 'model.status.not-tested');
+    const verified = checks.filter(
+      check => check.status.kind === 'verified'
+    ).length;
+    if (verified === checks.length) return byokT(t, 'model.status.verified');
+    if (verified === 0) return byokT(t, 'model.status.failed');
+    return byokT(t, 'model.status.partially-verified', {
+      verified,
+      total: checks.length,
+    });
+  };
+
+  return (
+    <>
+      <div className={styles.modelToolbar}>
+        <span className={styles.description}>
+          {byokT(t, 'models.description.order')}
+        </span>
+        <Button
+          variant="secondary"
+          onClick={() => {
+            setEditingIndex(null);
+            setEditorOpen(true);
+          }}
+        >
+          {byokT(t, 'action.add-model')}
+        </Button>
+      </div>
+      {models.length ? (
+        <ol className={styles.selectedModels}>
+          {models.map((model, index) => {
+            const catalogModel = catalog.find(
+              item => item.modelId === model.modelId
+            );
+            const selected = modelUseCases(model);
+            return (
+              <li
+                className={`${styles.selectedModel} ${
+                  model.enabled ? '' : styles.selectedModelDisabled
+                }`}
+                key={model.modelId}
+                onDragOver={event => event.preventDefault()}
+                onDrop={event => {
+                  event.preventDefault();
+                  drop(index);
+                }}
+              >
+                <div
+                  className={styles.modelDragHandle}
+                  draggable
+                  title={byokT(t, 'action.reorder')}
+                  onDragStart={() => setDraggingIndex(index)}
+                  onDragEnd={() => setDraggingIndex(null)}
+                >
+                  <DragHandle dragging={draggingIndex === index} />
+                </div>
+                <div className={styles.modelCopy}>
+                  <strong>{catalogModel?.displayName ?? model.modelId}</strong>
+                  {catalogModel?.displayName ? (
+                    <small>{model.modelId}</small>
+                  ) : null}
+                  <span className={styles.tags}>
+                    {selected.slice(0, 3).map(useCase => {
+                      const item = useCases.find(item => item.id === useCase);
+                      return item ? (
+                        <span className={styles.tag} key={useCase}>
+                          {byokT(t, item.labelKey)}
+                        </span>
+                      ) : null;
+                    })}
+                    {selected.length > 3 ? (
+                      <span className={styles.tag}>+{selected.length - 3}</span>
+                    ) : null}
+                  </span>
+                </div>
+                <span className={styles.modelStatus}>
+                  {model.enabled
+                    ? evidence(model.modelId)
+                    : byokT(t, 'model.status.disabled')}
+                </span>
+                <Switch
+                  checked={model.enabled}
+                  aria-label={byokT(
+                    t,
+                    model.enabled
+                      ? 'action.disable-model'
+                      : 'action.enable-model',
+                    { model: catalogModel?.displayName ?? model.modelId }
+                  )}
+                  onChange={enabled => update(index, { ...model, enabled })}
+                />
+                <Menu
+                  items={
+                    <>
+                      {customEndpoint ? (
+                        <MenuItem
+                          onSelect={() => {
+                            setEditingIndex(index);
+                            setEditorOpen(true);
+                          }}
+                        >
+                          {byokT(t, 'action.edit')}
+                        </MenuItem>
+                      ) : null}
+                      <MenuItem
+                        disabled={index === 0}
+                        onSelect={() => move(index, -1)}
+                      >
+                        {byokT(t, 'action.move-up')}
+                      </MenuItem>
+                      <MenuItem
+                        disabled={index === models.length - 1}
+                        onSelect={() => move(index, 1)}
+                      >
+                        {byokT(t, 'action.move-down')}
+                      </MenuItem>
+                      <MenuItem
+                        type="danger"
+                        onSelect={() =>
+                          onChange(models.filter((_, i) => i !== index))
+                        }
+                      >
+                        {byokT(t, 'action.remove')}
+                      </MenuItem>
+                    </>
+                  }
+                >
+                  <IconButton
+                    size="20"
+                    title={byokT(t, 'action.model-options', {
+                      model: catalogModel?.displayName ?? model.modelId,
+                    })}
+                    icon={<MoreHorizontalIcon />}
+                  />
+                </Menu>
+              </li>
+            );
+          })}
+        </ol>
+      ) : (
+        <div className={styles.modelEmpty}>{byokT(t, 'models.empty')}</div>
+      )}
+      <ModelEditorModal
+        open={editorOpen}
+        customEndpoint={customEndpoint}
+        catalog={catalog}
+        models={models}
+        editingModel={editingIndex === null ? null : models[editingIndex]}
+        onOpenChange={open => {
+          setEditorOpen(open);
+          if (!open) setEditingIndex(null);
+        }}
+        onSubmit={next => {
+          if (editingIndex === null) {
+            onChange([...models, ...next]);
+          } else if (next[0]) {
+            update(editingIndex, next[0]);
+          }
+          setEditingIndex(null);
+        }}
+      />
+    </>
+  );
+};

packages/frontend/core/src/desktop/dialogs/setting/workspace-setting/byok/model-utils.spec.ts

@@ -0,0 +1,46 @@
+import { describe, expect, test } from 'vitest';
+
+import {
+  capabilitiesForUseCases,
+  type ModelDeclaration,
+  modelUseCases,
+} from './model-utils';
+
+describe('BYOK model capabilities', () => {
+  test('maps richer catalog capabilities by minimum requirements', () => {
+    const model: ModelDeclaration = {
+      modelId: 'multimodal-tools',
+      enabled: true,
+      capabilities: [
+        {
+          input: ['text', 'image'],
+          output: ['text'],
+          features: ['tools'],
+          attachmentKinds: ['image'],
+          attachmentSources: ['url', 'data', 'bytes', 'file_handle'],
+        },
+      ],
+    };
+
+    expect(modelUseCases(model)).toEqual(['chat', 'actions', 'vision']);
+  });
+
+  test('preserves a rich capability when its represented uses stay selected', () => {
+    const capability = {
+      input: ['text', 'image'],
+      output: ['text'],
+      features: ['tools'],
+      attachmentKinds: ['image'],
+      attachmentSources: ['url', 'data', 'bytes', 'file_handle'],
+    };
+    const model: ModelDeclaration = {
+      modelId: 'multimodal-tools',
+      enabled: true,
+      capabilities: [capability],
+    };
+
+    expect(
+      capabilitiesForUseCases(model, ['chat', 'actions', 'vision'])
+    ).toEqual([capability]);
+  });
+});

packages/frontend/core/src/desktop/dialogs/setting/workspace-setting/byok/model-utils.ts

@@ -0,0 +1,150 @@
+import type { ByokProvider } from '@affine/graphql';
+
+import type { ByokDefinition, ByokSettings } from './types';
+
+export type ModelDeclaration = ByokDefinition['models'][number];
+type Capability = ModelDeclaration['capabilities'][number];
+export type UseCase =
+  | 'chat'
+  | 'actions'
+  | 'structured'
+  | 'vision'
+  | 'image'
+  | 'transcript'
+  | 'embedding'
+  | 'rerank';
+
+export const useCases: { id: UseCase; labelKey: string }[] = [
+  { id: 'chat', labelKey: 'model.use.chat' },
+  { id: 'actions', labelKey: 'model.use.actions' },
+  { id: 'structured', labelKey: 'model.use.structured' },
+  { id: 'vision', labelKey: 'model.use.vision' },
+  { id: 'image', labelKey: 'model.use.image' },
+  { id: 'transcript', labelKey: 'model.use.transcript' },
+  { id: 'embedding', labelKey: 'model.use.embedding' },
+  { id: 'rerank', labelKey: 'model.use.rerank' },
+];
+
+export function capabilityForUseCase(useCase: UseCase): Capability {
+  switch (useCase) {
+    case 'actions':
+      return modelCapability(['text'], ['text'], ['tools']);
+    case 'structured':
+      return modelCapability(['text'], ['structured']);
+    case 'vision':
+      return modelCapability(
+        ['text', 'image'],
+        ['text'],
+        [],
+        ['image'],
+        ['url', 'data', 'bytes', 'file_handle']
+      );
+    case 'image':
+      return modelCapability(['text'], ['image']);
+    case 'transcript':
+      return modelCapability(
+        ['audio'],
+        ['structured'],
+        [],
+        ['audio'],
+        ['url', 'data', 'bytes', 'file_handle']
+      );
+    case 'embedding':
+      return modelCapability(['text'], ['embedding']);
+    case 'rerank':
+      return modelCapability(['text'], ['rerank']);
+    default:
+      return modelCapability(['text'], ['text']);
+  }
+}
+
+function modelCapability(
+  input: string[],
+  output: string[],
+  features: string[] = [],
+  attachmentKinds: string[] = [],
+  attachmentSources: string[] = []
+): Capability {
+  return { input, output, features, attachmentKinds, attachmentSources };
+}
+
+function matchesCapability(value: Capability, useCase: UseCase) {
+  const expected = capabilityForUseCase(useCase);
+  const fields = [
+    'input',
+    'output',
+    'features',
+    'attachmentKinds',
+    'attachmentSources',
+  ] as const;
+  return fields.every(field =>
+    expected[field].every(item => value[field].includes(item))
+  );
+}
+
+export function modelUseCases(model: ModelDeclaration) {
+  return useCases
+    .filter(({ id }) =>
+      model.capabilities.some(item => matchesCapability(item, id))
+    )
+    .map(({ id }) => id);
+}
+
+export function capabilitiesForUseCases(
+  model: ModelDeclaration | null,
+  selectedUseCases: UseCase[]
+) {
+  const selected = new Set(selectedUseCases);
+  const capabilities = (model?.capabilities ?? []).filter(capability => {
+    const represented = useCases
+      .map(({ id }) => id)
+      .filter(useCase => matchesCapability(capability, useCase));
+    return (
+      represented.length > 0 &&
+      represented.every(useCase => selected.has(useCase))
+    );
+  });
+
+  for (const useCase of selectedUseCases) {
+    if (
+      !capabilities.some(capability => matchesCapability(capability, useCase))
+    ) {
+      capabilities.push(capabilityForUseCase(useCase));
+    }
+  }
+  return capabilities;
+}
+
+export function probeChecks(models: ModelDeclaration[], includeImage: boolean) {
+  return models
+    .filter(model => model.enabled)
+    .flatMap(model =>
+      modelUseCases(model)
+        .filter(
+          useCase =>
+            !['vision', 'transcript'].includes(useCase) &&
+            (useCase !== 'image' || includeImage)
+        )
+        .map(useCase => ({
+          modelId: model.modelId,
+          operation: useCase === 'actions' ? 'tools' : useCase,
+        }))
+    );
+}
+
+export function catalogModels(settings: ByokSettings, provider: ByokProvider) {
+  return (
+    settings.catalog.providers.find(item => item.provider === provider)
+      ?.models ?? []
+  );
+}
+
+export function defaultModels(settings: ByokSettings, provider: ByokProvider) {
+  const catalog = catalogModels(settings, provider);
+  const selected = catalog.filter(model => model.recommended);
+  return (selected.length ? selected : catalog.slice(0, 1)).map(model => ({
+    modelId: model.modelId,
+    enabled: true,
+    capabilities: model.capabilities,
+  }));
+}

packages/frontend/core/src/desktop/dialogs/setting/workspace-setting/byok/types.ts

@@ -1,76 +1,64 @@
 import {
-  type ByokKeyStorage,
-  type ByokKeyTestStatus,
   type ByokProvider,
   type GraphQLQuery,
   type QueryOptions,
   type QueryResponse,
+  type WorkspaceByokSettingsQuery,
 } from '@affine/graphql';
 
-export type ByokStorage = ByokKeyStorage;
+export const ByokStorage = {
+  server: 'server',
+  local: 'local',
+} as const;
+export type ByokStorage = (typeof ByokStorage)[keyof typeof ByokStorage];
 
-export type ByokKey = {
+export type ByokDefinition =
+  WorkspaceByokSettingsQuery['workspace']['byokSettings']['profiles'][number]['definition'];
+
+type ByokKeyBase = {
   id: string;
   provider: ByokProvider;
   name: string;
   description?: string | null;
-  storage: ByokStorage;
   configured: boolean;
   enabled: boolean;
-  endpoint?: string | null;
-  endpointEditable: boolean;
   sortOrder: number;
+  definition: ByokDefinition;
   capabilities: string[];
-  testStatus: ByokKeyTestStatus;
-  disabledReason?: string | null;
-  lastTestedAt?: string | null;
-  lastTestError?: string | null;
-  lastUsedAt?: string | null;
-  lastErrorAt?: string | null;
-  lastError?: string | null;
+  validation?: WorkspaceByokSettingsQuery['workspace']['byokSettings']['profiles'][number]['validation'];
 };
 
+export type ByokKey = ByokKeyBase &
+  (
+    | { storage: typeof ByokStorage.server; revision: number }
+    | { storage: typeof ByokStorage.local; revision?: never }
+  );
+
 export type LocalByokKeyInput = Pick<
   ByokKey,
   | 'id'
   | 'provider'
   | 'name'
   | 'description'
-  | 'endpoint'
   | 'sortOrder'
   | 'enabled'
+  | 'definition'
+> & { credential: string };
+
+export type ByokSettings = Omit<
+  WorkspaceByokSettingsQuery['workspace']['byokSettings'],
+  'profiles'
 > & {
-  apiKey: string;
-};
-
-export type ByokSettings = {
-  workspaceId: string;
-  entitled: boolean;
-  serverEntitled: boolean;
-  localEntitled: boolean;
-  entitlementRequired: string[];
   keys: ByokKey[];
-  allowedProviders: ByokProvider[];
   localStorageSupported: boolean;
-  customEndpointSupported: boolean;
-  privateEndpointSupported: boolean;
-  hasAiPlan: boolean;
-  warnings: Array<{
-    featureKind: string;
-    reason: string;
-    requiredProviders: ByokProvider[];
-  }>;
 };
 
-export type ByokUsagePoint = {
-  date: string;
-  featureKind: string;
-  totalTokens: number;
-};
+export type ByokUsagePoint =
+  WorkspaceByokSettingsQuery['workspace']['byokUsage'][number];
 
 export type ByokTestResult = {
   ok: boolean;
-  status: ByokKey['testStatus'];
+  status: string;
   message?: string | null;
 };
 
@@ -78,15 +66,6 @@ export type GqlFn = <Query extends GraphQLQuery>(
   input: QueryOptions<Query>
 ) => Promise<QueryResponse<Query>>;
 
-export type LocalByokPublicKey = {
-  id: string;
-  provider: ByokProvider;
-  name: string;
-  description?: string | null;
-  endpoint?: string | null;
-  endpointEditable?: boolean;
-  sortOrder?: number | null;
-  enabled?: boolean | null;
+export type LocalByokPublicKey = Omit<LocalByokKeyInput, 'credential'> & {
   configured?: boolean;
-  testStatus?: ByokKey['testStatus'];
 };

packages/frontend/core/src/desktop/pages/workspace/chat/index.tsx

@@ -19,9 +19,9 @@ import { useAISpecs } from '@affine/core/components/hooks/affine/use-ai-specs';
 import { useAISubscribe } from '@affine/core/components/hooks/affine/use-ai-subscribe';
 import {
   AIDraftService,
+  AIModelService,
   AIToolsConfigService,
 } from '@affine/core/modules/ai-button';
-import { AIModelService } from '@affine/core/modules/ai-button/services/models';
 import {
   EventSourceService,
   GraphQLService,
@@ -197,9 +197,9 @@ export const Component = () => {
       content.notificationService = notificationService;
       content.aiDraftService = framework.get(AIDraftService);
       content.aiToolsConfigService = framework.get(AIToolsConfigService);
+      content.aiModelService = framework.get(AIModelService);
       content.serverService = framework.get(ServerService);
       content.subscriptionService = framework.get(SubscriptionService);
-      content.aiModelService = framework.get(AIModelService);
       content.onAISubscribe = handleAISubscribe;
       content.onOpenDoc = onOpenDoc;
     },

packages/frontend/core/src/desktop/pages/workspace/detail-page/tabs/chat.tsx

@@ -23,9 +23,9 @@ import { useAISpecs } from '@affine/core/components/hooks/affine/use-ai-specs';
 import { useAISubscribe } from '@affine/core/components/hooks/affine/use-ai-subscribe';
 import {
   AIDraftService,
+  AIModelService,
   AIToolsConfigService,
 } from '@affine/core/modules/ai-button';
-import { AIModelService } from '@affine/core/modules/ai-button/services/models';
 import {
   EventSourceService,
   GraphQLService,
@@ -283,9 +283,9 @@ export const EditorChatPanel = ({
       content.notificationService = notificationService;
       content.aiDraftService = framework.get(AIDraftService);
       content.aiToolsConfigService = framework.get(AIToolsConfigService);
+      content.aiModelService = framework.get(AIModelService);
       content.peekViewService = framework.get(PeekViewService);
       content.subscriptionService = framework.get(SubscriptionService);
-      content.aiModelService = framework.get(AIModelService);
       content.onAISubscribe = handleAISubscribe;
       content.width = sidebarWidthSignal;
       content.onOpenDoc = (docId: string, sessionId?: string) => {
@@ -399,8 +399,8 @@ export const EditorChatPanel = ({
         .notificationService=${notificationService}
         .affineWorkspaceDialogService=${framework.get(WorkspaceDialogService)}
         .aiToolsConfigService=${framework.get(AIToolsConfigService)}
-        .subscriptionService=${framework.get(SubscriptionService)}
         .aiModelService=${framework.get(AIModelService)}
+        .subscriptionService=${framework.get(SubscriptionService)}
       ></playground-content>
     `;
 

packages/frontend/core/src/modules/ai-button/index.ts

@@ -1,6 +1,7 @@
 export { AIButtonProvider } from './provider/ai-button';
 export { AIButtonService } from './services/ai-button';
 export { AIDraftService } from './services/ai-draft';
+export { AIModelService } from './services/models';
 export {
   type AIToolsConfig,
   AIToolsConfigService,

packages/frontend/core/src/modules/ai-button/services/models.spec.ts

@@ -0,0 +1,67 @@
+import { Framework } from '@toeverything/infra';
+import { EMPTY } from 'rxjs';
+import { describe, expect, test, vi } from 'vitest';
+
+import { AIModelService } from './models';
+
+describe('AIModelService', () => {
+  test('clears the previous model while a new scope is loading', async () => {
+    let resolveSecond: ((value: unknown) => void) | undefined;
+    const gql = vi
+      .fn()
+      .mockResolvedValueOnce({
+        currentUser: {
+          copilot: {
+            routeOptions: {
+              choices: [
+                {
+                  id: 'model-a',
+                  displayName: 'Model A',
+                  available: true,
+                },
+              ],
+            },
+          },
+        },
+      })
+      .mockImplementationOnce(
+        () =>
+          new Promise(resolve => {
+            resolveSecond = resolve;
+          })
+      );
+    const stored = new Map<string, unknown>();
+    const framework = new Framework();
+    framework.service(
+      AIModelService,
+      () =>
+        new AIModelService(
+          {
+            globalState: {
+              get: (key: string) => stored.get(key),
+              set: (key: string, value: unknown) => stored.set(key, value),
+            },
+          } as never,
+          { gql } as never,
+          {
+            subscription: { ai$: EMPTY },
+          } as never
+        )
+    );
+    const service = framework.provider().get(AIModelService);
+
+    service.setScope('workspace-a', 'route-a');
+    await vi.waitFor(() => expect(service.models.value).toHaveLength(1));
+    service.setModel('model-a');
+
+    service.setScope('workspace-b', 'route-b');
+    expect(service.models.value).toEqual([]);
+    expect(service.modelId.value).toBeUndefined();
+    service.setModel('model-a');
+    expect(stored.has('AIManagedRouteTarget:workspace-b:route-b')).toBe(false);
+
+    resolveSecond?.({
+      currentUser: { copilot: { routeOptions: { choices: [] } } },
+    });
+  });
+});

packages/frontend/core/src/modules/ai-button/services/models.ts

@@ -1,112 +1,109 @@
-import { getPromptModelsQuery, SubscriptionStatus } from '@affine/graphql';
-import {
-  createSignalFromObservable,
-  type Signal,
-} from '@blocksuite/affine/shared/utils';
+import { getCopilotRouteOptionsQuery } from '@affine/graphql';
 import { signal } from '@preact/signals-core';
-import { LiveData, Service } from '@toeverything/infra';
+import { Service } from '@toeverything/infra';
 
 import type { GraphQLService, SubscriptionService } from '../../cloud';
 import type { GlobalStateService } from '../../storage';
 
-const AI_MODEL_ID_KEY = 'AIModelId';
+const ROUTE_TARGET_KEY = 'AIManagedRouteTarget';
 
 export interface AIModel {
-  name: string;
   id: string;
-  version: string;
+  name: string;
   category: string;
-  isPro: boolean;
-  isDefault: boolean;
+  version: string;
+  available: boolean;
 }
 
 export class AIModelService extends Service {
-  modelId: Signal<string | undefined>;
+  readonly modelId = signal<string | undefined>(undefined);
+  readonly models = signal<AIModel[]>([]);
 
-  models: Signal<AIModel[]> = signal([]);
-
-  private readonly modelId$ = LiveData.from(
-    this.globalStateService.globalState.watch<string>(AI_MODEL_ID_KEY),
-    undefined
-  );
+  private workspaceId: string | undefined;
+  private routeId: string | undefined;
+  private requestId = 0;
 
   constructor(
     private readonly globalStateService: GlobalStateService,
     private readonly gqlService: GraphQLService,
-    private readonly subscriptionService: SubscriptionService
+    subscriptionService: SubscriptionService
   ) {
     super();
-
-    const { signal: modelId, cleanup } = createSignalFromObservable<
-      string | undefined
-    >(this.modelId$, undefined);
-    this.modelId = modelId;
-    this.disposables.push(cleanup);
-
-    this.init().catch(err => {
-      console.error(err);
+    const subscription = subscriptionService.subscription.ai$.subscribe(() => {
+      if (this.workspaceId && this.routeId) {
+        this.load(this.workspaceId, this.routeId).catch(console.error);
+      }
     });
+    this.disposables.push(() => subscription.unsubscribe());
   }
 
-  resetModel = () => {
-    this.globalStateService.globalState.set(AI_MODEL_ID_KEY, undefined);
-  };
+  setScope(workspaceId: string, routeId: string) {
+    if (workspaceId === this.workspaceId && routeId === this.routeId) return;
+    this.workspaceId = workspaceId;
+    this.routeId = routeId;
+    this.models.value = [];
+    this.modelId.value = undefined;
+    this.load(workspaceId, routeId).catch(console.error);
+  }
 
-  setModel = (modelId: string) => {
-    const isSubscribed =
-      this.subscriptionService.subscription.ai$.value?.status ===
-      SubscriptionStatus.Active;
-    const model = this.models.value.find(model => model.id === modelId);
-    if (!isSubscribed && model?.isPro) {
+  resetModel() {
+    this.setModel(undefined);
+  }
+
+  setModel(modelId: string | undefined) {
+    if (!this.workspaceId || !this.routeId) return;
+    if (
+      modelId &&
+      !this.models.value.find(model => model.id === modelId)?.available
+    ) {
       return;
     }
-    this.globalStateService.globalState.set(AI_MODEL_ID_KEY, modelId);
-  };
-
-  private readonly init = async () => {
-    await this.initModels();
-
-    // subscribe to ai purchase status
-    const sub = this.subscriptionService.subscription.ai$.subscribe(
-      subscription => {
-        const isSubscribed = subscription?.status === SubscriptionStatus.Active;
-        const model = this.models.value.find(
-          model => model.id === this.modelId.value
-        );
-        if (!isSubscribed && model?.isPro) {
-          this.resetModel();
-        }
-      }
+    this.modelId.value = modelId;
+    this.globalStateService.globalState.set(
+      this.storageKey(this.workspaceId, this.routeId),
+      modelId
     );
-    this.disposables.push(() => sub.unsubscribe());
-  };
+  }
 
-  private readonly initModels = async (prompt?: string) => {
-    const promptName = prompt || 'Chat With AFFiNE AI';
-    const models = await this.getModelsByPrompt(promptName);
-    if (models) {
-      const { defaultModel, optionalModels, proModels } = models;
-      this.models.value = optionalModels.map(model => {
-        const [category] = model.name.split(' ');
-        const version = model.name.slice(category.length + 1);
-        return {
-          name: model.name,
-          id: model.id,
-          version,
-          category,
-          isPro: proModels.some(proModel => proModel.id === model.id),
-          isDefault: model.id === defaultModel,
-        };
-      });
+  private async load(workspaceId: string, routeId: string) {
+    const requestId = ++this.requestId;
+    const result = await this.gqlService.gql({
+      query: getCopilotRouteOptionsQuery,
+      variables: { promptName: routeId },
+    });
+    if (requestId !== this.requestId) return;
+    const options = result.currentUser?.copilot?.routeOptions;
+    if (!options) {
+      this.models.value = [];
+      this.modelId.value = undefined;
+      return;
     }
-  };
+    this.models.value = options.choices.map(choice => {
+      const [category] = choice.displayName.split(' ');
+      return {
+        id: choice.id,
+        name: choice.displayName,
+        category,
+        version: choice.displayName.slice(category.length + 1),
+        available: choice.available,
+      };
+    });
+    const selected = this.globalStateService.globalState.get<string>(
+      this.storageKey(workspaceId, routeId)
+    );
+    const selectedAvailable = this.models.value.some(
+      model => model.id === selected && model.available
+    );
+    this.modelId.value = selectedAvailable ? selected : undefined;
+    if (selected && !selectedAvailable) {
+      this.globalStateService.globalState.set(
+        this.storageKey(workspaceId, routeId),
+        undefined
+      );
+    }
+  }
 
-  private readonly getModelsByPrompt = async (promptName: string) => {
-    return this.gqlService
-      .gql({
-        query: getPromptModelsQuery,
-        variables: { promptName },
-      })
-      .then(res => res.currentUser?.copilot?.models);
-  };
+  private storageKey(workspaceId: string, routeId: string) {
+    return `${ROUTE_TARGET_KEY}:${workspaceId}:${routeId}`;
+  }
 }

packages/frontend/core/src/modules/media/entities/audio-transcription-job-store.spec.ts

@@ -69,10 +69,7 @@ describe('AudioTranscriptionJobStore transcript task API', () => {
       .mockResolvedValueOnce({ submitTranscriptTask: { id: 'task-1' } })
       .mockResolvedValueOnce({ retryTranscriptTask: { id: 'task-2' } })
       .mockResolvedValueOnce({ settleTranscriptTask: { id: 'task-2' } });
-    const store = createStore(gql, async () => ({
-      files: [file],
-      input: { strategy: 'gemini' },
-    }));
+    const store = createStore(gql, async () => ({ files: [file] }));
 
     await store.submitTranscriptTask();
     await store.retryTranscriptTask('task-1');
@@ -87,7 +84,7 @@ describe('AudioTranscriptionJobStore transcript task API', () => {
           workspaceId: 'workspace-1',
           blobId: 'blob-1',
           blobs: [file],
-          input: { strategy: 'gemini' },
+          input: undefined,
         },
       })
     );

packages/frontend/core/src/modules/peek-view/view/ai-chat-block-peek-view/index.tsx

@@ -6,9 +6,9 @@ import { useAIChatConfig } from '@affine/core/components/hooks/affine/use-ai-cha
 import { useAISubscribe } from '@affine/core/components/hooks/affine/use-ai-subscribe';
 import {
   AIDraftService,
+  AIModelService,
   AIToolsConfigService,
 } from '@affine/core/modules/ai-button';
-import { AIModelService } from '@affine/core/modules/ai-button/services/models';
 import { ServerService, SubscriptionService } from '@affine/core/modules/cloud';
 import { WorkspaceDialogService } from '@affine/core/modules/dialogs';
 import { FeatureFlagService } from '@affine/core/modules/feature-flag';
@@ -36,8 +36,8 @@ export const AIChatBlockPeekView = ({
   const affineWorkspaceDialogService = framework.get(WorkspaceDialogService);
   const aiDraftService = framework.get(AIDraftService);
   const aiToolsConfigService = framework.get(AIToolsConfigService);
-  const subscriptionService = framework.get(SubscriptionService);
   const aiModelService = framework.get(AIModelService);
+  const subscriptionService = framework.get(SubscriptionService);
   const handleAISubscribe = useAISubscribe();
 
   return useMemo(() => {
@@ -52,8 +52,8 @@ export const AIChatBlockPeekView = ({
       affineWorkspaceDialogService,
       aiDraftService,
       aiToolsConfigService,
-      subscriptionService,
       aiModelService,
+      subscriptionService,
       handleAISubscribe
     );
     return toReactNode(template);
@@ -68,8 +68,8 @@ export const AIChatBlockPeekView = ({
     affineWorkspaceDialogService,
     aiDraftService,
     aiToolsConfigService,
-    subscriptionService,
     aiModelService,
+    subscriptionService,
     handleAISubscribe,
   ]);
 };