mirror of
https://github.com/toeverything/AFFiNE.git
synced 2026-07-22 12:36:24 +08:00
feat(server): refactor for byok (#14911)
This commit is contained in:
DarkSky
@@ -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' });
|
||||
Reference in New Issue
Block a user