Jayden/AFFiNE-Mirror
0
0
Fork 0
mirror of https://github.com/toeverything/AFFiNE.git synced 2026-07-24 22:09:08 +08:00
Code Issues Packages Projects Releases Wiki Activity

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

Browse Source 
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
Download Patch File Download Diff File Expand all files Collapse all files
packages/backend/server/src/__tests__/copilot/byok-probe.spec.ts
+53
View File
@@ -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);
});
packages/backend/server/src/plugins/copilot/byok/probe.ts
+94
View File
@@ -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}.`;
}
}
packages/backend/server/src/plugins/copilot/byok/resolver.ts
+3
View File
@@ -117,6 +117,9 @@ class WorkspaceByokSettingsType {
@Field(() => Boolean)
customEndpointSupported!: boolean;
@Field(() => Boolean)
privateEndpointSupported!: boolean;
@Field(() => Boolean)
hasAiPlan!: boolean;
packages/backend/server/src/plugins/copilot/byok/service.ts
+19 -81
View File
@@ -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);
}
packages/backend/server/src/plugins/copilot/config.ts
+6
View File
@@ -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: [],
packages/backend/server/src/schema.gql
+1
View File
@@ -2796,6 +2796,7 @@ type WorkspaceByokSettingsType {
keys: [WorkspaceByokKeyConfigType!]!
localEntitled: Boolean!
localStorageSupported: Boolean!
privateEndpointSupported: Boolean!
serverEntitled: Boolean!
warnings: [WorkspaceByokCapabilityWarningType!]!
workspaceId: String!