feat(core): improve byok ux (#15303)

fix #15265

#### PR Dependency Tree


* **PR #15303** 👈

This tree was auto-generated by
[Charcoal](https://github.com/danerwilliams/charcoal)
This commit is contained in:
DarkSky
2026-07-21 19:52:34 +08:00
committed by GitHub
parent 7318ef1ed4
commit b6b7f1eeaf
22 changed files with 381 additions and 85 deletions
+5
View File
@@ -1163,6 +1163,11 @@
"description": "Whether workspace BYOK custom endpoints are accepted.\n@default false",
"default": false
},
"byok.allowPrivateEndpoint": {
"type": "boolean",
"description": "Whether workspace BYOK custom endpoints may resolve to private network targets. Enabling this allows workspace owners and admins to send provider probe requests to the private network.\n@default false",
"default": false
},
"providers.profiles": {
"type": "array",
"description": "The profile list for copilot providers.\n@default []",
@@ -0,0 +1,53 @@
import test from 'ava';
import Sinon from 'sinon';
import type { safeFetch } from '../../base';
import {
PROVIDER_PROBE_MAX_BYTES,
runProviderProbe,
} from '../../plugins/copilot/byok/probe';
import { ByokProvider } from '../../plugins/copilot/byok/types';
test('provider probe allows model responses and explicitly configured private targets', async t => {
const fetch = Sinon.stub<
Parameters<typeof safeFetch>,
ReturnType<typeof safeFetch>
>().resolves(new Response('{}', { status: 200 }));
await runProviderProbe(
fetch,
ByokProvider.openai,
'secret',
'http://provider.internal/v1',
true
);
t.is(fetch.firstCall.args[0], 'http://provider.internal/v1/models');
t.deepEqual(fetch.firstCall.args[2], {
timeoutMs: 10_000,
maxRedirects: 3,
maxBytes: PROVIDER_PROBE_MAX_BYTES,
allowedHeaders: ['Authorization'],
allowHttp: true,
allowPrivateTargetOrigin: true,
});
t.true(PROVIDER_PROBE_MAX_BYTES >= 64 * 1024);
});
test('provider probe keeps private targets blocked by default', async t => {
const fetch = Sinon.stub<
Parameters<typeof safeFetch>,
ReturnType<typeof safeFetch>
>().resolves(new Response('{}', { status: 200 }));
await runProviderProbe(
fetch,
ByokProvider.gemini,
'secret',
'https://provider.example/v1beta',
false
);
t.false(fetch.firstCall.args[2]?.allowHttp);
t.false(fetch.firstCall.args[2]?.allowPrivateTargetOrigin);
});
@@ -0,0 +1,94 @@
import { BadRequestException } from '@nestjs/common';
import type { safeFetch } from '../../../base';
import { ByokProvider } from './types';
const TEST_TIMEOUT_MS = 10_000;
export const PROVIDER_PROBE_MAX_BYTES = 1024 * 1024;
type ProbeFetch = typeof safeFetch;
export async function runProviderProbe(
probeFetch: ProbeFetch,
provider: ByokProvider,
apiKey: string,
endpoint: string | null,
allowPrivateEndpoint: boolean
) {
const request = buildProbeRequest(provider, apiKey, endpoint);
const response = await probeFetch(
request.url,
{
method: request.method,
headers: request.headers,
},
{
timeoutMs: TEST_TIMEOUT_MS,
maxRedirects: 3,
maxBytes: PROVIDER_PROBE_MAX_BYTES,
allowedHeaders: Object.keys(request.headers),
allowHttp: endpoint?.startsWith('http:') ?? false,
allowPrivateTargetOrigin: allowPrivateEndpoint,
}
);
if (!response.ok) {
throw new BadRequestException(providerProbeFailureMessage(response.status));
}
}
function buildProbeRequest(
provider: ByokProvider,
apiKey: string,
endpoint: string | null
): {
method: 'GET';
url: string;
headers: Record<string, string>;
} {
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}` },
};
}
}
function 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}.`;
}
}
@@ -117,6 +117,9 @@ class WorkspaceByokSettingsType {
@Field(() => Boolean)
customEndpointSupported!: boolean;
@Field(() => Boolean)
privateEndpointSupported!: boolean;
@Field(() => Boolean)
hasAiPlan!: boolean;
@@ -13,6 +13,7 @@ import {
import { Models } from '../../../models';
import type { CopilotProviderProfile } from '../config';
import { ByokEntitlementPolicy } from './policy';
import { runProviderProbe } from './probe';
import {
BYOK_ALLOWED_PROVIDERS,
type ByokFeatureKind,
@@ -27,7 +28,6 @@ import {
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;
@@ -71,6 +71,7 @@ export type ByokSettings = {
allowedProviders: ByokProvider[];
localStorageSupported: boolean;
customEndpointSupported: boolean;
privateEndpointSupported: boolean;
hasAiPlan: boolean;
warnings: Array<{
featureKind: string;
@@ -124,6 +125,13 @@ export class ByokService {
return env.selfhosted && this.config.copilot.byok.allowCustomEndpoint;
}
get privateEndpointSupported() {
return (
this.customEndpointSupported &&
this.config.copilot.byok.allowPrivateEndpoint
);
}
async getSettings(
workspaceId: string,
userId?: string
@@ -139,6 +147,7 @@ export class ByokService {
allowedProviders: [...BYOK_ALLOWED_PROVIDERS],
localStorageSupported: false,
customEndpointSupported: this.customEndpointSupported,
privateEndpointSupported: this.privateEndpointSupported,
hasAiPlan: await this.entitlement.hasAiPlan(userId),
warnings: [],
};
@@ -158,6 +167,7 @@ export class ByokService {
allowedProviders: [...BYOK_ALLOWED_PROVIDERS],
localStorageSupported: false,
customEndpointSupported: this.customEndpointSupported,
privateEndpointSupported: this.privateEndpointSupported,
hasAiPlan: await this.entitlement.hasAiPlan(userId),
warnings: [],
};
@@ -178,6 +188,7 @@ export class ByokService {
allowedProviders: [...BYOK_ALLOWED_PROVIDERS],
localStorageSupported: false,
customEndpointSupported: this.customEndpointSupported,
privateEndpointSupported: this.privateEndpointSupported,
hasAiPlan: await this.entitlement.hasAiPlan(userId),
warnings: this.buildWarnings(keys),
};
@@ -332,7 +343,13 @@ export class ByokService {
}
try {
await this.runProviderProbe(input.provider, apiKey, endpoint);
await runProviderProbe(
this.probeFetch,
input.provider,
apiKey,
endpoint,
this.privateEndpointSupported
);
if (input.configId && input.storage === ByokKeyStorage.server) {
await this.models.copilotWorkspaceByokConfig.markValidated(
input.workspaceId,
@@ -760,68 +777,6 @@ export class ByokService {
}
}
private async runProviderProbe(
provider: ByokProvider,
apiKey: string,
endpoint: string | null
) {
const request = this.buildProbeRequest(provider, apiKey, endpoint);
const response = await this.probeFetch(
request.url,
{
method: request.method,
headers: request.headers as unknown as Record<string, string>,
},
{
timeoutMs: TEST_TIMEOUT_MS,
maxRedirects: 3,
maxBytes: 1024,
allowedHeaders: Object.keys(request.headers),
}
);
if (!response.ok) {
throw new BadRequestException(
this.providerProbeFailureMessage(response.status)
);
}
}
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.';
@@ -832,23 +787,6 @@ export class ByokService {
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);
}
@@ -191,6 +191,7 @@ declare global {
Array<'openai' | 'anthropic' | 'gemini' | 'fal'>
>;
allowCustomEndpoint: ConfigItem<boolean>;
allowPrivateEndpoint: ConfigItem<boolean>;
};
unsplash: ConfigItem<{
key: string;
@@ -234,6 +235,11 @@ defineModuleConfig('copilot', {
default: false,
shape: z.boolean(),
},
'byok.allowPrivateEndpoint': {
desc: 'Whether workspace BYOK custom endpoints may resolve to private network targets. Enabling this allows workspace owners and admins to send provider probe requests to the private network.',
default: false,
shape: z.boolean(),
},
'providers.profiles': {
desc: 'The profile list for copilot providers.',
default: [],
+1
View File
@@ -2796,6 +2796,7 @@ type WorkspaceByokSettingsType {
keys: [WorkspaceByokKeyConfigType!]!
localEntitled: Boolean!
localStorageSupported: Boolean!
privateEndpointSupported: Boolean!
serverEntitled: Boolean!
warnings: [WorkspaceByokCapabilityWarningType!]!
workspaceId: String!
@@ -3060,6 +3060,7 @@ export const workspaceByokSettingsQuery = {
allowedProviders
localStorageSupported
customEndpointSupported
privateEndpointSupported
hasAiPlan
keys {
id
@@ -10,6 +10,7 @@ query workspaceByokSettings($id: String!, $from: DateTime!, $to: DateTime!) {
allowedProviders
localStorageSupported
customEndpointSupported
privateEndpointSupported
hasAiPlan
keys {
id
+2
View File
@@ -3569,6 +3569,7 @@ export interface WorkspaceByokSettingsType {
keys: Array<WorkspaceByokKeyConfigType>;
localEntitled: Scalars['Boolean']['output'];
localStorageSupported: Scalars['Boolean']['output'];
privateEndpointSupported: Scalars['Boolean']['output'];
serverEntitled: Scalars['Boolean']['output'];
warnings: Array<WorkspaceByokCapabilityWarningType>;
workspaceId: Scalars['String']['output'];
@@ -7784,6 +7785,7 @@ export type WorkspaceByokSettingsQuery = {
allowedProviders: Array<ByokProvider>;
localStorageSupported: boolean;
customEndpointSupported: boolean;
privateEndpointSupported: boolean;
hasAiPlan: boolean;
keys: Array<{
__typename?: 'WorkspaceByokKeyConfigType';
+4
View File
@@ -381,6 +381,10 @@
"type": "Boolean",
"desc": "Whether workspace BYOK custom endpoints are accepted."
},
"byok.allowPrivateEndpoint": {
"type": "Boolean",
"desc": "Whether workspace BYOK custom endpoints may resolve to private network targets. Enabling this allows workspace owners and admins to send provider probe requests to the private network."
},
"providers.profiles": {
"type": "Array",
"desc": "The profile list for copilot providers."
@@ -0,0 +1,86 @@
import { AuthSession } from '@affine/core/modules/cloud/entities/session';
import { AuthService } from '@affine/core/modules/cloud/services/auth';
import { FetchService } from '@affine/core/modules/cloud/services/fetch';
import { ServerService } from '@affine/core/modules/cloud/services/server';
import { AuthStore } from '@affine/core/modules/cloud/stores/auth';
import { GlobalDialogService } from '@affine/core/modules/dialogs/services/dialog';
import { NbstoreService } from '@affine/core/modules/storage';
import { UrlService } from '@affine/core/modules/url/services/url';
import { ServerDeploymentType } from '@affine/graphql';
import { Framework } from '@toeverything/infra';
import { of } from 'rxjs';
import { describe, expect, test, vi } from 'vitest';
describe('AuthService native password sign-in', () => {
test.each(['android', 'ios'])(
'waits for the %s session bootstrap before completing',
async () => {
let resolveSession!: () => void;
const revalidateOnce = vi.fn(
() =>
new Promise<void>(resolve => {
resolveSession = resolve;
})
);
const signInPassword = vi.fn().mockResolvedValue(undefined);
const framework = new Framework();
framework.entity(
AuthSession,
() =>
({
account$: of(null),
revalidate: vi.fn(),
revalidateOnce,
}) as any
);
framework.service(FetchService, { fetch: vi.fn() } as any);
framework.store(AuthStore, {
signInPassword,
setCachedSignInUser: vi.fn(),
} as any);
framework.service(UrlService, { getClientScheme: () => null } as any);
framework.service(GlobalDialogService, { open: vi.fn() } as any);
framework.service(NbstoreService, {
realtime: { subscribe: () => of() },
} as any);
framework.service(ServerService, {
server: {
['config$']: {
value: {
type: ServerDeploymentType.Selfhosted,
version: '0.27.2',
},
},
},
} as any);
framework.service(AuthService, [
FetchService,
AuthStore,
UrlService,
GlobalDialogService,
NbstoreService,
ServerService,
]);
const auth = framework.provider().get(AuthService);
let completed = false;
const signInPromise = auth
.signInPassword({
email: 'user@example.com',
password: 'password',
})
.then(() => {
completed = true;
});
await vi.waitFor(() => expect(revalidateOnce).toHaveBeenCalledOnce());
expect(completed).toBe(false);
resolveSession();
await signInPromise;
expect(signInPassword).toHaveBeenCalledOnce();
expect(completed).toBe(true);
}
);
});
@@ -11,7 +11,13 @@ import { useCallback, useEffect, useState } from 'react';
import { logByokError } from './errors';
import * as styles from './index.css';
import { readLocalKeys, upsertLocalKey } from './local-storage';
import { byokT, providerLabels, storageLabel } from './metadata';
import {
byokT,
endpointHintKey,
providerLabels,
shouldShowEndpoint,
storageLabel,
} from './metadata';
import type {
ByokKey,
ByokSettings,
@@ -32,6 +38,7 @@ export const AddKeyModal = ({
localStorageSupported,
canAddServerKey,
canAddLocalKey,
isSelfHosted,
gql,
}: {
workspaceId: string;
@@ -45,6 +52,7 @@ export const AddKeyModal = ({
localStorageSupported: boolean;
canAddServerKey: boolean;
canAddLocalKey: boolean;
isSelfHosted: boolean;
gql?: GqlFn;
}) => {
const t = useI18n();
@@ -61,6 +69,10 @@ export const AddKeyModal = ({
editingKey?.storage === ByokKeyStorage.server &&
editingKey.provider === provider;
const canTest = !!apiKey || canTestStoredConfig;
const endpointHint = endpointHintKey(
settings.customEndpointSupported,
settings.privateEndpointSupported
);
useEffect(() => {
if (!open) {
@@ -275,19 +287,23 @@ export const AddKeyModal = ({
}}
type="password"
/>
</label>
{settings.customEndpointSupported ? (
<label className={styles.field}>
</label>{' '}
{shouldShowEndpoint(isSelfHosted, settings.customEndpointSupported) ? (
<label className={styles.endpointField}>
<span className={styles.label}>{byokT(t, 'field.endpoint')}</span>
<input
className={styles.input}
value={endpoint}
disabled={!settings.customEndpointSupported}
onChange={event => {
setEndpoint(event.target.value);
setTestResult(null);
}}
placeholder="https://api.example.com/v1"
/>
{endpointHint ? (
<span className={styles.fieldHint}>{byokT(t, endpointHint)}</span>
) : null}
</label>
) : null}
<div className={styles.modalActions}>
@@ -186,11 +186,19 @@ export const field = style({
height: '3em',
});
export const endpointField = style([field, { height: 'auto' }]);
export const label = style({
fontSize: cssVar('fontXs'),
color: cssVarV2('text/secondary'),
});
export const fieldHint = style({
fontSize: cssVar('fontXs'),
lineHeight: '18px',
color: cssVarV2('text/secondary'),
});
export const input = style({
height: 32,
width: '100%',
@@ -47,6 +47,10 @@ const ByokKeyTestStatus = vi.hoisted(() => ({
passed: 'passed',
failed: 'failed',
}));
const ServerDeploymentType = vi.hoisted(() => ({
Affine: 'Affine',
Selfhosted: 'Selfhosted',
}));
const workspaceByokSettingsQuery = vi.hoisted(() =>
Symbol('workspaceByokSettingsQuery')
@@ -130,6 +134,7 @@ vi.mock('@affine/graphql', () => ({
ByokKeyStorage,
ByokKeyTestStatus,
ByokProvider,
ServerDeploymentType,
clearWorkspaceByokConfigsMutation,
deleteWorkspaceByokConfigMutation,
testWorkspaceByokConfigMutation,
@@ -220,6 +225,7 @@ vi.mock('@toeverything/infra', async importOriginal => {
if (token === WorkspaceServerServiceToken) {
return {
server: {
['config$']: { value: { type: ServerDeploymentType.Affine } },
gql: gqlMock,
},
};
@@ -10,6 +10,7 @@ import {
clearWorkspaceByokConfigsMutation as clearByokMutation,
deleteWorkspaceByokConfigMutation as deleteByokMutation,
type GraphQLQuery,
ServerDeploymentType,
workspaceByokSettingsQuery as byokSettingsQuery,
} from '@affine/graphql';
import { useI18n } from '@affine/i18n';
@@ -114,6 +115,9 @@ export const WorkspaceByokSetting = () => {
(settings?.localEntitled ?? false) &&
(settings?.localStorageSupported ?? false);
const canManageKeys = canAddServerKey || canAddLocalKey;
const isSelfHosted =
workspaceServer.server?.config$.value.type ===
ServerDeploymentType.Selfhosted;
const clearAll = useCallback(async () => {
if (!settings) {
@@ -356,6 +360,7 @@ export const WorkspaceByokSetting = () => {
localStorageSupported={settings.localStorageSupported}
canAddServerKey={canAddServerKey}
canAddLocalKey={canAddLocalKey}
isSelfHosted={isSelfHosted}
gql={workspaceServer.server?.gql as GqlFn | undefined}
/>
</>
@@ -0,0 +1,34 @@
import { describe, expect, test } from 'vitest';
import { endpointHintKey, shouldShowEndpoint } from './metadata';
describe('endpointHintKey', () => {
test.each([
[false, false, 'endpoint.custom-disabled'],
[true, false, 'endpoint.private-disabled'],
[true, true, null],
] as const)(
'maps custom=%s private=%s to %s',
(customEndpointSupported, privateEndpointSupported, expected) => {
expect(
endpointHintKey(customEndpointSupported, privateEndpointSupported)
).toBe(expected);
}
);
});
describe('shouldShowEndpoint', () => {
test.each([
[true, false, true],
[true, true, true],
[false, true, true],
[false, false, false],
])(
'self-hosted %s with custom endpoint support %s returns %s',
(isSelfHosted, customEndpointSupported, expected) => {
expect(shouldShowEndpoint(isSelfHosted, customEndpointSupported)).toBe(
expected
);
}
);
});
@@ -24,6 +24,26 @@ export function storageLabel(t: I18nInstance, storage: ByokStorage) {
: byokT(t, 'storage.server');
}
export function endpointHintKey(
customEndpointSupported: boolean,
privateEndpointSupported: boolean
) {
if (!customEndpointSupported) {
return 'endpoint.custom-disabled';
}
if (!privateEndpointSupported) {
return 'endpoint.private-disabled';
}
return null;
}
export function shouldShowEndpoint(
isSelfHosted: boolean,
customEndpointSupported: boolean
) {
return isSelfHosted || customEndpointSupported;
}
export function capabilitiesFor(provider: ByokProvider, storage: ByokStorage) {
switch (provider) {
case ByokProvider.openai:
@@ -53,6 +53,7 @@ export type ByokSettings = {
allowedProviders: ByokProvider[];
localStorageSupported: boolean;
customEndpointSupported: boolean;
privateEndpointSupported: boolean;
hasAiPlan: boolean;
warnings: Array<{
featureKind: string;
+8
View File
@@ -6140,6 +6140,14 @@ export function useAFFiNEI18N(): {
* `Endpoint`
*/
["com.affine.settings.workspace.byok.field.endpoint"](): string;
/**
* `Custom endpoints are disabled by the server administrator. In Self-hosted Admin, enable copilot.byok.allowCustomEndpoint.`
*/
["com.affine.settings.workspace.byok.endpoint.custom-disabled"](): string;
/**
* `Private network endpoints additionally require the server administrator to enable copilot.byok.allowPrivateEndpoint.`
*/
["com.affine.settings.workspace.byok.endpoint.private-disabled"](): string;
/**
* `Primary`
*/
@@ -1532,6 +1532,8 @@
"com.affine.settings.workspace.byok.field.storage": "Key storage",
"com.affine.settings.workspace.byok.field.api-key": "API key",
"com.affine.settings.workspace.byok.field.endpoint": "Endpoint",
"com.affine.settings.workspace.byok.endpoint.custom-disabled": "Custom endpoints are disabled by the server administrator. In Self-hosted Admin, enable copilot.byok.allowCustomEndpoint.",
"com.affine.settings.workspace.byok.endpoint.private-disabled": "Private network endpoints additionally require the server administrator to enable copilot.byok.allowPrivateEndpoint.",
"com.affine.settings.workspace.byok.placeholder.key-name": "Primary",
"com.affine.settings.workspace.byok.placeholder.description": "Workspace fallback key",
"com.affine.settings.workspace.byok.action.add-key": "Add key",
@@ -1532,6 +1532,8 @@
"com.affine.settings.workspace.byok.field.storage": "密钥存储",
"com.affine.settings.workspace.byok.field.api-key": "API key",
"com.affine.settings.workspace.byok.field.endpoint": "Endpoint",
"com.affine.settings.workspace.byok.endpoint.custom-disabled": "服务器管理员已禁用自定义端点。请在自托管管理后台启用 copilot.byok.allowCustomEndpoint。",
"com.affine.settings.workspace.byok.endpoint.private-disabled": "私有网络端点还需要服务器管理员启用 copilot.byok.allowPrivateEndpoint。",
"com.affine.settings.workspace.byok.placeholder.key-name": "Primary",
"com.affine.settings.workspace.byok.placeholder.description": "工作区回退密钥",
"com.affine.settings.workspace.byok.action.add-key": "添加密钥",