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

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

Browse Source
This commit is contained in:
DarkSky
2026-05-07 04:03:14 +08:00
committed by GitHub
parent 4e169ea5c7
commit eb9cc22502
115 changed files with 10369 additions and 1256 deletions
Download Patch File Download Diff File Expand all files Collapse all files
packages/frontend/core/src/blocksuite/ai/actions/types.ts
-6
View File
@@ -354,12 +354,6 @@ declare global {
files?: ContextMatchedFileChunk[];
docs?: ContextMatchedDocChunk[];
}>;
applyDocUpdates: (
workspaceId: string,
docId: string,
op: string,
updates: string
) => Promise<string>;
addContextBlob: (options: {
blobId: string;
contextId: string;
packages/frontend/core/src/blocksuite/ai/components/ai-tools/doc-edit.ts
+20 -105
View File
@@ -2,7 +2,6 @@ import track from '@affine/track';
import { WithDisposable } from '@blocksuite/affine/global/lit';
import { unsafeCSSVar, unsafeCSSVarV2 } from '@blocksuite/affine/shared/theme';
import { type EditorHost, ShadowlessElement } from '@blocksuite/affine/std';
import { LoadingIcon } from '@blocksuite/affine-components/icons';
import type { NotificationService } from '@blocksuite/affine-shared/services';
import {
CloseIcon,
@@ -17,8 +16,6 @@ import { css, html, nothing } from 'lit';
import { property, state } from 'lit/decorators.js';
import { repeat } from 'lit/directives/repeat.js';
import { AIProvider } from '../../provider';
import { BlockDiffProvider } from '../../services/block-diff';
import { diffMarkdown } from '../../utils/apply-model/markdown-diff';
import { copyText } from '../../utils/editor-actions';
import { AI_CHAT_AUTO_SCROLL_PAUSE_EVENT } from '../ai-chat-messages/auto-scroll';
@@ -218,61 +215,21 @@ export class DocEditTool extends WithDisposable(ShadowlessElement) {
@state()
accessor isCollapsed = false;
@state()
accessor applyingMap: Record<string, boolean> = {};
@state()
accessor acceptingMap: Record<string, boolean> = {};
get blockDiffService() {
return this.host?.std.getOptional(BlockDiffProvider);
}
get isBusy() {
return undefined;
}
isBusyForOp(op: string) {
return this.applyingMap[op] || this.acceptingMap[op];
}
private async _handleApply(op: string, updates: string) {
if (
!this.host ||
this.data.type !== 'tool-result' ||
this.isBusyForOp(op)
) {
private _handleApply(op: string) {
if (!this.host || this.data.type !== 'tool-result') {
return;
}
this.applyingMap = { ...this.applyingMap, [op]: true };
try {
const markdown = await AIProvider.context?.applyDocUpdates(
this.host.std.workspace.id,
this.data.args.doc_id,
op,
updates
);
if (!markdown) {
return;
}
track.applyModel.chat.$.apply({
instruction: this.data.args.instructions,
operation: op,
});
await this.blockDiffService?.apply(this.host.store, markdown);
} catch (error) {
this.notificationService.notify({
title: 'Failed to apply updates',
message: error instanceof Error ? error.message : 'Unknown error',
accent: 'error',
onClose: function (): void {},
});
} finally {
this.applyingMap = { ...this.applyingMap, [op]: false };
}
track.applyModel.chat.$.apply({
instruction: this.data.args.instructions,
operation: op,
});
}
private async _handleReject(op: string) {
private _handleReject(op: string) {
if (!this.host || this.data.type !== 'tool-result') {
return;
}
@@ -281,45 +238,16 @@ export class DocEditTool extends WithDisposable(ShadowlessElement) {
instruction: this.data.args.instructions,
operation: op,
});
this.blockDiffService?.setChangedMarkdown(null);
this.blockDiffService?.rejectAll();
}
private async _handleAccept(op: string, updates: string) {
if (
!this.host ||
this.data.type !== 'tool-result' ||
this.isBusyForOp(op)
) {
private _handleAccept(op: string) {
if (!this.host || this.data.type !== 'tool-result') {
return;
}
this.acceptingMap = { ...this.acceptingMap, [op]: true };
try {
const changedMarkdown = await AIProvider.context?.applyDocUpdates(
this.host.std.workspace.id,
this.data.args.doc_id,
op,
updates
);
if (!changedMarkdown) {
return;
}
track.applyModel.chat.$.accept({
instruction: this.data.args.instructions,
operation: op,
});
await this.blockDiffService?.apply(this.host.store, changedMarkdown);
await this.blockDiffService?.acceptAll(this.host.store);
} catch (error) {
this.notificationService.notify({
title: 'Failed to apply updates',
message: error instanceof Error ? error.message : 'Unknown error',
accent: 'error',
onClose: function (): void {},
});
} finally {
this.acceptingMap = { ...this.acceptingMap, [op]: false };
}
track.applyModel.chat.$.accept({
instruction: this.data.args.instructions,
operation: op,
});
}
private async _toggleCollapse() {
@@ -421,7 +349,7 @@ export class DocEditTool extends WithDisposable(ShadowlessElement) {
return repeat(
result.result,
change => change.op,
({ op, updates, originalContent, changedContent }) => {
({ op, originalContent, changedContent }) => {
const diffs = diffMarkdown(originalContent, changedContent);
return html`
<div class="doc-edit-tool-result-wrapper">
@@ -449,14 +377,7 @@ export class DocEditTool extends WithDisposable(ShadowlessElement) {
${CopyIcon()}
<affine-tooltip>Copy</affine-tooltip>
</button>
<button
@click=${() => this._handleApply(op, updates)}
?disabled=${this.isBusyForOp(op)}
>
${this.applyingMap[op]
? html`${LoadingIcon()} Applying`
: 'Apply'}
</button>
<button @click=${() => this._handleApply(op)}>Apply</button>
</div>
</div>
<div class="doc-edit-tool-result-card-content">
@@ -473,18 +394,12 @@ export class DocEditTool extends WithDisposable(ShadowlessElement) {
</button>
<button
class="doc-edit-tool-result-accept"
@click=${() => this._handleAccept(op, updates)}
?disabled=${this.isBusyForOp(op)}
style="${this.isBusyForOp(op)
? 'pointer-events: none; opacity: 0.6;'
: ''}"
@click=${() => this._handleAccept(op)}
>
${this.acceptingMap[op]
? html`${LoadingIcon()}`
: DoneIcon({
style: `color: ${unsafeCSSVarV2('icon/activated')}`,
})}
${this.acceptingMap[op] ? 'Accepting...' : 'Accept'}
${DoneIcon({
style: `color: ${unsafeCSSVarV2('icon/activated')}`,
})}
Accept
</button>
</div>
</div>
packages/frontend/core/src/blocksuite/ai/provider/copilot-client.ts
+6 -19
View File
@@ -6,7 +6,6 @@ import {
addContextCategoryMutation,
addContextDocMutation,
addContextFileMutation,
applyDocUpdatesMutation,
cleanupCopilotSessionMutation,
createCopilotContextMutation,
createCopilotMessageMutation,
@@ -473,6 +472,7 @@ export class CopilotClient {
actionVersion,
runId,
retry,
byokLeaseId,
}: {
sessionId: string;
messageId?: string;
@@ -483,6 +483,7 @@ export class CopilotClient {
actionVersion?: string;
runId?: string;
retry?: boolean;
byokLeaseId?: string;
},
endpoint = Endpoint.StreamObject
) {
@@ -499,6 +500,7 @@ export class CopilotClient {
actionVersion,
runId,
retry,
byokLeaseId,
});
if (queryString) {
url += `?${queryString}`;
@@ -511,12 +513,14 @@ export class CopilotClient {
sessionId: string,
messageId?: string,
seed?: string,
endpoint = Endpoint.Images
endpoint = Endpoint.Images,
byokLeaseId?: string
) {
let url = `/api/copilot/chat/${sessionId}/${endpoint}`;
const queryString = this.paramsToQueryString({
messageId,
seed,
byokLeaseId,
});
if (queryString) {
url += `?${queryString}`;
@@ -549,23 +553,6 @@ export class CopilotClient {
}).then(res => res.queryWorkspaceEmbeddingStatus);
}
applyDocUpdates(
workspaceId: string,
docId: string,
op: string,
updates: string
) {
return this.gql({
query: applyDocUpdatesMutation,
variables: {
workspaceId,
docId,
op,
updates,
},
}).then(res => res.applyDocUpdates);
}
addContextBlob(options: OptionsField<typeof addContextBlobMutation>) {
return this.gql({
query: addContextBlobMutation,
packages/frontend/core/src/blocksuite/ai/provider/request.spec.ts
+284
View File
@@ -0,0 +1,284 @@
/**
* @vitest-environment happy-dom
*/
import { UserFriendlyError } from '@affine/error';
import { beforeEach, describe, expect, test, vi } from 'vitest';
import { type CopilotClient, Endpoint } from './copilot-client';
import { textToText, toImage } from './request';
const electronApis = vi.hoisted(() => ({
byokStorage: undefined as
| {
isSupported: () => Promise<boolean>;
getWorkspaceLeaseProviders: (workspaceId: string) => Promise<
Array<{
provider: string;
name: string;
apiKey: string;
description?: string | null;
endpoint?: string | null;
sortOrder?: number | null;
enabled?: boolean | null;
}>
>;
}
| undefined,
}));
const createWorkspaceByokLocalLeaseMutation = vi.hoisted(() =>
Symbol('createWorkspaceByokLocalLeaseMutation')
);
vi.mock('@affine/electron-api', () => ({
apis: electronApis,
}));
vi.mock('@affine/graphql', () => ({
ByokProvider: {
openai: 'openai',
anthropic: 'anthropic',
gemini: 'gemini',
fal: 'fal',
},
createWorkspaceByokLocalLeaseMutation,
}));
function createClient(
overrides: Partial<
Pick<
CopilotClient,
'gql' | 'createMessage' | 'chatTextStream' | 'imagesStream'
>
> = {}
) {
return {
gql: vi.fn().mockResolvedValue({
createWorkspaceByokLocalLease: { leaseId: 'lease-1' },
}),
createMessage: vi.fn().mockResolvedValue('message-1'),
chatTextStream: vi.fn(),
imagesStream: vi.fn(),
...overrides,
} as unknown as CopilotClient;
}
async function drain(stream: AsyncIterable<unknown>) {
for await (const chunk of stream) {
void chunk;
}
}
describe('AI request BYOK local lease handling', () => {
beforeEach(() => {
vi.stubGlobal('BUILD_CONFIG', { isElectron: true });
electronApis.byokStorage = {
isSupported: vi.fn().mockResolvedValue(true),
getWorkspaceLeaseProviders: vi.fn().mockResolvedValue([
{
provider: 'openai',
name: 'OpenAI',
apiKey: 'sk-local',
},
]),
};
});
test('fails closed when local BYOK providers exist but lease creation fails', async () => {
const client = createClient({
gql: vi.fn().mockRejectedValue(new Error('mutation failed')),
});
const result = textToText({
client,
sessionId: 'session-1',
workspaceId: 'workspace-1',
content: 'hello',
}) as Promise<string>;
await expect(result).rejects.toThrow('mutation failed');
await expect(result).rejects.toBeInstanceOf(UserFriendlyError);
expect(client.chatTextStream).not.toHaveBeenCalled();
});
test('wraps local BYOK storage support failures as user friendly errors', async () => {
electronApis.byokStorage = {
isSupported: vi.fn().mockRejectedValue(new Error('support check failed')),
getWorkspaceLeaseProviders: vi.fn(),
};
const client = createClient();
const result = textToText({
client,
sessionId: 'session-1',
workspaceId: 'workspace-1',
content: 'hello',
}) as Promise<string>;
await expect(result).rejects.toThrow('support check failed');
await expect(result).rejects.toBeInstanceOf(UserFriendlyError);
expect(client.chatTextStream).not.toHaveBeenCalled();
});
test('wraps local BYOK provider loading failures as user friendly errors', async () => {
electronApis.byokStorage = {
isSupported: vi.fn().mockResolvedValue(true),
getWorkspaceLeaseProviders: vi
.fn()
.mockRejectedValue(new Error('provider load failed')),
};
const client = createClient();
const result = textToText({
client,
sessionId: 'session-1',
workspaceId: 'workspace-1',
content: 'hello',
}) as Promise<string>;
await expect(result).rejects.toThrow('provider load failed');
await expect(result).rejects.toBeInstanceOf(UserFriendlyError);
expect(client.chatTextStream).not.toHaveBeenCalled();
});
test('does not create local BYOK lease after cancellation', async () => {
const controller = new AbortController();
const client = createClient({
createMessage: vi.fn().mockImplementation(async () => {
controller.abort();
return 'message-1';
}),
});
await expect(
textToText({
client,
sessionId: 'session-1',
workspaceId: 'workspace-1',
content: 'hello',
signal: controller.signal,
}) as Promise<string>
).resolves.toBe('');
expect(client.gql).not.toHaveBeenCalled();
expect(client.chatTextStream).not.toHaveBeenCalled();
});
test('does not create stream local BYOK lease after cancellation', async () => {
const controller = new AbortController();
const client = createClient({
createMessage: vi.fn().mockImplementation(async () => {
controller.abort();
return 'message-1';
}),
});
await drain(
textToText({
client,
sessionId: 'session-1',
workspaceId: 'workspace-1',
content: 'hello',
stream: true,
signal: controller.signal,
}) as AsyncIterable<string>
);
expect(client.gql).not.toHaveBeenCalled();
expect(client.chatTextStream).not.toHaveBeenCalled();
});
test('does not create text stream when cancelled while creating local BYOK lease', async () => {
const controller = new AbortController();
const client = createClient({
gql: vi.fn().mockImplementation(async () => {
controller.abort();
return { createWorkspaceByokLocalLease: { leaseId: 'lease-1' } };
}),
});
await drain(
textToText({
client,
sessionId: 'session-1',
workspaceId: 'workspace-1',
content: 'hello',
stream: true,
signal: controller.signal,
}) as AsyncIterable<string>
);
expect(client.gql).toHaveBeenCalled();
expect(client.chatTextStream).not.toHaveBeenCalled();
});
test('does not create text request when cancelled while creating local BYOK lease', async () => {
const controller = new AbortController();
const client = createClient({
gql: vi.fn().mockImplementation(async () => {
controller.abort();
return { createWorkspaceByokLocalLease: { leaseId: 'lease-1' } };
}),
});
await expect(
textToText({
client,
sessionId: 'session-1',
workspaceId: 'workspace-1',
content: 'hello',
signal: controller.signal,
}) as Promise<string>
).resolves.toBe('');
expect(client.gql).toHaveBeenCalled();
expect(client.chatTextStream).not.toHaveBeenCalled();
});
test('does not create image local BYOK lease after cancellation', async () => {
const controller = new AbortController();
const client = createClient({
createMessage: vi.fn().mockImplementation(async () => {
controller.abort();
return 'message-1';
}),
});
await drain(
toImage({
client,
sessionId: 'session-1',
workspaceId: 'workspace-1',
content: 'image',
endpoint: Endpoint.Images,
signal: controller.signal,
}) as AsyncIterable<string>
);
expect(client.gql).not.toHaveBeenCalled();
expect(client.imagesStream).not.toHaveBeenCalled();
});
test('does not create image stream when cancelled while creating local BYOK lease', async () => {
const controller = new AbortController();
const client = createClient({
gql: vi.fn().mockImplementation(async () => {
controller.abort();
return { createWorkspaceByokLocalLease: { leaseId: 'lease-1' } };
}),
});
await drain(
toImage({
client,
sessionId: 'session-1',
workspaceId: 'workspace-1',
content: 'image',
endpoint: Endpoint.Images,
signal: controller.signal,
}) as AsyncIterable<string>
);
expect(client.gql).toHaveBeenCalled();
expect(client.imagesStream).not.toHaveBeenCalled();
});
});
packages/frontend/core/src/blocksuite/ai/provider/request.ts
+138 -1
View File
@@ -1,4 +1,10 @@
import type { AIToolsConfig } from '@affine/core/modules/ai-button';
import { apis, type ClientHandler } from '@affine/electron-api';
import { UserFriendlyError } from '@affine/error';
import {
ByokProvider,
createWorkspaceByokLocalLeaseMutation,
} from '@affine/graphql';
import { partition } from 'lodash-es';
import { AIProvider } from './ai-provider';
@@ -7,9 +13,99 @@ import { toTextStream } from './event-source';
const TIMEOUT = 50000;
function isElectronBuild() {
return typeof BUILD_CONFIG !== 'undefined' && BUILD_CONFIG.isElectron;
}
function byokStorageApi(): ClientHandler['byokStorage'] | undefined {
return isElectronBuild() ? apis?.byokStorage : undefined;
}
function toGraphqlByokProvider(provider: string): ByokProvider | null {
switch (provider) {
case ByokProvider.openai:
return ByokProvider.openai;
case ByokProvider.anthropic:
return ByokProvider.anthropic;
case ByokProvider.gemini:
return ByokProvider.gemini;
case ByokProvider.fal:
return ByokProvider.fal;
default:
return null;
}
}
function errorMetadata(error: unknown) {
if (!error || typeof error !== 'object') {
return { kind: typeof error };
}
const record = error as Record<string, unknown>;
return {
name: typeof record.name === 'string' ? record.name : undefined,
code: typeof record.code === 'string' ? record.code : undefined,
status:
typeof record.status === 'number' || typeof record.status === 'string'
? record.status
: undefined,
type: typeof record.type === 'string' ? record.type : undefined,
};
}
async function createWorkspaceByokLocalLease(
client: CopilotClient,
workspaceId?: string
) {
const storage = byokStorageApi();
if (!workspaceId || !storage) {
return undefined;
}
try {
if (!(await storage.isSupported())) return undefined;
const providers = await storage.getWorkspaceLeaseProviders(workspaceId);
if (!providers.length) return undefined;
const leaseProviders = providers.flatMap(provider => {
const gqlProvider = toGraphqlByokProvider(provider.provider);
return gqlProvider
? [
{
provider: gqlProvider,
name: provider.name,
description: provider.description ?? null,
apiKey: provider.apiKey,
endpoint: provider.endpoint ?? null,
sortOrder: provider.sortOrder ?? 0,
enabled: provider.enabled ?? true,
},
]
: [];
});
if (!leaseProviders.length) return undefined;
const result = await client.gql({
query: createWorkspaceByokLocalLeaseMutation,
variables: {
input: {
workspaceId,
providers: leaseProviders,
},
},
});
return result.createWorkspaceByokLocalLease.leaseId;
} catch (error) {
console.warn(
'Failed to create workspace BYOK local lease',
errorMetadata(error)
);
throw UserFriendlyError.fromAny(error);
}
}
export type TextToTextOptions = {
client: CopilotClient;
sessionId: string;
workspaceId?: string;
content?: string;
attachments?: (string | Blob | File)[];
params?: Record<string, any>;
@@ -114,6 +210,7 @@ async function createMessage({
export function textToText({
client,
sessionId,
workspaceId,
content,
attachments,
params,
@@ -145,6 +242,16 @@ export function textToText({
signal,
});
}
if (signal?.aborted) {
return;
}
const byokLeaseId = await createWorkspaceByokLocalLease(
client,
workspaceId
);
if (signal?.aborted) {
return;
}
const eventSource = client.chatTextStream(
{
sessionId,
@@ -156,6 +263,7 @@ export function textToText({
actionVersion,
runId,
retry,
byokLeaseId,
},
endpoint
);
@@ -203,6 +311,16 @@ export function textToText({
signal,
});
}
if (signal?.aborted) {
return '';
}
const byokLeaseId = await createWorkspaceByokLocalLease(
client,
workspaceId
);
if (signal?.aborted) {
return '';
}
const eventSource = client.chatTextStream(
{
sessionId,
@@ -214,6 +332,7 @@ export function textToText({
actionVersion,
runId,
retry,
byokLeaseId,
},
endpoint
);
@@ -258,6 +377,7 @@ export function textToText({
export function toImage({
content,
sessionId,
workspaceId,
attachments,
params,
seed,
@@ -284,6 +404,16 @@ export function toImage({
signal,
});
}
if (signal?.aborted) {
return;
}
const byokLeaseId = await createWorkspaceByokLocalLease(
client,
workspaceId
);
if (signal?.aborted) {
return;
}
const eventSource =
endpoint === Endpoint.Action
? client.chatTextStream(
@@ -294,10 +424,17 @@ export function toImage({
actionVersion,
runId,
retry,
byokLeaseId,
},
Endpoint.Action
)
: client.imagesStream(sessionId, messageId, seed, endpoint);
: client.imagesStream(
sessionId,
messageId,
seed,
endpoint,
byokLeaseId
);
AIProvider.LAST_ACTION_SESSIONID = sessionId;
for await (const event of toTextStream(eventSource, {
packages/frontend/core/src/blocksuite/ai/provider/setup-provider.tsx
-8
View File
@@ -722,14 +722,6 @@ Could you make a new website based on these notes and send back just the html fi
threshold
);
},
applyDocUpdates: async (
workspaceId: string,
docId: string,
op: string,
updates: string
) => {
return client.applyDocUpdates(workspaceId, docId, op, updates);
},
addContextBlob: async (options: { blobId: string; contextId: string }) => {
return client.addContextBlob({
contextId: options.contextId,