feat(server): refactor for byok (#14911)

2026-07-20 03:26:47 +08:00 · 2026-05-07 04:03:14 +08:00
parent 4e169ea5c7
commit eb9cc22502
115 changed files with 10369 additions and 1256 deletions
@@ -0,0 +1,44 @@
+import type { ByokFeatureKind } from '../byok/types';
+
+export type ByokSourceCoverage = {
+  local: boolean;
+  server: boolean;
+};
+
+export type CopilotFeatureAccessRule = ByokSourceCoverage & {
+  quotaMetered: boolean;
+};
+
+const DEFAULT_BYOK_COVERAGE: ByokSourceCoverage = {
+  local: true,
+  server: true,
+};
+
+const DEFAULT_FEATURE_ACCESS: CopilotFeatureAccessRule = {
+  ...DEFAULT_BYOK_COVERAGE,
+  quotaMetered: true,
+};
+
+const COPILOT_FEATURE_ACCESS: Partial<
+  Record<ByokFeatureKind, CopilotFeatureAccessRule>
+> = {
+  transcript: { local: false, server: true, quotaMetered: true },
+  embedding: { local: false, server: true, quotaMetered: false },
+  workspace_indexing: { local: false, server: true, quotaMetered: false },
+  rerank: { local: false, server: true, quotaMetered: false },
+};
+
+export function getByokSourceCoverage(
+  featureKind?: ByokFeatureKind
+): ByokSourceCoverage {
+  const access = getCopilotFeatureAccess(featureKind);
+  return { local: access.local, server: access.server };
+}
+
+export function getCopilotFeatureAccess(
+  featureKind?: ByokFeatureKind
+): CopilotFeatureAccessRule {
+  return featureKind
+    ? (COPILOT_FEATURE_ACCESS[featureKind] ?? DEFAULT_FEATURE_ACCESS)
+    : DEFAULT_FEATURE_ACCESS;
+}
@@ -0,0 +1,2 @@
+export * from './feature-coverage';
+export * from './policy';
@@ -0,0 +1,106 @@
+import { Injectable } from '@nestjs/common';
+
+import { CopilotQuotaExceeded } from '../../../base';
+import { ByokService } from '../byok/service';
+import type { ByokFeatureKind } from '../byok/types';
+import type { CopilotProviderProfile } from '../config';
+import { ConversationPolicy } from '../conversation/policy';
+import {
+  getByokSourceCoverage,
+  getCopilotFeatureAccess,
+} from './feature-coverage';
+
+export type CopilotAccessContext = {
+  userId?: string;
+  workspaceId?: string;
+  byokLeaseId?: string;
+  featureKind?: ByokFeatureKind;
+  quotaBackedRoutesAllowed?: boolean;
+};
+
+export type CopilotRouteAccess = {
+  byokProfiles: CopilotProviderProfile[];
+  quotaBackedRoutesAvailable: boolean;
+};
+
+export type CopilotTurnRouteAccess = {
+  byokProfiles: CopilotProviderProfile[];
+  quotaBackedRoutesAllowed?: boolean;
+};
+
+@Injectable()
+export class CopilotAccessPolicy {
+  constructor(
+    private readonly conversationPolicy: ConversationPolicy,
+    private readonly byok: ByokService
+  ) {}
+
+  async getByokProfiles(context: CopilotAccessContext = {}) {
+    const coverage = getByokSourceCoverage(context.featureKind);
+    return await this.byok.getProfiles(context, coverage);
+  }
+
+  async canUseQuotaBackedRoutes(context: CopilotAccessContext = {}) {
+    if (context.quotaBackedRoutesAllowed !== undefined) {
+      return context.quotaBackedRoutesAllowed;
+    }
+    if (!getCopilotFeatureAccess(context.featureKind).quotaMetered) {
+      return true;
+    }
+    if (!context.userId) {
+      return true;
+    }
+    return await this.conversationPolicy.hasQuota(context.userId);
+  }
+
+  async getQuota(userId: string) {
+    return await this.conversationPolicy.getQuota(userId);
+  }
+
+  async checkQuota(userId: string) {
+    await this.conversationPolicy.checkQuota(userId);
+  }
+
+  async resolveRouteAccess(
+    context: CopilotAccessContext = {}
+  ): Promise<CopilotRouteAccess> {
+    const [byokProfiles, quotaBackedRoutesAvailable] = await Promise.all([
+      this.getByokProfiles(context),
+      this.canUseQuotaBackedRoutes(context),
+    ]);
+
+    return { byokProfiles, quotaBackedRoutesAvailable };
+  }
+
+  async resolveTurnRouteAccess(
+    context: CopilotAccessContext
+  ): Promise<CopilotTurnRouteAccess> {
+    const byokProfiles = await this.getByokProfiles(context);
+    if (context.quotaBackedRoutesAllowed === false) {
+      return { byokProfiles, quotaBackedRoutesAllowed: false };
+    }
+    const featureAccess = getCopilotFeatureAccess(context.featureKind);
+    if (!byokProfiles.length && context.userId && featureAccess.quotaMetered) {
+      await this.conversationPolicy.checkQuota(context.userId);
+    }
+
+    const quotaBackedRoutesAllowed = byokProfiles.length
+      ? context.quotaBackedRoutesAllowed
+      : true;
+    return { byokProfiles, quotaBackedRoutesAllowed };
+  }
+
+  async assertQuotaOrByok(context: CopilotAccessContext) {
+    const byokProfiles = await this.getByokProfiles(context);
+    if (context.quotaBackedRoutesAllowed === false) {
+      if (!byokProfiles.length) {
+        throw new CopilotQuotaExceeded();
+      }
+      return;
+    }
+    const featureAccess = getCopilotFeatureAccess(context.featureKind);
+    if (!byokProfiles.length && context.userId && featureAccess.quotaMetered) {
+      await this.conversationPolicy.checkQuota(context.userId);
+    }
+  }
+}
@@ -0,0 +1,4 @@
+export { ByokEntitlementPolicy } from './policy';
+export { WorkspaceByokResolver } from './resolver';
+export { type ByokProviderRequestContext, ByokService } from './service';
+export * from './types';
@@ -0,0 +1,105 @@
+import { Injectable } from '@nestjs/common';
+
+import { ActionForbidden } from '../../../base';
+import { Models, WorkspaceRole } from '../../../models';
+
+@Injectable()
+export class ByokEntitlementPolicy {
+  constructor(private readonly models: Models) {}
+
+  private isUserPlanEntitled(features: string[]) {
+    return (
+      features.includes('pro_plan_v1') ||
+      features.includes('lifetime_pro_plan_v1') ||
+      features.includes('unlimited_copilot')
+    );
+  }
+
+  async hasAiPlan(userId?: string) {
+    if (!userId) return false;
+    const features = await this.models.userFeature.list(userId);
+    return this.isUserPlanEntitled(features);
+  }
+
+  async hasManagementAccess(workspaceId: string, userId?: string) {
+    if (!userId) return false;
+    const role = await this.models.workspaceUser.getActive(workspaceId, userId);
+    return (
+      role?.type === WorkspaceRole.Owner || role?.type === WorkspaceRole.Admin
+    );
+  }
+
+  async assertManagementAccess(workspaceId: string, userId?: string) {
+    if (!(await this.hasManagementAccess(workspaceId, userId))) {
+      throw new ActionForbidden(
+        'BYOK settings require workspace owner or admin.'
+      );
+    }
+  }
+
+  private async getWorkspaceOwnerId(workspaceId: string) {
+    const workspace = await this.models.workspace.get(workspaceId);
+    if (!workspace) {
+      return null;
+    }
+
+    try {
+      return (await this.models.workspaceUser.getOwner(workspaceId)).id;
+    } catch (error) {
+      if (
+        error instanceof Error &&
+        error.message === 'Workspace owner not found'
+      ) {
+        return null;
+      }
+      throw error;
+    }
+  }
+
+  async hasLocalEntitlement(workspaceId: string, userId?: string) {
+    if (env.selfhosted) return true;
+
+    if (await this.models.workspaceFeature.has(workspaceId, 'team_plan_v1')) {
+      return true;
+    }
+
+    const ownerId = await this.getWorkspaceOwnerId(workspaceId);
+    if (!ownerId) return false;
+
+    if (await this.hasAiPlan(userId)) return true;
+    return await this.hasAiPlan(ownerId);
+  }
+
+  async hasServerEntitlement(workspaceId: string) {
+    if (env.selfhosted) return true;
+
+    if (await this.models.workspaceFeature.has(workspaceId, 'team_plan_v1')) {
+      return true;
+    }
+
+    const ownerId = await this.getWorkspaceOwnerId(workspaceId);
+    if (!ownerId) return false;
+    return await this.hasAiPlan(ownerId);
+  }
+
+  async hasEntitlement(workspaceId: string, userId?: string) {
+    const [serverEntitled, localEntitled] = await Promise.all([
+      this.hasServerEntitlement(workspaceId),
+      this.hasLocalEntitlement(workspaceId, userId),
+    ]);
+
+    return [serverEntitled, localEntitled] as const;
+  }
+
+  async assertServerEntitled(workspaceId: string) {
+    if (!(await this.hasServerEntitlement(workspaceId))) {
+      throw new ActionForbidden('BYOK requires Pro, Team, or Believer.');
+    }
+  }
+
+  async assertLocalEntitled(workspaceId: string, userId?: string) {
+    if (!(await this.hasLocalEntitlement(workspaceId, userId))) {
+      throw new ActionForbidden('BYOK requires Pro, Team, or Believer.');
+    }
+  }
+}
@@ -0,0 +1,405 @@
+import {
+  Args,
+  Field,
+  ID,
+  InputType,
+  Mutation,
+  ObjectType,
+  Parent,
+  ResolveField,
+  Resolver,
+} from '@nestjs/graphql';
+import { SafeIntResolver } from 'graphql-scalars';
+
+import { Throttle } from '../../../base';
+import { CurrentUser } from '../../../core/auth';
+import { AccessController } from '../../../core/permission';
+import { WorkspaceType } from '../../../core/workspaces';
+import { ByokEntitlementPolicy } from './policy';
+import { ByokKeyConfig, ByokLocalLeaseProvider, ByokService } from './service';
+import { ByokKeyStorage, ByokKeyTestStatus, ByokProvider } from './types';
+
+@ObjectType()
+export class WorkspaceByokKeyConfigType implements ByokKeyConfig {
+  @Field(() => ID)
+  id!: string;
+
+  @Field(() => ByokProvider)
+  provider!: ByokProvider;
+
+  @Field(() => String)
+  name!: string;
+
+  @Field(() => String, { nullable: true })
+  description!: string | null;
+
+  @Field(() => ByokKeyStorage)
+  storage!: ByokKeyStorage;
+
+  @Field(() => Boolean)
+  configured!: boolean;
+
+  @Field(() => Boolean)
+  enabled!: boolean;
+
+  @Field(() => String, { nullable: true })
+  endpoint!: string | null;
+
+  @Field(() => Boolean)
+  endpointEditable!: boolean;
+
+  @Field(() => SafeIntResolver)
+  sortOrder!: number;
+
+  @Field(() => [String])
+  capabilities!: string[];
+
+  @Field(() => ByokKeyTestStatus)
+  testStatus!: ByokKeyTestStatus;
+
+  @Field(() => String, { nullable: true })
+  disabledReason!: string | null;
+
+  @Field(() => Date, { nullable: true })
+  lastTestedAt!: Date | null;
+
+  @Field(() => String, { nullable: true })
+  lastTestError!: string | null;
+
+  @Field(() => Date, { nullable: true })
+  lastUsedAt!: Date | null;
+
+  @Field(() => Date, { nullable: true })
+  lastErrorAt!: Date | null;
+
+  @Field(() => String, { nullable: true })
+  lastError!: string | null;
+}
+
+@ObjectType()
+class WorkspaceByokCapabilityWarningType {
+  @Field(() => String)
+  featureKind!: string;
+
+  @Field(() => String)
+  reason!: string;
+
+  @Field(() => [ByokProvider])
+  requiredProviders!: ByokProvider[];
+}
+
+@ObjectType()
+class WorkspaceByokSettingsType {
+  @Field(() => String)
+  workspaceId!: string;
+
+  @Field(() => Boolean)
+  entitled!: boolean;
+
+  @Field(() => Boolean)
+  serverEntitled!: boolean;
+
+  @Field(() => Boolean)
+  localEntitled!: boolean;
+
+  @Field(() => [String])
+  entitlementRequired!: string[];
+
+  @Field(() => [WorkspaceByokKeyConfigType])
+  keys!: WorkspaceByokKeyConfigType[];
+
+  @Field(() => [ByokProvider])
+  allowedProviders!: ByokProvider[];
+
+  @Field(() => Boolean)
+  localStorageSupported!: boolean;
+
+  @Field(() => Boolean)
+  customEndpointSupported!: boolean;
+
+  @Field(() => Boolean)
+  hasAiPlan!: boolean;
+
+  @Field(() => [WorkspaceByokCapabilityWarningType])
+  warnings!: WorkspaceByokCapabilityWarningType[];
+}
+
+@ObjectType()
+class WorkspaceByokUsagePointType {
+  @Field(() => Date)
+  date!: Date;
+
+  @Field(() => String)
+  featureKind!: string;
+
+  @Field(() => SafeIntResolver)
+  totalTokens!: number;
+}
+
+@ObjectType()
+class TestWorkspaceByokConfigResultType {
+  @Field(() => Boolean)
+  ok!: boolean;
+
+  @Field(() => ByokKeyTestStatus)
+  status!: ByokKeyTestStatus;
+
+  @Field(() => String, { nullable: true })
+  message!: string | null;
+}
+
+@ObjectType()
+class CreateWorkspaceByokLocalLeaseResultType {
+  @Field(() => String)
+  leaseId!: string;
+
+  @Field(() => Date)
+  expiresAt!: Date;
+}
+
+@InputType()
+class UpsertWorkspaceByokConfigInput {
+  @Field(() => ID, { nullable: true })
+  id?: string;
+
+  @Field(() => String)
+  workspaceId!: string;
+
+  @Field(() => ByokProvider)
+  provider!: ByokProvider;
+
+  @Field(() => String)
+  name!: string;
+
+  @Field(() => String, { nullable: true })
+  description?: string | null;
+
+  @Field(() => ByokKeyStorage)
+  storage!: ByokKeyStorage;
+
+  @Field(() => String, { nullable: true })
+  apiKey?: string | null;
+
+  @Field(() => String, { nullable: true })
+  endpoint?: string | null;
+
+  @Field(() => SafeIntResolver, { nullable: true })
+  sortOrder?: number | null;
+
+  @Field(() => Boolean, { nullable: true })
+  enabled?: boolean | null;
+}
+
+@InputType()
+class TestWorkspaceByokConfigInput {
+  @Field(() => String)
+  workspaceId!: string;
+
+  @Field(() => ByokProvider)
+  provider!: ByokProvider;
+
+  @Field(() => ByokKeyStorage)
+  storage!: ByokKeyStorage;
+
+  @Field(() => String, { nullable: true })
+  apiKey?: string | null;
+
+  @Field(() => String, { nullable: true })
+  endpoint?: string | null;
+
+  @Field(() => ID, { nullable: true })
+  configId?: string | null;
+}
+
+@InputType()
+class ReorderWorkspaceByokConfigsInput {
+  @Field(() => String)
+  workspaceId!: string;
+
+  @Field(() => ByokKeyStorage)
+  storage!: ByokKeyStorage;
+
+  @Field(() => [ID])
+  ids!: string[];
+}
+
+@InputType()
+class CreateWorkspaceByokLocalLeaseProviderInput implements ByokLocalLeaseProvider {
+  @Field(() => ByokProvider)
+  provider!: ByokProvider;
+
+  @Field(() => String)
+  name!: string;
+
+  @Field(() => String, { nullable: true })
+  description?: string | null;
+
+  @Field(() => String)
+  apiKey!: string;
+
+  @Field(() => String, { nullable: true })
+  endpoint?: string | null;
+
+  @Field(() => SafeIntResolver, { nullable: true })
+  sortOrder?: number | null;
+
+  @Field(() => Boolean, { nullable: true })
+  enabled?: boolean | null;
+}
+
+@InputType()
+class CreateWorkspaceByokLocalLeaseInput {
+  @Field(() => String)
+  workspaceId!: string;
+
+  @Field(() => [CreateWorkspaceByokLocalLeaseProviderInput])
+  providers!: CreateWorkspaceByokLocalLeaseProviderInput[];
+}
+
+@Resolver(() => WorkspaceType)
+export class WorkspaceByokResolver {
+  constructor(
+    private readonly ac: AccessController,
+    private readonly entitlement: ByokEntitlementPolicy,
+    private readonly byok: ByokService
+  ) {}
+
+  @ResolveField(() => WorkspaceByokSettingsType, {
+    name: 'byokSettings',
+    complexity: 2,
+  })
+  async settings(
+    @CurrentUser() user: CurrentUser,
+    @Parent() workspace: WorkspaceType
+  ) {
+    await this.ac
+      .user(user.id)
+      .workspace(workspace.id)
+      .allowLocal()
+      .assert('Workspace.Settings.Read');
+    await this.entitlement.assertManagementAccess(workspace.id, user.id);
+    return await this.byok.getSettings(workspace.id, user.id);
+  }
+
+  @ResolveField(() => [WorkspaceByokUsagePointType], {
+    name: 'byokUsage',
+    complexity: 2,
+  })
+  async usage(
+    @CurrentUser() user: CurrentUser,
+    @Parent() workspace: WorkspaceType,
+    @Args('from', { type: () => Date }) from: Date,
+    @Args('to', { type: () => Date }) to: Date
+  ) {
+    await this.ac
+      .user(user.id)
+      .workspace(workspace.id)
+      .allowLocal()
+      .assert('Workspace.Settings.Read');
+    await this.entitlement.assertManagementAccess(workspace.id, user.id);
+    return await this.byok.getUsage(workspace.id, from, to);
+  }
+
+  @Throttle('strict')
+  @Mutation(() => TestWorkspaceByokConfigResultType)
+  async testWorkspaceByokConfig(
+    @CurrentUser() user: CurrentUser,
+    @Args('input') input: TestWorkspaceByokConfigInput
+  ) {
+    await this.ac
+      .user(user.id)
+      .workspace(input.workspaceId)
+      .allowLocal()
+      .assert('Workspace.Settings.Update');
+    await this.entitlement.assertManagementAccess(input.workspaceId, user.id);
+    if (input.storage === ByokKeyStorage.server) {
+      await this.entitlement.assertServerEntitled(input.workspaceId);
+    } else {
+      await this.entitlement.assertLocalEntitled(input.workspaceId, user.id);
+    }
+    return await this.byok.testConfig({ ...input, userId: user.id });
+  }
+
+  @Mutation(() => WorkspaceByokKeyConfigType)
+  @Throttle('strict')
+  async upsertWorkspaceByokConfig(
+    @CurrentUser() user: CurrentUser,
+    @Args('input') input: UpsertWorkspaceByokConfigInput
+  ) {
+    await this.ac
+      .user(user.id)
+      .workspace(input.workspaceId)
+      .allowLocal()
+      .assert('Workspace.Settings.Update');
+    await this.entitlement.assertManagementAccess(input.workspaceId, user.id);
+    await this.entitlement.assertServerEntitled(input.workspaceId);
+    return await this.byok.upsertConfig({ ...input, userId: user.id });
+  }
+
+  @Mutation(() => [WorkspaceByokKeyConfigType])
+  @Throttle('strict')
+  async reorderWorkspaceByokConfigs(
+    @CurrentUser() user: CurrentUser,
+    @Args('input') input: ReorderWorkspaceByokConfigsInput
+  ) {
+    await this.ac
+      .user(user.id)
+      .workspace(input.workspaceId)
+      .allowLocal()
+      .assert('Workspace.Settings.Update');
+    await this.entitlement.assertManagementAccess(input.workspaceId, user.id);
+    await this.entitlement.assertServerEntitled(input.workspaceId);
+    return await this.byok.reorderConfigs({ ...input, userId: user.id });
+  }
+
+  @Mutation(() => Boolean)
+  @Throttle('strict')
+  async deleteWorkspaceByokConfig(
+    @CurrentUser() user: CurrentUser,
+    @Args('id', { type: () => ID }) id: string,
+    @Args('workspaceId', { type: () => String }) workspaceId: string
+  ) {
+    await this.ac
+      .user(user.id)
+      .workspace(workspaceId)
+      .allowLocal()
+      .assert('Workspace.Settings.Update');
+    await this.entitlement.assertManagementAccess(workspaceId, user.id);
+    await this.entitlement.assertServerEntitled(workspaceId);
+    return await this.byok.deleteConfig(workspaceId, id, user.id);
+  }
+
+  @Mutation(() => Boolean)
+  @Throttle('strict')
+  async clearWorkspaceByokConfigs(
+    @CurrentUser() user: CurrentUser,
+    @Args('workspaceId', { type: () => String }) workspaceId: string,
+    @Args('provider', { type: () => ByokProvider, nullable: true })
+    provider?: ByokProvider | null
+  ) {
+    await this.ac
+      .user(user.id)
+      .workspace(workspaceId)
+      .allowLocal()
+      .assert('Workspace.Settings.Update');
+    await this.entitlement.assertManagementAccess(workspaceId, user.id);
+    await this.entitlement.assertServerEntitled(workspaceId);
+    return await this.byok.clearConfigs(workspaceId, provider, user.id);
+  }
+
+  @Mutation(() => CreateWorkspaceByokLocalLeaseResultType)
+  @Throttle('strict')
+  async createWorkspaceByokLocalLease(
+    @CurrentUser() user: CurrentUser,
+    @Args('input') input: CreateWorkspaceByokLocalLeaseInput
+  ) {
+    await this.ac
+      .user(user.id)
+      .workspace(input.workspaceId)
+      .allowLocal()
+      .assert('Workspace.Copilot');
+    await this.entitlement.assertManagementAccess(input.workspaceId, user.id);
+    await this.entitlement.assertLocalEntitled(input.workspaceId, user.id);
+    return await this.byok.createLocalLease({ ...input, userId: user.id });
+  }
+}
@@ -0,0 +1,883 @@
+import { createHash, createHmac, randomUUID } from 'node:crypto';
+
+import { BadRequestException, Injectable } from '@nestjs/common';
+
+import { BadRequest, Cache, CryptoHelper, metrics } from '../../../base';
+import { Models } from '../../../models';
+import type { CopilotProviderProfile } from '../config';
+import { ByokEntitlementPolicy } from './policy';
+import {
+  BYOK_ALLOWED_PROVIDERS,
+  type ByokFeatureKind,
+  ByokKeyStorage,
+  ByokKeyTestStatus,
+  ByokProvider,
+  ByokProviderSource,
+  byokProviderToCopilotType,
+  isByokProvider,
+} from './types';
+
+const LOCAL_LEASE_TTL_MS = 10 * 60 * 1000;
+const BYOK_PROFILE_PRIORITY_BASE = 10_000;
+const SERVER_PROFILE_PRIORITY_OFFSET = 2_000;
+const TEST_TIMEOUT_MS = 10_000;
+
+export type ByokProviderRequestContext = {
+  userId?: string;
+  workspaceId?: string;
+  byokLeaseId?: string;
+};
+
+export type ByokProfileSourceFilter = {
+  local?: boolean;
+  server?: boolean;
+};
+
+export type ByokKeyConfig = {
+  id: string;
+  provider: ByokProvider;
+  name: string;
+  description: string | null;
+  storage: ByokKeyStorage;
+  configured: boolean;
+  enabled: boolean;
+  endpoint: string | null;
+  endpointEditable: boolean;
+  sortOrder: number;
+  capabilities: string[];
+  testStatus: ByokKeyTestStatus;
+  disabledReason: string | null;
+  lastTestedAt: Date | null;
+  lastTestError: string | null;
+  lastUsedAt: Date | null;
+  lastErrorAt: Date | null;
+  lastError: string | null;
+};
+
+export type ByokSettings = {
+  workspaceId: string;
+  entitled: boolean;
+  serverEntitled: boolean;
+  localEntitled: boolean;
+  entitlementRequired: string[];
+  keys: ByokKeyConfig[];
+  allowedProviders: ByokProvider[];
+  localStorageSupported: boolean;
+  customEndpointSupported: boolean;
+  hasAiPlan: boolean;
+  warnings: Array<{
+    featureKind: string;
+    reason: string;
+    requiredProviders: ByokProvider[];
+  }>;
+};
+
+export type ByokLocalLeaseProvider = {
+  provider: ByokProvider;
+  name: string;
+  description?: string | null;
+  apiKey: string;
+  endpoint?: string | null;
+  sortOrder?: number | null;
+  enabled?: boolean | null;
+};
+
+type LocalLeasePayload = {
+  workspaceId: string;
+  userId: string;
+  providers: Array<
+    Omit<ByokLocalLeaseProvider, 'apiKey'> & { encryptedApiKey: string }
+  >;
+};
+
+type LocalLeaseActive = {
+  leaseId: string;
+  expiresAt: string;
+};
+
+type ByokProfileMeta = {
+  source: ByokProviderSource.Server | ByokProviderSource.Local;
+  keyId?: string;
+  provider: ByokProvider;
+};
+
+@Injectable()
+export class ByokService {
+  constructor(
+    private readonly models: Models,
+    private readonly crypto: CryptoHelper,
+    private readonly cache: Cache,
+    private readonly entitlement: ByokEntitlementPolicy
+  ) {}
+
+  get customEndpointSupported() {
+    return env.selfhosted;
+  }
+
+  async getSettings(
+    workspaceId: string,
+    userId?: string
+  ): Promise<ByokSettings> {
+    if (!(await this.entitlement.hasManagementAccess(workspaceId, userId))) {
+      return {
+        workspaceId,
+        entitled: false,
+        serverEntitled: false,
+        localEntitled: false,
+        entitlementRequired: ['Workspace owner or admin'],
+        keys: [],
+        allowedProviders: [...BYOK_ALLOWED_PROVIDERS],
+        localStorageSupported: false,
+        customEndpointSupported: this.customEndpointSupported,
+        hasAiPlan: await this.entitlement.hasAiPlan(userId),
+        warnings: [],
+      };
+    }
+
+    const [serverEntitled, localEntitled] =
+      await this.entitlement.hasEntitlement(workspaceId, userId);
+    const entitled = serverEntitled || localEntitled;
+    if (!entitled) {
+      return {
+        workspaceId,
+        entitled: false,
+        serverEntitled: false,
+        localEntitled: false,
+        entitlementRequired: ['Pro', 'Team', 'Believer'],
+        keys: [],
+        allowedProviders: [...BYOK_ALLOWED_PROVIDERS],
+        localStorageSupported: false,
+        customEndpointSupported: this.customEndpointSupported,
+        hasAiPlan: await this.entitlement.hasAiPlan(userId),
+        warnings: [],
+      };
+    }
+
+    const rows = serverEntitled
+      ? await this.models.copilotWorkspaceByokConfig.list(workspaceId)
+      : [];
+    const keys = rows.map(row => this.toKeyConfig(row));
+
+    return {
+      workspaceId,
+      entitled: true,
+      serverEntitled,
+      localEntitled,
+      entitlementRequired: ['Pro', 'Team', 'Believer'],
+      keys,
+      allowedProviders: [...BYOK_ALLOWED_PROVIDERS],
+      localStorageSupported: false,
+      customEndpointSupported: this.customEndpointSupported,
+      hasAiPlan: await this.entitlement.hasAiPlan(userId),
+      warnings: this.buildWarnings(keys),
+    };
+  }
+
+  async upsertConfig(input: {
+    id?: string | null;
+    workspaceId: string;
+    provider: ByokProvider;
+    name: string;
+    description?: string | null;
+    storage: ByokKeyStorage;
+    apiKey?: string | null;
+    endpoint?: string | null;
+    sortOrder?: number | null;
+    enabled?: boolean | null;
+    userId?: string;
+  }): Promise<ByokKeyConfig> {
+    await this.entitlement.assertManagementAccess(
+      input.workspaceId,
+      input.userId
+    );
+    await this.entitlement.assertServerEntitled(input.workspaceId);
+    this.assertProvider(input.provider);
+    if (input.storage !== ByokKeyStorage.server) {
+      throw new BadRequestException('Only server BYOK keys are persisted.');
+    }
+    const existing = input.id
+      ? await this.models.copilotWorkspaceByokConfig.get(input.id)
+      : null;
+    if (input.id && (!existing || existing.workspaceId !== input.workspaceId)) {
+      throw new BadRequest('BYOK config not found.');
+    }
+    const encryptedApiKey = input.apiKey
+      ? this.crypto.encrypt(input.apiKey)
+      : undefined;
+
+    if (!input.id && !encryptedApiKey) {
+      throw new BadRequestException('apiKey is required.');
+    }
+
+    const description =
+      input.description !== undefined
+        ? input.description?.trim() || null
+        : (existing?.description ?? null);
+    const endpoint =
+      input.endpoint !== undefined
+        ? this.normalizeEndpoint(input.endpoint)
+        : (existing?.endpoint ?? null);
+    const sortOrder = input.sortOrder ?? existing?.sortOrder ?? 0;
+    const enabled = input.enabled ?? existing?.enabled ?? true;
+
+    const row = await this.models.copilotWorkspaceByokConfig.upsert({
+      id: input.id,
+      workspaceId: input.workspaceId,
+      provider: input.provider,
+      name: input.name.trim(),
+      description,
+      encryptedApiKey,
+      endpoint,
+      sortOrder,
+      enabled,
+      userId: input.userId,
+    });
+
+    return this.toKeyConfig(row);
+  }
+
+  async reorderConfigs(input: {
+    workspaceId: string;
+    storage: ByokKeyStorage;
+    ids: string[];
+    userId?: string;
+  }) {
+    await this.entitlement.assertManagementAccess(
+      input.workspaceId,
+      input.userId
+    );
+    await this.entitlement.assertServerEntitled(input.workspaceId);
+    if (input.storage !== ByokKeyStorage.server) {
+      throw new BadRequestException('Only server BYOK keys are persisted.');
+    }
+    await this.models.copilotWorkspaceByokConfig.reorder(
+      input.workspaceId,
+      input.ids,
+      input.userId
+    );
+    return (await this.getSettings(input.workspaceId, input.userId)).keys;
+  }
+
+  async deleteConfig(workspaceId: string, id: string, _userId?: string) {
+    await this.entitlement.assertManagementAccess(workspaceId, _userId);
+    await this.entitlement.assertServerEntitled(workspaceId);
+    await this.models.copilotWorkspaceByokConfig.delete(workspaceId, id);
+    return true;
+  }
+
+  async clearConfigs(
+    workspaceId: string,
+    provider: ByokProvider | null | undefined,
+    _userId?: string
+  ) {
+    await this.entitlement.assertManagementAccess(workspaceId, _userId);
+    await this.entitlement.assertServerEntitled(workspaceId);
+    await this.models.copilotWorkspaceByokConfig.clear(workspaceId, provider);
+    return true;
+  }
+
+  async testConfig(input: {
+    workspaceId: string;
+    provider: ByokProvider;
+    storage: ByokKeyStorage;
+    apiKey?: string | null;
+    endpoint?: string | null;
+    configId?: string | null;
+    userId?: string;
+  }) {
+    await this.entitlement.assertManagementAccess(
+      input.workspaceId,
+      input.userId
+    );
+    if (input.storage === ByokKeyStorage.server) {
+      await this.entitlement.assertServerEntitled(input.workspaceId);
+    } else {
+      await this.entitlement.assertLocalEntitled(
+        input.workspaceId,
+        input.userId
+      );
+    }
+    this.assertProvider(input.provider);
+    let apiKey = input.apiKey;
+    let endpoint = this.normalizeEndpoint(input.endpoint);
+    if (!apiKey && input.configId && input.storage === ByokKeyStorage.server) {
+      const config = await this.models.copilotWorkspaceByokConfig.get(
+        input.configId
+      );
+      if (
+        !config ||
+        config.workspaceId !== input.workspaceId ||
+        config.provider !== input.provider
+      ) {
+        throw new BadRequestException('BYOK config not found.');
+      }
+      apiKey = this.crypto.decrypt(config.encryptedApiKey);
+      endpoint =
+        input.endpoint !== undefined
+          ? endpoint
+          : this.normalizeEndpoint(config.endpoint);
+    }
+    if (!apiKey) {
+      throw new BadRequestException('apiKey is required.');
+    }
+
+    try {
+      await this.runProviderProbe(input.provider, apiKey, endpoint);
+      if (input.configId && input.storage === ByokKeyStorage.server) {
+        await this.models.copilotWorkspaceByokConfig.markValidated(
+          input.workspaceId,
+          input.configId,
+          input.userId
+        );
+      }
+      metrics.ai.counter('byok_test_key').add(1, {
+        workspace: input.workspaceId,
+        provider: input.provider,
+        storage: input.storage,
+        result: 'passed',
+      });
+      return { ok: true, status: ByokKeyTestStatus.passed, message: null };
+    } catch (error) {
+      const message = this.sanitizeError(error);
+      if (input.configId && input.storage === ByokKeyStorage.server) {
+        await this.models.copilotWorkspaceByokConfig.markFailure(
+          input.workspaceId,
+          input.configId,
+          message
+        );
+      }
+      metrics.ai.counter('byok_test_key').add(1, {
+        workspace: input.workspaceId,
+        provider: input.provider,
+        storage: input.storage,
+        result: 'failed',
+      });
+      return { ok: false, status: ByokKeyTestStatus.failed, message };
+    }
+  }
+
+  async createLocalLease(input: {
+    workspaceId: string;
+    providers: ByokLocalLeaseProvider[];
+    userId: string;
+  }) {
+    await this.entitlement.assertManagementAccess(
+      input.workspaceId,
+      input.userId
+    );
+    await this.entitlement.assertLocalEntitled(input.workspaceId, input.userId);
+    const providers = input.providers.map(provider => {
+      this.assertProvider(provider.provider);
+      const endpoint = this.normalizeEndpoint(provider.endpoint);
+      return { ...provider, endpoint };
+    });
+    const activeCacheKey = this.localLeaseActiveCacheKey({
+      ...input,
+      providers,
+    });
+    const activeLease = await this.getActiveLocalLease(activeCacheKey);
+    if (activeLease) return activeLease;
+
+    const leaseId = randomUUID();
+    const expiresAt = new Date(Date.now() + LOCAL_LEASE_TTL_MS);
+    const payload: LocalLeasePayload = {
+      workspaceId: input.workspaceId,
+      userId: input.userId,
+      providers: providers.map(provider => ({
+        provider: provider.provider,
+        name: provider.name,
+        description: provider.description,
+        encryptedApiKey: this.crypto.encrypt(provider.apiKey),
+        endpoint: provider.endpoint,
+        sortOrder: provider.sortOrder,
+        enabled: provider.enabled,
+      })),
+    };
+    await this.cache.set(this.leaseCacheKey(leaseId), payload, {
+      ttl: LOCAL_LEASE_TTL_MS,
+    });
+    const registered = await this.cache.setnx<LocalLeaseActive>(
+      activeCacheKey,
+      { leaseId, expiresAt: expiresAt.toISOString() },
+      { ttl: LOCAL_LEASE_TTL_MS }
+    );
+    if (!registered) {
+      const current = await this.getActiveLocalLease(activeCacheKey);
+      if (current) {
+        await this.cache.delete(this.leaseCacheKey(leaseId));
+        return current;
+      }
+    }
+    return { leaseId, expiresAt };
+  }
+
+  async getProfiles(
+    context: ByokProviderRequestContext = {},
+    sources: ByokProfileSourceFilter = { local: true, server: true }
+  ): Promise<CopilotProviderProfile[]> {
+    if (!context.workspaceId) {
+      return [];
+    }
+    const [localEntitled, serverEntitled] = await Promise.all([
+      this.entitlement.hasLocalEntitlement(context.workspaceId, context.userId),
+      this.entitlement.hasServerEntitlement(context.workspaceId),
+    ]);
+    const [localProfiles, serverProfiles] = await Promise.all([
+      sources.local && localEntitled
+        ? this.getLocalProfiles(context)
+        : Promise.resolve([]),
+      sources.server && serverEntitled
+        ? this.getServerProfiles(context.workspaceId)
+        : Promise.resolve([]),
+    ]);
+
+    return [...localProfiles, ...serverProfiles];
+  }
+
+  async recordUsage(input: {
+    workspaceId?: string;
+    userId?: string;
+    providerId?: string;
+    model?: string | null;
+    featureKind: ByokFeatureKind;
+    sessionId?: string;
+    taskId?: string;
+    actionId?: string;
+    billingUnitId?: string;
+    usage?: {
+      prompt_tokens?: number;
+      completion_tokens?: number;
+      total_tokens?: number;
+      cached_tokens?: number;
+    };
+  }) {
+    if (!input.workspaceId || !input.providerId) return;
+    const meta = this.parseProfileMeta(input.providerId, input.workspaceId);
+    if (!meta) return;
+
+    metrics.ai.counter('byok_usage').add(1, {
+      workspace: input.workspaceId,
+      provider: meta.provider,
+      source: meta.source,
+      feature: input.featureKind,
+    });
+    await this.models.copilotUsage.create({
+      workspaceId: input.workspaceId,
+      userId: input.userId,
+      provider: meta.provider,
+      providerSource: meta.source,
+      featureKind: input.featureKind,
+      model: input.model ?? null,
+      sessionId: input.sessionId,
+      taskId: input.taskId,
+      actionId: input.actionId,
+      billingUnitId: input.billingUnitId,
+      promptTokens: input.usage?.prompt_tokens ?? 0,
+      completionTokens: input.usage?.completion_tokens ?? 0,
+      totalTokens: input.usage?.total_tokens ?? 0,
+      cachedTokens: input.usage?.cached_tokens ?? 0,
+    });
+    if (meta.source === ByokProviderSource.Server && meta.keyId) {
+      await this.models.copilotWorkspaceByokConfig.touchUsed(
+        input.workspaceId,
+        meta.keyId
+      );
+    }
+  }
+
+  async recordProviderFailure(input: {
+    workspaceId?: string;
+    providerId?: string;
+    featureKind: ByokFeatureKind;
+    error: unknown;
+  }) {
+    if (!input.workspaceId || !input.providerId) return;
+    const meta = this.parseProfileMeta(input.providerId, input.workspaceId);
+    if (!meta) return;
+
+    const message = this.sanitizeError(input.error);
+    metrics.ai.counter('byok_route_failure').add(1, {
+      workspace: input.workspaceId,
+      provider: meta.provider,
+      source: meta.source,
+      feature: input.featureKind,
+    });
+    if (meta.source === ByokProviderSource.Server && meta.keyId) {
+      await this.models.copilotWorkspaceByokConfig.markFailure(
+        input.workspaceId,
+        meta.keyId,
+        message
+      );
+    }
+  }
+
+  async getUsage(workspaceId: string, from: Date, to: Date) {
+    return await this.models.copilotUsage.aggregateByDay({
+      workspaceId,
+      from,
+      to,
+      providerSources: [ByokProviderSource.Server, ByokProviderSource.Local],
+    });
+  }
+
+  private async getServerProfiles(workspaceId: string) {
+    const rows =
+      await this.models.copilotWorkspaceByokConfig.listEnabled(workspaceId);
+
+    return rows
+      .filter(row => isByokProvider(row.provider))
+      .map((row, index): CopilotProviderProfile => {
+        const provider = row.provider as ByokProvider;
+        return {
+          id: this.profileId(workspaceId, provider, row.id, 'server'),
+          type: byokProviderToCopilotType(provider),
+          priority:
+            BYOK_PROFILE_PRIORITY_BASE - SERVER_PROFILE_PRIORITY_OFFSET - index,
+          config: this.providerConfig(
+            provider,
+            row.encryptedApiKey,
+            row.endpoint
+          ),
+        } as CopilotProviderProfile;
+      });
+  }
+
+  private async getLocalProfiles(context: ByokProviderRequestContext) {
+    if (!context.byokLeaseId || !context.workspaceId || !context.userId) {
+      return [];
+    }
+    if (
+      !(await this.entitlement.hasManagementAccess(
+        context.workspaceId,
+        context.userId
+      ))
+    ) {
+      return [];
+    }
+    const lease = await this.cache.get<LocalLeasePayload>(
+      this.leaseCacheKey(context.byokLeaseId)
+    );
+    if (
+      !lease ||
+      lease.workspaceId !== context.workspaceId ||
+      lease.userId !== context.userId
+    ) {
+      return [];
+    }
+    return lease.providers
+      .filter(provider => provider.enabled !== false)
+      .map((provider, index): CopilotProviderProfile => {
+        return {
+          id: this.profileId(
+            context.workspaceId ?? lease.workspaceId,
+            provider.provider,
+            `${index}`,
+            'local'
+          ),
+          type: byokProviderToCopilotType(provider.provider),
+          priority: BYOK_PROFILE_PRIORITY_BASE - index,
+          config: this.providerConfig(
+            provider.provider,
+            provider.encryptedApiKey,
+            provider.endpoint ?? null
+          ),
+        } as CopilotProviderProfile;
+      });
+  }
+
+  private providerConfig(
+    provider: ByokProvider,
+    encryptedApiKey: string,
+    endpoint: string | null
+  ) {
+    const apiKey = this.crypto.decrypt(encryptedApiKey);
+    switch (provider) {
+      case ByokProvider.openai:
+      case ByokProvider.gemini:
+      case ByokProvider.anthropic:
+        return { apiKey, ...(endpoint ? { baseURL: endpoint } : {}) };
+      case ByokProvider.fal:
+        return { apiKey };
+    }
+  }
+
+  private profileId(
+    workspaceId: string,
+    provider: ByokProvider,
+    keyId: string,
+    storage: 'server' | 'local'
+  ) {
+    const hash = this.workspaceHash(workspaceId);
+    const sanitizedKeyId = keyId.replaceAll(/[^a-zA-Z0-9-_]/g, '');
+    return storage === 'local'
+      ? `byok-${hash}-${provider}-local-${sanitizedKeyId}`
+      : `byok-${hash}-${provider}-${sanitizedKeyId}`;
+  }
+
+  parseProfileMeta(
+    providerId: string,
+    workspaceId?: string
+  ): ByokProfileMeta | null {
+    const match =
+      /^byok-([a-f0-9]{12})-(openai|anthropic|gemini|fal)-(.+)$/.exec(
+        providerId
+      );
+    if (!match) return null;
+    if (workspaceId && match[1] !== this.workspaceHash(workspaceId)) {
+      return null;
+    }
+
+    const keyId = match[3];
+    return {
+      provider: match[2] as ByokProvider,
+      source: keyId.startsWith('local-')
+        ? ByokProviderSource.Local
+        : ByokProviderSource.Server,
+      keyId: keyId.startsWith('local-') ? undefined : keyId,
+    };
+  }
+
+  private toKeyConfig(row: {
+    id: string;
+    provider: string;
+    name: string;
+    description: string | null;
+    endpoint: string | null;
+    sortOrder: number;
+    enabled: boolean;
+    disabledReason: string | null;
+    lastValidatedAt: Date | null;
+    lastValidationError: string | null;
+    lastUsedAt: Date | null;
+    lastErrorAt: Date | null;
+    lastError: string | null;
+  }): ByokKeyConfig {
+    const provider = row.provider as ByokProvider;
+    return {
+      id: row.id,
+      provider,
+      name: row.name,
+      description: row.description,
+      storage: ByokKeyStorage.server,
+      configured: true,
+      enabled: row.enabled,
+      endpoint: row.endpoint,
+      endpointEditable: this.customEndpointSupported,
+      sortOrder: row.sortOrder,
+      capabilities: this.capabilities(provider, 'server'),
+      testStatus: row.lastValidationError
+        ? ByokKeyTestStatus.failed
+        : row.lastValidatedAt
+          ? ByokKeyTestStatus.passed
+          : ByokKeyTestStatus.untested,
+      disabledReason: row.disabledReason,
+      lastTestedAt: row.lastValidatedAt,
+      lastTestError: row.lastValidationError,
+      lastUsedAt: row.lastUsedAt,
+      lastErrorAt: row.lastErrorAt,
+      lastError: row.lastError,
+    };
+  }
+
+  private capabilities(provider: ByokProvider, storage: 'server' | 'local') {
+    switch (provider) {
+      case ByokProvider.openai:
+        return ['Text', 'Image input', 'Actions', 'Image generate'];
+      case ByokProvider.anthropic:
+        return ['Text', 'Image input'];
+      case ByokProvider.gemini:
+        return storage === 'server'
+          ? [
+              'Text',
+              'Image input',
+              'Actions',
+              'Image generate',
+              'Transcript',
+              'Indexing',
+            ]
+          : ['Text', 'Image input', 'Actions', 'Image generate'];
+      case ByokProvider.fal:
+        return ['Image generate'];
+    }
+  }
+
+  private buildWarnings(keys: ByokKeyConfig[]) {
+    const activeServerGemini = keys.some(
+      key =>
+        key.provider === ByokProvider.gemini &&
+        key.storage === ByokKeyStorage.server &&
+        key.enabled
+    );
+    if (activeServerGemini) {
+      return [];
+    }
+    return [
+      {
+        featureKind: 'transcript',
+        reason:
+          'Transcript and workspace indexing require a server Gemini BYOK key or AFFiNE AI plan fallback.',
+        requiredProviders: [ByokProvider.gemini],
+      },
+      {
+        featureKind: 'workspace_indexing',
+        reason:
+          'Workspace indexing requires a server Gemini BYOK key or AFFiNE AI plan fallback.',
+        requiredProviders: [ByokProvider.gemini],
+      },
+    ];
+  }
+
+  private normalizeEndpoint(endpoint?: string | null) {
+    if (!endpoint) return null;
+    if (!this.customEndpointSupported) {
+      throw new BadRequestException('Custom BYOK endpoint is not supported.');
+    }
+    let parsed: URL;
+    try {
+      parsed = new URL(endpoint);
+    } catch {
+      throw new BadRequestException('Invalid BYOK endpoint.');
+    }
+    if (!['https:', 'http:'].includes(parsed.protocol)) {
+      throw new BadRequestException('BYOK endpoint must use HTTP or HTTPS.');
+    }
+    return parsed.toString().replace(/\/$/, '');
+  }
+
+  private assertProvider(provider: ByokProvider) {
+    if (!BYOK_ALLOWED_PROVIDERS.includes(provider)) {
+      throw new BadRequestException('Unsupported BYOK provider.');
+    }
+  }
+
+  private async runProviderProbe(
+    provider: ByokProvider,
+    apiKey: string,
+    endpoint: string | null
+  ) {
+    const controller = new AbortController();
+    const timeout = setTimeout(() => controller.abort(), TEST_TIMEOUT_MS);
+    try {
+      const request = this.buildProbeRequest(provider, apiKey, endpoint);
+      const response = await fetch(request.url, {
+        method: request.method,
+        headers: request.headers as unknown as Record<string, string>,
+        signal: controller.signal,
+      });
+      if (!response.ok) {
+        throw new BadRequestException(
+          this.providerProbeFailureMessage(response.status)
+        );
+      }
+    } finally {
+      clearTimeout(timeout);
+    }
+  }
+
+  private buildProbeRequest(
+    provider: ByokProvider,
+    apiKey: string,
+    endpoint: string | null
+  ) {
+    switch (provider) {
+      case ByokProvider.openai:
+        return {
+          method: 'GET',
+          url: `${endpoint ?? 'https://api.openai.com/v1'}/models`,
+          headers: { Authorization: `Bearer ${apiKey}` },
+        };
+      case ByokProvider.anthropic:
+        return {
+          method: 'GET',
+          url: `${endpoint ?? 'https://api.anthropic.com/v1'}/models`,
+          headers: {
+            'x-api-key': apiKey,
+            'anthropic-version': '2023-06-01',
+          },
+        };
+      case ByokProvider.gemini:
+        return {
+          method: 'GET',
+          url: `${endpoint ?? 'https://generativelanguage.googleapis.com/v1beta'}/models`,
+          headers: { 'x-goog-api-key': apiKey },
+        };
+      case ByokProvider.fal:
+        return {
+          method: 'GET',
+          url: 'https://api.fal.ai/v1/models?limit=10',
+          headers: { Authorization: `Key ${apiKey}` },
+        };
+    }
+  }
+
+  private sanitizeError(error: unknown) {
+    if (error instanceof Error && error.name === 'AbortError') {
+      return 'Provider key test timed out.';
+    }
+    if (error instanceof BadRequestException && error.message) {
+      return error.message.slice(0, 300);
+    }
+    return 'Provider request failed.';
+  }
+
+  private providerProbeFailureMessage(status: number) {
+    switch (status) {
+      case 401:
+        return 'Provider rejected the BYOK key.';
+      case 403:
+        return 'Provider rejected the BYOK key permissions.';
+      case 404:
+        return 'Provider probe endpoint was not found.';
+      case 429:
+        return 'Provider rate limit exceeded while testing the key.';
+      default:
+        return status >= 500
+          ? 'Provider service is unavailable.'
+          : `Provider key test failed with HTTP ${status}.`;
+    }
+  }
+
+  private workspaceHash(workspaceId: string) {
+    return createHash('sha256').update(workspaceId).digest('hex').slice(0, 12);
+  }
+
+  private leaseCacheKey(leaseId: string) {
+    return `copilot:byok:lease:${leaseId}`;
+  }
+
+  private async getActiveLocalLease(activeCacheKey: string) {
+    const active = await this.cache.get<LocalLeaseActive>(activeCacheKey);
+    if (!active) return null;
+    if (await this.cache.has(this.leaseCacheKey(active.leaseId))) {
+      return { leaseId: active.leaseId, expiresAt: new Date(active.expiresAt) };
+    }
+    await this.cache.delete(activeCacheKey);
+    return null;
+  }
+
+  private localLeaseActiveCacheKey(input: {
+    workspaceId: string;
+    userId: string;
+    providers: ByokLocalLeaseProvider[];
+  }) {
+    const fingerprint = createHmac(
+      'sha256',
+      this.crypto.keyPair.sha256.privateKey
+    )
+      .update(
+        JSON.stringify(
+          input.providers.map(provider => ({
+            provider: provider.provider,
+            name: provider.name,
+            description: provider.description ?? null,
+            apiKey: provider.apiKey,
+            endpoint: provider.endpoint ?? null,
+            sortOrder: provider.sortOrder ?? 0,
+            enabled: provider.enabled ?? true,
+          }))
+        )
+      )
+      .digest('hex');
+    return `copilot:byok:lease:active:${input.workspaceId}:${input.userId}:${fingerprint}`;
+  }
+}
@@ -0,0 +1,79 @@
+import { registerEnumType } from '@nestjs/graphql';
+
+import { CopilotProviderType } from '../providers/types';
+
+export enum ByokProvider {
+  openai = 'openai',
+  anthropic = 'anthropic',
+  gemini = 'gemini',
+  fal = 'fal',
+}
+
+export enum ByokKeyStorage {
+  server = 'server',
+  local = 'local',
+}
+
+export enum ByokKeyTestStatus {
+  untested = 'untested',
+  passed = 'passed',
+  failed = 'failed',
+}
+
+export enum ByokProviderSource {
+  Server = 'byok_server',
+  Local = 'byok_local',
+  AffinePlan = 'affine_plan',
+}
+
+export type ByokFeatureKind =
+  | 'chat'
+  | 'action'
+  | 'image'
+  | 'embedding'
+  | 'rerank'
+  | 'transcript'
+  | 'workspace_indexing';
+
+export const BYOK_ALLOWED_PROVIDERS = [
+  ByokProvider.openai,
+  ByokProvider.anthropic,
+  ByokProvider.gemini,
+  ByokProvider.fal,
+] as const;
+
+export function byokProviderToCopilotType(provider: ByokProvider) {
+  switch (provider) {
+    case ByokProvider.openai:
+      return CopilotProviderType.OpenAI;
+    case ByokProvider.anthropic:
+      return CopilotProviderType.Anthropic;
+    case ByokProvider.gemini:
+      return CopilotProviderType.Gemini;
+    case ByokProvider.fal:
+      return CopilotProviderType.FAL;
+  }
+}
+
+export function copilotTypeToByokProvider(type: CopilotProviderType) {
+  switch (type) {
+    case CopilotProviderType.OpenAI:
+      return ByokProvider.openai;
+    case CopilotProviderType.Anthropic:
+      return ByokProvider.anthropic;
+    case CopilotProviderType.Gemini:
+      return ByokProvider.gemini;
+    case CopilotProviderType.FAL:
+      return ByokProvider.fal;
+    default:
+      return null;
+  }
+}
+
+export function isByokProvider(value: string): value is ByokProvider {
+  return (BYOK_ALLOWED_PROVIDERS as readonly string[]).includes(value);
+}
+
+registerEnumType(ByokProvider, { name: 'ByokProvider' });
+registerEnumType(ByokKeyStorage, { name: 'ByokKeyStorage' });
+registerEnumType(ByokKeyTestStatus, { name: 'ByokKeyTestStatus' });
@@ -12,9 +12,7 @@ import {
 import { CloudflareWorkersAIConfig } from './providers/cloudflare';
 import type { FalConfig } from './providers/fal';
 import { GeminiGenerativeConfig, GeminiVertexConfig } from './providers/gemini';
-import { MorphConfig } from './providers/morph';
 import { OpenAIConfig } from './providers/openai';
-import { PerplexityConfig } from './providers/perplexity';
 import {
  CopilotProviderType,
  ModelOutputType,
@@ -27,10 +25,8 @@ export type CopilotProviderConfigMap = {
  [CopilotProviderType.FAL]: FalConfig;
  [CopilotProviderType.Gemini]: GeminiGenerativeConfig;
  [CopilotProviderType.GeminiVertex]: GeminiVertexConfig;
-  [CopilotProviderType.Perplexity]: PerplexityConfig;
  [CopilotProviderType.Anthropic]: AnthropicOfficialConfig;
  [CopilotProviderType.AnthropicVertex]: AnthropicVertexConfig;
-  [CopilotProviderType.Morph]: MorphConfig;
 };

 export type ProviderSpecificConfig =
@@ -138,20 +134,11 @@ const VertexProviderConfigShape = z.object({
  fetch: z.any().optional(),
 });

-const PerplexityConfigShape = z.object({
-  apiKey: z.string(),
-  endpoint: z.string().optional(),
-});
-
 const AnthropicOfficialConfigShape = z.object({
  apiKey: z.string(),
  baseURL: z.string().optional(),
 });

-const MorphConfigShape = z.object({
-  apiKey: z.string().optional(),
-});
-
 const CopilotProviderProfileShape = z.discriminatedUnion('type', [
  CopilotProviderProfileBaseShape.extend({
    type: z.literal(CopilotProviderType.OpenAI),
@@ -173,10 +160,6 @@ const CopilotProviderProfileShape = z.discriminatedUnion('type', [
    type: z.literal(CopilotProviderType.GeminiVertex),
    config: VertexProviderConfigShape,
  }),
-  CopilotProviderProfileBaseShape.extend({
-    type: z.literal(CopilotProviderType.Perplexity),
-    config: PerplexityConfigShape,
-  }),
  CopilotProviderProfileBaseShape.extend({
    type: z.literal(CopilotProviderType.Anthropic),
    config: AnthropicOfficialConfigShape,
@@ -185,10 +168,6 @@ const CopilotProviderProfileShape = z.discriminatedUnion('type', [
    type: z.literal(CopilotProviderType.AnthropicVertex),
    config: VertexProviderConfigShape,
  }),
-  CopilotProviderProfileBaseShape.extend({
-    type: z.literal(CopilotProviderType.Morph),
-    config: MorphConfigShape,
-  }),
 ]);

 const CopilotProviderDefaultsShape = z.object({
@@ -205,6 +184,13 @@ declare global {
  interface AppConfigSchema {
    copilot: {
      enabled: boolean;
+      byok: {
+        enabled: ConfigItem<boolean>;
+        allowedProviders: ConfigItem<
+          Array<'openai' | 'anthropic' | 'gemini' | 'fal'>
+        >;
+        allowCustomEndpoint: ConfigItem<boolean>;
+      };
      unsplash: ConfigItem<{
        key: string;
      }>;
@@ -220,10 +206,8 @@ declare global {
        fal: ConfigItem<FalConfig>;
        gemini: ConfigItem<GeminiGenerativeConfig>;
        geminiVertex: ConfigItem<GeminiVertexConfig>;
-        perplexity: ConfigItem<PerplexityConfig>;
        anthropic: ConfigItem<AnthropicOfficialConfig>;
        anthropicVertex: ConfigItem<AnthropicVertexConfig>;
-        morph: ConfigItem<MorphConfig>;
      };
    };
  }
@@ -234,6 +218,21 @@ defineModuleConfig('copilot', {
    desc: 'Whether to enable the copilot plugin. <br> Document: <a href="https://docs.affine.pro/self-host-affine/administer/ai" target="_blank">https://docs.affine.pro/self-host-affine/administer/ai</a>',
    default: false,
  },
+  'byok.enabled': {
+    desc: 'Whether to enable workspace BYOK.',
+    default: true,
+    shape: z.boolean(),
+  },
+  'byok.allowedProviders': {
+    desc: 'The allowlist for workspace BYOK providers.',
+    default: ['openai', 'anthropic', 'gemini', 'fal'],
+    shape: z.array(z.enum(['openai', 'anthropic', 'gemini', 'fal'])),
+  },
+  'byok.allowCustomEndpoint': {
+    desc: 'Whether workspace BYOK custom endpoints are accepted.',
+    default: false,
+    shape: z.boolean(),
+  },
  'providers.profiles': {
    desc: 'The profile list for copilot providers.',
    default: [],
@@ -277,12 +276,6 @@ defineModuleConfig('copilot', {
    default: {},
    schema: VertexSchema,
  },
-  'providers.perplexity': {
-    desc: 'The config for the perplexity provider.',
-    default: {
-      apiKey: '',
-    },
-  },
  'providers.anthropic': {
    desc: 'The config for the anthropic provider.',
    default: {
@@ -295,10 +288,6 @@ defineModuleConfig('copilot', {
    default: {},
    schema: VertexSchema,
  },
-  'providers.morph': {
-    desc: 'The config for the morph provider.',
-    default: {},
-  },
  unsplash: {
    desc: 'The config for the unsplash key.',
    default: {
@@ -15,7 +15,11 @@ import {
  Models,
 } from '../../../models';
 import { CopilotEmbeddingClientService } from '../embedding/client';
-import type { EmbeddingClient } from '../embedding/types';
+import type {
+  EmbeddingCallOptions,
+  EmbeddingClient,
+  EmbeddingRouteContext,
+} from '../embedding/types';
 import { ContextSession } from './session';

 const CONTEXT_SESSION_KEY = 'context-session';
@@ -62,6 +66,14 @@ export class CopilotContextService implements OnApplicationBootstrap {
    return this.client ?? this.embeddingClients.getClient();
  }

+  private embeddingOptions(
+    workspaceId: string,
+    signal?: AbortSignal,
+    routeContext: EmbeddingRouteContext = {}
+  ): EmbeddingCallOptions {
+    return { workspaceId, signal, ...routeContext, featureKind: 'embedding' };
+  }
+
  private async saveConfig(
    contextId: string,
    config: ContextConfig,
@@ -172,11 +184,13 @@ export class CopilotContextService implements OnApplicationBootstrap {
    content: string,
    topK: number = 5,
    signal?: AbortSignal,
-    threshold: number = 0.5
+    threshold: number = 0.5,
+    routeContext?: EmbeddingRouteContext
  ) {
    const client = this.embeddingClient;
    if (!client) return [];
-    const embedding = await client.getEmbedding(content, signal);
+    const options = this.embeddingOptions(workspaceId, signal, routeContext);
+    const embedding = await client.getEmbedding(content, options);
    if (!embedding) return [];

    const blobChunks = await this.models.copilotWorkspace.matchBlobEmbedding(
@@ -187,7 +201,7 @@ export class CopilotContextService implements OnApplicationBootstrap {
    );
    if (!blobChunks.length) return [];

-    return await client.reRank(content, blobChunks, topK, signal);
+    return await client.reRank(content, blobChunks, topK, options);
  }

  async matchWorkspaceFiles(
@@ -195,11 +209,13 @@ export class CopilotContextService implements OnApplicationBootstrap {
    content: string,
    topK: number = 5,
    signal?: AbortSignal,
-    threshold: number = 0.5
+    threshold: number = 0.5,
+    routeContext?: EmbeddingRouteContext
  ) {
    const client = this.embeddingClient;
    if (!client) return [];
-    const embedding = await client.getEmbedding(content, signal);
+    const options = this.embeddingOptions(workspaceId, signal, routeContext);
+    const embedding = await client.getEmbedding(content, options);
    if (!embedding) return [];

    const fileChunks = await this.models.copilotWorkspace.matchFileEmbedding(
@@ -210,7 +226,7 @@ export class CopilotContextService implements OnApplicationBootstrap {
    );
    if (!fileChunks.length) return [];

-    return await client.reRank(content, fileChunks, topK, signal);
+    return await client.reRank(content, fileChunks, topK, options);
  }

  async matchWorkspaceDocs(
@@ -218,11 +234,13 @@ export class CopilotContextService implements OnApplicationBootstrap {
    content: string,
    topK: number = 5,
    signal?: AbortSignal,
-    threshold: number = 0.5
+    threshold: number = 0.5,
+    routeContext?: EmbeddingRouteContext
  ) {
    const client = this.embeddingClient;
    if (!client) return [];
-    const embedding = await client.getEmbedding(content, signal);
+    const options = this.embeddingOptions(workspaceId, signal, routeContext);
+    const embedding = await client.getEmbedding(content, options);
    if (!embedding) return [];

    const workspaceChunks =
@@ -234,7 +252,7 @@ export class CopilotContextService implements OnApplicationBootstrap {
      );
    if (!workspaceChunks.length) return [];

-    return await client.reRank(content, workspaceChunks, topK, signal);
+    return await client.reRank(content, workspaceChunks, topK, options);
  }

  async matchWorkspaceAll(
@@ -244,11 +262,13 @@ export class CopilotContextService implements OnApplicationBootstrap {
    signal?: AbortSignal,
    threshold: number = 0.8,
    docIds?: string[],
-    scopedThreshold: number = 0.85
+    scopedThreshold: number = 0.85,
+    routeContext?: EmbeddingRouteContext
  ) {
    const client = this.embeddingClient;
    if (!client) return [];
-    const embedding = await client.getEmbedding(content, signal);
+    const options = this.embeddingOptions(workspaceId, signal, routeContext);
+    const embedding = await client.getEmbedding(content, options);
    if (!embedding) return [];

    const [fileChunks, blobChunks, workspaceChunks, scopedWorkspaceChunks] =
@@ -300,7 +320,7 @@ export class CopilotContextService implements OnApplicationBootstrap {
        ...(scopedWorkspaceChunks || []),
      ],
      topK,
-      signal
+      options
    );
  }

@@ -11,7 +11,11 @@ import {
  FileChunkSimilarity,
  Models,
 } from '../../../models';
-import { EmbeddingClient } from '../embedding/types';
+import type {
+  EmbeddingCallOptions,
+  EmbeddingClient,
+  EmbeddingRouteContext,
+} from '../embedding/types';

 export class ContextSession implements AsyncDisposable {
  constructor(
@@ -69,6 +73,18 @@ export class ContextSession implements AsyncDisposable {
    );
  }

+  private embeddingOptions(
+    signal?: AbortSignal,
+    routeContext: EmbeddingRouteContext = {}
+  ): EmbeddingCallOptions {
+    return {
+      workspaceId: this.workspaceId,
+      signal,
+      ...routeContext,
+      featureKind: 'embedding',
+    };
+  }
+
  async addCategoryRecord(type: ContextCategories, id: string, docs: string[]) {
    const category = this.config.categories.find(
      c => c.type === type && c.id === id
@@ -269,10 +285,12 @@ export class ContextSession implements AsyncDisposable {
    topK: number = 5,
    signal?: AbortSignal,
    scopedThreshold: number = 0.85,
-    threshold: number = 0.5
+    threshold: number = 0.5,
+    routeContext?: EmbeddingRouteContext
  ): Promise<FileChunkSimilarity[]> {
    if (!this.client) return [];
-    const embedding = await this.client.getEmbedding(content, signal);
+    const options = this.embeddingOptions(signal, routeContext);
+    const embedding = await this.client.getEmbedding(content, options);
    if (!embedding) return [];

    const [context, workspace] = await Promise.all([
@@ -305,7 +323,7 @@ export class ContextSession implements AsyncDisposable {
        ...workspace,
      ],
      topK,
-      signal
+      options
    );
  }

@@ -322,10 +340,12 @@ export class ContextSession implements AsyncDisposable {
    topK: number = 5,
    signal?: AbortSignal,
    scopedThreshold: number = 0.85,
-    threshold: number = 0.5
+    threshold: number = 0.5,
+    routeContext?: EmbeddingRouteContext
  ) {
    if (!this.client) return [];
-    const embedding = await this.client.getEmbedding(content, signal);
+    const options = this.embeddingOptions(signal, routeContext);
+    const embedding = await this.client.getEmbedding(content, options);
    if (!embedding) return [];

    const docIds = this.docIds;
@@ -349,7 +369,7 @@ export class ContextSession implements AsyncDisposable {
      content,
      [...inContext, ...workspace],
      topK,
-      signal
+      options
    );

    // sort result, doc recorded in context first
@@ -1,7 +1,7 @@
 import { Injectable } from '@nestjs/common';

 import { CopilotQuotaExceeded } from '../../../base';
-import { QuotaService } from '../../../core/quota';
+import { QuotaService } from '../../../core/quota/service';
 import { Models } from '../../../models';
 import type { Turn } from '../core';
 import type { ResolvedPrompt } from '../prompt';
@@ -31,12 +31,16 @@ export class ConversationPolicy {
  }

  async checkQuota(userId: string) {
-    const { limit, used } = await this.getQuota(userId);
-    if (limit && Number.isFinite(limit) && used >= limit) {
+    if (!(await this.hasQuota(userId))) {
      throw new CopilotQuotaExceeded();
    }
  }

+  async hasQuota(userId: string) {
+    const { limit, used } = await this.getQuota(userId);
+    return !(limit !== undefined && Number.isFinite(limit) && used >= limit);
+  }
+
  shouldScheduleTitle(prompt: Pick<ResolvedPrompt, 'action'>) {
    return !prompt.action;
  }
@@ -11,7 +11,12 @@ import {
 import { type CopilotRerankRequest } from '../providers/types';
 import { CapabilityRuntime } from '../runtime/capability-runtime';
 import { TaskPolicy } from '../runtime/task-policy';
-import { EmbeddingClient, type ReRankResult } from './types';
+import {
+  type EmbeddingCallOptionsInput,
+  EmbeddingClient,
+  normalizeEmbeddingCallOptions,
+  type ReRankResult,
+} from './types';

 class ProductionEmbeddingClient extends EmbeddingClient {
  private readonly logger = new Logger(ProductionEmbeddingClient.name);
@@ -35,10 +40,19 @@ class ProductionEmbeddingClient extends EmbeddingClient {
    return result;
  }

-  async getEmbeddings(input: string[]): Promise<Embedding[]> {
+  async getEmbeddings(
+    input: string[],
+    options?: EmbeddingCallOptionsInput
+  ): Promise<Embedding[]> {
+    const normalizedOptions = normalizeEmbeddingCallOptions(options);
    const modelId = this.taskPolicy.resolveEmbeddingModelId();
    const embeddings = await this.runtime.embed(modelId, input, {
      dimensions: EMBEDDING_DIMENSIONS,
+      signal: normalizedOptions.signal,
+      user: normalizedOptions.userId,
+      workspace: normalizedOptions.workspaceId,
+      byokLeaseId: normalizedOptions.byokLeaseId,
+      featureKind: normalizedOptions.featureKind ?? 'embedding',
    });
    if (embeddings.length !== input.length) {
      throw new CopilotFailedToGenerateEmbedding({
@@ -67,8 +81,9 @@ class ProductionEmbeddingClient extends EmbeddingClient {
  >(
    query: string,
    embeddings: Chunk[],
-    signal?: AbortSignal
+    options?: EmbeddingCallOptionsInput
  ): Promise<ReRankResult> {
+    const normalizedOptions = normalizeEmbeddingCallOptions(options);
    if (!embeddings.length) return [];

    const rerankRequest: CopilotRerankRequest = {
@@ -82,7 +97,13 @@ class ProductionEmbeddingClient extends EmbeddingClient {
    const ranks = await this.runtime.rerank(
      this.taskPolicy.resolveRerankModelId(),
      rerankRequest,
-      { signal }
+      {
+        signal: normalizedOptions.signal,
+        user: normalizedOptions.userId,
+        workspace: normalizedOptions.workspaceId,
+        byokLeaseId: normalizedOptions.byokLeaseId,
+        featureKind: 'rerank',
+      }
    );

    try {
@@ -105,8 +126,9 @@ class ProductionEmbeddingClient extends EmbeddingClient {
    query: string,
    embeddings: Chunk[],
    topK: number,
-    signal?: AbortSignal
+    options?: EmbeddingCallOptionsInput
  ): Promise<Chunk[]> {
+    const normalizedOptions = normalizeEmbeddingCallOptions(options);
    // search in context and workspace may find same chunks, de-duplicate them
    const { deduped: dedupedEmbeddings } = embeddings.reduce(
      (acc, e) => {
@@ -138,14 +160,19 @@ class ProductionEmbeddingClient extends EmbeddingClient {
      const ranks = await this.getEmbeddingRelevance(
        query,
        sortedEmbeddings,
-        signal
+        normalizedOptions
      );
      if (sortedEmbeddings.length !== ranks.length) {
        // llm return wrong result, fallback to default sorting
        this.logger.warn(
          `Batch size mismatch: expected ${sortedEmbeddings.length}, got ${ranks.length}`
        );
-        return await super.reRank(query, dedupedEmbeddings, topK, signal);
+        return await super.reRank(
+          query,
+          dedupedEmbeddings,
+          topK,
+          normalizedOptions
+        );
      }

      const highConfidenceChunks = ranks
@@ -164,7 +191,12 @@ class ProductionEmbeddingClient extends EmbeddingClient {
      return highConfidenceChunks.slice(0, topK);
    } catch (error) {
      this.logger.warn('ReRank failed, falling back to default sorting', error);
-      return await super.reRank(query, dedupedEmbeddings, topK, signal);
+      return await super.reRank(
+        query,
+        dedupedEmbeddings,
+        topK,
+        normalizedOptions
+      );
    }
  }
 }
@@ -180,7 +212,8 @@ export class CopilotEmbeddingClientService {

  async refresh() {
    const client = new ProductionEmbeddingClient(this.taskPolicy, this.runtime);
-    this.client = (await client.configured()) ? client : undefined;
+    await client.configured();
+    this.client = client;
    return this.client;
  }

@@ -18,7 +18,7 @@ import { Models } from '../../../models';
 import { CopilotStorage } from '../storage';
 import { readStream } from '../utils';
 import { CopilotEmbeddingClientService } from './client';
-import type { Chunk, DocFragment } from './types';
+import type { Chunk, DocFragment, EmbeddingCallOptions } from './types';
 import { EmbeddingClient } from './types';

@Injectable()
@@ -242,6 +242,19 @@ export class CopilotEmbeddingJob {
    return new File([buffer], fileName);
  }

+  private workspaceIndexingOptions(
+    workspaceId: string,
+    signal?: AbortSignal,
+    userId?: string
+  ): EmbeddingCallOptions {
+    return {
+      workspaceId,
+      userId,
+      signal,
+      featureKind: 'workspace_indexing',
+    };
+  }
+
  @OnJob('copilot.embedding.files')
  async embedPendingFile({
    userId,
@@ -266,7 +279,10 @@ export class CopilotEmbeddingJob {
      const total = chunks.reduce((acc, c) => acc + c.length, 0);

      for (const chunk of chunks) {
-        const embeddings = await this.embeddingClient.generateEmbeddings(chunk);
+        const embeddings = await this.embeddingClient.generateEmbeddings(
+          chunk,
+          this.workspaceIndexingOptions(workspaceId, undefined, userId)
+        );
        if (contextId) {
          // for context files
          await this.models.copilotContext.insertFileEmbedding(
@@ -320,7 +336,10 @@ export class CopilotEmbeddingJob {
      const total = chunks.reduce((acc, c) => acc + c.length, 0);

      for (const chunk of chunks) {
-        const embeddings = await this.embeddingClient.generateEmbeddings(chunk);
+        const embeddings = await this.embeddingClient.generateEmbeddings(
+          chunk,
+          this.workspaceIndexingOptions(workspaceId)
+        );
        await this.models.copilotWorkspace.insertBlobEmbeddings(
          workspaceId,
          blobId,
@@ -462,7 +481,7 @@ export class CopilotEmbeddingJob {
                `${fragment.title || 'Untitled'}.md`
              ),
              chunks => this.formatDocChunks(chunks, fragment),
-              signal
+              this.workspaceIndexingOptions(workspaceId, signal)
            );

            for (const chunks of embeddings) {
@@ -6,6 +6,7 @@ import { CopilotContextFileNotSupported } from '../../../base';
 import type { PageDocContent } from '../../../core/utils/blocksuite';
 import { ChunkSimilarity, Embedding } from '../../../models';
 import { parseDoc } from '../../../native';
+import type { ByokFeatureKind } from '../byok/types';

 declare global {
  interface Events {
@@ -103,6 +104,35 @@ export type Chunk = {
  content: string;
 };

+export type EmbeddingCallOptions = {
+  signal?: AbortSignal;
+  userId?: string;
+  workspaceId?: string;
+  byokLeaseId?: string;
+  featureKind?: Extract<
+    ByokFeatureKind,
+    'embedding' | 'workspace_indexing' | 'rerank'
+  >;
+};
+
+export type EmbeddingCallOptionsInput = AbortSignal | EmbeddingCallOptions;
+export type EmbeddingRouteContext = Pick<
+  EmbeddingCallOptions,
+  'userId' | 'byokLeaseId'
+>;
+
+export function normalizeEmbeddingCallOptions(
+  options?: EmbeddingCallOptionsInput
+): EmbeddingCallOptions {
+  if (!options) {
+    return {};
+  }
+  if ('aborted' in options && 'addEventListener' in options) {
+    return { signal: options };
+  }
+  return options;
+}
+
 export abstract class EmbeddingClient {
  async configured() {
    return true;
@@ -111,11 +141,14 @@ export abstract class EmbeddingClient {
  async getFileEmbeddings(
    file: File,
    chunkMapper: (chunk: Chunk[]) => Chunk[],
-    signal?: AbortSignal
+    options?: EmbeddingCallOptionsInput
  ): Promise<Embedding[][]> {
-    const chunks = await this.getFileChunks(file, signal);
+    const normalizedOptions = normalizeEmbeddingCallOptions(options);
+    const chunks = await this.getFileChunks(file, normalizedOptions.signal);
    const chunkedEmbeddings = await Promise.all(
-      chunks.map(chunk => this.generateEmbeddings(chunkMapper(chunk)))
+      chunks.map(chunk =>
+        this.generateEmbeddings(chunkMapper(chunk), normalizedOptions)
+      )
    );
    return chunkedEmbeddings;
  }
@@ -154,8 +187,9 @@ export abstract class EmbeddingClient {

  async generateEmbeddings(
    chunks: Chunk[],
-    signal?: AbortSignal
+    options?: EmbeddingCallOptionsInput
  ): Promise<Embedding[]> {
+    const normalizedOptions = normalizeEmbeddingCallOptions(options);
    const retry = 3;

    let embeddings: Embedding[] = [];
@@ -164,7 +198,7 @@ export abstract class EmbeddingClient {
      try {
        embeddings = await this.getEmbeddings(
          chunks.map(c => c.content),
-          signal
+          normalizedOptions
        );
        break;
      } catch (e) {
@@ -181,7 +215,7 @@ export abstract class EmbeddingClient {
    _query: string,
    embeddings: Chunk[],
    topK: number,
-    _signal?: AbortSignal
+    _options?: EmbeddingCallOptionsInput
  ): Promise<Chunk[]> {
    // sort by distance with ascending order
    return embeddings
@@ -189,14 +223,14 @@ export abstract class EmbeddingClient {
      .slice(0, topK);
  }

-  async getEmbedding(query: string, signal?: AbortSignal) {
-    const embedding = await this.getEmbeddings([query], signal);
+  async getEmbedding(query: string, options?: EmbeddingCallOptionsInput) {
+    const embedding = await this.getEmbeddings([query], options);
    return embedding?.[0]?.embedding;
  }

  abstract getEmbeddings(
    input: string[],
-    signal?: AbortSignal
+    options?: EmbeddingCallOptionsInput
  ): Promise<Embedding[]>;
 }

@@ -1,3 +1,9 @@
+import { CopilotAccessPolicy } from './access';
+import {
+  ByokEntitlementPolicy,
+  ByokService,
+  WorkspaceByokResolver,
+} from './byok';
 import { HistoryAttachmentUrlProjector } from './compat/history-attachment-url-projector';
 import { CompatHistoryProjector } from './compat/history-projector';
 import { HistoryPromptPreloadProjector } from './compat/history-prompt-preload-projector';
@@ -64,10 +70,13 @@ export const COPILOT_PROVIDER_PROVIDERS = [
 ];

 export const COPILOT_RUNTIME_PROVIDERS = [
+  ByokEntitlementPolicy,
+  ByokService,
  ChatSessionService,
  ConversationStore,
  ConversationInboxService,
  ConversationPolicy,
+  CopilotAccessPolicy,
  HistoryAttachmentUrlProjector,
  CompatHistoryProjector,
  HistoryPromptPreloadProjector,
@@ -114,6 +123,7 @@ export const COPILOT_RESOLVER_PROVIDERS = [
  CopilotResolver,
  UserCopilotResolver,
  CopilotContextRootResolver,
+  WorkspaceByokResolver,
 ];

 export const COPILOT_JOB_PROVIDERS = [CopilotEmbeddingJob, CopilotCronJobs];
@@ -1,12 +1,11 @@
 import { CopilotProviderSideError, UserFriendlyError } from '../../../base';
 import { type LlmBackendConfig } from '../../../native';
-import type { CopilotTool } from '../tools';
 import { CopilotProvider } from './provider';
 import {
  type CopilotProviderExecution,
  type ProviderDriverSpec,
 } from './provider-runtime-contract';
-import { type CopilotChatTools, CopilotProviderType } from './types';
+import { CopilotProviderType } from './types';

 export type CloudflareWorkersAIConfig = {
  apiToken: string;
@@ -25,16 +24,6 @@ export class CloudflareWorkersAIProvider extends CopilotProvider<CloudflareWorke
    const config = this.getConfig(execution);
    return !!config.apiToken && (!!config.accountId || !!config.baseURL);
  }
-  override getProviderSpecificTools(
-    toolName: CopilotChatTools,
-    _model: string
-  ): [string, CopilotTool?] | undefined {
-    if (toolName === 'docEdit') {
-      return ['doc_edit', undefined];
-    }
-    return;
-  }
-
  private handleError(e: any) {
    if (e instanceof UserFriendlyError) {
      return e;
@@ -1,10 +1,14 @@
 import { Injectable, Logger } from '@nestjs/common';

+import { CopilotQuotaExceeded } from '../../../base';
 import { ServerFeature, ServerService } from '../../../core';
+import { type CopilotAccessContext, CopilotAccessPolicy } from '../access';
 import type { RequiredStructuredOutputContract } from '../runtime/contracts';
 import { getProviderRuntimeHost } from '../runtime/provider-runtime-context';
 import type { CopilotProvider } from './provider';
 import {
+  buildProviderRegistry,
+  type CopilotProviderRegistry,
  type NormalizedCopilotProviderProfile,
  resolveModel,
  stripProviderPrefix,
@@ -57,11 +61,18 @@ type RoutePreparationResult = Partial<
  >
 >;

+type EffectiveProviderRegistry = {
+  byokRegistry: CopilotProviderRegistry;
+  quotaBackedRegistry: CopilotProviderRegistry;
+  quotaBackedRoutesAvailable: boolean;
+};
+
@Injectable()
 export class CopilotProviderFactory {
  constructor(
    private readonly server: ServerService,
-    private readonly registries: CopilotProviderRegistryService
+    private readonly registries: CopilotProviderRegistryService,
+    private readonly access: CopilotAccessPolicy
  ) {}

  private readonly logger = new Logger(CopilotProviderFactory.name);
@@ -73,20 +84,84 @@ export class CopilotProviderFactory {
    return this.registries.getRegistry();
  }

-  private getPreferredProviderIds(type?: CopilotProviderType) {
+  private getProviderByProfile(
+    providerId: string,
+    profile: NormalizedCopilotProviderProfile
+  ) {
+    return (
+      this.#providers.get(providerId) ??
+      Array.from(this.#providerIdsByType.get(profile.type) ?? [])
+        .map(id => this.#providers.get(id))
+        .find((provider): provider is CopilotProvider => !!provider)
+    );
+  }
+
+  private providerAvailable(
+    providerId: string,
+    profile: NormalizedCopilotProviderProfile
+  ) {
+    return !!this.getProviderByProfile(providerId, profile);
+  }
+
+  private getAvailableProviderIds(registry: CopilotProviderRegistry) {
+    return Array.from(registry.profiles.entries())
+      .filter(([providerId, profile]) =>
+        this.providerAvailable(providerId, profile)
+      )
+      .map(([providerId]) => providerId);
+  }
+
+  private getPreferredProviderIds(
+    registry: CopilotProviderRegistry,
+    type?: CopilotProviderType
+  ) {
    if (!type) return undefined;
-    return this.#providerIdsByType.get(type);
+    return registry.byType.get(type)?.filter(providerId => {
+      const profile = registry.profiles.get(providerId);
+      return profile ? this.providerAvailable(providerId, profile) : false;
+    });
  }

  private normalizeCond(
+    registry: CopilotProviderRegistry,
    providerId: string,
    cond: ModelFullConditions
  ): ModelFullConditions {
-    const registry = this.getRegistry();
    const modelId = stripProviderPrefix(registry, providerId, cond.modelId);
    return { ...cond, modelId };
  }

+  private async getEffectiveRegistry(
+    context: CopilotAccessContext = {}
+  ): Promise<EffectiveProviderRegistry> {
+    const quotaBackedRegistry = this.getRegistry();
+    const routeAccess = await this.access.resolveRouteAccess(context);
+
+    return {
+      byokRegistry: buildProviderRegistry({
+        profiles: routeAccess.byokProfiles,
+        defaults: {},
+      }),
+      quotaBackedRegistry,
+      quotaBackedRoutesAvailable: routeAccess.quotaBackedRoutesAvailable,
+    };
+  }
+
+  private getRequestContext(
+    options?:
+      | CopilotChatOptions
+      | CopilotStructuredOptions
+      | CopilotImageOptions
+  ): CopilotAccessContext {
+    return {
+      userId: options?.user,
+      workspaceId: options?.workspace,
+      byokLeaseId: options?.byokLeaseId,
+      featureKind: options?.featureKind,
+      quotaBackedRoutesAllowed: options?.quotaBackedRoutesAllowed,
+    };
+  }
+
  private filterPreparedRoutes(routes: Array<ResolvedCopilotProvider | null>) {
    return routes.filter(
      (route): route is ResolvedCopilotProvider => route !== null
@@ -113,36 +188,89 @@ export class CopilotProviderFactory {
    cond: ModelFullConditions,
    filter: {
      prefer?: CopilotProviderType;
-    } = {}
+    } = {},
+    context: CopilotAccessContext = {}
  ): Promise<ResolvedCopilotProvider | null> {
-    return (await this.resolveRoutes(cond, filter))[0] ?? null;
+    return (await this.resolveRoutes(cond, filter, context))[0] ?? null;
  }

  async resolveRoutes(
    cond: ModelFullConditions,
    filter: {
      prefer?: CopilotProviderType;
-    } = {}
+    } = {},
+    context: CopilotAccessContext = {}
  ): Promise<ResolvedCopilotProvider[]> {
    this.logger.debug(
      `Resolving copilot provider for output type: ${cond.outputType}`
    );
-    const registry = this.getRegistry();
+    const { byokRegistry, quotaBackedRegistry, quotaBackedRoutesAvailable } =
+      await this.getEffectiveRegistry(context);
+    const byokRoutes = await this.resolveRoutesFromRegistry(
+      byokRegistry,
+      cond,
+      filter
+    );
+    const resolved = byokRoutes.length
+      ? byokRoutes
+      : quotaBackedRoutesAvailable
+        ? await this.resolveRoutesFromRegistry(
+            quotaBackedRegistry,
+            cond,
+            filter
+          )
+        : [];
+    for (const route of resolved) {
+      this.logger.debug(
+        `Copilot provider candidate found: ${route.provider.type} (${route.providerId})`
+      );
+    }
+
+    if (
+      !resolved.length &&
+      !quotaBackedRoutesAvailable &&
+      context.quotaBackedRoutesAllowed !== false
+    ) {
+      const quotaBackedRoutes = await this.resolveRoutesFromRegistry(
+        quotaBackedRegistry,
+        cond,
+        filter
+      );
+      if (quotaBackedRoutes.length) {
+        throw new CopilotQuotaExceeded();
+      }
+    }
+
+    return resolved;
+  }
+
+  private async resolveRoutesFromRegistry(
+    registry: CopilotProviderRegistry,
+    cond: ModelFullConditions,
+    filter: {
+      prefer?: CopilotProviderType;
+    } = {}
+  ): Promise<ResolvedCopilotProvider[]> {
    const route = resolveModel({
      registry,
      modelId: cond.modelId,
      outputType: cond.outputType,
-      availableProviderIds: this.#providers.keys(),
-      preferredProviderIds: this.getPreferredProviderIds(filter.prefer),
+      availableProviderIds: this.getAvailableProviderIds(registry),
+      preferredProviderIds: this.getPreferredProviderIds(
+        registry,
+        filter.prefer
+      ),
    });

    const resolved: ResolvedCopilotProvider[] = [];
    for (const providerId of route.candidateProviderIds) {
-      const provider = this.#providers.get(providerId);
      const profile = registry.profiles.get(providerId);
+      const provider = profile
+        ? this.getProviderByProfile(providerId, profile)
+        : undefined;
      if (!provider || !profile) continue;

-      const normalizedCond = this.normalizeCond(providerId, cond);
+      const normalizedCond = this.normalizeCond(registry, providerId, cond);
      if (
        normalizedCond.modelId &&
        profile.models?.length &&
@@ -155,9 +283,6 @@ export class CopilotProviderFactory {
      const matched = await provider.match(normalizedCond, execution);
      if (!matched) continue;

-      this.logger.debug(
-        `Copilot provider candidate found: ${provider.type} (${providerId})`
-      );
      resolved.push({
        providerId,
        provider,
@@ -181,7 +306,11 @@ export class CopilotProviderFactory {
      prefer?: CopilotProviderType;
    } = {}
  ): Promise<ResolvedCopilotProvider[]> {
-    const routes = await this.resolveRoutes(cond, filter);
+    const routes = await this.resolveRoutes(
+      cond,
+      filter,
+      this.getRequestContext(options)
+    );
    return await this.prepareResolvedRoutes(routes, async route => {
      const prepared = await getProviderRuntimeHost(
        route.provider
@@ -213,7 +342,11 @@ export class CopilotProviderFactory {
    } = {},
    responseContract?: RequiredStructuredOutputContract
  ): Promise<ResolvedCopilotProvider[]> {
-    const routes = await this.resolveRoutes(cond, filter);
+    const routes = await this.resolveRoutes(
+      cond,
+      filter,
+      this.getRequestContext(options)
+    );
    return await this.prepareResolvedRoutes(routes, async route => {
      const preparedStructured =
        (await getProviderRuntimeHost(route.provider).prepare.structured(
@@ -239,10 +372,14 @@ export class CopilotProviderFactory {
    input: string | string[],
    options: CopilotEmbeddingOptions = {}
  ): Promise<ResolvedCopilotProvider[]> {
-    const routes = await this.resolveRoutes({
-      modelId,
-      outputType: ModelOutputType.Embedding,
-    });
+    const routes = await this.resolveRoutes(
+      { modelId, outputType: ModelOutputType.Embedding },
+      {},
+      {
+        ...this.getRequestContext(options),
+        featureKind: options?.featureKind ?? 'embedding',
+      }
+    );
    return await this.prepareResolvedRoutes(routes, async route => {
      const preparedEmbedding =
        (await getProviderRuntimeHost(route.provider).prepare.embedding(
@@ -267,10 +404,14 @@ export class CopilotProviderFactory {
    request: CopilotRerankRequest,
    options: CopilotChatOptions = {}
  ): Promise<ResolvedCopilotProvider[]> {
-    const routes = await this.resolveRoutes({
-      modelId,
-      outputType: ModelOutputType.Rerank,
-    });
+    const routes = await this.resolveRoutes(
+      {
+        modelId,
+        outputType: ModelOutputType.Rerank,
+      },
+      {},
+      { ...this.getRequestContext(options), featureKind: 'rerank' }
+    );
    return await this.prepareResolvedRoutes(routes, async route => {
      const preparedRerank =
        (await getProviderRuntimeHost(route.provider).prepare.rerank(
@@ -298,7 +439,10 @@ export class CopilotProviderFactory {
      prefer?: CopilotProviderType;
    } = {}
  ): Promise<ResolvedCopilotProvider[]> {
-    const routes = await this.resolveRoutes(cond, filter);
+    const routes = await this.resolveRoutes(cond, filter, {
+      ...this.getRequestContext(options),
+      featureKind: options?.featureKind ?? 'image',
+    });
    return await this.prepareResolvedRoutes(routes, async route => {
      const preparedImage =
        (await getProviderRuntimeHost(route.provider).prepare.image(
@@ -8,7 +8,6 @@ export { FalProvider } from './fal';
 export { GeminiGenerativeProvider, GeminiVertexProvider } from './gemini';
 export { CopilotProviderLifecycleService } from './lifecycle-service';
 export { OpenAIProvider } from './openai';
-export { PerplexityProvider } from './perplexity';
 export type { CopilotProvider } from './provider';
 export { CopilotProviders } from './provider-tokens';
 export { CopilotProviderRegistryService } from './registry-service';
@@ -1,60 +0,0 @@
-import { CopilotProviderSideError, UserFriendlyError } from '../../../base';
-import { type LlmBackendConfig } from '../../../native';
-import { CopilotProvider } from './provider';
-import {
-  type CopilotProviderExecution,
-  type ProviderDriverSpec,
-} from './provider-runtime-contract';
-import { CopilotProviderType, ModelOutputType } from './types';
-
-export const DEFAULT_DIMENSIONS = 256;
-
-export type MorphConfig = {
-  apiKey?: string;
-};
-
-export class MorphProvider extends CopilotProvider<MorphConfig> {
-  readonly type = CopilotProviderType.Morph;
-
-  protected resolveModelBackendKind() {
-    return 'morph' as const;
-  }
-
-  override configured(execution?: CopilotProviderExecution): boolean {
-    return !!this.getConfig(execution).apiKey;
-  }
-
-  private handleError(e: any) {
-    if (e instanceof UserFriendlyError) {
-      return e;
-    }
-    return new CopilotProviderSideError({
-      provider: this.type,
-      kind: 'unexpected_response',
-      message: e?.message || 'Unexpected morph response',
-    });
-  }
-
-  private createNativeConfig(
-    execution?: CopilotProviderExecution
-  ): LlmBackendConfig {
-    return {
-      base_url: 'https://api.morphllm.com',
-      auth_token: this.getConfig(execution).apiKey ?? '',
-    };
-  }
-
-  override getDriverSpec(): ProviderDriverSpec {
-    return {
-      createBackendConfig: execution => this.createNativeConfig(execution),
-      mapError: error => this.handleError(error),
-      chat: {
-        resolveOutputType: kind =>
-          kind === 'streamObject' ? null : ModelOutputType.Text,
-      },
-      structured: false,
-      embedding: false,
-      rerank: false,
-    };
-  }
-}
@@ -14,7 +14,6 @@ import {
  AttachmentAdmissionHost,
 } from '../runtime/hosts/attachment-admission';
 import { AttachmentMaterializer } from '../runtime/hosts/attachment-materializer';
-import type { CopilotTool } from '../tools';
 import { CopilotProvider } from './provider';
 import { hasProviderModelBehaviorFlag } from './provider-model-runtime';
 import type {
@@ -22,7 +21,6 @@ import type {
  ProviderDriverSpec,
 } from './provider-runtime-contract';
 import {
-  CopilotChatTools,
  CopilotProviderType,
  type PromptAttachment,
  type PromptMessage,
@@ -64,16 +62,6 @@ export class OpenAIProvider extends CopilotProvider<OpenAIConfig> {
    });
  }

-  override getProviderSpecificTools(
-    toolName: CopilotChatTools,
-    _model: string
-  ): [string, CopilotTool?] | undefined {
-    if (toolName === 'docEdit') {
-      return ['doc_edit', undefined];
-    }
-    return;
-  }
-
  protected createNativeConfig(
    execution?: CopilotProviderExecution
  ): LlmBackendConfig {
@@ -1,75 +0,0 @@
-import { CopilotProviderSideError } from '../../../base';
-import { type LlmBackendConfig } from '../../../native';
-import { CopilotProvider } from './provider';
-import { hasProviderModelBehaviorFlag } from './provider-model-runtime';
-import {
-  type CopilotProviderExecution,
-  type ProviderDriverSpec,
-} from './provider-runtime-contract';
-import { CopilotProviderType, ModelOutputType } from './types';
-
-export type PerplexityConfig = {
-  apiKey: string;
-  endpoint?: string;
-};
-
-export class PerplexityProvider extends CopilotProvider<PerplexityConfig> {
-  readonly type = CopilotProviderType.Perplexity;
-
-  protected resolveModelBackendKind() {
-    return 'perplexity' as const;
-  }
-
-  override configured(execution?: CopilotProviderExecution): boolean {
-    return !!this.getConfig(execution).apiKey;
-  }
-
-  override getDriverSpec(): ProviderDriverSpec {
-    return {
-      createBackendConfig: execution => this.createNativeConfig(execution),
-      mapError: error => this.handleError(error),
-      chat: {
-        resolveOutputType: kind =>
-          kind === 'streamObject' ? null : ModelOutputType.Text,
-        withAttachment: false,
-        resolveRequestOptions: async context => ({
-          withAttachment: !hasProviderModelBehaviorFlag(
-            context.model,
-            'no_attachments'
-          ),
-          include: hasProviderModelBehaviorFlag(
-            context.model,
-            'citations_include'
-          )
-            ? ['citations']
-            : undefined,
-        }),
-      },
-      structured: false,
-      embedding: false,
-      rerank: false,
-    };
-  }
-
-  private createNativeConfig(
-    execution?: CopilotProviderExecution
-  ): LlmBackendConfig {
-    const config = this.getConfig(execution);
-    const baseUrl = config.endpoint || 'https://api.perplexity.ai';
-    return {
-      base_url: baseUrl.replace(/\/v1\/?$/, ''),
-      auth_token: config.apiKey,
-    };
-  }
-
-  private handleError(e: any) {
-    if (e instanceof CopilotProviderSideError) {
-      return e;
-    }
-    return new CopilotProviderSideError({
-      provider: this.type,
-      kind: 'unexpected_response',
-      message: e?.message || 'Unexpected perplexity response',
-    });
-  }
-}
@@ -21,18 +21,6 @@ const DEFAULT_MIDDLEWARE_BY_TYPE: Record<
  [CopilotProviderType.AnthropicVertex]: {
    node: { text: DEFAULT_NODE_TEXT_MIDDLEWARE },
  },
-  [CopilotProviderType.Morph]: {
-    rust: {
-      request: ['clamp_max_tokens'],
-    },
-    node: { text: DEFAULT_NODE_TEXT_MIDDLEWARE },
-  },
-  [CopilotProviderType.Perplexity]: {
-    rust: {
-      request: ['clamp_max_tokens'],
-    },
-    node: { text: DEFAULT_NODE_TEXT_MIDDLEWARE },
-  },
  [CopilotProviderType.Gemini]: {
    node: { text: DEFAULT_NODE_TEXT_MIDDLEWARE },
  },
@@ -15,10 +15,8 @@ const LEGACY_PROVIDER_ORDER: CopilotProviderType[] = [
  CopilotProviderType.FAL,
  CopilotProviderType.Gemini,
  CopilotProviderType.GeminiVertex,
-  CopilotProviderType.Perplexity,
  CopilotProviderType.Anthropic,
  CopilotProviderType.AnthropicVertex,
-  CopilotProviderType.Morph,
 ];

 const LEGACY_PROVIDER_PRIORITY = LEGACY_PROVIDER_ORDER.reduce(
@@ -5,9 +5,7 @@ import {
 import { CloudflareWorkersAIProvider } from './cloudflare';
 import { FalProvider } from './fal';
 import { GeminiGenerativeProvider, GeminiVertexProvider } from './gemini';
-import { MorphProvider } from './morph';
 import { OpenAIProvider } from './openai';
-import { PerplexityProvider } from './perplexity';

 export const CopilotProviders = [
  OpenAIProvider,
@@ -15,8 +13,6 @@ export const CopilotProviders = [
  FalProvider,
  GeminiGenerativeProvider,
  GeminiVertexProvider,
-  PerplexityProvider,
  AnthropicOfficialProvider,
  AnthropicVertexProvider,
-  MorphProvider,
 ];
@@ -30,8 +30,6 @@ export enum CopilotProviderType {
  Gemini = 'gemini',
  GeminiVertex = 'geminiVertex',
  OpenAI = 'openai',
-  Perplexity = 'perplexity',
-  Morph = 'morph',
 }

 export const CopilotProviderSchema = z.object({
@@ -80,8 +78,6 @@ export const PromptToolsSchema = z
    'blobRead',
    'codeArtifact',
    'conversationSummary',
-    // work with morph
-    'docEdit',
    // work with indexer
    'docRead',
    'docCreate',
@@ -268,6 +264,22 @@ const CopilotProviderOptionsSchema = z.object({
  user: z.string().optional(),
  session: z.string().optional(),
  workspace: z.string().optional(),
+  byokLeaseId: z.string().optional(),
+  billingUnitId: z.string().optional(),
+  taskId: z.string().optional(),
+  actionId: z.string().optional(),
+  quotaBackedRoutesAllowed: z.boolean().optional(),
+  featureKind: z
+    .enum([
+      'chat',
+      'action',
+      'image',
+      'embedding',
+      'workspace_indexing',
+      'rerank',
+      'transcript',
+    ])
+    .optional(),
 });

 export const CopilotChatOptionsSchema = CopilotProviderOptionsSchema.merge(
@@ -164,11 +164,6 @@ export function toError(error: unknown): Error {
  }
 }

-type DocEditFootnote = {
-  intent: string;
-  result: string;
-};
-
 function asRecord(value: unknown): Record<string, unknown> | null {
  if (value && typeof value === 'object' && !Array.isArray(value)) {
    return value as Record<string, unknown>;
@@ -184,8 +179,6 @@ export class TextStreamParser {

  private prefix: string | null = this.CALLOUT_PREFIX;

-  private readonly docEditFootnotes: DocEditFootnote[] = [];
-
  public parse(chunk: CopilotTextStreamPart) {
    let result = '';
    switch (chunk.type) {
@@ -233,13 +226,6 @@ export class TextStreamParser {
            result += `\nWriting document "${chunk.input.title}"\n`;
            break;
          }
-          case 'doc_edit': {
-            this.docEditFootnotes.push({
-              intent: String(chunk.input.instructions ?? ''),
-              result: '',
-            });
-            break;
-          }
        }
        result = this.markAsCallout(result);
        break;
@@ -250,22 +236,6 @@ export class TextStreamParser {
        );
        result = this.addPrefix(result);
        switch (chunk.toolName) {
-          case 'doc_edit': {
-            const output = asRecord(chunk.output);
-            const array = output?.result;
-            if (Array.isArray(array)) {
-              result += array
-                .map(item => {
-                  return `\n${String(asRecord(item)?.changedContent ?? '')}\n`;
-                })
-                .join('');
-              this.docEditFootnotes[this.docEditFootnotes.length - 1].result =
-                result;
-            } else {
-              this.docEditFootnotes.pop();
-            }
-            break;
-          }
          case 'doc_semantic_search': {
            const output = chunk.output;
            if (Array.isArray(output)) {
@@ -319,10 +289,7 @@ export class TextStreamParser {
  }

  public end() {
-    const footnotes = this.docEditFootnotes.map((footnote, index) => {
-      return `[^edit${index + 1}]: ${JSON.stringify({ type: 'doc-edit', ...footnote })}`;
-    });
-    return footnotes.join('\n');
+    return '';
  }

  private addPrefix(text: string) {
@@ -1,4 +1,4 @@
-import { BadRequestException, NotFoundException } from '@nestjs/common';
+import { NotFoundException } from '@nestjs/common';
 import {
  Args,
  Field,
@@ -7,7 +7,6 @@ import {
  Mutation,
  ObjectType,
  Parent,
-  Query,
  registerEnumType,
  ResolveField,
  Resolver,
@@ -19,7 +18,6 @@ import {
  CallMetric,
  CopilotDocNotFound,
  CopilotFailedToCreateMessage,
-  CopilotProviderSideError,
  CopilotSessionNotFound,
  type FileUpload,
  paginate,
@@ -28,10 +26,8 @@ import {
  RequestMutex,
  Throttle,
  TooManyRequest,
-  UserFriendlyError,
 } from '../../base';
 import { CurrentUser } from '../../core/auth';
-import { DocReader } from '../../core/doc';
 import { AccessController, DocAction } from '../../core/permission';
 import { UserType } from '../../core/user';
 import type { ListSessionOptions, UpdateChatSession } from '../../models';
@@ -40,7 +36,6 @@ import { ConversationInboxService } from './conversation/inbox';
 import { PromptService } from './prompt/service';
 import { CopilotProviderFactory } from './providers/factory';
 import { ModelOutputType, type StreamObject } from './providers/types';
-import { CapabilityRuntime } from './runtime/capability-runtime';
 import { ChatSessionService } from './session';
 import { type ChatHistory, type ChatMessage, SubmittedMessage } from './types';

@@ -376,9 +371,7 @@ export class CopilotResolver {
    private readonly chatSession: ChatSessionService,
    private readonly historyProjector: CompatHistoryProjector,
    private readonly inbox: ConversationInboxService,
-    private readonly docReader: DocReader,
-    private readonly providerFactory: CopilotProviderFactory,
-    private readonly runtime: CapabilityRuntime
+    private readonly providerFactory: CopilotProviderFactory
  ) {}

  @ResolveField(() => CopilotQuotaType, {
@@ -641,8 +634,6 @@ export class CopilotResolver {
      throw new TooManyRequest('Server is busy');
    }

-    await this.chatSession.checkQuota(user.id);
-
    return await this.chatSession.create({
      ...options,
      pinned: options.pinned ?? false,
@@ -724,7 +715,6 @@ export class CopilotResolver {
      throw new TooManyRequest('Server is busy');
    }

-    await this.chatSession.checkQuota(user.id);
    return await this.chatSession.update({
      ...options,
      userId: user.id,
@@ -752,8 +742,6 @@ export class CopilotResolver {
      throw new CopilotDocNotFound({ docId: options.docId });
    }

-    await this.chatSession.checkQuota(user.id);
-
    return await this.chatSession.fork({
      ...options,
      userId: user.id,
@@ -819,96 +807,6 @@ export class CopilotResolver {
    }
  }

-  @Query(() => String, {
-    description:
-      'Apply updates to a doc using LLM and return the merged markdown.',
-    deprecationReason: 'use Mutation.applyDocUpdates',
-  })
-  async applyDocUpdates(
-    @CurrentUser() user: CurrentUser,
-    @Args({ name: 'workspaceId', type: () => String })
-    workspaceId: string,
-    @Args({ name: 'docId', type: () => String })
-    docId: string,
-    @Args({ name: 'op', type: () => String })
-    op: string,
-    @Args({ name: 'updates', type: () => String })
-    updates: string
-  ): Promise<string> {
-    return this.applyDocUpdatesInternal(user, workspaceId, docId, op, updates);
-  }
-
-  @Mutation(() => String, {
-    description:
-      'Apply updates to a doc using LLM and return the merged markdown.',
-    name: 'applyDocUpdates',
-  })
-  async applyDocUpdatesMutation(
-    @CurrentUser() user: CurrentUser,
-    @Args({ name: 'workspaceId', type: () => String })
-    workspaceId: string,
-    @Args({ name: 'docId', type: () => String })
-    docId: string,
-    @Args({ name: 'op', type: () => String })
-    op: string,
-    @Args({ name: 'updates', type: () => String })
-    updates: string
-  ): Promise<string> {
-    return this.applyDocUpdatesInternal(user, workspaceId, docId, op, updates);
-  }
-
-  private async applyDocUpdatesInternal(
-    user: CurrentUser,
-    workspaceId: string,
-    docId: string,
-    op: string,
-    updates: string
-  ): Promise<string> {
-    await this.assertPermission(user, { workspaceId, docId });
-
-    const docContent = await this.docReader.getDocMarkdown(
-      workspaceId,
-      docId,
-      true
-    );
-    if (!docContent || !docContent.markdown) {
-      throw new NotFoundException('Doc not found or empty');
-    }
-
-    const markdown = docContent.markdown.trim();
-
-    const resolved = await this.providerFactory.resolveProvider({
-      modelId: 'morph-v3-large',
-      outputType: ModelOutputType.Text,
-    });
-    if (!resolved) {
-      throw new BadRequestException('No LLM provider available');
-    }
-
-    try {
-      return await this.runtime.text(
-        { modelId: 'morph-v3-large' },
-        [
-          {
-            role: 'user',
-            content: `<instruction>${op}</instruction>\n<code>${markdown}</code>\n<update>${updates}</update>`,
-          },
-        ],
-        { reasoning: false }
-      );
-    } catch (e: any) {
-      if (e instanceof UserFriendlyError) {
-        throw e;
-      } else {
-        throw new CopilotProviderSideError({
-          provider: resolved.provider.type,
-          kind: 'unexpected_response',
-          message: e?.message || 'Unexpected apply response',
-        });
-      }
-    }
-  }
-
  private transformToSessionType(session: Omit<ChatHistory, 'messages'>) {
    return { id: session.sessionId, ...session };
  }
@@ -269,19 +269,20 @@ export class ActionRuntimeBridge {
      attempt,
    });

+    const inputWithBillingUnit = this.withBillingUnit(input, run.id);
    let finalEvent: NativeActionEvent | undefined;
    const attachments: unknown[] = [];
    try {
      const nativeInput = await this.prepareNativeInput({
-        ...input,
+        ...inputWithBillingUnit,
      });
      for await (const event of this.runNativeStream(
        {
          ...nativeInput,
-          recipeId: input.actionId,
-          recipeVersion: input.actionVersion,
+          recipeId: inputWithBillingUnit.actionId,
+          recipeVersion: inputWithBillingUnit.actionVersion,
        },
-        input.signal
+        inputWithBillingUnit.signal
      )) {
        finalEvent = event;
        let projectedEvent = event;
@@ -343,4 +344,40 @@ export class ActionRuntimeBridge {
      });
    }
  }
+
+  private withBillingUnit(
+    input: ActionRuntimeBridgeInput,
+    billingUnitId: string
+  ): ActionRuntimeBridgeInput {
+    return {
+      ...input,
+      prepareStructuredRoutes: input.prepareStructuredRoutes
+        ? {
+            ...input.prepareStructuredRoutes,
+            options: {
+              ...input.prepareStructuredRoutes.options,
+              actionId:
+                input.prepareStructuredRoutes.options?.actionId ??
+                input.actionId,
+              billingUnitId:
+                input.prepareStructuredRoutes.options?.billingUnitId ??
+                billingUnitId,
+            },
+          }
+        : undefined,
+      prepareImageRoutes: input.prepareImageRoutes
+        ? {
+            ...input.prepareImageRoutes,
+            options: {
+              ...input.prepareImageRoutes.options,
+              actionId:
+                input.prepareImageRoutes.options?.actionId ?? input.actionId,
+              billingUnitId:
+                input.prepareImageRoutes.options?.billingUnitId ??
+                billingUnitId,
+            },
+          }
+        : undefined,
+    };
+  }
 }
@@ -411,6 +411,7 @@ function stripHostOnlyOptions<TOptions extends object | undefined>(
    user: _user,
    session: _session,
    workspace: _workspace,
+    quotaBackedRoutesAllowed: _quotaBackedRoutesAllowed,
    ...serializable
  } = options as Record<string, unknown>;

@@ -90,6 +90,8 @@ export class ActionStreamHost {
      prepared.session,
      params,
      userId,
+      parsedQuery.byokLeaseId,
+      prepared.quotaBackedRoutesAllowed,
      signal
    );
    const runStream = this.bridge.runStream({
@@ -130,6 +132,9 @@ export class ActionStreamHost {
              user: userId,
              workspace: prepared.session.config.workspaceId,
              session: sessionId,
+              byokLeaseId: parsedQuery.byokLeaseId,
+              quotaBackedRoutesAllowed: prepared.quotaBackedRoutesAllowed,
+              featureKind: 'action',
            },
          },
      prepareImageRoutes: imageRoutes
@@ -177,6 +182,8 @@ export class ActionStreamHost {
    session: ChatSession,
    params: Record<string, unknown>,
    userId: string,
+    byokLeaseId?: string,
+    quotaBackedRoutesAllowed?: boolean,
    signal?: AbortSignal
  ): Promise<ImageActionRoutePreparation | undefined> {
    if (!isImageAction(actionId)) {
@@ -201,6 +208,9 @@ export class ActionStreamHost {
        user: userId,
        workspace: session.config.workspaceId,
        session: session.config.sessionId,
+        byokLeaseId,
+        quotaBackedRoutesAllowed,
+        featureKind: 'image',
      },
    };
  }
@@ -18,6 +18,10 @@ export type ChatSelectionOptions = {
  reasoning?: boolean;
  webSearch?: boolean;
  toolsConfig?: ToolsConfig;
+  byokLeaseId?: string;
+  billingUnitId?: string;
+  featureKind?: 'chat' | 'action' | 'image';
+  quotaBackedRoutesAllowed?: boolean;
 };

 type ResolvePolicyModelInput = ResolveModelInput & {
@@ -97,6 +101,10 @@ export class CapabilityPolicyHost {
        user: session.config.userId,
        session: session.config.sessionId,
        workspace: session.config.workspaceId,
+        byokLeaseId: options.byokLeaseId,
+        billingUnitId: options.billingUnitId,
+        featureKind: options.featureKind ?? 'chat',
+        quotaBackedRoutesAllowed: options.quotaBackedRoutesAllowed,
        reasoning: options.reasoning,
        webSearch: options.webSearch,
        tools,
@@ -5,6 +5,7 @@ import {
  CopilotSessionNotFound,
  Mutex,
 } from '../../../../base';
+import { CopilotAccessPolicy } from '../../access';
 import { CompatSubmissionStore } from '../../compat/submission-store';
 import {
  canonicalizeTurnTrace,
@@ -20,6 +21,12 @@ export type PreparedConversationTurn = {
  params: Record<string, string>;
  session: ChatSession;
  latestTurn?: Turn;
+  quotaBackedRoutesAllowed?: boolean;
+};
+
+type AppendedSessionMessage = {
+  turn?: Turn;
+  quotaBackedRoutesAllowed?: boolean;
 };

@Injectable()
@@ -27,7 +34,8 @@ export class ConversationHost {
  constructor(
    private readonly sessions: ChatSessionService,
    private readonly submissions: CompatSubmissionStore,
-    private readonly mutex: Mutex
+    private readonly mutex: Mutex,
+    private readonly access: CopilotAccessPolicy
  ) {}

  private async loadAcceptedTurn(
@@ -101,12 +109,32 @@ export class ConversationHost {
    session: ChatSession,
    sessionId: string,
    messageId?: string,
-    retry = false
-  ): Promise<Turn | undefined> {
+    retry = false,
+    byokLeaseId?: string
+  ): Promise<AppendedSessionMessage> {
+    const resolveChatRouteAccess = () =>
+      this.access.resolveTurnRouteAccess({
+        userId,
+        workspaceId: session.config.workspaceId,
+        byokLeaseId,
+        featureKind: 'chat',
+      });
+
    if (!messageId) {
      await this.sessions.revertLatestMessage(sessionId, false);
      session.revertLatestMessage(false);
-      return session.latestUserTurn;
+      if (!session.latestUserTurn) {
+        const routeAccess = await resolveChatRouteAccess();
+        return {
+          turn: session.latestUserTurn,
+          quotaBackedRoutesAllowed: routeAccess.quotaBackedRoutesAllowed,
+        };
+      }
+      const routeAccess = await resolveChatRouteAccess();
+      return {
+        turn: session.latestUserTurn,
+        quotaBackedRoutesAllowed: routeAccess.quotaBackedRoutesAllowed,
+      };
    }

    const acceptedTurn = await this.loadAcceptedTurn(
@@ -116,7 +144,7 @@ export class ConversationHost {
      retry
    );
    if (acceptedTurn) {
-      return acceptedTurn;
+      return { turn: acceptedTurn, quotaBackedRoutesAllowed: true };
    }

    await using lock = await this.mutex.acquire(
@@ -132,7 +160,9 @@ export class ConversationHost {
      messageId,
      retry
    );
-    if (acceptedAfterLock) return acceptedAfterLock;
+    if (acceptedAfterLock) {
+      return { turn: acceptedAfterLock, quotaBackedRoutesAllowed: true };
+    }

    const durableTurn = await this.loadDurableTurn(
      session,
@@ -140,9 +170,14 @@ export class ConversationHost {
      messageId,
      retry
    );
-    if (durableTurn) return durableTurn;
+    if (durableTurn) {
+      return {
+        turn: durableTurn,
+        quotaBackedRoutesAllowed: true,
+      };
+    }

-    await this.sessions.checkQuota(userId);
+    const routeAccess = await resolveChatRouteAccess();

    const submission = await this.submissions.get(messageId);
    if (!submission || submission.sessionId !== sessionId) {
@@ -176,7 +211,10 @@ export class ConversationHost {
      turnId: turn.id ?? '',
    });
    session.pushPersistedTurn(turn);
-    return turn;
+    return {
+      turn,
+      quotaBackedRoutesAllowed: routeAccess.quotaBackedRoutesAllowed,
+    };
  }

  async prepareTurn(
@@ -184,27 +222,30 @@ export class ConversationHost {
    sessionId: string,
    query: Record<string, string | string[]>
  ): Promise<PreparedConversationTurn> {
-    const { messageId, retry, params } = ChatQuerySchema.parse(query);
+    const { messageId, retry, params, byokLeaseId } =
+      ChatQuerySchema.parse(query);
    const session = await this.sessions.get(sessionId);
    if (!session || session.config.userId !== userId) {
      throw new CopilotSessionNotFound();
    }
-    const latestMessage = await this.appendSessionMessage(
+    const appended = await this.appendSessionMessage(
      userId,
      session,
      sessionId,
      messageId,
-      retry
+      retry,
+      byokLeaseId
    );
    const currentUserMessage =
      session.stashTurns.findLast(turn => turn.role === 'user') ??
-      latestMessage;
+      appended.turn;

    return {
      messageId,
      params,
      session,
      latestTurn: currentUserMessage,
+      quotaBackedRoutesAllowed: appended.quotaBackedRoutesAllowed,
    };
  }

@@ -1,4 +1,4 @@
-import { Injectable } from '@nestjs/common';
+import { Injectable, Logger } from '@nestjs/common';

 import { NoCopilotProviderAvailable } from '../../../base';
 import {
@@ -13,6 +13,7 @@ import {
  llmValidateJsonSchema,
  parseNativeStructuredOutput,
 } from '../../../native';
+import { type ByokFeatureKind, ByokService } from '../byok';
 import { type StreamObject } from '../providers/types';
 import { CopilotExecutionMetrics } from './execution-metrics';
 import {
@@ -25,8 +26,11 @@ import { mapNativeSemanticError } from './native-errors';
 import {
  createNativeToolLoopAdapter,
  NativeProviderAdapter,
+  type NativeProviderAdapterOptions,
 } from './tool/native-adapter';

+const logger = new Logger('NativeExecutionEngine');
+
 function modelIdForError(modelId?: string) {
  return modelId ?? 'auto';
 }
@@ -60,6 +64,83 @@ function extractTextResponse(response: LlmDispatchResponse) {
    .trim();
 }

+function getUsageContext(plan: ExecutionPlan) {
+  const options = 'options' in plan.request ? plan.request.options : undefined;
+  const requestFeatureKind =
+    plan.request.kind === 'text' ||
+    plan.request.kind === 'streamText' ||
+    plan.request.kind === 'streamObject'
+      ? 'chat'
+      : plan.request.kind;
+  return {
+    workspaceId: options?.workspace,
+    userId: options?.user,
+    sessionId: options?.session,
+    taskId: options?.taskId,
+    actionId: options?.actionId,
+    billingUnitId: options?.billingUnitId,
+    featureKind: options?.featureKind ?? requestFeatureKind,
+  };
+}
+
+async function recordByokUsage(
+  byok: ByokService,
+  plan: ExecutionPlan,
+  input: {
+    providerId?: string;
+    model?: string | null;
+    usage?: LlmDispatchResponse['usage'];
+  }
+) {
+  const context = getUsageContext(plan);
+  try {
+    await byok.recordUsage({
+      workspaceId: context.workspaceId,
+      userId: context.userId,
+      sessionId: context.sessionId,
+      taskId: context.taskId,
+      actionId: context.actionId,
+      billingUnitId: context.billingUnitId,
+      featureKind: context.featureKind as ByokFeatureKind,
+      providerId: input.providerId,
+      model: input.model,
+      usage: input.usage,
+    });
+  } catch (error) {
+    logger.warn(
+      `Failed to record BYOK usage: ${
+        error instanceof Error ? error.message : String(error)
+      }`
+    );
+  }
+}
+
+async function recordSingleByokRouteFailure(
+  byok: ByokService,
+  plan: ExecutionPlan,
+  error: unknown
+) {
+  const [providerId] = plan.routePolicy.fallbackOrder;
+  if (plan.routePolicy.fallbackOrder.length !== 1 || !providerId) {
+    return;
+  }
+  const context = getUsageContext(plan);
+  try {
+    await byok.recordProviderFailure({
+      workspaceId: context.workspaceId,
+      providerId,
+      featureKind: context.featureKind as ByokFeatureKind,
+      error,
+    });
+  } catch (recordError) {
+    logger.warn(
+      `Failed to record BYOK provider failure: ${
+        recordError instanceof Error ? recordError.message : String(recordError)
+      }`
+    );
+  }
+}
+
 function recordPreparedDispatch(
  executionMetrics: CopilotExecutionMetrics | undefined,
  plan: ExecutionPlan,
@@ -72,7 +153,12 @@ function recordPreparedDispatch(
  );
 }

-function createNativeChatAdapter(dispatch: NativeChatDispatchPlan) {
+function createNativeChatAdapter(
+  dispatch: NativeChatDispatchPlan,
+  options?: {
+    onUsage?: NativeProviderAdapterOptions['onUsage'];
+  }
+) {
  if (dispatch.hasTools) {
    return createNativeToolLoopAdapter(
      { preparedRoutes: dispatch.routes },
@@ -80,6 +166,7 @@ function createNativeChatAdapter(dispatch: NativeChatDispatchPlan) {
      {
        maxSteps: dispatch.prepared.maxSteps,
        nodeTextMiddleware: dispatch.prepared.postprocess?.nodeTextMiddleware,
+        onUsage: options?.onUsage,
      }
    );
  }
@@ -95,6 +182,7 @@ function createNativeChatAdapter(dispatch: NativeChatDispatchPlan) {

  return new NativeProviderAdapter(nativeDispatch, {
    nodeTextMiddleware: dispatch.prepared.postprocess?.nodeTextMiddleware,
+    onUsage: options?.onUsage,
  });
 }

@@ -102,30 +190,38 @@ async function runPreparedValuePlan<TResult>(
  plan: ExecutionPlan,
  routeCount: number,
  executionMetrics: CopilotExecutionMetrics | undefined,
-  run: () => Promise<TResult>
+  run: () => Promise<TResult>,
+  byok: ByokService
 ) {
  recordPreparedDispatch(executionMetrics, plan, routeCount);
  try {
    return await run();
  } catch (error) {
-    throw mapNativeSemanticError(error);
+    const mapped = mapNativeSemanticError(error);
+    await recordSingleByokRouteFailure(byok, plan, mapped);
+    throw mapped;
  }
 }

 async function* mapPreparedStreamErrors<T>(
-  source: AsyncIterable<T>
+  source: AsyncIterable<T>,
+  plan: ExecutionPlan,
+  byok: ByokService
 ): AsyncIterableIterator<T> {
  try {
    yield* source;
  } catch (error) {
-    throw mapNativeSemanticError(error);
+    const mapped = mapNativeSemanticError(error);
+    await recordSingleByokRouteFailure(byok, plan, mapped);
+    throw mapped;
  }
 }

 async function runChatValuePlan(
  plan: ExecutionPlan,
  dispatch: NativeChatDispatchPlan,
-  executionMetrics?: CopilotExecutionMetrics
+  executionMetrics: CopilotExecutionMetrics | undefined,
+  byok: ByokService
 ) {
  const adapter = createNativeChatAdapter(dispatch);
  return await runPreparedValuePlan(
@@ -140,6 +236,11 @@ async function runChatValuePlan(
        const result = await llmDispatchPlan({
          preparedRoutes: dispatch.routes,
        });
+        await recordByokUsage(byok, plan, {
+          providerId: result.provider_id,
+          model: result.response.model,
+          usage: result.response.usage,
+        });
        return extractTextResponse(result.response);
      }

@@ -152,16 +253,26 @@ async function runChatValuePlan(
        plan.hostContext.signal,
        plan.request.messages
      );
-    }
+    },
+    byok
  );
 }

 async function* runChatStreamPlan(
  plan: ExecutionPlan,
  dispatch: NativeChatDispatchPlan,
-  executionMetrics?: CopilotExecutionMetrics
+  executionMetrics: CopilotExecutionMetrics | undefined,
+  byok: ByokService
 ): AsyncIterableIterator<string | StreamObject> {
-  const adapter = createNativeChatAdapter(dispatch);
+  const adapter = createNativeChatAdapter(dispatch, {
+    onUsage: async usage => {
+      await recordByokUsage(byok, plan, {
+        providerId: usage.providerId,
+        model: usage.model,
+        usage: usage.usage,
+      });
+    },
+  });
  recordPreparedDispatch(executionMetrics, plan, dispatch.routes.length);

  if (plan.request.kind === 'streamText') {
@@ -170,7 +281,9 @@ async function* runChatStreamPlan(
        dispatch.prepared.request,
        plan.hostContext.signal,
        plan.request.messages
-      )
+      ),
+      plan,
+      byok
    );
    return;
  }
@@ -181,7 +294,9 @@ async function* runChatStreamPlan(
        dispatch.prepared.request,
        plan.hostContext.signal,
        plan.request.messages
-      )
+      ),
+      plan,
+      byok
    );
    return;
  }
@@ -192,7 +307,8 @@ async function* runChatStreamPlan(
 async function* runPreparedImageArtifactPlan(
  dispatch: NativeImageDispatchPlan,
  plan: ExecutionPlan,
-  executionMetrics?: CopilotExecutionMetrics
+  executionMetrics: CopilotExecutionMetrics | undefined,
+  byok: ByokService
 ): AsyncIterableIterator<NativeImageArtifact> {
  if (plan.request.kind !== 'image') {
    throw new Error('image dispatch requires image plan');
@@ -204,8 +320,21 @@ async function* runPreparedImageArtifactPlan(
    result = await llmImageDispatchPlan({
      preparedRoutes: dispatch.routes,
    });
+    await recordByokUsage(byok, plan, {
+      providerId: result.provider_id,
+      model: dispatch.prepared.route.model,
+      usage: result.response.usage
+        ? {
+            prompt_tokens: result.response.usage.input_tokens ?? 0,
+            completion_tokens: result.response.usage.output_tokens ?? 0,
+            total_tokens: result.response.usage.total_tokens ?? 0,
+          }
+        : undefined,
+    });
  } catch (error) {
-    throw mapNativeSemanticError(error);
+    const mapped = mapNativeSemanticError(error);
+    await recordSingleByokRouteFailure(byok, plan, mapped);
+    throw mapped;
  }
  for (const artifact of result.response.images) {
    yield artifact;
@@ -214,13 +343,14 @@ async function* runPreparedImageArtifactPlan(

 async function executePreparedPlan(
  plan: ExecutionPlan,
-  executionMetrics?: CopilotExecutionMetrics
+  executionMetrics: CopilotExecutionMetrics | undefined,
+  byok: ByokService
 ): Promise<string | number[][] | number[] | null> {
  switch (plan.request.kind) {
    case 'text': {
      const dispatch = plan.nativeDispatch?.chat;
      return dispatch
-        ? await runChatValuePlan(plan, dispatch, executionMetrics)
+        ? await runChatValuePlan(plan, dispatch, executionMetrics, byok)
        : null;
    }
    case 'structured': {
@@ -236,13 +366,19 @@ async function executePreparedPlan(
          const result = await llmStructuredDispatchPlan({
            preparedRoutes: dispatch.routes,
          });
+          await recordByokUsage(byok, plan, {
+            providerId: result.provider_id,
+            model: result.response.model,
+            usage: result.response.usage,
+          });
          const parsed = parseNativeStructuredOutput(result.response);
          const validated = llmValidateJsonSchema(
            dispatch.prepared.request.schema,
            parsed
          );
          return JSON.stringify(validated);
-        }
+        },
+        byok
      );
    }
    case 'embedding': {
@@ -258,8 +394,20 @@ async function executePreparedPlan(
          const result = await llmEmbeddingDispatchPlan({
            preparedRoutes: dispatch.routes,
          });
+          await recordByokUsage(byok, plan, {
+            providerId: result.provider_id,
+            model: result.response.model,
+            usage: result.response.usage
+              ? {
+                  prompt_tokens: result.response.usage.prompt_tokens,
+                  completion_tokens: 0,
+                  total_tokens: result.response.usage.total_tokens,
+                }
+              : undefined,
+          });
          return result.response.embeddings;
-        }
+        },
+        byok
      );
    }
    case 'rerank': {
@@ -275,8 +423,13 @@ async function executePreparedPlan(
          const result = await llmRerankDispatchPlan({
            preparedRoutes: dispatch.routes,
          });
+          await recordByokUsage(byok, plan, {
+            providerId: result.provider_id,
+            model: result.response.model,
+          });
          return result.response.scores;
-        }
+        },
+        byok
      );
    }
    default:
@@ -286,14 +439,15 @@ async function executePreparedPlan(

 function executePreparedStreamPlan(
  plan: ExecutionPlan,
-  executionMetrics?: CopilotExecutionMetrics
+  executionMetrics: CopilotExecutionMetrics | undefined,
+  byok: ByokService
 ): AsyncIterableIterator<string | StreamObject> | null {
  switch (plan.request.kind) {
    case 'streamText':
    case 'streamObject': {
      const dispatch = plan.nativeDispatch?.chat;
      return dispatch
-        ? runChatStreamPlan(plan, dispatch, executionMetrics)
+        ? runChatStreamPlan(plan, dispatch, executionMetrics, byok)
        : null;
    }
    default:
@@ -312,7 +466,10 @@ function noRouteStream<T>(plan: ExecutionPlan) {

@Injectable()
 export class NativeExecutionEngine {
-  constructor(private readonly executionMetrics?: CopilotExecutionMetrics) {}
+  constructor(
+    private readonly byok: ByokService,
+    private readonly executionMetrics?: CopilotExecutionMetrics
+  ) {}

  private noRoute(plan: ExecutionPlan): never {
    throw new NoCopilotProviderAvailable({
@@ -328,7 +485,11 @@ export class NativeExecutionEngine {
  async execute(
    plan: ExecutionPlanForKind<ValueExecutionKind>
  ): Promise<string | number[][] | number[]> {
-    const result = await executePreparedPlan(plan, this.executionMetrics);
+    const result = await executePreparedPlan(
+      plan,
+      this.executionMetrics,
+      this.byok
+    );
    if (result === null) {
      return this.noRoute(plan);
    }
@@ -345,7 +506,11 @@ export class NativeExecutionEngine {
  executeStream(
    plan: ExecutionPlanForKind<StreamExecutionKind>
  ): AsyncIterableIterator<string | StreamObject> {
-    const result = executePreparedStreamPlan(plan, this.executionMetrics);
+    const result = executePreparedStreamPlan(
+      plan,
+      this.executionMetrics,
+      this.byok
+    );
    if (result) {
      return result;
    }
@@ -361,7 +526,8 @@ export class NativeExecutionEngine {
      return runPreparedImageArtifactPlan(
        dispatch,
        plan,
-        this.executionMetrics
+        this.executionMetrics,
+        this.byok
      );
    }

@@ -13,7 +13,6 @@ import {
 } from '../providers/types';
 import {
  buildBlobContentGetter,
-  buildContentGetter,
  buildDocContentGetter,
  buildDocCreateHandler,
  buildDocKeywordSearchGetter,
@@ -27,7 +26,6 @@ import {
  createConversationSummaryTool,
  createDocComposeTool,
  createDocCreateTool,
-  createDocEditTool,
  createDocKeywordSearchTool,
  createDocReadTool,
  createDocSemanticSearchTool,
@@ -68,6 +66,21 @@ export class ToolRuntime {
    if (!options?.tools?.length) {
      return tools;
    }
+    const runPromptText = (
+      promptName: string,
+      params: Record<string, unknown>
+    ) =>
+      this.promptRuntime.runText(promptName, params, {
+        providerOptions: {
+          user: options.user,
+          session: options.session,
+          workspace: options.workspace,
+          byokLeaseId: options.byokLeaseId,
+          billingUnitId: options.billingUnitId,
+          quotaBackedRoutesAllowed: options.quotaBackedRoutesAllowed,
+          featureKind: options.featureKind,
+        },
+      });

    for (const tool of options.tools) {
      const toolDef = resolveProviderSpecificTool?.(tool, model);
@@ -97,23 +110,13 @@ export class ToolRuntime {
          break;
        }
        case 'codeArtifact': {
-          tools.code_artifact = createCodeArtifactTool(
-            this.promptRuntime.runText.bind(this.promptRuntime)
-          );
+          tools.code_artifact = createCodeArtifactTool(runPromptText);
          break;
        }
        case 'conversationSummary': {
          tools.conversation_summary = createConversationSummaryTool(
            options.session,
-            this.promptRuntime.runText.bind(this.promptRuntime)
-          );
-          break;
-        }
-        case 'docEdit': {
-          const getDocContent = buildContentGetter(this.ac, this.docReader);
-          tools.doc_edit = createDocEditTool(
-            this.promptRuntime.runText.bind(this.promptRuntime),
-            getDocContent.bind(null, options)
+            runPromptText
          );
          break;
        }
@@ -177,15 +180,11 @@ export class ToolRuntime {
          break;
        }
        case 'docCompose': {
-          tools.doc_compose = createDocComposeTool(
-            this.promptRuntime.runText.bind(this.promptRuntime)
-          );
+          tools.doc_compose = createDocComposeTool(runPromptText);
          break;
        }
        case 'sectionEdit': {
-          tools.section_edit = createSectionEditTool(
-            this.promptRuntime.runText.bind(this.promptRuntime)
-          );
+          tools.section_edit = createSectionEditTool(runPromptText);
          break;
        }
      }
@@ -1,3 +1,5 @@
+import { Logger } from '@nestjs/common';
+
 import type { LlmRequest, LlmToolLoopStreamEvent } from '../../../../native';
 import type { NodeTextMiddleware } from '../../config';
 import type { PromptMessage, StreamObject } from '../../providers/types';
@@ -20,9 +22,14 @@ type AttachmentFootnote = {
  fileType: string;
 };

-type NativeProviderAdapterOptions = {
+export type NativeProviderAdapterOptions = {
  maxSteps?: number;
  nodeTextMiddleware?: NodeTextMiddleware[];
+  onUsage?: (input: {
+    providerId: string;
+    model?: string;
+    usage?: Extract<LlmToolLoopStreamEvent, { type: 'usage' }>['usage'];
+  }) => void | Promise<void>;
 };

 type NativeStreamDispatch = ConstructorParameters<
@@ -103,9 +110,11 @@ function formatAttachmentFootnotes(
 }

 export class NativeProviderAdapter {
+  readonly logger = new Logger(NativeProviderAdapter.name);
  readonly #runtime: NativeRuntimeAdapter;
  readonly #enableCallout: boolean;
  readonly #enableCitationFootnote: boolean;
+  readonly #onUsage?: NativeProviderAdapterOptions['onUsage'];

  constructor(
    dispatchWithTools: NativeStreamDispatch,
@@ -120,6 +129,36 @@ export class NativeProviderAdapter {
      enabledNodeTextMiddlewares.has('thinking_format');
    this.#enableCitationFootnote =
      enabledNodeTextMiddlewares.has('citation_footnote');
+    this.#onUsage = options.onUsage;
+  }
+
+  async #recordUsageOnProviderSelected(
+    event: { type: string; [key: string]: unknown },
+    state: {
+      model?: string;
+      usage?: Extract<LlmToolLoopStreamEvent, { type: 'usage' }>['usage'];
+    }
+  ) {
+    if (
+      event.type !== 'provider_selected' ||
+      typeof event.provider_id !== 'string'
+    ) {
+      return;
+    }
+    try {
+      await this.#onUsage?.({
+        providerId: event.provider_id,
+        model: state.model,
+        usage: state.usage,
+      });
+    } catch (error) {
+      this.logger.warn(
+        `Provider usage callback failed: ${
+          error instanceof Error ? error.message : String(error)
+        }`
+      );
+    }
+    state.usage = undefined;
  }

  async text(
@@ -144,6 +183,10 @@ export class NativeProviderAdapter {
      ? new CitationFootnoteFormatter()
      : null;
    let streamPartId = 0;
+    const usageState: {
+      model?: string;
+      usage?: Extract<LlmToolLoopStreamEvent, { type: 'usage' }>['usage'];
+    } = {};

    for await (const event of this.#runtime.streamEvents(
      request,
@@ -151,6 +194,22 @@ export class NativeProviderAdapter {
      messages
    )) {
      switch (event.type) {
+        case 'message_start': {
+          const startEvent = event as Extract<
+            LlmToolLoopStreamEvent,
+            { type: 'message_start' }
+          >;
+          usageState.model = startEvent.model;
+          break;
+        }
+        case 'usage': {
+          const usageEvent = event as Extract<
+            LlmToolLoopStreamEvent,
+            { type: 'usage' }
+          >;
+          usageState.usage = usageEvent.usage;
+          break;
+        }
        case 'text_delta': {
          const textEvent = event as unknown as { text: string };
          if (textParser) {
@@ -216,6 +275,11 @@ export class NativeProviderAdapter {
          break;
        }
        case 'done': {
+          const doneEvent = event as Extract<
+            LlmToolLoopStreamEvent,
+            { type: 'done' }
+          >;
+          usageState.usage = doneEvent.usage ?? usageState.usage;
          const footnotes = textParser?.end() ?? '';
          const citations = citationFormatter?.end() ?? '';
          const tails = [citations, footnotes].filter(Boolean).join('\n');
@@ -224,6 +288,9 @@ export class NativeProviderAdapter {
          }
          break;
        }
+        case 'provider_selected':
+          await this.#recordUsageOnProviderSelected(event, usageState);
+          break;
        case 'error':
          throw new Error(
            typeof event.message === 'string'
@@ -246,6 +313,10 @@ export class NativeProviderAdapter {
      : null;
    const fallbackAttachmentFootnotes = new Map<string, AttachmentFootnote>();
    let hasFootnoteReference = false;
+    const usageState: {
+      model?: string;
+      usage?: Extract<LlmToolLoopStreamEvent, { type: 'usage' }>['usage'];
+    } = {};

    for await (const event of this.#runtime.streamEvents(
      request,
@@ -253,6 +324,22 @@ export class NativeProviderAdapter {
      messages
    )) {
      switch (event.type) {
+        case 'message_start': {
+          const startEvent = event as Extract<
+            LlmToolLoopStreamEvent,
+            { type: 'message_start' }
+          >;
+          usageState.model = startEvent.model;
+          break;
+        }
+        case 'usage': {
+          const usageEvent = event as Extract<
+            LlmToolLoopStreamEvent,
+            { type: 'usage' }
+          >;
+          usageState.usage = usageEvent.usage;
+          break;
+        }
        case 'text_delta': {
          const textEvent = event as unknown as { text: string };
          if (textEvent.text.includes('[^')) {
@@ -302,6 +389,11 @@ export class NativeProviderAdapter {
          break;
        }
        case 'done': {
+          const doneEvent = event as Extract<
+            LlmToolLoopStreamEvent,
+            { type: 'done' }
+          >;
+          usageState.usage = doneEvent.usage ?? usageState.usage;
          const citations = citationFormatter?.end() ?? '';
          if (citations) {
            hasFootnoteReference = true;
@@ -318,6 +410,9 @@ export class NativeProviderAdapter {
          }
          break;
        }
+        case 'provider_selected':
+          await this.#recordUsageOnProviderSelected(event, usageState);
+          break;
        case 'error':
          throw new Error(
            typeof event.message === 'string'
@@ -62,7 +62,7 @@ export class TurnOrchestrator {
      sessionId,
      query
    );
-    const { modelId, reasoning, webSearch, toolsConfig } =
+    const { modelId, reasoning, webSearch, toolsConfig, byokLeaseId } =
      ChatQuerySchema.parse(query);
    const promptParams = await this.buildPromptParams(sessionId, {
      latestTurn: prepared.latestTurn,
@@ -82,6 +82,15 @@ export class TurnOrchestrator {
        reasoning,
        webSearch,
        toolsConfig,
+        byokLeaseId,
+        billingUnitId: prepared.latestTurn?.id,
+        quotaBackedRoutesAllowed: prepared.quotaBackedRoutesAllowed,
+        featureKind:
+          selection.responseMode === 'image'
+            ? 'image'
+            : selection.responseMode === 'object'
+              ? 'action'
+              : 'chat',
      }),
    };
  }
@@ -20,6 +20,7 @@ import {
  type UpdateChatSession,
  UpdateChatSessionOptions,
 } from '../../models';
+import { CopilotAccessPolicy } from './access';
 import { ConversationPolicy } from './conversation/policy';
 import { ConversationStore } from './conversation/store';
 import { type Conversation, promptMessageFromTurn, type Turn } from './core';
@@ -186,6 +187,7 @@ export class ChatSessionService {
    private readonly models: Models,
    private readonly jobs: JobQueue,
    private readonly store: ConversationStore,
+    private readonly access: CopilotAccessPolicy,
    private readonly conversationPolicy: ConversationPolicy,
    private readonly prompts: PromptService,
    private readonly promptRuntime: PromptRuntime
@@ -298,11 +300,11 @@ export class ChatSessionService {
  }

  async getQuota(userId: string) {
-    return await this.conversationPolicy.getQuota(userId);
+    return await this.access.getQuota(userId);
  }

  async checkQuota(userId: string) {
-    await this.conversationPolicy.checkQuota(userId);
+    await this.access.checkQuota(userId);
  }

  async create(options: ChatSessionOptions): Promise<string> {
@@ -1,213 +0,0 @@
-import { z } from 'zod';
-
-import { DocReader } from '../../../core/doc';
-import { AccessController } from '../../../core/permission';
-import { defineTool } from './tool';
-import type { CopilotChatOptions } from './types';
-
-type RunPromptText = (
-  promptName: string,
-  params: Record<string, unknown>
-) => Promise<string>;
-
-const CodeEditSchema = z
-  .array(
-    z.object({
-      op: z
-        .string()
-        .describe(
-          'A short description of the change, such as "Bold intro name"'
-        ),
-      updates: z
-        .string()
-        .describe(
-          'Markdown block fragments that represent the change, including the block_id and type'
-        ),
-    })
-  )
-  .describe(
-    'An array of independent semantic changes to apply to the document.'
-  );
-
-export const buildContentGetter = (ac: AccessController, doc: DocReader) => {
-  const getDocContent = async (options: CopilotChatOptions, docId?: string) => {
-    if (!options || !docId || !options.user || !options.workspace) {
-      return undefined;
-    }
-    const canAccess = await ac
-      .user(options.user)
-      .workspace(options.workspace)
-      .doc(docId)
-      .can('Doc.Read');
-    if (!canAccess) return undefined;
-    const content = await doc.getDocMarkdown(options.workspace, docId, true);
-    return content?.markdown.trim() || undefined;
-  };
-  return getDocContent;
-};
-
-export const createDocEditTool = (
-  prompt: RunPromptText,
-  getContent: (targetId?: string) => Promise<string | undefined>
-) => {
-  return defineTool({
-    description: `
-Use this tool to propose an edit to a structured Markdown document with identifiable blocks. 
-Each block begins with a comment like <!-- block_id=... -->, and represents a unit of editable content such as a heading, paragraph, list, or code snippet.
-This will be read by a less intelligent model, which will quickly apply the edit. You should make it clear what the edit is, while also minimizing the unchanged code you write.
-
-If you receive a markdown without block_id comments, you should call \`doc_read\` tool to get the content.
-
-Your task is to return a list of block-level changes needed to fulfill the user's intent. **Each change in code_edit must be completely independent: each code_edit entry should only perform a single, isolated change, and must not include the effects of other changes. For example, the updates for a delete operation should only show the context related to the deletion, and must not include any content modified by other operations (such as bolding or insertion). This ensures that each change can be applied independently and in any order.**
-
-Each change should correspond to a specific user instruction and be represented by one of the following operations:
-
-replace: Replace the content of a block with updated Markdown.
-
-delete: Remove a block entirely.
-
-insert: Add a new block, and specify its block_id and content.
-
-Important Instructions:
- Use the existing block structure as-is. Do not reformat or reorder blocks unless explicitly asked.
- When inserting, follow the same format as a replacement, but ensure the new block_id does not conflict with existing IDs.
- When replacing content, always keep the original block_id unchanged.
- When deleting content, only use the format <!-- delete block_id=xxx -->, and only for valid block_id present in the original <code> content.
- Each top-level list item should be a block. Like this:
-  \`\`\`markdown
-  <!-- block_id=001 flavour=affine:list -->
-  * Item 1
-    * SubItem 1
-  <!-- block_id=002 flavour=affine:list -->
-  1. Item 1
-    1. SubItem 1
-  \`\`\`
- Your task is to return a list of block-level changes needed to fulfill the user's intent. 
- **Each change in code_edit must be completely independent: each code_edit entry should only perform a single, isolated change, and must not include the effects of other changes. For example, the updates for a delete operation should only show the context related to the deletion, and must not include any content modified by other operations (such as bolding or insertion). This ensures that each change can be applied independently and in any order.**
-
-Original Content:
-\`\`\`markdown
-<!-- block_id=001 flavour=paragraph -->
-# Andriy Shevchenko
-
-<!-- block_id=002 flavour=paragraph -->
-## Player Profile
-
-<!-- block_id=003 flavour=paragraph -->
-Andriy Shevchenko is a legendary Ukrainian striker, best known for his time at AC Milan and Dynamo Kyiv. He won the Ballon d'Or in 2004.
-
-<!-- block_id=004 flavour=paragraph -->
-## Career Overview
-
-<!-- block_id=005 flavour=list -->
- Born in 1976 in Ukraine.
-<!-- block_id=006 flavour=list -->
- Rose to fame at Dynamo Kyiv in the 1990s.
-<!-- block_id=007 flavour=list -->
- Starred at AC Milan (1999–2006), scoring over 170 goals.
-<!-- block_id=008 flavour=list -->
- Played for Chelsea (2006–2009) before returning to Kyiv.
-<!-- block_id=009 flavour=list -->
- Coached Ukraine national team, reaching Euro 2020 quarter-finals.
-\`\`\`
-
-User Request：
-\`\`\`
-Bold the player’s name in the intro, add a summary section at the end, and remove the career overview.
-\`\`\`
-
-Example response:
-\`\`\`json
-[
-  {
-    "op": "Bold the player's name in the introduction",
-    "updates": "
-      <!-- block_id=003 flavour=paragraph -->
-      **Andriy Shevchenko** is a legendary Ukrainian striker, best known for his time at AC Milan and Dynamo Kyiv. He won the Ballon d'Or in 2004.
-    "
-  },
-  {
-    "op": "Add a summary section at the end",
-    "updates": "
-      <!-- block_id=new-abc123 flavour=paragraph -->
-      ## Summary
-      <!-- block_id=new-def456 flavour=paragraph -->
-      Shevchenko is celebrated as one of the greatest Ukrainian footballers of all time. Known for his composure, strength, and goal-scoring instinct, he left a lasting legacy both on and off the pitch.
-    "
-  },
-  {
-    "op": "Delete the career overview section",
-    "updates": "
-      <!-- delete block_id=004 -->
-      <!-- delete block_id=005 -->
-      <!-- delete block_id=006 -->
-      <!-- delete block_id=007 -->
-      <!-- delete block_id=008 -->
-      <!-- delete block_id=009 -->
-    "
-  }
-]
-\`\`\`
-You should specify the following arguments before the others: [doc_id], [origin_content]
-
-    `,
-    inputSchema: z.object({
-      doc_id: z
-        .string()
-        .describe(
-          'The unique ID of the document being edited. Required when editing an existing document stored in the system. If you are editing ad-hoc Markdown content instead, leave this empty and use origin_content.'
-        )
-        .optional(),
-
-      origin_content: z
-        .string()
-        .describe(
-          'The full original Markdown content, including all block_id comments (e.g., <!-- block_id=block-001 type=paragraph -->). Required when doc_id is not provided. This content will be parsed into discrete editable blocks.'
-        )
-        .optional(),
-
-      instructions: z
-        .string()
-        .describe(
-          'A short, first-person description of the intended edit, clearly summarizing what I will change. For example: "I will translate the steps into English and delete the paragraph explaining the delay." This helps the downstream system understand the purpose of the changes.'
-        ),
-
-      code_edit: z.preprocess(val => {
-        // BACKGROUND: LLM sometimes returns a JSON string instead of an array.
-        if (typeof val === 'string') {
-          return JSON.parse(val);
-        }
-        return val;
-      }, CodeEditSchema) as unknown as typeof CodeEditSchema,
-    }),
-    execute: async ({ doc_id, origin_content, code_edit }) => {
-      try {
-        const content = origin_content || (await getContent(doc_id));
-        if (!content) {
-          return 'Doc not found or doc is empty';
-        }
-
-        const changedContents = await Promise.all(
-          code_edit.map(async edit => {
-            return await prompt('Apply Updates', {
-              content,
-              op: edit.op,
-              updates: edit.updates,
-            });
-          })
-        );
-
-        return {
-          result: changedContents.map((changedContent, index) => ({
-            op: code_edit[index].op,
-            updates: code_edit[index].updates,
-            originalContent: content,
-            changedContent,
-          })),
-        };
-      } catch {
-        return 'Failed to apply edit to the doc';
-      }
-    },
-  });
-};
@@ -13,6 +13,11 @@ import { toolError } from './error';
 import { defineTool } from './tool';
 import type { CopilotChatOptions } from './types';

+const getEmbeddingRouteContext = (options: CopilotChatOptions) => ({
+  userId: options?.user,
+  byokLeaseId: options?.byokLeaseId,
+});
+
 export const buildDocSearchGetter = (
  ac: AccessController,
  context: CopilotContextService,
@@ -43,12 +48,32 @@ export const buildDocSearchGetter = (
        'Doc Semantic Search Failed',
        'You do not have permission to access this workspace.'
      );
+    const routeContext = getEmbeddingRouteContext(options);
    const [chunks, contextChunks] = await Promise.all([
-      context.matchWorkspaceAll(options.workspace, query, 10, signal),
+      context.matchWorkspaceAll(
+        options.workspace,
+        query,
+        10,
+        signal,
+        0.8,
+        undefined,
+        0.85,
+        routeContext
+      ),
      sessionId
        ? context
            .getBySessionId(sessionId)
-            .then(current => current?.matchFiles(query, 10, signal) ?? [])
+            .then(
+              current =>
+                current?.matchFiles(
+                  query,
+                  10,
+                  signal,
+                  0.85,
+                  0.5,
+                  routeContext
+                ) ?? []
+            )
        : [],
    ]);

@@ -2,7 +2,6 @@ export * from './blob-read';
 export * from './code-artifact';
 export * from './conversation-summary';
 export * from './doc-compose';
-export * from './doc-edit';
 export * from './doc-keyword-search';
 export * from './doc-read';
 export * from './doc-semantic-search';
@@ -10,7 +10,7 @@ import {
  sniffMime,
 } from '../../../base';
 import { Models } from '../../../models';
-import { ConversationPolicy } from '../conversation/policy';
+import { CopilotAccessPolicy } from '../access';
 import { PromptService } from '../prompt';
 import { CopilotProviderType } from '../providers/types';
 import { ActionRuntimeBridge } from '../runtime/action-runtime-bridge';
@@ -62,7 +62,7 @@ export class CopilotTranscriptionService {
    private readonly tasks: TaskPolicy,
    private readonly prompts: PromptService,
    private readonly actionBridge: ActionRuntimeBridge,
-    @Optional() private readonly conversationPolicy?: ConversationPolicy
+    @Optional() private readonly access?: CopilotAccessPolicy
  ) {}

  private parseTaskPayload(payload: unknown): TranscriptionPayloadV2 {
@@ -223,6 +223,12 @@ export class CopilotTranscriptionService {
      throw new CopilotTranscriptionJobExists();
    }

+    await this.access?.assertQuotaOrByok({
+      userId,
+      workspaceId,
+      featureKind: 'transcript',
+    });
+
    const { model, strategy } = await this.resolveTranscriptStrategy(
      userId,
      input?.strategy ?? undefined
@@ -270,6 +276,12 @@ export class CopilotTranscriptionService {
      );
    }

+    await this.access?.assertQuotaOrByok({
+      userId,
+      workspaceId,
+      featureKind: 'transcript',
+    });
+
    const payload = this.parseTaskPayload(task.protectedResult);
    const { model } = await this.resolveTranscriptStrategy(
      userId,
@@ -307,13 +319,17 @@ export class CopilotTranscriptionService {
      return null;
    }

-    const settled =
-      task.status === 'settled'
-        ? task
-        : await (async () => {
-            await this.conversationPolicy?.checkQuota(userId);
-            return await this.models.copilotTranscriptTask.settle(task.id);
-          })();
+    if (task.status === 'settled') {
+      return this.taskToJob(task);
+    }
+
+    await this.access?.assertQuotaOrByok({
+      userId,
+      workspaceId,
+      featureKind: 'transcript',
+    });
+
+    const settled = await this.models.copilotTranscriptTask.settle(task.id);
    return this.taskToJob(settled);
  }

@@ -378,6 +394,13 @@ export class CopilotTranscriptionService {
          stepId: 'transcribe',
          modelId,
          messages,
+          options: {
+            user: task.userId,
+            workspace: task.workspaceId,
+            taskId,
+            billingUnitId: taskId,
+            featureKind: 'transcript',
+          },
          prefer: CopilotProviderType.Gemini,
          responseContract: TranscriptActionResultContract,
        },
@@ -37,6 +37,7 @@ export const ChatQuerySchema = z
  .object({
    messageId: zMaybeString,
    modelId: zMaybeString,
+    byokLeaseId: zMaybeString,
    retry: zBool,
    reasoning: zBool,
    webSearch: zBool,
@@ -47,6 +48,7 @@ export const ChatQuerySchema = z
    ({
      messageId,
      modelId,
+      byokLeaseId,
      retry,
      reasoning,
      webSearch,
@@ -55,6 +57,7 @@ export const ChatQuerySchema = z
    }) => ({
      messageId,
      modelId,
+      byokLeaseId,
      retry,
      reasoning,
      webSearch,