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

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
@@ -526,17 +526,6 @@ Generated by [AVA](https://avajs.dev).
remoteAttachmentRequests: [],
}
## PerplexityProvider should ignore attachments during text model matching
> Snapshot 1
[
{
text: 'summarize this',
type: 'text',
},
]
## GeminiVertexProvider should prefetch bearer token for native config
> Snapshot 1
File diff suppressed because it is too large Load Diff
@@ -771,7 +771,7 @@ function actionRunRecord(
};
}
function installActionSessionMock(
async function installActionSessionMock(
t: ExecutionContext<Tester>,
{
actionId,
@@ -786,8 +786,12 @@ function installActionSessionMock(
const { models, session } = t.context;
const sandbox = Sinon.createSandbox();
const sessionId = `copilot-provider-action-${actionId}-${randomUUID()}`;
const userId = `copilot-provider-user-${randomUUID()}`;
const workspaceId = `copilot-provider-action-${actionId}`;
const user = await models.user.create({
email: `copilot-provider-user-${randomUUID()}@affine.test`,
});
const userId = user.id;
const workspace = await models.workspace.create(userId);
const workspaceId = workspace.id;
const docId = `copilot-provider-action-${actionId}-doc`;
const savedTurns: Array<{ role: string }> = [];
const userTurn = {
@@ -904,7 +908,11 @@ for (const { actionId, content, verifier } of actionRecipeCases) {
}
const { sandbox, sessionId, userId, savedTurns } =
installActionSessionMock(t, { actionId, actionPrompt, content });
await installActionSessionMock(t, {
actionId,
actionPrompt,
content,
});
let result = '';
try {
@@ -976,8 +984,10 @@ for (const testCase of TRANSCRIPT_AUDIO_CASES) {
runIfCopilotConfigured,
async t => {
const { models, transcript } = t.context;
const userId = `copilot-provider-transcript-user-${randomUUID()}`;
const workspaceId = `copilot-provider-transcript-workspace-${randomUUID()}`;
const user = await models.user.create({
email: `copilot-provider-transcript-${randomUUID()}@affine.pro`,
});
const workspace = await models.workspace.create(user.id);
const blobId = `copilot-provider-transcript-blob-${randomUUID()}`;
const payload = TranscriptPayloadSchema.parse({
sourceAudio: { blobId, mimeType: testCase.mimeType },
@@ -990,8 +1000,8 @@ for (const testCase of TRANSCRIPT_AUDIO_CASES) {
],
});
const task = await models.copilotTranscriptTask.create({
userId,
workspaceId,
userId: user.id,
workspaceId: workspace.id,
blobId,
strategy: 'gemini',
recipeId: 'transcript.audio.gemini',
@@ -139,9 +139,6 @@ test.before(async t => {
fal: {
apiKey: process.env.COPILOT_FAL_API_KEY ?? '1',
},
perplexity: {
apiKey: process.env.COPILOT_PERPLEXITY_API_KEY ?? '1',
},
anthropic: {
apiKey: process.env.COPILOT_ANTHROPIC_API_KEY ?? '1',
},
@@ -1,11 +1,15 @@
import test from 'ava';
import Sinon from 'sinon';
import type { Models } from '../../models';
import { type Models } from '../../models';
import { CopilotAccessPolicy } from '../../plugins/copilot/access';
import type { ByokFeatureKind } from '../../plugins/copilot/byok/types';
import { HistoryAttachmentUrlProjector } from '../../plugins/copilot/compat/history-attachment-url-projector';
import { CompatHistoryProjector } from '../../plugins/copilot/compat/history-projector';
import { HistoryPromptPreloadProjector } from '../../plugins/copilot/compat/history-prompt-preload-projector';
import { HistoryVisibilityPolicy } from '../../plugins/copilot/compat/history-visibility-policy';
import { ConversationPolicy } from '../../plugins/copilot/conversation/policy';
import type { Turn } from '../../plugins/copilot/core';
import { CopilotEmbeddingClientService } from '../../plugins/copilot/embedding/client';
import { CopilotProviderType } from '../../plugins/copilot/providers/types';
import {
@@ -29,9 +33,11 @@ import {
AttachmentMaterializer,
resolveAttachmentFetchUrl,
} from '../../plugins/copilot/runtime/hosts/attachment-materializer';
import { ConversationHost } from '../../plugins/copilot/runtime/hosts/conversation-host';
import { ImageResultHost } from '../../plugins/copilot/runtime/hosts/image-result-host';
import { ResponsePostprocessor } from '../../plugins/copilot/runtime/hosts/response-postprocessor';
import { TurnPersistence } from '../../plugins/copilot/runtime/hosts/turn-persistence';
import { ToolRuntime } from '../../plugins/copilot/runtime/tool-runtime';
function stubTurnPersistence(
persistProjectedResult: Sinon.SinonStub = Sinon.stub().resolves(null)
@@ -41,6 +47,367 @@ function stubTurnPersistence(
} as unknown as TurnPersistence;
}
function stubConversationSession(latestUserTurn?: unknown) {
return {
config: {
sessionId: 'session-1',
userId: 'user-1',
workspaceId: 'workspace-1',
},
model: 'gpt-4o-mini',
stashTurns: latestUserTurn ? [latestUserTurn] : [],
latestUserTurn,
revertLatestMessage: Sinon.stub(),
};
}
test('ConversationPolicy should treat zero quota limit as exhausted', async t => {
const policy = new ConversationPolicy(
{
userFeature: { has: Sinon.stub().resolves(false) },
copilotSession: { countUserMessages: Sinon.stub().resolves(0) },
} as any,
{
getUserQuota: Sinon.stub().resolves({ copilotActionLimit: 0 }),
} as any
);
t.false(await policy.hasQuota('user-1'));
await t.throwsAsync(policy.checkQuota('user-1'));
});
type TurnRouteAccessCase = {
name: string;
profiles: Array<{ id: string }>;
featureKind?: 'embedding' | 'rerank' | 'workspace_indexing';
byokLeaseId?: string;
quotaBackedRoutesAllowed?: boolean;
expectedQuotaCalls: number;
expectedError?: string;
expectedQuotaBackedRoutesAllowed?: boolean;
};
const turnRouteAccessCases: TurnRouteAccessCase[] = [
{
name: 'checks quota when BYOK does not cover the route',
profiles: [],
expectedQuotaCalls: 1,
expectedError: 'quota exceeded',
},
{
name: 'skips quota when BYOK covers the route',
profiles: [{ id: 'profile-1' }],
byokLeaseId: 'lease-1',
expectedQuotaCalls: 0,
expectedQuotaBackedRoutesAllowed: undefined,
},
{
name: 'preserves explicit quota-backed route disable override',
profiles: [],
quotaBackedRoutesAllowed: false,
expectedQuotaCalls: 0,
expectedQuotaBackedRoutesAllowed: false,
},
{
name: 'does not check user quota for unmetered service features',
profiles: [],
featureKind: 'rerank',
expectedQuotaCalls: 0,
expectedQuotaBackedRoutesAllowed: true,
},
];
for (const matrixCase of turnRouteAccessCases) {
test(`CopilotAccessPolicy resolve turn route access: ${matrixCase.name}`, async t => {
const checkQuota = Sinon.stub().rejects(new Error('quota exceeded'));
const getProfiles = Sinon.stub().resolves(matrixCase.profiles);
const access = new CopilotAccessPolicy(
{ checkQuota } as any,
{ getProfiles } as any
);
const promise = access.resolveTurnRouteAccess({
userId: 'user-1',
workspaceId: 'workspace-1',
byokLeaseId: matrixCase.byokLeaseId,
featureKind: matrixCase.featureKind,
quotaBackedRoutesAllowed: matrixCase.quotaBackedRoutesAllowed,
});
if (matrixCase.expectedError) {
await t.throwsAsync(promise, { message: matrixCase.expectedError });
} else {
const routeAccess = await promise;
t.is(
routeAccess.quotaBackedRoutesAllowed,
matrixCase.expectedQuotaBackedRoutesAllowed
);
}
t.is(checkQuota.callCount, matrixCase.expectedQuotaCalls);
if (matrixCase.expectedQuotaCalls) {
Sinon.assert.calledWithExactly(checkQuota, 'user-1');
}
if (matrixCase.byokLeaseId) {
Sinon.assert.calledWithMatch(getProfiles, {
byokLeaseId: matrixCase.byokLeaseId,
});
}
});
}
type ByokCoverageCase = {
featureKind?: ByokFeatureKind;
expected: { local: boolean; server: boolean };
};
const byokCoverageCases: ByokCoverageCase[] = [
{ featureKind: 'chat', expected: { local: true, server: true } },
{ featureKind: 'action', expected: { local: true, server: true } },
{ featureKind: 'image', expected: { local: true, server: true } },
{ featureKind: 'transcript', expected: { local: false, server: true } },
{ featureKind: 'embedding', expected: { local: false, server: true } },
{
featureKind: 'workspace_indexing',
expected: { local: false, server: true },
},
{ featureKind: 'rerank', expected: { local: false, server: true } },
{ expected: { local: true, server: true } },
];
for (const matrixCase of byokCoverageCases) {
test(`CopilotAccessPolicy should resolve BYOK coverage for ${matrixCase.featureKind ?? 'default'}`, async t => {
const getProfiles = Sinon.stub().resolves([]);
const access = new CopilotAccessPolicy(
{ hasQuota: Sinon.stub().resolves(true) } as any,
{ getProfiles } as any
);
await access.getByokProfiles({
userId: 'user-1',
workspaceId: 'workspace-1',
featureKind: matrixCase.featureKind,
});
t.like(getProfiles.firstCall.args[0], {
userId: 'user-1',
workspaceId: 'workspace-1',
});
t.is(getProfiles.firstCall.args[0].featureKind, matrixCase.featureKind);
t.deepEqual(getProfiles.firstCall.args[1], matrixCase.expected);
});
}
test('CopilotAccessPolicy assertQuotaOrByok should honor quota-backed route disable', async t => {
const checkQuota = Sinon.stub().resolves(undefined);
const access = new CopilotAccessPolicy(
{ checkQuota } as any,
{ getProfiles: Sinon.stub().resolves([]) } as any
);
await t.throwsAsync(
access.assertQuotaOrByok({
userId: 'user-1',
workspaceId: 'workspace-1',
featureKind: 'transcript',
quotaBackedRoutesAllowed: false,
})
);
Sinon.assert.notCalled(checkQuota);
});
test('ConversationHost should delegate empty no-message stream access', async t => {
const session = stubConversationSession();
const resolveTurnRouteAccess = Sinon.stub().rejects(
new Error('quota exceeded')
);
const host = new ConversationHost(
{
get: Sinon.stub().resolves(session),
revertLatestMessage: Sinon.stub().resolves(undefined),
} as any,
{} as any,
{} as any,
{ resolveTurnRouteAccess } as any
);
await t.throwsAsync(host.prepareTurn('user-1', 'session-1', {}), {
message: 'quota exceeded',
});
Sinon.assert.calledOnceWithMatch(resolveTurnRouteAccess, {
userId: 'user-1',
workspaceId: 'workspace-1',
});
});
test('ConversationHost should return access decision for empty no-message stream', async t => {
const session = stubConversationSession();
const resolveTurnRouteAccess = Sinon.stub().resolves({
byokProfiles: [{ id: 'profile-1' }],
quotaBackedRoutesAllowed: undefined,
});
const host = new ConversationHost(
{
get: Sinon.stub().resolves(session),
revertLatestMessage: Sinon.stub().resolves(undefined),
} as any,
{} as any,
{} as any,
{ resolveTurnRouteAccess } as any
);
const prepared = await host.prepareTurn('user-1', 'session-1', {});
t.is(prepared.latestTurn, undefined);
t.is(prepared.quotaBackedRoutesAllowed, undefined);
Sinon.assert.calledOnce(resolveTurnRouteAccess);
});
test('ConversationHost should replay accepted tokens without rechecking quota', async t => {
const acceptedTurn: Turn = {
id: 'turn-1',
conversationId: 'session-1',
role: 'user',
content: 'hello',
attachments: [],
metadata: {},
renderTrace: [],
toolEvents: [],
createdAt: new Date(),
};
const session = {
...stubConversationSession(acceptedTurn),
findTurn: Sinon.stub().withArgs('turn-1').returns(acceptedTurn),
};
const resolveTurnRouteAccess = Sinon.stub().rejects(
new Error('quota exceeded')
);
const host = new ConversationHost(
{
get: Sinon.stub().resolves(session),
revertLatestMessage: Sinon.stub().resolves(undefined),
} as any,
{
getAccepted: Sinon.stub().resolves({
sessionId: 'session-1',
turnId: 'turn-1',
}),
} as any,
{} as any,
{ resolveTurnRouteAccess } as any
);
const prepared = await host.prepareTurn('user-1', 'session-1', {
messageId: 'message-1',
});
t.is(prepared.latestTurn, acceptedTurn);
t.true(prepared.quotaBackedRoutesAllowed);
Sinon.assert.notCalled(resolveTurnRouteAccess);
});
test('ConversationHost should replay durable tokens without rechecking quota', async t => {
const durableTurn: Turn = {
id: 'turn-1',
conversationId: 'session-1',
role: 'user',
content: 'hello',
attachments: [],
metadata: {},
renderTrace: [],
toolEvents: [],
createdAt: new Date(),
};
const session = {
...stubConversationSession(durableTurn),
findTurn: Sinon.stub().withArgs('turn-1').returns(durableTurn),
pushPersistedTurn: Sinon.stub(),
};
const resolveTurnRouteAccess = Sinon.stub().rejects(
new Error('quota exceeded')
);
const markAccepted = Sinon.stub().resolves(undefined);
const host = new ConversationHost(
{
get: Sinon.stub().resolves(session),
findTurnByCompatSubmissionId: Sinon.stub().resolves(durableTurn),
revertLatestMessage: Sinon.stub().resolves(undefined),
} as any,
{
getAccepted: Sinon.stub().resolves(undefined),
markAccepted,
} as any,
{
acquire: Sinon.stub().resolves({
[Symbol.asyncDispose]: Sinon.stub().resolves(undefined),
}),
} as any,
{ resolveTurnRouteAccess } as any
);
const prepared = await host.prepareTurn('user-1', 'session-1', {
messageId: 'message-1',
});
t.is(prepared.latestTurn, durableTurn);
t.true(prepared.quotaBackedRoutesAllowed);
Sinon.assert.calledOnceWithMatch(markAccepted, 'message-1', {
sessionId: 'session-1',
turnId: 'turn-1',
});
Sinon.assert.notCalled(resolveTurnRouteAccess);
});
test('ToolRuntime should pass route context into prompt-backed tools', async t => {
const promptRuntime = {
runText: Sinon.stub().resolves('<html><body>done</body></html>'),
};
const runtime = new ToolRuntime(
{} as any,
{} as any,
{} as any,
{} as any,
{} as any,
{} as any,
promptRuntime as any,
{} as any
);
const tools = await runtime.getTools(
{
tools: ['codeArtifact'],
user: 'user-1',
session: 'session-1',
workspace: 'workspace-1',
byokLeaseId: 'lease-1',
featureKind: 'chat',
quotaBackedRoutesAllowed: false,
},
'gpt-4o-mini'
);
const result = await tools.code_artifact.execute?.(
{ title: 'Demo', userPrompt: 'build a page' },
{}
);
t.like(result as object, { title: 'Demo' });
Sinon.assert.calledOnceWithMatch(
promptRuntime.runText,
'Code Artifact',
{ content: 'build a page' },
{
providerOptions: {
user: 'user-1',
session: 'session-1',
workspace: 'workspace-1',
byokLeaseId: 'lease-1',
featureKind: 'chat',
quotaBackedRoutesAllowed: false,
},
}
);
});
test('ResponsePostprocessor should build text, object and image assistant turns', t => {
const postprocessor = new ResponsePostprocessor();
@@ -267,7 +634,7 @@ test('action result projection should map image result url to assistant attachme
t.deepEqual(turn?.attachments, ['https://example.com/final.png']);
});
test('CopilotEmbeddingClientService should refresh configured client and clear unavailable client', async t => {
test('CopilotEmbeddingClientService should keep dispatch client across global config refreshes', async t => {
const taskPolicy = {
resolveEmbeddingModelId: () => 'text-embedding-3-large',
};
@@ -288,8 +655,8 @@ test('CopilotEmbeddingClientService should refresh configured client and clear u
t.truthy(service.getClient());
const second = await service.refresh();
t.is(second, undefined);
t.is(service.getClient(), undefined);
t.truthy(second);
t.is(service.getClient(), second);
Sinon.assert.calledTwice(runtime.embeddingConfigured);
Sinon.assert.alwaysCalledWithExactly(
runtime.embeddingConfigured,
@@ -297,6 +664,95 @@ test('CopilotEmbeddingClientService should refresh configured client and clear u
);
});
test('CopilotEmbeddingClientService should keep workspace-routed embedding client without global provider', async t => {
const taskPolicy = {
resolveEmbeddingModelId: () => 'gemini-embedding-001',
resolveRerankModelId: () => 'gpt-4o-mini',
};
const runtime = {
embeddingConfigured: Sinon.stub().resolves(false),
};
const service = new CopilotEmbeddingClientService(
taskPolicy as any,
runtime as any
);
const client = await service.refresh();
t.truthy(client);
t.is(service.getClient(), client);
Sinon.assert.calledOnceWithExactly(
runtime.embeddingConfigured,
'gemini-embedding-001'
);
});
test('CopilotEmbeddingClientService should pass workspace context into embedding routes', async t => {
const signal = new AbortController().signal;
const taskPolicy = {
resolveEmbeddingModelId: () => 'gemini-embedding-001',
resolveRerankModelId: () => 'gpt-4o-mini',
};
const runtime = {
embeddingConfigured: Sinon.stub().resolves(true),
embed: Sinon.stub().resolves([[0.1]]),
rerank: Sinon.stub().resolves([0.8]),
};
const service = new CopilotEmbeddingClientService(
taskPolicy as any,
runtime as any
);
const client = await service.refresh();
t.truthy(client);
await client?.getEmbeddings(['hello'], {
workspaceId: 'workspace-1',
userId: 'user-1',
featureKind: 'workspace_indexing',
signal,
});
Sinon.assert.calledOnceWithMatch(
runtime.embed,
'gemini-embedding-001',
['hello'],
{
dimensions: Sinon.match.number,
workspace: 'workspace-1',
user: 'user-1',
featureKind: 'workspace_indexing',
signal,
}
);
await client?.reRank(
'hello',
[{ chunk: 0, content: 'hello', distance: 0.2 }],
1,
{
workspaceId: 'workspace-1',
userId: 'user-1',
featureKind: 'workspace_indexing',
signal,
}
);
Sinon.assert.calledOnceWithMatch(
runtime.rerank,
'gpt-4o-mini',
{
query: 'hello',
candidates: [{ id: '0', text: 'hello' }],
},
{
workspace: 'workspace-1',
user: 'user-1',
featureKind: 'rerank',
signal,
}
);
});
test('CompatHistoryProjector should compose visibility, prompt preload and attachment url projection', t => {
const projector = new CompatHistoryProjector(
new HistoryVisibilityPolicy(),
@@ -20,7 +20,6 @@ import {
import { GeminiProvider } from '../../plugins/copilot/providers/gemini/gemini';
import { GeminiVertexProvider } from '../../plugins/copilot/providers/gemini/vertex';
import { OpenAIProvider } from '../../plugins/copilot/providers/openai';
import { PerplexityProvider } from '../../plugins/copilot/providers/perplexity';
import {
CopilotProviderType,
type PromptMessage,
@@ -589,16 +588,6 @@ class TestOpenAIProvider extends OpenAIProvider {
}
}
class TestPerplexityProvider extends PerplexityProvider {
override get config() {
return { apiKey: 'perplexity-key' };
}
override configured() {
return true;
}
}
test('NativeProviderAdapter should append citation and attachment footnotes', async t => {
const dispatch = () =>
(async function* (): AsyncIterableIterator<LlmToolLoopStreamEvent> {
@@ -818,6 +807,91 @@ test('NativeProviderAdapter streamObject should map tool and text events', async
t.snapshot(events);
});
test('NativeProviderAdapter streamObject should finalize usage with selected provider', async t => {
const usageEvents: Array<{
providerId: string;
model?: string;
usage?: {
prompt_tokens: number;
completion_tokens: number;
total_tokens: number;
cached_tokens?: number;
};
}> = [];
const adapter = new NativeProviderAdapter(
() =>
stream(() => [
{ type: 'message_start', model: 'gpt-5-mini' },
{ type: 'text_delta', text: 'ok' },
{
type: 'done',
finish_reason: 'stop',
usage: { prompt_tokens: 2, completion_tokens: 3, total_tokens: 5 },
},
{
type: 'provider_selected',
provider_id: 'byok-aaaaaaaaaaaa-openai-server-key1',
},
]),
{
onUsage: input => {
usageEvents.push(input);
},
}
);
const events = await collectChunks(
adapter.streamObject({
model: 'gpt-5-mini',
stream: true,
messages: nativeMessages(nativeUserText('hi')),
})
);
t.deepEqual(events, [{ type: 'text-delta', textDelta: 'ok' }]);
t.deepEqual(usageEvents, [
{
providerId: 'byok-aaaaaaaaaaaa-openai-server-key1',
model: 'gpt-5-mini',
usage: { prompt_tokens: 2, completion_tokens: 3, total_tokens: 5 },
},
]);
});
test('NativeProviderAdapter streamObject should keep streaming when usage callback fails', async t => {
const adapter = new NativeProviderAdapter(
() =>
stream(() => [
{ type: 'message_start', model: 'gpt-5-mini' },
{ type: 'text_delta', text: 'ok' },
{
type: 'done',
finish_reason: 'stop',
usage: { prompt_tokens: 2, completion_tokens: 3, total_tokens: 5 },
},
{
type: 'provider_selected',
provider_id: 'byok-aaaaaaaaaaaa-openai-server-key1',
},
]),
{
onUsage: () => {
throw new Error('usage callback failed');
},
}
);
const events = await collectChunks(
adapter.streamObject({
model: 'gpt-5-mini',
stream: true,
messages: nativeMessages(nativeUserText('hi')),
})
);
t.deepEqual(events, [{ type: 'text-delta', textDelta: 'ok' }]);
});
test('NativeRuntimeAdapter streamObject should keep raw runtime stream objects only', async t => {
const adapter = new NativeRuntimeAdapter(
createTestToolLoopBridge(mockDispatch, {}, 3)
@@ -1653,36 +1727,6 @@ test('GeminiProvider should not pass materialized inline attachment URL to nativ
t.false('url' in (attachmentPart?.source ?? {}));
});
test('PerplexityProvider should ignore attachments during text model matching', async t => {
const provider = new TestPerplexityProvider();
let capturedRequest: LlmRequest | undefined;
(provider as any).getActiveProviderMiddleware = () => ({});
(provider as any).getTools = async () => ({});
(provider as any).createNativeAdapter = () => ({
text: async (request: LlmRequest) => {
capturedRequest = request;
return 'ok';
},
});
const result = await getProviderRuntimeHost(provider).run.text(
{ modelId: 'sonar' },
[
{
role: 'user',
content: 'summarize this',
attachments: ['https://example.com/a.pdf'],
params: { mimetype: 'application/pdf' },
},
],
{}
);
t.is(result, 'ok');
t.snapshot(capturedRequest?.messages[0]?.content);
});
test('GeminiProvider should reject unsupported attachment schemes at input validation', async t => {
const provider = new TestGeminiProvider();
@@ -3,7 +3,11 @@ import test from 'ava';
import Sinon from 'sinon';
import { z } from 'zod';
import { CopilotPromptInvalid, NoCopilotProviderAvailable } from '../../base';
import {
CopilotPromptInvalid,
CopilotQuotaExceeded,
NoCopilotProviderAvailable,
} from '../../base';
import {
type LlmBackendConfig,
type LlmEmbeddingRequest,
@@ -20,9 +24,11 @@ import {
llmResolveRequestedModelMatch,
type LlmStructuredRequest,
} from '../../native';
import type { ProviderMiddlewareConfig } from '../../plugins/copilot/config';
import type {
CopilotProviderProfile,
ProviderMiddlewareConfig,
} from '../../plugins/copilot/config';
import { CopilotProviderFactory } from '../../plugins/copilot/providers/factory';
import { MorphProvider } from '../../plugins/copilot/providers/morph';
import { OpenAIProvider } from '../../plugins/copilot/providers/openai';
import { CopilotProvider } from '../../plugins/copilot/providers/provider';
import { buildProviderRegistry } from '../../plugins/copilot/providers/provider-registry';
@@ -62,6 +68,13 @@ import {
userPrompt,
} from './prompt-test-helper';
function createNativeExecutionEngine() {
return new NativeExecutionEngine({
recordUsage: Sinon.stub().resolves(),
recordProviderFailure: Sinon.stub().resolves(),
} as never);
}
function structuredOptions(
schema: z.ZodTypeAny,
extra?: Record<string, unknown>
@@ -541,7 +554,7 @@ test('CapabilityRuntime should defer no-route embedding plans to native engine',
});
test('NativeExecutionEngine should expose execute/executeStream as the single plan entrypoints', async t => {
const engine = new NativeExecutionEngine();
const engine = createNativeExecutionEngine();
let dispatchCalls = 0;
let streamCalls = 0;
@@ -660,6 +673,279 @@ test('NativeExecutionEngine should expose execute/executeStream as the single pl
t.is(streamCalls, 1);
});
test('NativeExecutionEngine should record BYOK usage when stream finalizes with selected provider', async t => {
const byok = {
recordUsage: Sinon.stub().resolves(),
};
const engine = new NativeExecutionEngine(byok as never);
const providerId = 'byok-aaaaaaaaaaaa-openai-server-key1';
const originalStream = (serverNativeModule as any).llmDispatchPreparedStream;
(serverNativeModule as any).llmDispatchPreparedStream = (
_routesJson: string,
callback: (error: Error | null, arg: string) => void
) => {
callback(
null,
JSON.stringify({
type: 'message_start',
model: 'gpt-5-mini',
})
);
callback(null, JSON.stringify({ type: 'text_delta', text: 'ok' }));
callback(
null,
JSON.stringify({
type: 'done',
finish_reason: 'stop',
usage: {
prompt_tokens: 2,
completion_tokens: 3,
total_tokens: 5,
},
})
);
callback(
null,
JSON.stringify({
type: 'provider_selected',
provider_id: providerId,
})
);
callback(null, '__AFFINE_LLM_STREAM_END__');
return { abort() {} };
};
t.teardown(() => {
(serverNativeModule as any).llmDispatchPreparedStream = originalStream;
});
const chunks = await collectAsync(
engine.executeStream({
nativeDispatch: {
chat: {
routes: [
nativeRoute({
providerId,
authToken: 'byok-key',
request: nativeTextRequest('hello'),
}),
],
prepared: {
route: preparedRoute({
providerId,
authToken: 'byok-key',
}),
request: nativeTextRequest('hello'),
tools: {},
postprocess: { nodeTextMiddleware: [] },
},
hasTools: false,
},
},
request: {
kind: 'streamText',
cond: { modelId: 'gpt-5-mini' },
messages: singleUserPromptMessages('hello'),
options: {
workspace: 'workspace-1',
user: 'user-1',
session: 'session-1',
featureKind: 'chat',
},
},
routePolicy: { fallbackOrder: [providerId] },
runtimePolicy: {},
attachmentPolicy: { materializeRemoteAttachments: true },
responsePostprocess: { mode: 'streamText' },
hostPersistence: { persistAssistantTurn: true, outputKind: 'streamText' },
hostContext: {},
})
);
t.deepEqual(chunks, ['ok']);
Sinon.assert.calledOnceWithMatch(byok.recordUsage, {
workspaceId: 'workspace-1',
userId: 'user-1',
sessionId: 'session-1',
featureKind: 'chat',
providerId,
model: 'gpt-5-mini',
usage: {
prompt_tokens: 2,
completion_tokens: 3,
total_tokens: 5,
},
});
});
test('NativeExecutionEngine should record plain text BYOK usage as chat by default', async t => {
const byok = {
recordUsage: Sinon.stub().resolves(),
recordProviderFailure: Sinon.stub().resolves(),
};
const engine = new NativeExecutionEngine(byok as never);
const providerId = 'byok-aaaaaaaaaaaa-openai-server-key1';
const originalDispatch = (serverNativeModule as any).llmDispatchPrepared;
(serverNativeModule as any).llmDispatchPrepared = () => {
return JSON.stringify({
provider_id: providerId,
response: {
id: 'chat_execute',
model: 'gpt-5-mini',
message: {
role: 'assistant',
content: [{ type: 'text', text: 'execute-ok' }],
},
usage: {
prompt_tokens: 1,
completion_tokens: 2,
total_tokens: 3,
},
finish_reason: 'stop',
},
});
};
t.teardown(() => {
(serverNativeModule as any).llmDispatchPrepared = originalDispatch;
});
const text = await engine.execute({
nativeDispatch: {
chat: {
routes: [
nativeRoute({
providerId,
authToken: 'byok-key',
request: nativeTextRequest('hello'),
}),
],
prepared: {
route: preparedRoute({
providerId,
authToken: 'byok-key',
}),
request: nativeTextRequest('hello'),
tools: {},
postprocess: { nodeTextMiddleware: [] },
},
hasTools: false,
},
},
request: {
kind: 'text',
cond: { modelId: 'gpt-5-mini' },
messages: singleUserPromptMessages('hello'),
options: {
workspace: 'workspace-1',
user: 'user-1',
session: 'session-1',
},
},
routePolicy: { fallbackOrder: [providerId] },
runtimePolicy: {},
attachmentPolicy: { materializeRemoteAttachments: true },
responsePostprocess: { mode: 'text' },
hostPersistence: { persistAssistantTurn: true, outputKind: 'text' },
hostContext: {},
});
t.is(text, 'execute-ok');
Sinon.assert.calledOnceWithMatch(byok.recordUsage, {
workspaceId: 'workspace-1',
userId: 'user-1',
sessionId: 'session-1',
featureKind: 'chat',
providerId,
model: 'gpt-5-mini',
usage: {
prompt_tokens: 1,
completion_tokens: 2,
total_tokens: 3,
},
});
});
test('NativeExecutionEngine should not fail stream when BYOK usage recording fails', async t => {
const byok = {
recordUsage: Sinon.stub().rejects(new Error('usage db down')),
};
const engine = new NativeExecutionEngine(byok as never);
const providerId = 'byok-aaaaaaaaaaaa-openai-server-key1';
const originalStream = (serverNativeModule as any).llmDispatchPreparedStream;
(serverNativeModule as any).llmDispatchPreparedStream = (
_routesJson: string,
callback: (error: Error | null, arg: string) => void
) => {
callback(
null,
JSON.stringify({ type: 'message_start', model: 'gpt-5-mini' })
);
callback(null, JSON.stringify({ type: 'text_delta', text: 'ok' }));
callback(
null,
JSON.stringify({
type: 'done',
finish_reason: 'stop',
usage: { prompt_tokens: 2, completion_tokens: 3, total_tokens: 5 },
})
);
callback(
null,
JSON.stringify({ type: 'provider_selected', provider_id: providerId })
);
callback(null, '__AFFINE_LLM_STREAM_END__');
return { abort() {} };
};
t.teardown(() => {
(serverNativeModule as any).llmDispatchPreparedStream = originalStream;
});
const chunks = await collectAsync(
engine.executeStream({
nativeDispatch: {
chat: {
routes: [
nativeRoute({
providerId,
authToken: 'byok-key',
request: nativeTextRequest('hello'),
}),
],
prepared: {
route: preparedRoute({ providerId, authToken: 'byok-key' }),
request: nativeTextRequest('hello'),
tools: {},
postprocess: { nodeTextMiddleware: [] },
},
hasTools: false,
},
},
request: {
kind: 'streamText',
cond: { modelId: 'gpt-5-mini' },
messages: singleUserPromptMessages('hello'),
options: {
workspace: 'workspace-1',
user: 'user-1',
session: 'session-1',
featureKind: 'chat',
},
},
routePolicy: { fallbackOrder: [providerId] },
runtimePolicy: {},
attachmentPolicy: { materializeRemoteAttachments: true },
responsePostprocess: { mode: 'streamText' },
hostPersistence: { persistAssistantTurn: true, outputKind: 'streamText' },
hostContext: {},
})
);
t.deepEqual(chunks, ['ok']);
Sinon.assert.calledOnce(byok.recordUsage);
});
test('CopilotProviderFactory should return no prepared routes when native prepare returns null', async t => {
const provider = new DriverOnlyProvider();
(provider as any).AFFiNEConfig = { copilot: { providers: { openai: {} } } };
@@ -693,9 +979,16 @@ test('CopilotProviderFactory should return no prepared routes when native prepar
enableFeature: Sinon.stub(),
disableFeature: Sinon.stub(),
};
const access = {
resolveRouteAccess: Sinon.stub().resolves({
byokProfiles: [],
quotaBackedRoutesAvailable: true,
}),
};
const factory = new CopilotProviderFactory(
server as never,
registryService as never
registryService as never,
access as never
);
factory.register('openai-main', provider);
@@ -923,52 +1216,6 @@ test('driver-only provider should require explicit structured response contracts
t.is(capturedRequest, undefined);
});
test('MorphProvider should reuse the base native chat driver template', async t => {
const provider = new MorphProvider();
(provider as any).AFFiNEConfig = {
copilot: { providers: { morph: { apiKey: 'test-key' } } },
};
(provider as any).toolExecutorHost = {
createNativeAdapter: () => ({
text: async () => 'morph text',
streamText: async function* () {
yield 'morph stream';
},
streamObject: async function* () {
yield { type: 'text-delta', textDelta: 'unused' };
},
}),
getTools: async () => ({}),
};
t.is(
await getProviderRuntimeHost(provider).run.text(
{ modelId: 'morph-v3-fast' },
promptMessages(userPrompt('hello'))
),
'morph text'
);
t.deepEqual(
await collectAsync(
getProviderRuntimeHost(provider).run.streamText(
{ modelId: 'morph-v3-fast' },
promptMessages(userPrompt('hello'))
)
),
['morph stream']
);
t.is(
await getProviderRuntimeHost(provider).prepare.chat(
'streamObject',
{
modelId: 'morph-v3-fast',
},
promptMessages(userPrompt('hello'))
),
null
);
});
test('getActiveProviderMiddleware should merge defaults with profile override', t => {
const provider = createProvider({
rust: { request: ['clamp_max_tokens'] },
@@ -1231,9 +1478,16 @@ test('CopilotProviderFactory should resolve legacy model ids through native regi
enableFeature: Sinon.stub(),
disableFeature: Sinon.stub(),
};
const access = {
resolveRouteAccess: Sinon.stub().resolves({
byokProfiles: [],
quotaBackedRoutesAvailable: true,
}),
};
const factory = new CopilotProviderFactory(
server as never,
registryService as never
registryService as never,
access as never
);
factory.register('openai-main', provider);
@@ -1242,6 +1496,230 @@ test('CopilotProviderFactory should resolve legacy model ids through native regi
t.is(provider.resolveModel('gpt-5-2025-08-07')?.id, 'gpt-5');
});
const BYOK_OPENAI_PROFILE: CopilotProviderProfile = {
id: 'byok-aaaaaaaaaaaa-openai-server-key1',
type: CopilotProviderType.OpenAI,
priority: 10_000,
config: { apiKey: 'byok-key' },
};
const BYOK_FAL_PROFILE: CopilotProviderProfile = {
id: 'byok-aaaaaaaaaaaa-fal-server-key1',
type: CopilotProviderType.FAL,
priority: 10_000,
config: { apiKey: 'byok-key' },
};
function createProviderFactoryWithByokRoutes({
byokProfiles = [BYOK_OPENAI_PROFILE],
hasQuota = true,
}: {
byokProfiles?: CopilotProviderProfile[];
hasQuota?: boolean;
} = {}) {
const provider = createProvider();
const registryService = {
getRegistry: () =>
buildProviderRegistry({
profiles: [
{
id: 'openai-main',
type: CopilotProviderType.OpenAI,
priority: 1,
config: { apiKey: 'test-key' },
},
],
defaults: {},
}),
};
const server = {
enableFeature: Sinon.stub(),
disableFeature: Sinon.stub(),
};
const byok = {
getProfiles: Sinon.stub().resolves(byokProfiles),
};
const access = {
resolveRouteAccess: Sinon.stub().callsFake(async context => ({
byokProfiles: await byok.getProfiles(context),
quotaBackedRoutesAvailable: context.quotaBackedRoutesAllowed ?? hasQuota,
})),
};
const factory = new CopilotProviderFactory(
server as never,
registryService as never,
access as never
);
factory.register('openai-main', provider);
return { factory, byok };
}
test('CopilotProviderFactory should use matching BYOK routes before quota-backed routes', async t => {
const { factory } = createProviderFactoryWithByokRoutes();
const routes = await factory.resolveRoutes(
{ modelId: 'gpt-5-mini', outputType: ModelOutputType.Text },
{},
{ userId: 'user-1', workspaceId: 'workspace-1' }
);
t.deepEqual(
routes.map(route => route.providerId),
['byok-aaaaaaaaaaaa-openai-server-key1']
);
});
test('CopilotProviderFactory should skip unsupported BYOK profiles and use quota-backed fallback', async t => {
const { factory } = createProviderFactoryWithByokRoutes({
byokProfiles: [BYOK_FAL_PROFILE],
});
const routes = await factory.resolveRoutes(
{ modelId: 'gpt-5-mini', outputType: ModelOutputType.Text },
{},
{ userId: 'user-1', workspaceId: 'workspace-1' }
);
t.deepEqual(
routes.map(route => route.providerId),
['openai-main']
);
});
test('CopilotProviderFactory should resolve BYOK embedding routes with workspace context', async t => {
const { factory, byok } = createProviderFactoryWithByokRoutes();
const routes = await factory.resolveRoutes(
{
modelId: 'text-embedding-3-small',
outputType: ModelOutputType.Embedding,
},
{},
{ workspaceId: 'workspace-1', featureKind: 'workspace_indexing' }
);
t.deepEqual(
routes.map(route => route.providerId),
['byok-aaaaaaaaaaaa-openai-server-key1']
);
Sinon.assert.calledOnceWithMatch(byok.getProfiles, {
workspaceId: 'workspace-1',
featureKind: 'workspace_indexing',
});
});
test('CopilotProviderFactory should treat embedding preparation as embedding feature by default', async t => {
const { factory, byok } = createProviderFactoryWithByokRoutes();
await factory.prepareEmbeddingRoutes('text-embedding-3-small', 'hello', {
workspace: 'workspace-1',
});
t.true(byok.getProfiles.calledOnce);
Sinon.assert.calledOnceWithMatch(byok.getProfiles, {
workspaceId: 'workspace-1',
featureKind: 'embedding',
});
});
test('CopilotProviderFactory should resolve BYOK rerank routes before quota-backed routes', async t => {
const { factory, byok } = createProviderFactoryWithByokRoutes();
const preparedRoutes = await factory.prepareRerankRoutes(
'gpt-4o-mini',
{
query: 'programming',
candidates: [{ text: 'React is a UI library.' }],
},
{ workspace: 'workspace-1' }
);
const resolvedRoutes = await factory.resolveRoutes(
{ modelId: 'gpt-4o-mini', outputType: ModelOutputType.Rerank },
{},
{ workspaceId: 'workspace-1', featureKind: 'rerank' }
);
t.deepEqual(
preparedRoutes.map(route => route.providerId),
[]
);
t.deepEqual(
resolvedRoutes.map(route => route.providerId),
['byok-aaaaaaaaaaaa-openai-server-key1']
);
Sinon.assert.calledWithMatch(byok.getProfiles, {
workspaceId: 'workspace-1',
featureKind: 'rerank',
});
});
test('CopilotProviderFactory should treat image preparation as image feature by default', async t => {
const { factory, byok } = createProviderFactoryWithByokRoutes();
await factory.prepareImageRoutes(
{ modelId: 'gpt-image-1', outputType: ModelOutputType.Image },
singleUserPromptMessages('draw a cat'),
{ workspace: 'workspace-1' }
);
t.true(byok.getProfiles.calledOnce);
Sinon.assert.calledOnceWithMatch(byok.getProfiles, {
workspaceId: 'workspace-1',
featureKind: 'image',
});
});
test('CopilotProviderFactory should omit quota-backed routes when quota is exhausted', async t => {
const { factory } = createProviderFactoryWithByokRoutes({ hasQuota: false });
const routes = await factory.resolveRoutes(
{ modelId: 'gpt-5-mini', outputType: ModelOutputType.Text },
{},
{ userId: 'user-1', workspaceId: 'workspace-1' }
);
t.deepEqual(
routes.map(route => route.providerId),
['byok-aaaaaaaaaaaa-openai-server-key1']
);
});
test('CopilotProviderFactory should raise quota exceeded when only quota-backed routes match', async t => {
const { factory } = createProviderFactoryWithByokRoutes({
byokProfiles: [],
hasQuota: false,
});
await t.throwsAsync(
factory.resolveRoutes(
{ modelId: 'gpt-5-mini', outputType: ModelOutputType.Text },
{},
{ userId: 'user-1', workspaceId: 'workspace-1' }
),
{ instanceOf: CopilotQuotaExceeded }
);
});
test('CopilotProviderFactory should not report quota exhausted when quota-backed routes are disabled', async t => {
const { factory } = createProviderFactoryWithByokRoutes({
byokProfiles: [],
hasQuota: true,
});
const routes = await factory.resolveRoutes(
{ modelId: 'gpt-5-mini', outputType: ModelOutputType.Text },
{},
{
userId: 'user-1',
workspaceId: 'workspace-1',
quotaBackedRoutesAllowed: false,
}
);
t.deepEqual(routes, []);
});
test('selectModel should reject unknown models without online fallback', t => {
const provider = new TestOpenAIProvider();
t.is(provider.resolveModel('online-preview'), undefined);
@@ -1476,7 +1954,7 @@ test('ProviderDriverSpec should freeze declarative driver shape', t => {
});
test('NativeExecutionEngine should dispatch prepared text routes through native fallback', async t => {
const engine = new NativeExecutionEngine();
const engine = createNativeExecutionEngine();
const registry = buildProviderRegistry({
profiles: [
{
@@ -1575,8 +2053,78 @@ test('NativeExecutionEngine should dispatch prepared text routes through native
t.snapshot(summarizePreparedDispatchRoutes(capturedRoutes));
});
test('NativeExecutionEngine should record single BYOK route dispatch failure', async t => {
const byok = {
recordProviderFailure: Sinon.stub().resolves(),
recordUsage: Sinon.stub().resolves(),
};
const engine = new NativeExecutionEngine(byok as never);
const providerId = 'byok-aaaaaaaaaaaa-openai-server-key1';
const original = (serverNativeModule as any).llmDispatchPrepared;
(serverNativeModule as any).llmDispatchPrepared = () => {
throw new Error('401 invalid sk-test-primary');
};
t.teardown(() => {
(serverNativeModule as any).llmDispatchPrepared = original;
});
const error = await t.throwsAsync(
engine.execute({
nativeDispatch: {
chat: {
routes: [
nativeRoute({
providerId,
authToken: 'primary-key',
request: nativeTextRequest('hello'),
}),
],
prepared: {
route: preparedRoute({
providerId,
authToken: 'primary-key',
}),
request: nativeTextRequest('hello'),
tools: {},
postprocess: { nodeTextMiddleware: [] },
},
hasTools: false,
},
},
request: {
kind: 'text',
cond: { modelId: 'gpt-5-mini' },
messages: singleUserPromptMessages('hello'),
options: {
workspace: 'workspace-1',
user: 'user-1',
session: 'session-1',
featureKind: 'chat',
},
},
routePolicy: { fallbackOrder: [providerId] },
runtimePolicy: {},
attachmentPolicy: { materializeRemoteAttachments: true },
responsePostprocess: { mode: 'text' },
hostPersistence: { persistAssistantTurn: true, outputKind: 'text' },
hostContext: {
currentMessages: singleUserPromptMessages('hello'),
},
})
);
t.truthy(error);
Sinon.assert.calledOnceWithMatch(byok.recordProviderFailure, {
workspaceId: 'workspace-1',
providerId,
featureKind: 'chat',
});
Sinon.assert.notCalled(byok.recordUsage);
});
test('NativeExecutionEngine should reject single-route plans when no native route is prepared', async t => {
const engine = new NativeExecutionEngine();
const engine = createNativeExecutionEngine();
const error = await t.throwsAsync(
engine.execute({
@@ -1604,7 +2152,7 @@ test('NativeExecutionEngine should reject single-route plans when no native rout
});
test('NativeExecutionEngine should prefer prepared native fallback dispatch for explicit routes', async t => {
const engine = new NativeExecutionEngine();
const engine = createNativeExecutionEngine();
let capturedRoutes: unknown;
let called = false;
@@ -1683,7 +2231,7 @@ test('NativeExecutionEngine should prefer prepared native fallback dispatch for
});
test('NativeExecutionEngine should stream through prepared native fallback dispatch', async t => {
const engine = new NativeExecutionEngine();
const engine = createNativeExecutionEngine();
let called = false;
const original = (serverNativeModule as any).llmDispatchPreparedStream;
@@ -1912,7 +2460,7 @@ test('ExecutionPlanBuilder should keep single-route tool chat plans on prepared_
});
test('NativeExecutionEngine should route tool-loop chat prepared routes through native dispatch', async t => {
const engine = new NativeExecutionEngine();
const engine = createNativeExecutionEngine();
let capturedRoutes: unknown;
let called = false;
let toolCallbackCount = 0;
@@ -2262,7 +2810,7 @@ test('ExecutionPlanBuilder should build native prepared routes for structured, i
});
test('NativeExecutionEngine should dispatch structured prepared routes through native execution', async t => {
const engine = new NativeExecutionEngine();
const engine = createNativeExecutionEngine();
let capturedRoutes: unknown;
let called = false;
@@ -2350,7 +2898,7 @@ test('NativeExecutionEngine should dispatch structured prepared routes through n
});
test('NativeExecutionEngine should dispatch embedding prepared routes through native execution', async t => {
const engine = new NativeExecutionEngine();
const engine = createNativeExecutionEngine();
let capturedRoutes: unknown;
let called = false;
@@ -2424,7 +2972,7 @@ test('NativeExecutionEngine should dispatch embedding prepared routes through na
});
test('NativeExecutionEngine should dispatch rerank prepared routes through native execution', async t => {
const engine = new NativeExecutionEngine();
const engine = createNativeExecutionEngine();
let capturedRoutes: unknown;
let called = false;
@@ -2507,7 +3055,7 @@ test('NativeExecutionEngine should dispatch rerank prepared routes through nativ
});
test('NativeExecutionEngine should dispatch image plans through prepared native routes', async t => {
const engine = new NativeExecutionEngine();
const engine = createNativeExecutionEngine();
let capturedRoutes: unknown;
const original = (serverNativeModule as any).llmImageDispatchPrepared;
(serverNativeModule as any).llmImageDispatchPrepared = (
@@ -2587,8 +3135,90 @@ test('NativeExecutionEngine should dispatch image plans through prepared native
t.snapshot(summarizePreparedDispatchRoutes(capturedRoutes));
});
test('NativeExecutionEngine should record zero-token BYOK image usage without provider usage', async t => {
const byok = {
recordUsage: Sinon.stub().resolves(),
};
const engine = new NativeExecutionEngine(byok as never);
const providerId = 'byok-aaaaaaaaaaaa-fal-server-key1';
const original = (serverNativeModule as any).llmImageDispatchPrepared;
(serverNativeModule as any).llmImageDispatchPrepared = () => {
return JSON.stringify({
provider_id: providerId,
response: {
images: [
{
url: 'https://cdn.example.com/image.png',
media_type: 'image/png',
},
],
},
});
};
t.teardown(() => {
(serverNativeModule as any).llmImageDispatchPrepared = original;
});
const request = nativeImageRequest('draw a cat');
const imageArtifacts = await collectAsync(
engine.executeImageArtifacts({
nativeDispatch: {
image: {
routes: [
nativeRoute({
providerId,
authToken: 'image-key',
protocol: 'fal_image',
model: 'fal-ai/fast-sdxl',
request,
}),
],
prepared: {
route: preparedRoute({
providerId,
authToken: 'image-key',
protocol: 'fal_image',
model: 'fal-ai/fast-sdxl',
}),
request,
},
},
},
request: {
kind: 'image',
cond: { modelId: 'fal-ai/fast-sdxl' },
messages: singleUserPromptMessages('draw a cat'),
options: {
workspace: 'workspace-1',
user: 'user-1',
session: 'session-1',
featureKind: 'image',
},
},
routePolicy: { fallbackOrder: [providerId] },
runtimePolicy: {},
attachmentPolicy: { materializeRemoteAttachments: true },
responsePostprocess: { mode: 'image' },
hostPersistence: { persistAssistantTurn: true, outputKind: 'image' },
hostContext: {},
})
);
t.is(imageArtifacts.length, 1);
Sinon.assert.calledOnceWithMatch(byok.recordUsage, {
workspaceId: 'workspace-1',
userId: 'user-1',
sessionId: 'session-1',
featureKind: 'image',
providerId,
model: 'fal-ai/fast-sdxl',
usage: undefined,
});
});
test('NativeExecutionEngine should reject image plans without native dispatch', async t => {
const engine = new NativeExecutionEngine();
const engine = createNativeExecutionEngine();
await t.throwsAsync(
collectAsync(
@@ -534,6 +534,58 @@ test('doc_semantic_search should return empty array when nothing matches', async
t.deepEqual(result, []);
});
test('doc_semantic_search should pass BYOK route context into embedding matches', async t => {
const ac = {
user: () => ({
workspace: () => ({
can: async () => true,
docs: async () => [],
}),
}),
} as unknown as AccessController;
const models = {
workspace: {
get: async () => ({ id: 'workspace-1' }),
},
} as unknown as Models;
let workspaceRouteContext: unknown;
let sessionRouteContext: unknown;
const contextService = {
matchWorkspaceAll: async (...args: unknown[]) => {
workspaceRouteContext = args[7];
return [];
},
getBySessionId: async () => ({
matchFiles: async (...args: unknown[]) => {
sessionRouteContext = args[5];
return [];
},
}),
} as unknown as Parameters<typeof buildDocSearchGetter>[1];
const semanticTool = createDocSemanticSearchTool(
buildDocSearchGetter(ac, contextService, 'session-1', models).bind(null, {
user: 'user-1',
workspace: 'workspace-1',
byokLeaseId: 'lease-1',
})
);
const result = await semanticTool.execute?.({ query: 'hello' }, {});
t.deepEqual(result, []);
t.deepEqual(workspaceRouteContext, {
userId: 'user-1',
byokLeaseId: 'lease-1',
});
t.deepEqual(sessionRouteContext, {
userId: 'user-1',
byokLeaseId: 'lease-1',
});
});
test('blob_read should return explicit error when attachment context is missing', async t => {
const ac = {
user: () => ({
@@ -215,7 +215,7 @@ test('settleTask checks copilot quota before unlocking ready task', async t => {
status: 'settled',
protectedResult: payload,
});
const checkQuota = Sinon.stub().rejects(new Error('quota exceeded'));
const assertQuotaOrByok = Sinon.stub().rejects(new Error('quota exceeded'));
const service = new CopilotTranscriptionService(
{
copilotTranscriptTask: {
@@ -232,14 +232,18 @@ test('settleTask checks copilot quota before unlocking ready task', async t => {
{} as never,
{} as never,
{} as never,
{ checkQuota } as never
{ assertQuotaOrByok } as never
);
await t.throwsAsync(
() => service.settleTask('user-1', 'workspace-1', 'task-1'),
{ message: /quota exceeded/ }
);
Sinon.assert.calledOnceWithExactly(checkQuota, 'user-1');
Sinon.assert.calledOnceWithMatch(assertQuotaOrByok, {
userId: 'user-1',
workspaceId: 'workspace-1',
featureKind: 'transcript',
});
Sinon.assert.notCalled(settle);
});
@@ -341,6 +345,48 @@ test('retryTask reuses failed task and queues a new action attempt', async t =>
Sinon.assert.calledOnceWithExactly(markRunning, 'task-1');
});
test('retryTask prechecks quota or BYOK before queueing provider work', async t => {
const add = Sinon.stub().resolves(undefined);
const markRunning = Sinon.stub().resolves({ id: 'task-1' });
const assertQuotaOrByok = Sinon.stub().rejects(new Error('quota exceeded'));
const payload = TranscriptPayloadSchema.parse({
normalizedTranscript: '00:00:05 A: Kickoff',
});
const service = new CopilotTranscriptionService(
{
copilotTranscriptTask: {
getWithUser: Sinon.stub().resolves({
id: 'task-1',
status: 'failed',
strategy: 'gemini',
protectedResult: payload,
}),
markRunning,
},
} as never,
{ add } as never,
{} as never,
{
resolveTranscriptionModel: Sinon.stub().resolves('gemini-2.5-flash'),
} as never,
{} as never,
{} as never,
{ assertQuotaOrByok } as never
);
await t.throwsAsync(
() => service.retryTask('user-1', 'workspace-1', 'task-1'),
{ message: /quota exceeded/ }
);
Sinon.assert.calledOnceWithMatch(assertQuotaOrByok, {
userId: 'user-1',
workspaceId: 'workspace-1',
featureKind: 'transcript',
});
Sinon.assert.notCalled(add);
Sinon.assert.notCalled(markRunning);
});
for (const status of ['ready', 'settled']) {
test(`submitTask allows a new task for the same blob after ${status} task`, async t => {
const createdTasks: unknown[] = [];
@@ -390,6 +436,37 @@ for (const status of ['ready', 'settled']) {
});
}
test('submitTask prechecks quota or BYOK before persisting uploads', async t => {
const assertQuotaOrByok = Sinon.stub().rejects(new Error('quota exceeded'));
const resolveTranscriptionModel = Sinon.stub().resolves('gemini-2.5-flash');
const service = new CopilotTranscriptionService(
{
copilotTranscriptTask: {
getWithUser: Sinon.stub().resolves(null),
},
} as never,
{} as never,
{} as never,
{
resolveTranscriptionModel,
} as never,
{} as never,
{} as never,
{ assertQuotaOrByok } as never
);
await t.throwsAsync(
() => service.submitTask('user-1', 'workspace-1', 'blob-1', []),
{ message: /quota exceeded/ }
);
Sinon.assert.calledOnceWithMatch(assertQuotaOrByok, {
userId: 'user-1',
workspaceId: 'workspace-1',
featureKind: 'transcript',
});
Sinon.assert.notCalled(resolveTranscriptionModel);
});
test('submitTask rejects unavailable transcript strategy', async t => {
const service = new CopilotTranscriptionService(
{
@@ -1145,6 +1145,110 @@ test('should count action runs without double-counting legacy action sessions',
t.truthy(legacyAction.sessionId);
});
test('should exclude BYOK provider usage from copilot quota cost', async t => {
const { copilotSession, db, models } = t.context;
await createTestPrompts(copilotSession, db);
const regular = await createTestSession(t);
const firstMessage = await copilotSession.appendMessage({
sessionId: regular.sessionId,
userId: user.id,
prompt: { model: 'gpt-5-mini' },
message: {
role: 'user',
content: 'regular message',
createdAt: new Date(),
},
});
const secondMessage = await copilotSession.appendMessage({
sessionId: regular.sessionId,
userId: user.id,
prompt: { model: 'gpt-5-mini' },
message: {
role: 'user',
content: 'second BYOK message',
createdAt: new Date(),
},
});
await copilotSession.appendMessage({
sessionId: regular.sessionId,
userId: user.id,
prompt: { model: 'gpt-5-mini' },
message: {
role: 'user',
content: 'quota-backed message',
createdAt: new Date(),
},
});
const failedRun = await models.copilotActionRun.create({
userId: user.id,
workspaceId: workspace.id,
actionId: 'mindmap.generate',
actionVersion: 'v1',
});
await models.copilotActionRun.complete(failedRun.id, {
status: 'failed',
errorCode: 'test_failed',
});
const pendingTranscriptTask = await models.copilotTranscriptTask.create({
userId: user.id,
workspaceId: workspace.id,
blobId: 'pending-audio',
strategy: 'gemini',
recipeId: 'transcript.audio.gemini',
recipeVersion: 'v1',
});
await models.copilotUsage.create({
workspaceId: workspace.id,
userId: user.id,
provider: 'openai',
providerSource: 'byok_server',
featureKind: 'chat',
billingUnitId: firstMessage.id,
});
await models.copilotUsage.create({
workspaceId: workspace.id,
userId: user.id,
provider: 'openai',
providerSource: 'byok_server',
featureKind: 'chat',
billingUnitId: firstMessage.id,
});
await models.copilotUsage.create({
workspaceId: workspace.id,
userId: user.id,
provider: 'fal',
providerSource: 'byok_server',
featureKind: 'image',
billingUnitId: secondMessage.id,
});
await models.copilotUsage.create({
workspaceId: workspace.id,
userId: user.id,
provider: 'fal',
providerSource: 'byok_server',
featureKind: 'image',
});
await models.copilotUsage.create({
workspaceId: workspace.id,
userId: user.id,
provider: 'openai',
providerSource: 'byok_server',
featureKind: 'action',
billingUnitId: failedRun.id,
});
await models.copilotUsage.create({
workspaceId: workspace.id,
userId: user.id,
provider: 'gemini',
providerSource: 'byok_server',
featureKind: 'transcript',
billingUnitId: pendingTranscriptTask.id,
});
t.is(await copilotSession.countUserMessages(user.id), 1);
});
test('should get sessions for title generation correctly', async t => {
const { copilotSession, db } = t.context;
await createTestPrompts(copilotSession, db);
@@ -5,8 +5,9 @@ import ava, { ExecutionContext, TestFn } from 'ava';
import Sinon from 'sinon';
import { Doc as YDoc } from 'yjs';
import { MockEventBus } from '../../../__tests__/mocks';
import { createTestingApp, type TestingApp } from '../../../__tests__/utils';
import { ConfigFactory } from '../../../base';
import { ConfigFactory, EventBus } from '../../../base';
import { Flavor } from '../../../env';
import { Models } from '../../../models';
import { DocReader, PgWorkspaceDocStorageAdapter } from '../../doc';
@@ -16,6 +17,7 @@ interface Context {
app: TestingApp;
adapter: PgWorkspaceDocStorageAdapter;
docReader: DocReader;
recordDocView: Sinon.SinonStub;
}
const test = ava as TestFn<Context>;
@@ -23,7 +25,9 @@ const test = ava as TestFn<Context>;
test.before(async t => {
// @ts-expect-error testing
env.FLAVOR = Flavor.Renderer;
const app = await createTestingApp();
const app = await createTestingApp({
tapModule: m => m.overrideProvider(EventBus).useClass(MockEventBus),
});
t.context.models = app.get(Models);
t.context.adapter = app.get(PgWorkspaceDocStorageAdapter);
@@ -45,6 +49,14 @@ test.beforeEach(async t => {
email: 'test@affine.pro',
});
workspace = await t.context.models.workspace.create(user.id);
t.context.recordDocView = Sinon.stub(
t.context.models.workspaceAnalytics,
'recordDocView'
).resolves();
});
test.afterEach.always(t => {
t.context.recordDocView?.restore();
});
test.after.always(async t => {
@@ -88,10 +100,7 @@ test('should record page view when rendering shared page', async t => {
title: 'analytics-doc',
summary: 'summary',
});
const record = Sinon.stub(
models.workspaceAnalytics,
'recordDocView'
).resolves();
const record = t.context.recordDocView;
await app.GET(`/workspace/${workspace.id}/${docId}`).expect(200);
@@ -103,7 +112,6 @@ test('should record page view when rendering shared page', async t => {
});
docContent.restore();
record.restore();
});
const policyCases: Array<{
@@ -146,10 +154,7 @@ const policyCases: Array<{
unknownBlocks: [],
}),
docContent: Sinon.stub(docReader, 'getDocContent'),
record: Sinon.stub(
models.workspaceAnalytics,
'recordDocView'
).resolves(),
record: models.workspaceAnalytics.recordDocView as Sinon.SinonStub,
};
},
request: (app, docId) =>
@@ -0,0 +1,139 @@
import { Injectable } from '@nestjs/common';
import { Transactional } from '@nestjs-cls/transactional';
import { BaseModel } from './base';
export type UpsertAiWorkspaceByokConfigInput = {
id?: string | null;
workspaceId: string;
provider: string;
name: string;
description: string | null;
encryptedApiKey?: string;
endpoint: string | null;
sortOrder: number;
enabled: boolean;
userId?: string;
};
@Injectable()
export class CopilotWorkspaceByokConfigModel extends BaseModel {
async list(workspaceId: string) {
return await this.db.aiWorkspaceByokConfig.findMany({
where: { workspaceId },
orderBy: [{ sortOrder: 'asc' }, { createdAt: 'asc' }],
});
}
async listEnabled(workspaceId: string) {
return await this.db.aiWorkspaceByokConfig.findMany({
where: { workspaceId, enabled: true },
orderBy: [{ sortOrder: 'asc' }, { createdAt: 'asc' }],
});
}
async get(id: string) {
return await this.db.aiWorkspaceByokConfig.findUnique({
where: { id },
});
}
@Transactional()
async upsert(input: UpsertAiWorkspaceByokConfigInput) {
const data = {
provider: input.provider,
name: input.name,
description: input.description,
endpoint: input.endpoint,
sortOrder: input.sortOrder,
enabled: input.enabled,
updatedBy: input.userId,
...(input.encryptedApiKey
? {
encryptedApiKey: input.encryptedApiKey,
lastValidatedAt: new Date(),
lastValidationError: null,
disabledReason: null,
lastError: null,
lastErrorAt: null,
}
: {}),
};
return input.id
? await this.db.aiWorkspaceByokConfig.update({
where: { id: input.id, workspaceId: input.workspaceId },
data,
})
: await this.db.aiWorkspaceByokConfig.create({
data: {
...data,
encryptedApiKey: input.encryptedApiKey ?? '',
workspaceId: input.workspaceId,
createdBy: input.userId,
},
});
}
@Transactional()
async reorder(workspaceId: string, ids: string[], userId?: string) {
await Promise.all(
ids.map((id, sortOrder) =>
this.db.aiWorkspaceByokConfig.update({
where: { id, workspaceId },
data: { sortOrder, updatedBy: userId },
})
)
);
}
@Transactional()
async delete(workspaceId: string, id: string) {
await this.db.aiWorkspaceByokConfig.delete({ where: { id, workspaceId } });
}
@Transactional()
async clear(workspaceId: string, provider?: string | null) {
await this.db.aiWorkspaceByokConfig.deleteMany({
where: { workspaceId, ...(provider ? { provider } : {}) },
});
}
@Transactional()
async markValidated(workspaceId: string, id: string, userId?: string) {
await this.db.aiWorkspaceByokConfig.update({
where: { id, workspaceId },
data: {
enabled: true,
disabledReason: null,
lastValidatedAt: new Date(),
lastValidationError: null,
lastError: null,
lastErrorAt: null,
updatedBy: userId,
},
});
}
@Transactional()
async markFailure(workspaceId: string, id: string, message: string) {
await this.db.aiWorkspaceByokConfig.update({
where: { id, workspaceId },
data: {
enabled: false,
disabledReason: 'recent_failure',
lastValidationError: message,
lastError: message,
lastErrorAt: new Date(),
},
});
}
@Transactional()
async touchUsed(workspaceId: string, id: string) {
await this.db.aiWorkspaceByokConfig.updateMany({
where: { id, workspaceId },
data: { lastUsedAt: new Date() },
});
}
}
@@ -1005,20 +1005,26 @@ export class CopilotSessionModel extends BaseModel {
.filter(({ promptAction }) => !promptAction)
.map(({ messageCost }) => messageCost)
.reduce((prev, cost) => prev + cost, 0);
const [actionRunCost, legacyActionSessionCost, transcriptSettlementCost] =
await Promise.all([
this.models.copilotActionRun.countSucceededByUser(userId),
this.models.copilotActionRun.countLegacyPromptActionSessionsWithoutRun(
userId
),
this.models.copilotTranscriptTask.countSettledByUser(userId),
]);
return (
const [
actionRunCost,
legacyActionSessionCost,
transcriptSettlementCost,
byokQuotaExemptCost,
] = await Promise.all([
this.models.copilotActionRun.countSucceededByUser(userId),
this.models.copilotActionRun.countLegacyPromptActionSessionsWithoutRun(
userId
),
this.models.copilotTranscriptTask.countSettledByUser(userId),
this.models.copilotUsage.countQuotaExemptByokUsage(userId),
]);
const quotaBackedCost =
regularMessageCost +
actionRunCost +
legacyActionSessionCost +
transcriptSettlementCost
);
transcriptSettlementCost -
byokQuotaExemptCost;
return Math.max(0, quotaBackedCost);
}
async cleanupEmptySessions(earlyThen: Date) {
@@ -0,0 +1,142 @@
import { Injectable } from '@nestjs/common';
import { Transactional } from '@nestjs-cls/transactional';
import { Prisma } from '@prisma/client';
import { BaseModel } from './base';
type CreateAiUsageEventInput = {
workspaceId: string;
userId?: string;
provider: string;
providerSource: string;
featureKind: string;
model?: string | null;
sessionId?: string;
taskId?: string;
actionId?: string;
billingUnitId?: string;
promptTokens?: number;
completionTokens?: number;
totalTokens?: number;
cachedTokens?: number;
};
type UsageAggregateRow = {
date: string;
featureKind: string;
totalTokens: number | bigint | null;
};
type CountRow = {
count: number | bigint;
};
const BYOK_PROVIDER_SOURCES = ['byok_server', 'byok_local'];
const QUOTA_EXEMPT_BYOK_FEATURES = ['chat', 'action', 'image', 'transcript'];
@Injectable()
export class CopilotUsageModel extends BaseModel {
@Transactional()
async create(input: CreateAiUsageEventInput) {
await this.db.aiUsageEvent.create({
data: {
workspaceId: input.workspaceId,
userId: input.userId,
provider: input.provider,
providerSource: input.providerSource,
featureKind: input.featureKind,
model: input.model ?? null,
sessionId: input.sessionId,
taskId: input.taskId,
actionId: input.actionId,
billingUnitId: input.billingUnitId,
promptTokens: input.promptTokens ?? 0,
completionTokens: input.completionTokens ?? 0,
totalTokens: input.totalTokens ?? 0,
cachedTokens: input.cachedTokens ?? 0,
},
});
}
async countQuotaExemptByokUsage(userId: string) {
const rows = await this.db.$queryRaw<CountRow[]>(Prisma.sql`
WITH "byok_usage" AS (
SELECT "billing_unit_id", "feature_kind"
FROM "ai_usage_events"
WHERE "user_id" = ${userId}
AND "provider_source" IN (${Prisma.join(BYOK_PROVIDER_SOURCES)})
AND "feature_kind" IN (${Prisma.join(QUOTA_EXEMPT_BYOK_FEATURES)})
AND "billing_unit_id" IS NOT NULL
),
"message_units" AS (
SELECT DISTINCT "usage"."billing_unit_id"
FROM "byok_usage" AS "usage"
JOIN "ai_sessions_messages" AS "message"
ON "message"."id" = "usage"."billing_unit_id"
JOIN "ai_sessions_metadata" AS "session"
ON "session"."id" = "message"."session_id"
WHERE "usage"."feature_kind" IN ('chat', 'action', 'image')
AND "message"."role" = 'user'
AND "session"."user_id" = ${userId}
AND ("session"."prompt_action" IS NULL OR "session"."prompt_action" = '')
),
"action_units" AS (
SELECT DISTINCT "usage"."billing_unit_id"
FROM "byok_usage" AS "usage"
JOIN "ai_action_runs" AS "run"
ON "run"."id" = "usage"."billing_unit_id"
WHERE "usage"."feature_kind" IN ('action', 'image')
AND "run"."user_id" = ${userId}
AND "run"."status" = 'succeeded'
AND "run"."action_id" NOT LIKE 'transcript.audio.%'
),
"transcript_units" AS (
SELECT DISTINCT "usage"."billing_unit_id"
FROM "byok_usage" AS "usage"
JOIN "ai_transcript_tasks" AS "task"
ON "task"."id" = "usage"."billing_unit_id"
WHERE "usage"."feature_kind" = 'transcript'
AND "task"."user_id" = ${userId}
AND "task"."status" = 'settled'
)
SELECT (
(SELECT COUNT(*) FROM "message_units") +
(SELECT COUNT(*) FROM "action_units") +
(SELECT COUNT(*) FROM "transcript_units")
) AS "count"
`);
const count = rows[0]?.count ?? 0;
return typeof count === 'bigint' ? Number(count) : count;
}
async aggregateByDay(input: {
workspaceId: string;
from: Date;
to: Date;
providerSources: string[];
}) {
if (!input.providerSources.length) return [];
const rows = await this.db.$queryRaw<UsageAggregateRow[]>(Prisma.sql`
SELECT
to_char(date_trunc('day', "created_at" AT TIME ZONE 'UTC'), 'YYYY-MM-DD') AS "date",
"feature_kind" AS "featureKind",
COALESCE(SUM("total_tokens"), 0)::bigint AS "totalTokens"
FROM "ai_usage_events"
WHERE "workspace_id" = ${input.workspaceId}
AND "provider_source" IN (${Prisma.join(input.providerSources)})
AND "created_at" >= ${input.from}
AND "created_at" < ${input.to}
GROUP BY 1, 2
ORDER BY 1 ASC, 2 ASC
`);
return rows.map(row => {
return {
date: new Date(`${row.date}T00:00:00.000Z`),
featureKind: row.featureKind,
totalTokens: Number(row.totalTokens ?? 0),
};
});
}
}
@@ -17,10 +17,12 @@ import { CommentModel } from './comment';
import { CommentAttachmentModel } from './comment-attachment';
import { AppConfigModel } from './config';
import { CopilotActionRunModel } from './copilot-action-run';
import { CopilotWorkspaceByokConfigModel } from './copilot-byok';
import { CopilotContextModel } from './copilot-context';
import { CopilotJobModel } from './copilot-job';
import { CopilotSessionModel } from './copilot-session';
import { CopilotTranscriptTaskModel } from './copilot-transcript-task';
import { CopilotUsageModel } from './copilot-usage';
import { CopilotWorkspaceConfigModel } from './copilot-workspace';
import { DocModel } from './doc';
import { DocUserModel } from './doc-user';
@@ -58,10 +60,12 @@ const MODELS = {
notification: NotificationModel,
userSettings: UserSettingsModel,
copilotSession: CopilotSessionModel,
copilotUsage: CopilotUsageModel,
copilotTranscriptTask: CopilotTranscriptTaskModel,
copilotActionRun: CopilotActionRunModel,
copilotContext: CopilotContextModel,
copilotWorkspace: CopilotWorkspaceConfigModel,
copilotWorkspaceByokConfig: CopilotWorkspaceByokConfigModel,
copilotJob: CopilotJobModel,
appConfig: AppConfigModel,
comment: CommentModel,
@@ -133,10 +137,12 @@ export * from './calendar-subscription';
export * from './comment';
export * from './comment-attachment';
export * from './common';
export * from './copilot-byok';
export * from './copilot-context';
export * from './copilot-job';
export * from './copilot-session';
export * from './copilot-transcript-task';
export * from './copilot-usage';
export * from './copilot-workspace';
export * from './doc';
export * from './doc-user';
+9 -1
View File
@@ -458,6 +458,7 @@ type LlmRerankResponse = {
export type LlmToolLoopStreamEvent =
| { type: 'message_start'; id?: string; model?: string }
| { type: 'provider_selected'; provider_id: string }
| { type: 'text_delta'; text: string }
| { type: 'reasoning_delta'; text: string }
| {
@@ -537,7 +538,14 @@ function parseLlmEventJson(eventJson: string): LlmStreamEvent {
function parseLlmToolLoopStreamEvent(
eventJson: string
): LlmToolLoopStreamEvent {
return parseToolLoopStreamEvent(parseLlmEventJson(eventJson));
const event = parseLlmEventJson(eventJson);
if (
event.type === 'provider_selected' &&
typeof event.provider_id === 'string'
) {
return event;
}
return parseToolLoopStreamEvent(event);
}
export function llmMatchModelCapabilities(
@@ -0,0 +1,44 @@
import type { ByokFeatureKind } from '../byok/types';
export type ByokSourceCoverage = {
local: boolean;
server: boolean;
};
export type CopilotFeatureAccessRule = ByokSourceCoverage & {
quotaMetered: boolean;
};
const DEFAULT_BYOK_COVERAGE: ByokSourceCoverage = {
local: true,
server: true,
};
const DEFAULT_FEATURE_ACCESS: CopilotFeatureAccessRule = {
...DEFAULT_BYOK_COVERAGE,
quotaMetered: true,
};
const COPILOT_FEATURE_ACCESS: Partial<
Record<ByokFeatureKind, CopilotFeatureAccessRule>
> = {
transcript: { local: false, server: true, quotaMetered: true },
embedding: { local: false, server: true, quotaMetered: false },
workspace_indexing: { local: false, server: true, quotaMetered: false },
rerank: { local: false, server: true, quotaMetered: false },
};
export function getByokSourceCoverage(
featureKind?: ByokFeatureKind
): ByokSourceCoverage {
const access = getCopilotFeatureAccess(featureKind);
return { local: access.local, server: access.server };
}
export function getCopilotFeatureAccess(
featureKind?: ByokFeatureKind
): CopilotFeatureAccessRule {
return featureKind
? (COPILOT_FEATURE_ACCESS[featureKind] ?? DEFAULT_FEATURE_ACCESS)
: DEFAULT_FEATURE_ACCESS;
}
@@ -0,0 +1,2 @@
export * from './feature-coverage';
export * from './policy';
@@ -0,0 +1,106 @@
import { Injectable } from '@nestjs/common';
import { CopilotQuotaExceeded } from '../../../base';
import { ByokService } from '../byok/service';
import type { ByokFeatureKind } from '../byok/types';
import type { CopilotProviderProfile } from '../config';
import { ConversationPolicy } from '../conversation/policy';
import {
getByokSourceCoverage,
getCopilotFeatureAccess,
} from './feature-coverage';
export type CopilotAccessContext = {
userId?: string;
workspaceId?: string;
byokLeaseId?: string;
featureKind?: ByokFeatureKind;
quotaBackedRoutesAllowed?: boolean;
};
export type CopilotRouteAccess = {
byokProfiles: CopilotProviderProfile[];
quotaBackedRoutesAvailable: boolean;
};
export type CopilotTurnRouteAccess = {
byokProfiles: CopilotProviderProfile[];
quotaBackedRoutesAllowed?: boolean;
};
@Injectable()
export class CopilotAccessPolicy {
constructor(
private readonly conversationPolicy: ConversationPolicy,
private readonly byok: ByokService
) {}
async getByokProfiles(context: CopilotAccessContext = {}) {
const coverage = getByokSourceCoverage(context.featureKind);
return await this.byok.getProfiles(context, coverage);
}
async canUseQuotaBackedRoutes(context: CopilotAccessContext = {}) {
if (context.quotaBackedRoutesAllowed !== undefined) {
return context.quotaBackedRoutesAllowed;
}
if (!getCopilotFeatureAccess(context.featureKind).quotaMetered) {
return true;
}
if (!context.userId) {
return true;
}
return await this.conversationPolicy.hasQuota(context.userId);
}
async getQuota(userId: string) {
return await this.conversationPolicy.getQuota(userId);
}
async checkQuota(userId: string) {
await this.conversationPolicy.checkQuota(userId);
}
async resolveRouteAccess(
context: CopilotAccessContext = {}
): Promise<CopilotRouteAccess> {
const [byokProfiles, quotaBackedRoutesAvailable] = await Promise.all([
this.getByokProfiles(context),
this.canUseQuotaBackedRoutes(context),
]);
return { byokProfiles, quotaBackedRoutesAvailable };
}
async resolveTurnRouteAccess(
context: CopilotAccessContext
): Promise<CopilotTurnRouteAccess> {
const byokProfiles = await this.getByokProfiles(context);
if (context.quotaBackedRoutesAllowed === false) {
return { byokProfiles, quotaBackedRoutesAllowed: false };
}
const featureAccess = getCopilotFeatureAccess(context.featureKind);
if (!byokProfiles.length && context.userId && featureAccess.quotaMetered) {
await this.conversationPolicy.checkQuota(context.userId);
}
const quotaBackedRoutesAllowed = byokProfiles.length
? context.quotaBackedRoutesAllowed
: true;
return { byokProfiles, quotaBackedRoutesAllowed };
}
async assertQuotaOrByok(context: CopilotAccessContext) {
const byokProfiles = await this.getByokProfiles(context);
if (context.quotaBackedRoutesAllowed === false) {
if (!byokProfiles.length) {
throw new CopilotQuotaExceeded();
}
return;
}
const featureAccess = getCopilotFeatureAccess(context.featureKind);
if (!byokProfiles.length && context.userId && featureAccess.quotaMetered) {
await this.conversationPolicy.checkQuota(context.userId);
}
}
}
@@ -0,0 +1,4 @@
export { ByokEntitlementPolicy } from './policy';
export { WorkspaceByokResolver } from './resolver';
export { type ByokProviderRequestContext, ByokService } from './service';
export * from './types';
@@ -0,0 +1,105 @@
import { Injectable } from '@nestjs/common';
import { ActionForbidden } from '../../../base';
import { Models, WorkspaceRole } from '../../../models';
@Injectable()
export class ByokEntitlementPolicy {
constructor(private readonly models: Models) {}
private isUserPlanEntitled(features: string[]) {
return (
features.includes('pro_plan_v1') ||
features.includes('lifetime_pro_plan_v1') ||
features.includes('unlimited_copilot')
);
}
async hasAiPlan(userId?: string) {
if (!userId) return false;
const features = await this.models.userFeature.list(userId);
return this.isUserPlanEntitled(features);
}
async hasManagementAccess(workspaceId: string, userId?: string) {
if (!userId) return false;
const role = await this.models.workspaceUser.getActive(workspaceId, userId);
return (
role?.type === WorkspaceRole.Owner || role?.type === WorkspaceRole.Admin
);
}
async assertManagementAccess(workspaceId: string, userId?: string) {
if (!(await this.hasManagementAccess(workspaceId, userId))) {
throw new ActionForbidden(
'BYOK settings require workspace owner or admin.'
);
}
}
private async getWorkspaceOwnerId(workspaceId: string) {
const workspace = await this.models.workspace.get(workspaceId);
if (!workspace) {
return null;
}
try {
return (await this.models.workspaceUser.getOwner(workspaceId)).id;
} catch (error) {
if (
error instanceof Error &&
error.message === 'Workspace owner not found'
) {
return null;
}
throw error;
}
}
async hasLocalEntitlement(workspaceId: string, userId?: string) {
if (env.selfhosted) return true;
if (await this.models.workspaceFeature.has(workspaceId, 'team_plan_v1')) {
return true;
}
const ownerId = await this.getWorkspaceOwnerId(workspaceId);
if (!ownerId) return false;
if (await this.hasAiPlan(userId)) return true;
return await this.hasAiPlan(ownerId);
}
async hasServerEntitlement(workspaceId: string) {
if (env.selfhosted) return true;
if (await this.models.workspaceFeature.has(workspaceId, 'team_plan_v1')) {
return true;
}
const ownerId = await this.getWorkspaceOwnerId(workspaceId);
if (!ownerId) return false;
return await this.hasAiPlan(ownerId);
}
async hasEntitlement(workspaceId: string, userId?: string) {
const [serverEntitled, localEntitled] = await Promise.all([
this.hasServerEntitlement(workspaceId),
this.hasLocalEntitlement(workspaceId, userId),
]);
return [serverEntitled, localEntitled] as const;
}
async assertServerEntitled(workspaceId: string) {
if (!(await this.hasServerEntitlement(workspaceId))) {
throw new ActionForbidden('BYOK requires Pro, Team, or Believer.');
}
}
async assertLocalEntitled(workspaceId: string, userId?: string) {
if (!(await this.hasLocalEntitlement(workspaceId, userId))) {
throw new ActionForbidden('BYOK requires Pro, Team, or Believer.');
}
}
}
@@ -0,0 +1,405 @@
import {
Args,
Field,
ID,
InputType,
Mutation,
ObjectType,
Parent,
ResolveField,
Resolver,
} from '@nestjs/graphql';
import { SafeIntResolver } from 'graphql-scalars';
import { Throttle } from '../../../base';
import { CurrentUser } from '../../../core/auth';
import { AccessController } from '../../../core/permission';
import { WorkspaceType } from '../../../core/workspaces';
import { ByokEntitlementPolicy } from './policy';
import { ByokKeyConfig, ByokLocalLeaseProvider, ByokService } from './service';
import { ByokKeyStorage, ByokKeyTestStatus, ByokProvider } from './types';
@ObjectType()
export class WorkspaceByokKeyConfigType implements ByokKeyConfig {
@Field(() => ID)
id!: string;
@Field(() => ByokProvider)
provider!: ByokProvider;
@Field(() => String)
name!: string;
@Field(() => String, { nullable: true })
description!: string | null;
@Field(() => ByokKeyStorage)
storage!: ByokKeyStorage;
@Field(() => Boolean)
configured!: boolean;
@Field(() => Boolean)
enabled!: boolean;
@Field(() => String, { nullable: true })
endpoint!: string | null;
@Field(() => Boolean)
endpointEditable!: boolean;
@Field(() => SafeIntResolver)
sortOrder!: number;
@Field(() => [String])
capabilities!: string[];
@Field(() => ByokKeyTestStatus)
testStatus!: ByokKeyTestStatus;
@Field(() => String, { nullable: true })
disabledReason!: string | null;
@Field(() => Date, { nullable: true })
lastTestedAt!: Date | null;
@Field(() => String, { nullable: true })
lastTestError!: string | null;
@Field(() => Date, { nullable: true })
lastUsedAt!: Date | null;
@Field(() => Date, { nullable: true })
lastErrorAt!: Date | null;
@Field(() => String, { nullable: true })
lastError!: string | null;
}
@ObjectType()
class WorkspaceByokCapabilityWarningType {
@Field(() => String)
featureKind!: string;
@Field(() => String)
reason!: string;
@Field(() => [ByokProvider])
requiredProviders!: ByokProvider[];
}
@ObjectType()
class WorkspaceByokSettingsType {
@Field(() => String)
workspaceId!: string;
@Field(() => Boolean)
entitled!: boolean;
@Field(() => Boolean)
serverEntitled!: boolean;
@Field(() => Boolean)
localEntitled!: boolean;
@Field(() => [String])
entitlementRequired!: string[];
@Field(() => [WorkspaceByokKeyConfigType])
keys!: WorkspaceByokKeyConfigType[];
@Field(() => [ByokProvider])
allowedProviders!: ByokProvider[];
@Field(() => Boolean)
localStorageSupported!: boolean;
@Field(() => Boolean)
customEndpointSupported!: boolean;
@Field(() => Boolean)
hasAiPlan!: boolean;
@Field(() => [WorkspaceByokCapabilityWarningType])
warnings!: WorkspaceByokCapabilityWarningType[];
}
@ObjectType()
class WorkspaceByokUsagePointType {
@Field(() => Date)
date!: Date;
@Field(() => String)
featureKind!: string;
@Field(() => SafeIntResolver)
totalTokens!: number;
}
@ObjectType()
class TestWorkspaceByokConfigResultType {
@Field(() => Boolean)
ok!: boolean;
@Field(() => ByokKeyTestStatus)
status!: ByokKeyTestStatus;
@Field(() => String, { nullable: true })
message!: string | null;
}
@ObjectType()
class CreateWorkspaceByokLocalLeaseResultType {
@Field(() => String)
leaseId!: string;
@Field(() => Date)
expiresAt!: Date;
}
@InputType()
class UpsertWorkspaceByokConfigInput {
@Field(() => ID, { nullable: true })
id?: string;
@Field(() => String)
workspaceId!: string;
@Field(() => ByokProvider)
provider!: ByokProvider;
@Field(() => String)
name!: string;
@Field(() => String, { nullable: true })
description?: string | null;
@Field(() => ByokKeyStorage)
storage!: ByokKeyStorage;
@Field(() => String, { nullable: true })
apiKey?: string | null;
@Field(() => String, { nullable: true })
endpoint?: string | null;
@Field(() => SafeIntResolver, { nullable: true })
sortOrder?: number | null;
@Field(() => Boolean, { nullable: true })
enabled?: boolean | null;
}
@InputType()
class TestWorkspaceByokConfigInput {
@Field(() => String)
workspaceId!: string;
@Field(() => ByokProvider)
provider!: ByokProvider;
@Field(() => ByokKeyStorage)
storage!: ByokKeyStorage;
@Field(() => String, { nullable: true })
apiKey?: string | null;
@Field(() => String, { nullable: true })
endpoint?: string | null;
@Field(() => ID, { nullable: true })
configId?: string | null;
}
@InputType()
class ReorderWorkspaceByokConfigsInput {
@Field(() => String)
workspaceId!: string;
@Field(() => ByokKeyStorage)
storage!: ByokKeyStorage;
@Field(() => [ID])
ids!: string[];
}
@InputType()
class CreateWorkspaceByokLocalLeaseProviderInput implements ByokLocalLeaseProvider {
@Field(() => ByokProvider)
provider!: ByokProvider;
@Field(() => String)
name!: string;
@Field(() => String, { nullable: true })
description?: string | null;
@Field(() => String)
apiKey!: string;
@Field(() => String, { nullable: true })
endpoint?: string | null;
@Field(() => SafeIntResolver, { nullable: true })
sortOrder?: number | null;
@Field(() => Boolean, { nullable: true })
enabled?: boolean | null;
}
@InputType()
class CreateWorkspaceByokLocalLeaseInput {
@Field(() => String)
workspaceId!: string;
@Field(() => [CreateWorkspaceByokLocalLeaseProviderInput])
providers!: CreateWorkspaceByokLocalLeaseProviderInput[];
}
@Resolver(() => WorkspaceType)
export class WorkspaceByokResolver {
constructor(
private readonly ac: AccessController,
private readonly entitlement: ByokEntitlementPolicy,
private readonly byok: ByokService
) {}
@ResolveField(() => WorkspaceByokSettingsType, {
name: 'byokSettings',
complexity: 2,
})
async settings(
@CurrentUser() user: CurrentUser,
@Parent() workspace: WorkspaceType
) {
await this.ac
.user(user.id)
.workspace(workspace.id)
.allowLocal()
.assert('Workspace.Settings.Read');
await this.entitlement.assertManagementAccess(workspace.id, user.id);
return await this.byok.getSettings(workspace.id, user.id);
}
@ResolveField(() => [WorkspaceByokUsagePointType], {
name: 'byokUsage',
complexity: 2,
})
async usage(
@CurrentUser() user: CurrentUser,
@Parent() workspace: WorkspaceType,
@Args('from', { type: () => Date }) from: Date,
@Args('to', { type: () => Date }) to: Date
) {
await this.ac
.user(user.id)
.workspace(workspace.id)
.allowLocal()
.assert('Workspace.Settings.Read');
await this.entitlement.assertManagementAccess(workspace.id, user.id);
return await this.byok.getUsage(workspace.id, from, to);
}
@Throttle('strict')
@Mutation(() => TestWorkspaceByokConfigResultType)
async testWorkspaceByokConfig(
@CurrentUser() user: CurrentUser,
@Args('input') input: TestWorkspaceByokConfigInput
) {
await this.ac
.user(user.id)
.workspace(input.workspaceId)
.allowLocal()
.assert('Workspace.Settings.Update');
await this.entitlement.assertManagementAccess(input.workspaceId, user.id);
if (input.storage === ByokKeyStorage.server) {
await this.entitlement.assertServerEntitled(input.workspaceId);
} else {
await this.entitlement.assertLocalEntitled(input.workspaceId, user.id);
}
return await this.byok.testConfig({ ...input, userId: user.id });
}
@Mutation(() => WorkspaceByokKeyConfigType)
@Throttle('strict')
async upsertWorkspaceByokConfig(
@CurrentUser() user: CurrentUser,
@Args('input') input: UpsertWorkspaceByokConfigInput
) {
await this.ac
.user(user.id)
.workspace(input.workspaceId)
.allowLocal()
.assert('Workspace.Settings.Update');
await this.entitlement.assertManagementAccess(input.workspaceId, user.id);
await this.entitlement.assertServerEntitled(input.workspaceId);
return await this.byok.upsertConfig({ ...input, userId: user.id });
}
@Mutation(() => [WorkspaceByokKeyConfigType])
@Throttle('strict')
async reorderWorkspaceByokConfigs(
@CurrentUser() user: CurrentUser,
@Args('input') input: ReorderWorkspaceByokConfigsInput
) {
await this.ac
.user(user.id)
.workspace(input.workspaceId)
.allowLocal()
.assert('Workspace.Settings.Update');
await this.entitlement.assertManagementAccess(input.workspaceId, user.id);
await this.entitlement.assertServerEntitled(input.workspaceId);
return await this.byok.reorderConfigs({ ...input, userId: user.id });
}
@Mutation(() => Boolean)
@Throttle('strict')
async deleteWorkspaceByokConfig(
@CurrentUser() user: CurrentUser,
@Args('id', { type: () => ID }) id: string,
@Args('workspaceId', { type: () => String }) workspaceId: string
) {
await this.ac
.user(user.id)
.workspace(workspaceId)
.allowLocal()
.assert('Workspace.Settings.Update');
await this.entitlement.assertManagementAccess(workspaceId, user.id);
await this.entitlement.assertServerEntitled(workspaceId);
return await this.byok.deleteConfig(workspaceId, id, user.id);
}
@Mutation(() => Boolean)
@Throttle('strict')
async clearWorkspaceByokConfigs(
@CurrentUser() user: CurrentUser,
@Args('workspaceId', { type: () => String }) workspaceId: string,
@Args('provider', { type: () => ByokProvider, nullable: true })
provider?: ByokProvider | null
) {
await this.ac
.user(user.id)
.workspace(workspaceId)
.allowLocal()
.assert('Workspace.Settings.Update');
await this.entitlement.assertManagementAccess(workspaceId, user.id);
await this.entitlement.assertServerEntitled(workspaceId);
return await this.byok.clearConfigs(workspaceId, provider, user.id);
}
@Mutation(() => CreateWorkspaceByokLocalLeaseResultType)
@Throttle('strict')
async createWorkspaceByokLocalLease(
@CurrentUser() user: CurrentUser,
@Args('input') input: CreateWorkspaceByokLocalLeaseInput
) {
await this.ac
.user(user.id)
.workspace(input.workspaceId)
.allowLocal()
.assert('Workspace.Copilot');
await this.entitlement.assertManagementAccess(input.workspaceId, user.id);
await this.entitlement.assertLocalEntitled(input.workspaceId, user.id);
return await this.byok.createLocalLease({ ...input, userId: user.id });
}
}
@@ -0,0 +1,883 @@
import { createHash, createHmac, randomUUID } from 'node:crypto';
import { BadRequestException, Injectable } from '@nestjs/common';
import { BadRequest, Cache, CryptoHelper, metrics } from '../../../base';
import { Models } from '../../../models';
import type { CopilotProviderProfile } from '../config';
import { ByokEntitlementPolicy } from './policy';
import {
BYOK_ALLOWED_PROVIDERS,
type ByokFeatureKind,
ByokKeyStorage,
ByokKeyTestStatus,
ByokProvider,
ByokProviderSource,
byokProviderToCopilotType,
isByokProvider,
} from './types';
const LOCAL_LEASE_TTL_MS = 10 * 60 * 1000;
const BYOK_PROFILE_PRIORITY_BASE = 10_000;
const SERVER_PROFILE_PRIORITY_OFFSET = 2_000;
const TEST_TIMEOUT_MS = 10_000;
export type ByokProviderRequestContext = {
userId?: string;
workspaceId?: string;
byokLeaseId?: string;
};
export type ByokProfileSourceFilter = {
local?: boolean;
server?: boolean;
};
export type ByokKeyConfig = {
id: string;
provider: ByokProvider;
name: string;
description: string | null;
storage: ByokKeyStorage;
configured: boolean;
enabled: boolean;
endpoint: string | null;
endpointEditable: boolean;
sortOrder: number;
capabilities: string[];
testStatus: ByokKeyTestStatus;
disabledReason: string | null;
lastTestedAt: Date | null;
lastTestError: string | null;
lastUsedAt: Date | null;
lastErrorAt: Date | null;
lastError: string | null;
};
export type ByokSettings = {
workspaceId: string;
entitled: boolean;
serverEntitled: boolean;
localEntitled: boolean;
entitlementRequired: string[];
keys: ByokKeyConfig[];
allowedProviders: ByokProvider[];
localStorageSupported: boolean;
customEndpointSupported: boolean;
hasAiPlan: boolean;
warnings: Array<{
featureKind: string;
reason: string;
requiredProviders: ByokProvider[];
}>;
};
export type ByokLocalLeaseProvider = {
provider: ByokProvider;
name: string;
description?: string | null;
apiKey: string;
endpoint?: string | null;
sortOrder?: number | null;
enabled?: boolean | null;
};
type LocalLeasePayload = {
workspaceId: string;
userId: string;
providers: Array<
Omit<ByokLocalLeaseProvider, 'apiKey'> & { encryptedApiKey: string }
>;
};
type LocalLeaseActive = {
leaseId: string;
expiresAt: string;
};
type ByokProfileMeta = {
source: ByokProviderSource.Server | ByokProviderSource.Local;
keyId?: string;
provider: ByokProvider;
};
@Injectable()
export class ByokService {
constructor(
private readonly models: Models,
private readonly crypto: CryptoHelper,
private readonly cache: Cache,
private readonly entitlement: ByokEntitlementPolicy
) {}
get customEndpointSupported() {
return env.selfhosted;
}
async getSettings(
workspaceId: string,
userId?: string
): Promise<ByokSettings> {
if (!(await this.entitlement.hasManagementAccess(workspaceId, userId))) {
return {
workspaceId,
entitled: false,
serverEntitled: false,
localEntitled: false,
entitlementRequired: ['Workspace owner or admin'],
keys: [],
allowedProviders: [...BYOK_ALLOWED_PROVIDERS],
localStorageSupported: false,
customEndpointSupported: this.customEndpointSupported,
hasAiPlan: await this.entitlement.hasAiPlan(userId),
warnings: [],
};
}
const [serverEntitled, localEntitled] =
await this.entitlement.hasEntitlement(workspaceId, userId);
const entitled = serverEntitled || localEntitled;
if (!entitled) {
return {
workspaceId,
entitled: false,
serverEntitled: false,
localEntitled: false,
entitlementRequired: ['Pro', 'Team', 'Believer'],
keys: [],
allowedProviders: [...BYOK_ALLOWED_PROVIDERS],
localStorageSupported: false,
customEndpointSupported: this.customEndpointSupported,
hasAiPlan: await this.entitlement.hasAiPlan(userId),
warnings: [],
};
}
const rows = serverEntitled
? await this.models.copilotWorkspaceByokConfig.list(workspaceId)
: [];
const keys = rows.map(row => this.toKeyConfig(row));
return {
workspaceId,
entitled: true,
serverEntitled,
localEntitled,
entitlementRequired: ['Pro', 'Team', 'Believer'],
keys,
allowedProviders: [...BYOK_ALLOWED_PROVIDERS],
localStorageSupported: false,
customEndpointSupported: this.customEndpointSupported,
hasAiPlan: await this.entitlement.hasAiPlan(userId),
warnings: this.buildWarnings(keys),
};
}
async upsertConfig(input: {
id?: string | null;
workspaceId: string;
provider: ByokProvider;
name: string;
description?: string | null;
storage: ByokKeyStorage;
apiKey?: string | null;
endpoint?: string | null;
sortOrder?: number | null;
enabled?: boolean | null;
userId?: string;
}): Promise<ByokKeyConfig> {
await this.entitlement.assertManagementAccess(
input.workspaceId,
input.userId
);
await this.entitlement.assertServerEntitled(input.workspaceId);
this.assertProvider(input.provider);
if (input.storage !== ByokKeyStorage.server) {
throw new BadRequestException('Only server BYOK keys are persisted.');
}
const existing = input.id
? await this.models.copilotWorkspaceByokConfig.get(input.id)
: null;
if (input.id && (!existing || existing.workspaceId !== input.workspaceId)) {
throw new BadRequest('BYOK config not found.');
}
const encryptedApiKey = input.apiKey
? this.crypto.encrypt(input.apiKey)
: undefined;
if (!input.id && !encryptedApiKey) {
throw new BadRequestException('apiKey is required.');
}
const description =
input.description !== undefined
? input.description?.trim() || null
: (existing?.description ?? null);
const endpoint =
input.endpoint !== undefined
? this.normalizeEndpoint(input.endpoint)
: (existing?.endpoint ?? null);
const sortOrder = input.sortOrder ?? existing?.sortOrder ?? 0;
const enabled = input.enabled ?? existing?.enabled ?? true;
const row = await this.models.copilotWorkspaceByokConfig.upsert({
id: input.id,
workspaceId: input.workspaceId,
provider: input.provider,
name: input.name.trim(),
description,
encryptedApiKey,
endpoint,
sortOrder,
enabled,
userId: input.userId,
});
return this.toKeyConfig(row);
}
async reorderConfigs(input: {
workspaceId: string;
storage: ByokKeyStorage;
ids: string[];
userId?: string;
}) {
await this.entitlement.assertManagementAccess(
input.workspaceId,
input.userId
);
await this.entitlement.assertServerEntitled(input.workspaceId);
if (input.storage !== ByokKeyStorage.server) {
throw new BadRequestException('Only server BYOK keys are persisted.');
}
await this.models.copilotWorkspaceByokConfig.reorder(
input.workspaceId,
input.ids,
input.userId
);
return (await this.getSettings(input.workspaceId, input.userId)).keys;
}
async deleteConfig(workspaceId: string, id: string, _userId?: string) {
await this.entitlement.assertManagementAccess(workspaceId, _userId);
await this.entitlement.assertServerEntitled(workspaceId);
await this.models.copilotWorkspaceByokConfig.delete(workspaceId, id);
return true;
}
async clearConfigs(
workspaceId: string,
provider: ByokProvider | null | undefined,
_userId?: string
) {
await this.entitlement.assertManagementAccess(workspaceId, _userId);
await this.entitlement.assertServerEntitled(workspaceId);
await this.models.copilotWorkspaceByokConfig.clear(workspaceId, provider);
return true;
}
async testConfig(input: {
workspaceId: string;
provider: ByokProvider;
storage: ByokKeyStorage;
apiKey?: string | null;
endpoint?: string | null;
configId?: string | null;
userId?: string;
}) {
await this.entitlement.assertManagementAccess(
input.workspaceId,
input.userId
);
if (input.storage === ByokKeyStorage.server) {
await this.entitlement.assertServerEntitled(input.workspaceId);
} else {
await this.entitlement.assertLocalEntitled(
input.workspaceId,
input.userId
);
}
this.assertProvider(input.provider);
let apiKey = input.apiKey;
let endpoint = this.normalizeEndpoint(input.endpoint);
if (!apiKey && input.configId && input.storage === ByokKeyStorage.server) {
const config = await this.models.copilotWorkspaceByokConfig.get(
input.configId
);
if (
!config ||
config.workspaceId !== input.workspaceId ||
config.provider !== input.provider
) {
throw new BadRequestException('BYOK config not found.');
}
apiKey = this.crypto.decrypt(config.encryptedApiKey);
endpoint =
input.endpoint !== undefined
? endpoint
: this.normalizeEndpoint(config.endpoint);
}
if (!apiKey) {
throw new BadRequestException('apiKey is required.');
}
try {
await this.runProviderProbe(input.provider, apiKey, endpoint);
if (input.configId && input.storage === ByokKeyStorage.server) {
await this.models.copilotWorkspaceByokConfig.markValidated(
input.workspaceId,
input.configId,
input.userId
);
}
metrics.ai.counter('byok_test_key').add(1, {
workspace: input.workspaceId,
provider: input.provider,
storage: input.storage,
result: 'passed',
});
return { ok: true, status: ByokKeyTestStatus.passed, message: null };
} catch (error) {
const message = this.sanitizeError(error);
if (input.configId && input.storage === ByokKeyStorage.server) {
await this.models.copilotWorkspaceByokConfig.markFailure(
input.workspaceId,
input.configId,
message
);
}
metrics.ai.counter('byok_test_key').add(1, {
workspace: input.workspaceId,
provider: input.provider,
storage: input.storage,
result: 'failed',
});
return { ok: false, status: ByokKeyTestStatus.failed, message };
}
}
async createLocalLease(input: {
workspaceId: string;
providers: ByokLocalLeaseProvider[];
userId: string;
}) {
await this.entitlement.assertManagementAccess(
input.workspaceId,
input.userId
);
await this.entitlement.assertLocalEntitled(input.workspaceId, input.userId);
const providers = input.providers.map(provider => {
this.assertProvider(provider.provider);
const endpoint = this.normalizeEndpoint(provider.endpoint);
return { ...provider, endpoint };
});
const activeCacheKey = this.localLeaseActiveCacheKey({
...input,
providers,
});
const activeLease = await this.getActiveLocalLease(activeCacheKey);
if (activeLease) return activeLease;
const leaseId = randomUUID();
const expiresAt = new Date(Date.now() + LOCAL_LEASE_TTL_MS);
const payload: LocalLeasePayload = {
workspaceId: input.workspaceId,
userId: input.userId,
providers: providers.map(provider => ({
provider: provider.provider,
name: provider.name,
description: provider.description,
encryptedApiKey: this.crypto.encrypt(provider.apiKey),
endpoint: provider.endpoint,
sortOrder: provider.sortOrder,
enabled: provider.enabled,
})),
};
await this.cache.set(this.leaseCacheKey(leaseId), payload, {
ttl: LOCAL_LEASE_TTL_MS,
});
const registered = await this.cache.setnx<LocalLeaseActive>(
activeCacheKey,
{ leaseId, expiresAt: expiresAt.toISOString() },
{ ttl: LOCAL_LEASE_TTL_MS }
);
if (!registered) {
const current = await this.getActiveLocalLease(activeCacheKey);
if (current) {
await this.cache.delete(this.leaseCacheKey(leaseId));
return current;
}
}
return { leaseId, expiresAt };
}
async getProfiles(
context: ByokProviderRequestContext = {},
sources: ByokProfileSourceFilter = { local: true, server: true }
): Promise<CopilotProviderProfile[]> {
if (!context.workspaceId) {
return [];
}
const [localEntitled, serverEntitled] = await Promise.all([
this.entitlement.hasLocalEntitlement(context.workspaceId, context.userId),
this.entitlement.hasServerEntitlement(context.workspaceId),
]);
const [localProfiles, serverProfiles] = await Promise.all([
sources.local && localEntitled
? this.getLocalProfiles(context)
: Promise.resolve([]),
sources.server && serverEntitled
? this.getServerProfiles(context.workspaceId)
: Promise.resolve([]),
]);
return [...localProfiles, ...serverProfiles];
}
async recordUsage(input: {
workspaceId?: string;
userId?: string;
providerId?: string;
model?: string | null;
featureKind: ByokFeatureKind;
sessionId?: string;
taskId?: string;
actionId?: string;
billingUnitId?: string;
usage?: {
prompt_tokens?: number;
completion_tokens?: number;
total_tokens?: number;
cached_tokens?: number;
};
}) {
if (!input.workspaceId || !input.providerId) return;
const meta = this.parseProfileMeta(input.providerId, input.workspaceId);
if (!meta) return;
metrics.ai.counter('byok_usage').add(1, {
workspace: input.workspaceId,
provider: meta.provider,
source: meta.source,
feature: input.featureKind,
});
await this.models.copilotUsage.create({
workspaceId: input.workspaceId,
userId: input.userId,
provider: meta.provider,
providerSource: meta.source,
featureKind: input.featureKind,
model: input.model ?? null,
sessionId: input.sessionId,
taskId: input.taskId,
actionId: input.actionId,
billingUnitId: input.billingUnitId,
promptTokens: input.usage?.prompt_tokens ?? 0,
completionTokens: input.usage?.completion_tokens ?? 0,
totalTokens: input.usage?.total_tokens ?? 0,
cachedTokens: input.usage?.cached_tokens ?? 0,
});
if (meta.source === ByokProviderSource.Server && meta.keyId) {
await this.models.copilotWorkspaceByokConfig.touchUsed(
input.workspaceId,
meta.keyId
);
}
}
async recordProviderFailure(input: {
workspaceId?: string;
providerId?: string;
featureKind: ByokFeatureKind;
error: unknown;
}) {
if (!input.workspaceId || !input.providerId) return;
const meta = this.parseProfileMeta(input.providerId, input.workspaceId);
if (!meta) return;
const message = this.sanitizeError(input.error);
metrics.ai.counter('byok_route_failure').add(1, {
workspace: input.workspaceId,
provider: meta.provider,
source: meta.source,
feature: input.featureKind,
});
if (meta.source === ByokProviderSource.Server && meta.keyId) {
await this.models.copilotWorkspaceByokConfig.markFailure(
input.workspaceId,
meta.keyId,
message
);
}
}
async getUsage(workspaceId: string, from: Date, to: Date) {
return await this.models.copilotUsage.aggregateByDay({
workspaceId,
from,
to,
providerSources: [ByokProviderSource.Server, ByokProviderSource.Local],
});
}
private async getServerProfiles(workspaceId: string) {
const rows =
await this.models.copilotWorkspaceByokConfig.listEnabled(workspaceId);
return rows
.filter(row => isByokProvider(row.provider))
.map((row, index): CopilotProviderProfile => {
const provider = row.provider as ByokProvider;
return {
id: this.profileId(workspaceId, provider, row.id, 'server'),
type: byokProviderToCopilotType(provider),
priority:
BYOK_PROFILE_PRIORITY_BASE - SERVER_PROFILE_PRIORITY_OFFSET - index,
config: this.providerConfig(
provider,
row.encryptedApiKey,
row.endpoint
),
} as CopilotProviderProfile;
});
}
private async getLocalProfiles(context: ByokProviderRequestContext) {
if (!context.byokLeaseId || !context.workspaceId || !context.userId) {
return [];
}
if (
!(await this.entitlement.hasManagementAccess(
context.workspaceId,
context.userId
))
) {
return [];
}
const lease = await this.cache.get<LocalLeasePayload>(
this.leaseCacheKey(context.byokLeaseId)
);
if (
!lease ||
lease.workspaceId !== context.workspaceId ||
lease.userId !== context.userId
) {
return [];
}
return lease.providers
.filter(provider => provider.enabled !== false)
.map((provider, index): CopilotProviderProfile => {
return {
id: this.profileId(
context.workspaceId ?? lease.workspaceId,
provider.provider,
`${index}`,
'local'
),
type: byokProviderToCopilotType(provider.provider),
priority: BYOK_PROFILE_PRIORITY_BASE - index,
config: this.providerConfig(
provider.provider,
provider.encryptedApiKey,
provider.endpoint ?? null
),
} as CopilotProviderProfile;
});
}
private providerConfig(
provider: ByokProvider,
encryptedApiKey: string,
endpoint: string | null
) {
const apiKey = this.crypto.decrypt(encryptedApiKey);
switch (provider) {
case ByokProvider.openai:
case ByokProvider.gemini:
case ByokProvider.anthropic:
return { apiKey, ...(endpoint ? { baseURL: endpoint } : {}) };
case ByokProvider.fal:
return { apiKey };
}
}
private profileId(
workspaceId: string,
provider: ByokProvider,
keyId: string,
storage: 'server' | 'local'
) {
const hash = this.workspaceHash(workspaceId);
const sanitizedKeyId = keyId.replaceAll(/[^a-zA-Z0-9-_]/g, '');
return storage === 'local'
? `byok-${hash}-${provider}-local-${sanitizedKeyId}`
: `byok-${hash}-${provider}-${sanitizedKeyId}`;
}
parseProfileMeta(
providerId: string,
workspaceId?: string
): ByokProfileMeta | null {
const match =
/^byok-([a-f0-9]{12})-(openai|anthropic|gemini|fal)-(.+)$/.exec(
providerId
);
if (!match) return null;
if (workspaceId && match[1] !== this.workspaceHash(workspaceId)) {
return null;
}
const keyId = match[3];
return {
provider: match[2] as ByokProvider,
source: keyId.startsWith('local-')
? ByokProviderSource.Local
: ByokProviderSource.Server,
keyId: keyId.startsWith('local-') ? undefined : keyId,
};
}
private toKeyConfig(row: {
id: string;
provider: string;
name: string;
description: string | null;
endpoint: string | null;
sortOrder: number;
enabled: boolean;
disabledReason: string | null;
lastValidatedAt: Date | null;
lastValidationError: string | null;
lastUsedAt: Date | null;
lastErrorAt: Date | null;
lastError: string | null;
}): ByokKeyConfig {
const provider = row.provider as ByokProvider;
return {
id: row.id,
provider,
name: row.name,
description: row.description,
storage: ByokKeyStorage.server,
configured: true,
enabled: row.enabled,
endpoint: row.endpoint,
endpointEditable: this.customEndpointSupported,
sortOrder: row.sortOrder,
capabilities: this.capabilities(provider, 'server'),
testStatus: row.lastValidationError
? ByokKeyTestStatus.failed
: row.lastValidatedAt
? ByokKeyTestStatus.passed
: ByokKeyTestStatus.untested,
disabledReason: row.disabledReason,
lastTestedAt: row.lastValidatedAt,
lastTestError: row.lastValidationError,
lastUsedAt: row.lastUsedAt,
lastErrorAt: row.lastErrorAt,
lastError: row.lastError,
};
}
private capabilities(provider: ByokProvider, storage: 'server' | 'local') {
switch (provider) {
case ByokProvider.openai:
return ['Text', 'Image input', 'Actions', 'Image generate'];
case ByokProvider.anthropic:
return ['Text', 'Image input'];
case ByokProvider.gemini:
return storage === 'server'
? [
'Text',
'Image input',
'Actions',
'Image generate',
'Transcript',
'Indexing',
]
: ['Text', 'Image input', 'Actions', 'Image generate'];
case ByokProvider.fal:
return ['Image generate'];
}
}
private buildWarnings(keys: ByokKeyConfig[]) {
const activeServerGemini = keys.some(
key =>
key.provider === ByokProvider.gemini &&
key.storage === ByokKeyStorage.server &&
key.enabled
);
if (activeServerGemini) {
return [];
}
return [
{
featureKind: 'transcript',
reason:
'Transcript and workspace indexing require a server Gemini BYOK key or AFFiNE AI plan fallback.',
requiredProviders: [ByokProvider.gemini],
},
{
featureKind: 'workspace_indexing',
reason:
'Workspace indexing requires a server Gemini BYOK key or AFFiNE AI plan fallback.',
requiredProviders: [ByokProvider.gemini],
},
];
}
private normalizeEndpoint(endpoint?: string | null) {
if (!endpoint) return null;
if (!this.customEndpointSupported) {
throw new BadRequestException('Custom BYOK endpoint is not supported.');
}
let parsed: URL;
try {
parsed = new URL(endpoint);
} catch {
throw new BadRequestException('Invalid BYOK endpoint.');
}
if (!['https:', 'http:'].includes(parsed.protocol)) {
throw new BadRequestException('BYOK endpoint must use HTTP or HTTPS.');
}
return parsed.toString().replace(/\/$/, '');
}
private assertProvider(provider: ByokProvider) {
if (!BYOK_ALLOWED_PROVIDERS.includes(provider)) {
throw new BadRequestException('Unsupported BYOK provider.');
}
}
private async runProviderProbe(
provider: ByokProvider,
apiKey: string,
endpoint: string | null
) {
const controller = new AbortController();
const timeout = setTimeout(() => controller.abort(), TEST_TIMEOUT_MS);
try {
const request = this.buildProbeRequest(provider, apiKey, endpoint);
const response = await fetch(request.url, {
method: request.method,
headers: request.headers as unknown as Record<string, string>,
signal: controller.signal,
});
if (!response.ok) {
throw new BadRequestException(
this.providerProbeFailureMessage(response.status)
);
}
} finally {
clearTimeout(timeout);
}
}
private buildProbeRequest(
provider: ByokProvider,
apiKey: string,
endpoint: string | null
) {
switch (provider) {
case ByokProvider.openai:
return {
method: 'GET',
url: `${endpoint ?? 'https://api.openai.com/v1'}/models`,
headers: { Authorization: `Bearer ${apiKey}` },
};
case ByokProvider.anthropic:
return {
method: 'GET',
url: `${endpoint ?? 'https://api.anthropic.com/v1'}/models`,
headers: {
'x-api-key': apiKey,
'anthropic-version': '2023-06-01',
},
};
case ByokProvider.gemini:
return {
method: 'GET',
url: `${endpoint ?? 'https://generativelanguage.googleapis.com/v1beta'}/models`,
headers: { 'x-goog-api-key': apiKey },
};
case ByokProvider.fal:
return {
method: 'GET',
url: 'https://api.fal.ai/v1/models?limit=10',
headers: { Authorization: `Key ${apiKey}` },
};
}
}
private sanitizeError(error: unknown) {
if (error instanceof Error && error.name === 'AbortError') {
return 'Provider key test timed out.';
}
if (error instanceof BadRequestException && error.message) {
return error.message.slice(0, 300);
}
return 'Provider request failed.';
}
private providerProbeFailureMessage(status: number) {
switch (status) {
case 401:
return 'Provider rejected the BYOK key.';
case 403:
return 'Provider rejected the BYOK key permissions.';
case 404:
return 'Provider probe endpoint was not found.';
case 429:
return 'Provider rate limit exceeded while testing the key.';
default:
return status >= 500
? 'Provider service is unavailable.'
: `Provider key test failed with HTTP ${status}.`;
}
}
private workspaceHash(workspaceId: string) {
return createHash('sha256').update(workspaceId).digest('hex').slice(0, 12);
}
private leaseCacheKey(leaseId: string) {
return `copilot:byok:lease:${leaseId}`;
}
private async getActiveLocalLease(activeCacheKey: string) {
const active = await this.cache.get<LocalLeaseActive>(activeCacheKey);
if (!active) return null;
if (await this.cache.has(this.leaseCacheKey(active.leaseId))) {
return { leaseId: active.leaseId, expiresAt: new Date(active.expiresAt) };
}
await this.cache.delete(activeCacheKey);
return null;
}
private localLeaseActiveCacheKey(input: {
workspaceId: string;
userId: string;
providers: ByokLocalLeaseProvider[];
}) {
const fingerprint = createHmac(
'sha256',
this.crypto.keyPair.sha256.privateKey
)
.update(
JSON.stringify(
input.providers.map(provider => ({
provider: provider.provider,
name: provider.name,
description: provider.description ?? null,
apiKey: provider.apiKey,
endpoint: provider.endpoint ?? null,
sortOrder: provider.sortOrder ?? 0,
enabled: provider.enabled ?? true,
}))
)
)
.digest('hex');
return `copilot:byok:lease:active:${input.workspaceId}:${input.userId}:${fingerprint}`;
}
}
@@ -0,0 +1,79 @@
import { registerEnumType } from '@nestjs/graphql';
import { CopilotProviderType } from '../providers/types';
export enum ByokProvider {
openai = 'openai',
anthropic = 'anthropic',
gemini = 'gemini',
fal = 'fal',
}
export enum ByokKeyStorage {
server = 'server',
local = 'local',
}
export enum ByokKeyTestStatus {
untested = 'untested',
passed = 'passed',
failed = 'failed',
}
export enum ByokProviderSource {
Server = 'byok_server',
Local = 'byok_local',
AffinePlan = 'affine_plan',
}
export type ByokFeatureKind =
| 'chat'
| 'action'
| 'image'
| 'embedding'
| 'rerank'
| 'transcript'
| 'workspace_indexing';
export const BYOK_ALLOWED_PROVIDERS = [
ByokProvider.openai,
ByokProvider.anthropic,
ByokProvider.gemini,
ByokProvider.fal,
] as const;
export function byokProviderToCopilotType(provider: ByokProvider) {
switch (provider) {
case ByokProvider.openai:
return CopilotProviderType.OpenAI;
case ByokProvider.anthropic:
return CopilotProviderType.Anthropic;
case ByokProvider.gemini:
return CopilotProviderType.Gemini;
case ByokProvider.fal:
return CopilotProviderType.FAL;
}
}
export function copilotTypeToByokProvider(type: CopilotProviderType) {
switch (type) {
case CopilotProviderType.OpenAI:
return ByokProvider.openai;
case CopilotProviderType.Anthropic:
return ByokProvider.anthropic;
case CopilotProviderType.Gemini:
return ByokProvider.gemini;
case CopilotProviderType.FAL:
return ByokProvider.fal;
default:
return null;
}
}
export function isByokProvider(value: string): value is ByokProvider {
return (BYOK_ALLOWED_PROVIDERS as readonly string[]).includes(value);
}
registerEnumType(ByokProvider, { name: 'ByokProvider' });
registerEnumType(ByokKeyStorage, { name: 'ByokKeyStorage' });
registerEnumType(ByokKeyTestStatus, { name: 'ByokKeyTestStatus' });
@@ -12,9 +12,7 @@ import {
import { CloudflareWorkersAIConfig } from './providers/cloudflare';
import type { FalConfig } from './providers/fal';
import { GeminiGenerativeConfig, GeminiVertexConfig } from './providers/gemini';
import { MorphConfig } from './providers/morph';
import { OpenAIConfig } from './providers/openai';
import { PerplexityConfig } from './providers/perplexity';
import {
CopilotProviderType,
ModelOutputType,
@@ -27,10 +25,8 @@ export type CopilotProviderConfigMap = {
[CopilotProviderType.FAL]: FalConfig;
[CopilotProviderType.Gemini]: GeminiGenerativeConfig;
[CopilotProviderType.GeminiVertex]: GeminiVertexConfig;
[CopilotProviderType.Perplexity]: PerplexityConfig;
[CopilotProviderType.Anthropic]: AnthropicOfficialConfig;
[CopilotProviderType.AnthropicVertex]: AnthropicVertexConfig;
[CopilotProviderType.Morph]: MorphConfig;
};
export type ProviderSpecificConfig =
@@ -138,20 +134,11 @@ const VertexProviderConfigShape = z.object({
fetch: z.any().optional(),
});
const PerplexityConfigShape = z.object({
apiKey: z.string(),
endpoint: z.string().optional(),
});
const AnthropicOfficialConfigShape = z.object({
apiKey: z.string(),
baseURL: z.string().optional(),
});
const MorphConfigShape = z.object({
apiKey: z.string().optional(),
});
const CopilotProviderProfileShape = z.discriminatedUnion('type', [
CopilotProviderProfileBaseShape.extend({
type: z.literal(CopilotProviderType.OpenAI),
@@ -173,10 +160,6 @@ const CopilotProviderProfileShape = z.discriminatedUnion('type', [
type: z.literal(CopilotProviderType.GeminiVertex),
config: VertexProviderConfigShape,
}),
CopilotProviderProfileBaseShape.extend({
type: z.literal(CopilotProviderType.Perplexity),
config: PerplexityConfigShape,
}),
CopilotProviderProfileBaseShape.extend({
type: z.literal(CopilotProviderType.Anthropic),
config: AnthropicOfficialConfigShape,
@@ -185,10 +168,6 @@ const CopilotProviderProfileShape = z.discriminatedUnion('type', [
type: z.literal(CopilotProviderType.AnthropicVertex),
config: VertexProviderConfigShape,
}),
CopilotProviderProfileBaseShape.extend({
type: z.literal(CopilotProviderType.Morph),
config: MorphConfigShape,
}),
]);
const CopilotProviderDefaultsShape = z.object({
@@ -205,6 +184,13 @@ declare global {
interface AppConfigSchema {
copilot: {
enabled: boolean;
byok: {
enabled: ConfigItem<boolean>;
allowedProviders: ConfigItem<
Array<'openai' | 'anthropic' | 'gemini' | 'fal'>
>;
allowCustomEndpoint: ConfigItem<boolean>;
};
unsplash: ConfigItem<{
key: string;
}>;
@@ -220,10 +206,8 @@ declare global {
fal: ConfigItem<FalConfig>;
gemini: ConfigItem<GeminiGenerativeConfig>;
geminiVertex: ConfigItem<GeminiVertexConfig>;
perplexity: ConfigItem<PerplexityConfig>;
anthropic: ConfigItem<AnthropicOfficialConfig>;
anthropicVertex: ConfigItem<AnthropicVertexConfig>;
morph: ConfigItem<MorphConfig>;
};
};
}
@@ -234,6 +218,21 @@ defineModuleConfig('copilot', {
desc: 'Whether to enable the copilot plugin. <br> Document: <a href="https://docs.affine.pro/self-host-affine/administer/ai" target="_blank">https://docs.affine.pro/self-host-affine/administer/ai</a>',
default: false,
},
'byok.enabled': {
desc: 'Whether to enable workspace BYOK.',
default: true,
shape: z.boolean(),
},
'byok.allowedProviders': {
desc: 'The allowlist for workspace BYOK providers.',
default: ['openai', 'anthropic', 'gemini', 'fal'],
shape: z.array(z.enum(['openai', 'anthropic', 'gemini', 'fal'])),
},
'byok.allowCustomEndpoint': {
desc: 'Whether workspace BYOK custom endpoints are accepted.',
default: false,
shape: z.boolean(),
},
'providers.profiles': {
desc: 'The profile list for copilot providers.',
default: [],
@@ -277,12 +276,6 @@ defineModuleConfig('copilot', {
default: {},
schema: VertexSchema,
},
'providers.perplexity': {
desc: 'The config for the perplexity provider.',
default: {
apiKey: '',
},
},
'providers.anthropic': {
desc: 'The config for the anthropic provider.',
default: {
@@ -295,10 +288,6 @@ defineModuleConfig('copilot', {
default: {},
schema: VertexSchema,
},
'providers.morph': {
desc: 'The config for the morph provider.',
default: {},
},
unsplash: {
desc: 'The config for the unsplash key.',
default: {
@@ -15,7 +15,11 @@ import {
Models,
} from '../../../models';
import { CopilotEmbeddingClientService } from '../embedding/client';
import type { EmbeddingClient } from '../embedding/types';
import type {
EmbeddingCallOptions,
EmbeddingClient,
EmbeddingRouteContext,
} from '../embedding/types';
import { ContextSession } from './session';
const CONTEXT_SESSION_KEY = 'context-session';
@@ -62,6 +66,14 @@ export class CopilotContextService implements OnApplicationBootstrap {
return this.client ?? this.embeddingClients.getClient();
}
private embeddingOptions(
workspaceId: string,
signal?: AbortSignal,
routeContext: EmbeddingRouteContext = {}
): EmbeddingCallOptions {
return { workspaceId, signal, ...routeContext, featureKind: 'embedding' };
}
private async saveConfig(
contextId: string,
config: ContextConfig,
@@ -172,11 +184,13 @@ export class CopilotContextService implements OnApplicationBootstrap {
content: string,
topK: number = 5,
signal?: AbortSignal,
threshold: number = 0.5
threshold: number = 0.5,
routeContext?: EmbeddingRouteContext
) {
const client = this.embeddingClient;
if (!client) return [];
const embedding = await client.getEmbedding(content, signal);
const options = this.embeddingOptions(workspaceId, signal, routeContext);
const embedding = await client.getEmbedding(content, options);
if (!embedding) return [];
const blobChunks = await this.models.copilotWorkspace.matchBlobEmbedding(
@@ -187,7 +201,7 @@ export class CopilotContextService implements OnApplicationBootstrap {
);
if (!blobChunks.length) return [];
return await client.reRank(content, blobChunks, topK, signal);
return await client.reRank(content, blobChunks, topK, options);
}
async matchWorkspaceFiles(
@@ -195,11 +209,13 @@ export class CopilotContextService implements OnApplicationBootstrap {
content: string,
topK: number = 5,
signal?: AbortSignal,
threshold: number = 0.5
threshold: number = 0.5,
routeContext?: EmbeddingRouteContext
) {
const client = this.embeddingClient;
if (!client) return [];
const embedding = await client.getEmbedding(content, signal);
const options = this.embeddingOptions(workspaceId, signal, routeContext);
const embedding = await client.getEmbedding(content, options);
if (!embedding) return [];
const fileChunks = await this.models.copilotWorkspace.matchFileEmbedding(
@@ -210,7 +226,7 @@ export class CopilotContextService implements OnApplicationBootstrap {
);
if (!fileChunks.length) return [];
return await client.reRank(content, fileChunks, topK, signal);
return await client.reRank(content, fileChunks, topK, options);
}
async matchWorkspaceDocs(
@@ -218,11 +234,13 @@ export class CopilotContextService implements OnApplicationBootstrap {
content: string,
topK: number = 5,
signal?: AbortSignal,
threshold: number = 0.5
threshold: number = 0.5,
routeContext?: EmbeddingRouteContext
) {
const client = this.embeddingClient;
if (!client) return [];
const embedding = await client.getEmbedding(content, signal);
const options = this.embeddingOptions(workspaceId, signal, routeContext);
const embedding = await client.getEmbedding(content, options);
if (!embedding) return [];
const workspaceChunks =
@@ -234,7 +252,7 @@ export class CopilotContextService implements OnApplicationBootstrap {
);
if (!workspaceChunks.length) return [];
return await client.reRank(content, workspaceChunks, topK, signal);
return await client.reRank(content, workspaceChunks, topK, options);
}
async matchWorkspaceAll(
@@ -244,11 +262,13 @@ export class CopilotContextService implements OnApplicationBootstrap {
signal?: AbortSignal,
threshold: number = 0.8,
docIds?: string[],
scopedThreshold: number = 0.85
scopedThreshold: number = 0.85,
routeContext?: EmbeddingRouteContext
) {
const client = this.embeddingClient;
if (!client) return [];
const embedding = await client.getEmbedding(content, signal);
const options = this.embeddingOptions(workspaceId, signal, routeContext);
const embedding = await client.getEmbedding(content, options);
if (!embedding) return [];
const [fileChunks, blobChunks, workspaceChunks, scopedWorkspaceChunks] =
@@ -300,7 +320,7 @@ export class CopilotContextService implements OnApplicationBootstrap {
...(scopedWorkspaceChunks || []),
],
topK,
signal
options
);
}
@@ -11,7 +11,11 @@ import {
FileChunkSimilarity,
Models,
} from '../../../models';
import { EmbeddingClient } from '../embedding/types';
import type {
EmbeddingCallOptions,
EmbeddingClient,
EmbeddingRouteContext,
} from '../embedding/types';
export class ContextSession implements AsyncDisposable {
constructor(
@@ -69,6 +73,18 @@ export class ContextSession implements AsyncDisposable {
);
}
private embeddingOptions(
signal?: AbortSignal,
routeContext: EmbeddingRouteContext = {}
): EmbeddingCallOptions {
return {
workspaceId: this.workspaceId,
signal,
...routeContext,
featureKind: 'embedding',
};
}
async addCategoryRecord(type: ContextCategories, id: string, docs: string[]) {
const category = this.config.categories.find(
c => c.type === type && c.id === id
@@ -269,10 +285,12 @@ export class ContextSession implements AsyncDisposable {
topK: number = 5,
signal?: AbortSignal,
scopedThreshold: number = 0.85,
threshold: number = 0.5
threshold: number = 0.5,
routeContext?: EmbeddingRouteContext
): Promise<FileChunkSimilarity[]> {
if (!this.client) return [];
const embedding = await this.client.getEmbedding(content, signal);
const options = this.embeddingOptions(signal, routeContext);
const embedding = await this.client.getEmbedding(content, options);
if (!embedding) return [];
const [context, workspace] = await Promise.all([
@@ -305,7 +323,7 @@ export class ContextSession implements AsyncDisposable {
...workspace,
],
topK,
signal
options
);
}
@@ -322,10 +340,12 @@ export class ContextSession implements AsyncDisposable {
topK: number = 5,
signal?: AbortSignal,
scopedThreshold: number = 0.85,
threshold: number = 0.5
threshold: number = 0.5,
routeContext?: EmbeddingRouteContext
) {
if (!this.client) return [];
const embedding = await this.client.getEmbedding(content, signal);
const options = this.embeddingOptions(signal, routeContext);
const embedding = await this.client.getEmbedding(content, options);
if (!embedding) return [];
const docIds = this.docIds;
@@ -349,7 +369,7 @@ export class ContextSession implements AsyncDisposable {
content,
[...inContext, ...workspace],
topK,
signal
options
);
// sort result, doc recorded in context first
@@ -1,7 +1,7 @@
import { Injectable } from '@nestjs/common';
import { CopilotQuotaExceeded } from '../../../base';
import { QuotaService } from '../../../core/quota';
import { QuotaService } from '../../../core/quota/service';
import { Models } from '../../../models';
import type { Turn } from '../core';
import type { ResolvedPrompt } from '../prompt';
@@ -31,12 +31,16 @@ export class ConversationPolicy {
}
async checkQuota(userId: string) {
const { limit, used } = await this.getQuota(userId);
if (limit && Number.isFinite(limit) && used >= limit) {
if (!(await this.hasQuota(userId))) {
throw new CopilotQuotaExceeded();
}
}
async hasQuota(userId: string) {
const { limit, used } = await this.getQuota(userId);
return !(limit !== undefined && Number.isFinite(limit) && used >= limit);
}
shouldScheduleTitle(prompt: Pick<ResolvedPrompt, 'action'>) {
return !prompt.action;
}
@@ -11,7 +11,12 @@ import {
import { type CopilotRerankRequest } from '../providers/types';
import { CapabilityRuntime } from '../runtime/capability-runtime';
import { TaskPolicy } from '../runtime/task-policy';
import { EmbeddingClient, type ReRankResult } from './types';
import {
type EmbeddingCallOptionsInput,
EmbeddingClient,
normalizeEmbeddingCallOptions,
type ReRankResult,
} from './types';
class ProductionEmbeddingClient extends EmbeddingClient {
private readonly logger = new Logger(ProductionEmbeddingClient.name);
@@ -35,10 +40,19 @@ class ProductionEmbeddingClient extends EmbeddingClient {
return result;
}
async getEmbeddings(input: string[]): Promise<Embedding[]> {
async getEmbeddings(
input: string[],
options?: EmbeddingCallOptionsInput
): Promise<Embedding[]> {
const normalizedOptions = normalizeEmbeddingCallOptions(options);
const modelId = this.taskPolicy.resolveEmbeddingModelId();
const embeddings = await this.runtime.embed(modelId, input, {
dimensions: EMBEDDING_DIMENSIONS,
signal: normalizedOptions.signal,
user: normalizedOptions.userId,
workspace: normalizedOptions.workspaceId,
byokLeaseId: normalizedOptions.byokLeaseId,
featureKind: normalizedOptions.featureKind ?? 'embedding',
});
if (embeddings.length !== input.length) {
throw new CopilotFailedToGenerateEmbedding({
@@ -67,8 +81,9 @@ class ProductionEmbeddingClient extends EmbeddingClient {
>(
query: string,
embeddings: Chunk[],
signal?: AbortSignal
options?: EmbeddingCallOptionsInput
): Promise<ReRankResult> {
const normalizedOptions = normalizeEmbeddingCallOptions(options);
if (!embeddings.length) return [];
const rerankRequest: CopilotRerankRequest = {
@@ -82,7 +97,13 @@ class ProductionEmbeddingClient extends EmbeddingClient {
const ranks = await this.runtime.rerank(
this.taskPolicy.resolveRerankModelId(),
rerankRequest,
{ signal }
{
signal: normalizedOptions.signal,
user: normalizedOptions.userId,
workspace: normalizedOptions.workspaceId,
byokLeaseId: normalizedOptions.byokLeaseId,
featureKind: 'rerank',
}
);
try {
@@ -105,8 +126,9 @@ class ProductionEmbeddingClient extends EmbeddingClient {
query: string,
embeddings: Chunk[],
topK: number,
signal?: AbortSignal
options?: EmbeddingCallOptionsInput
): Promise<Chunk[]> {
const normalizedOptions = normalizeEmbeddingCallOptions(options);
// search in context and workspace may find same chunks, de-duplicate them
const { deduped: dedupedEmbeddings } = embeddings.reduce(
(acc, e) => {
@@ -138,14 +160,19 @@ class ProductionEmbeddingClient extends EmbeddingClient {
const ranks = await this.getEmbeddingRelevance(
query,
sortedEmbeddings,
signal
normalizedOptions
);
if (sortedEmbeddings.length !== ranks.length) {
// llm return wrong result, fallback to default sorting
this.logger.warn(
`Batch size mismatch: expected ${sortedEmbeddings.length}, got ${ranks.length}`
);
return await super.reRank(query, dedupedEmbeddings, topK, signal);
return await super.reRank(
query,
dedupedEmbeddings,
topK,
normalizedOptions
);
}
const highConfidenceChunks = ranks
@@ -164,7 +191,12 @@ class ProductionEmbeddingClient extends EmbeddingClient {
return highConfidenceChunks.slice(0, topK);
} catch (error) {
this.logger.warn('ReRank failed, falling back to default sorting', error);
return await super.reRank(query, dedupedEmbeddings, topK, signal);
return await super.reRank(
query,
dedupedEmbeddings,
topK,
normalizedOptions
);
}
}
}
@@ -180,7 +212,8 @@ export class CopilotEmbeddingClientService {
async refresh() {
const client = new ProductionEmbeddingClient(this.taskPolicy, this.runtime);
this.client = (await client.configured()) ? client : undefined;
await client.configured();
this.client = client;
return this.client;
}
@@ -18,7 +18,7 @@ import { Models } from '../../../models';
import { CopilotStorage } from '../storage';
import { readStream } from '../utils';
import { CopilotEmbeddingClientService } from './client';
import type { Chunk, DocFragment } from './types';
import type { Chunk, DocFragment, EmbeddingCallOptions } from './types';
import { EmbeddingClient } from './types';
@Injectable()
@@ -242,6 +242,19 @@ export class CopilotEmbeddingJob {
return new File([buffer], fileName);
}
private workspaceIndexingOptions(
workspaceId: string,
signal?: AbortSignal,
userId?: string
): EmbeddingCallOptions {
return {
workspaceId,
userId,
signal,
featureKind: 'workspace_indexing',
};
}
@OnJob('copilot.embedding.files')
async embedPendingFile({
userId,
@@ -266,7 +279,10 @@ export class CopilotEmbeddingJob {
const total = chunks.reduce((acc, c) => acc + c.length, 0);
for (const chunk of chunks) {
const embeddings = await this.embeddingClient.generateEmbeddings(chunk);
const embeddings = await this.embeddingClient.generateEmbeddings(
chunk,
this.workspaceIndexingOptions(workspaceId, undefined, userId)
);
if (contextId) {
// for context files
await this.models.copilotContext.insertFileEmbedding(
@@ -320,7 +336,10 @@ export class CopilotEmbeddingJob {
const total = chunks.reduce((acc, c) => acc + c.length, 0);
for (const chunk of chunks) {
const embeddings = await this.embeddingClient.generateEmbeddings(chunk);
const embeddings = await this.embeddingClient.generateEmbeddings(
chunk,
this.workspaceIndexingOptions(workspaceId)
);
await this.models.copilotWorkspace.insertBlobEmbeddings(
workspaceId,
blobId,
@@ -462,7 +481,7 @@ export class CopilotEmbeddingJob {
`${fragment.title || 'Untitled'}.md`
),
chunks => this.formatDocChunks(chunks, fragment),
signal
this.workspaceIndexingOptions(workspaceId, signal)
);
for (const chunks of embeddings) {
@@ -6,6 +6,7 @@ import { CopilotContextFileNotSupported } from '../../../base';
import type { PageDocContent } from '../../../core/utils/blocksuite';
import { ChunkSimilarity, Embedding } from '../../../models';
import { parseDoc } from '../../../native';
import type { ByokFeatureKind } from '../byok/types';
declare global {
interface Events {
@@ -103,6 +104,35 @@ export type Chunk = {
content: string;
};
export type EmbeddingCallOptions = {
signal?: AbortSignal;
userId?: string;
workspaceId?: string;
byokLeaseId?: string;
featureKind?: Extract<
ByokFeatureKind,
'embedding' | 'workspace_indexing' | 'rerank'
>;
};
export type EmbeddingCallOptionsInput = AbortSignal | EmbeddingCallOptions;
export type EmbeddingRouteContext = Pick<
EmbeddingCallOptions,
'userId' | 'byokLeaseId'
>;
export function normalizeEmbeddingCallOptions(
options?: EmbeddingCallOptionsInput
): EmbeddingCallOptions {
if (!options) {
return {};
}
if ('aborted' in options && 'addEventListener' in options) {
return { signal: options };
}
return options;
}
export abstract class EmbeddingClient {
async configured() {
return true;
@@ -111,11 +141,14 @@ export abstract class EmbeddingClient {
async getFileEmbeddings(
file: File,
chunkMapper: (chunk: Chunk[]) => Chunk[],
signal?: AbortSignal
options?: EmbeddingCallOptionsInput
): Promise<Embedding[][]> {
const chunks = await this.getFileChunks(file, signal);
const normalizedOptions = normalizeEmbeddingCallOptions(options);
const chunks = await this.getFileChunks(file, normalizedOptions.signal);
const chunkedEmbeddings = await Promise.all(
chunks.map(chunk => this.generateEmbeddings(chunkMapper(chunk)))
chunks.map(chunk =>
this.generateEmbeddings(chunkMapper(chunk), normalizedOptions)
)
);
return chunkedEmbeddings;
}
@@ -154,8 +187,9 @@ export abstract class EmbeddingClient {
async generateEmbeddings(
chunks: Chunk[],
signal?: AbortSignal
options?: EmbeddingCallOptionsInput
): Promise<Embedding[]> {
const normalizedOptions = normalizeEmbeddingCallOptions(options);
const retry = 3;
let embeddings: Embedding[] = [];
@@ -164,7 +198,7 @@ export abstract class EmbeddingClient {
try {
embeddings = await this.getEmbeddings(
chunks.map(c => c.content),
signal
normalizedOptions
);
break;
} catch (e) {
@@ -181,7 +215,7 @@ export abstract class EmbeddingClient {
_query: string,
embeddings: Chunk[],
topK: number,
_signal?: AbortSignal
_options?: EmbeddingCallOptionsInput
): Promise<Chunk[]> {
// sort by distance with ascending order
return embeddings
@@ -189,14 +223,14 @@ export abstract class EmbeddingClient {
.slice(0, topK);
}
async getEmbedding(query: string, signal?: AbortSignal) {
const embedding = await this.getEmbeddings([query], signal);
async getEmbedding(query: string, options?: EmbeddingCallOptionsInput) {
const embedding = await this.getEmbeddings([query], options);
return embedding?.[0]?.embedding;
}
abstract getEmbeddings(
input: string[],
signal?: AbortSignal
options?: EmbeddingCallOptionsInput
): Promise<Embedding[]>;
}
@@ -1,3 +1,9 @@
import { CopilotAccessPolicy } from './access';
import {
ByokEntitlementPolicy,
ByokService,
WorkspaceByokResolver,
} from './byok';
import { HistoryAttachmentUrlProjector } from './compat/history-attachment-url-projector';
import { CompatHistoryProjector } from './compat/history-projector';
import { HistoryPromptPreloadProjector } from './compat/history-prompt-preload-projector';
@@ -64,10 +70,13 @@ export const COPILOT_PROVIDER_PROVIDERS = [
];
export const COPILOT_RUNTIME_PROVIDERS = [
ByokEntitlementPolicy,
ByokService,
ChatSessionService,
ConversationStore,
ConversationInboxService,
ConversationPolicy,
CopilotAccessPolicy,
HistoryAttachmentUrlProjector,
CompatHistoryProjector,
HistoryPromptPreloadProjector,
@@ -114,6 +123,7 @@ export const COPILOT_RESOLVER_PROVIDERS = [
CopilotResolver,
UserCopilotResolver,
CopilotContextRootResolver,
WorkspaceByokResolver,
];
export const COPILOT_JOB_PROVIDERS = [CopilotEmbeddingJob, CopilotCronJobs];
@@ -1,12 +1,11 @@
import { CopilotProviderSideError, UserFriendlyError } from '../../../base';
import { type LlmBackendConfig } from '../../../native';
import type { CopilotTool } from '../tools';
import { CopilotProvider } from './provider';
import {
type CopilotProviderExecution,
type ProviderDriverSpec,
} from './provider-runtime-contract';
import { type CopilotChatTools, CopilotProviderType } from './types';
import { CopilotProviderType } from './types';
export type CloudflareWorkersAIConfig = {
apiToken: string;
@@ -25,16 +24,6 @@ export class CloudflareWorkersAIProvider extends CopilotProvider<CloudflareWorke
const config = this.getConfig(execution);
return !!config.apiToken && (!!config.accountId || !!config.baseURL);
}
override getProviderSpecificTools(
toolName: CopilotChatTools,
_model: string
): [string, CopilotTool?] | undefined {
if (toolName === 'docEdit') {
return ['doc_edit', undefined];
}
return;
}
private handleError(e: any) {
if (e instanceof UserFriendlyError) {
return e;
@@ -1,10 +1,14 @@
import { Injectable, Logger } from '@nestjs/common';
import { CopilotQuotaExceeded } from '../../../base';
import { ServerFeature, ServerService } from '../../../core';
import { type CopilotAccessContext, CopilotAccessPolicy } from '../access';
import type { RequiredStructuredOutputContract } from '../runtime/contracts';
import { getProviderRuntimeHost } from '../runtime/provider-runtime-context';
import type { CopilotProvider } from './provider';
import {
buildProviderRegistry,
type CopilotProviderRegistry,
type NormalizedCopilotProviderProfile,
resolveModel,
stripProviderPrefix,
@@ -57,11 +61,18 @@ type RoutePreparationResult = Partial<
>
>;
type EffectiveProviderRegistry = {
byokRegistry: CopilotProviderRegistry;
quotaBackedRegistry: CopilotProviderRegistry;
quotaBackedRoutesAvailable: boolean;
};
@Injectable()
export class CopilotProviderFactory {
constructor(
private readonly server: ServerService,
private readonly registries: CopilotProviderRegistryService
private readonly registries: CopilotProviderRegistryService,
private readonly access: CopilotAccessPolicy
) {}
private readonly logger = new Logger(CopilotProviderFactory.name);
@@ -73,20 +84,84 @@ export class CopilotProviderFactory {
return this.registries.getRegistry();
}
private getPreferredProviderIds(type?: CopilotProviderType) {
private getProviderByProfile(
providerId: string,
profile: NormalizedCopilotProviderProfile
) {
return (
this.#providers.get(providerId) ??
Array.from(this.#providerIdsByType.get(profile.type) ?? [])
.map(id => this.#providers.get(id))
.find((provider): provider is CopilotProvider => !!provider)
);
}
private providerAvailable(
providerId: string,
profile: NormalizedCopilotProviderProfile
) {
return !!this.getProviderByProfile(providerId, profile);
}
private getAvailableProviderIds(registry: CopilotProviderRegistry) {
return Array.from(registry.profiles.entries())
.filter(([providerId, profile]) =>
this.providerAvailable(providerId, profile)
)
.map(([providerId]) => providerId);
}
private getPreferredProviderIds(
registry: CopilotProviderRegistry,
type?: CopilotProviderType
) {
if (!type) return undefined;
return this.#providerIdsByType.get(type);
return registry.byType.get(type)?.filter(providerId => {
const profile = registry.profiles.get(providerId);
return profile ? this.providerAvailable(providerId, profile) : false;
});
}
private normalizeCond(
registry: CopilotProviderRegistry,
providerId: string,
cond: ModelFullConditions
): ModelFullConditions {
const registry = this.getRegistry();
const modelId = stripProviderPrefix(registry, providerId, cond.modelId);
return { ...cond, modelId };
}
private async getEffectiveRegistry(
context: CopilotAccessContext = {}
): Promise<EffectiveProviderRegistry> {
const quotaBackedRegistry = this.getRegistry();
const routeAccess = await this.access.resolveRouteAccess(context);
return {
byokRegistry: buildProviderRegistry({
profiles: routeAccess.byokProfiles,
defaults: {},
}),
quotaBackedRegistry,
quotaBackedRoutesAvailable: routeAccess.quotaBackedRoutesAvailable,
};
}
private getRequestContext(
options?:
| CopilotChatOptions
| CopilotStructuredOptions
| CopilotImageOptions
): CopilotAccessContext {
return {
userId: options?.user,
workspaceId: options?.workspace,
byokLeaseId: options?.byokLeaseId,
featureKind: options?.featureKind,
quotaBackedRoutesAllowed: options?.quotaBackedRoutesAllowed,
};
}
private filterPreparedRoutes(routes: Array<ResolvedCopilotProvider | null>) {
return routes.filter(
(route): route is ResolvedCopilotProvider => route !== null
@@ -113,36 +188,89 @@ export class CopilotProviderFactory {
cond: ModelFullConditions,
filter: {
prefer?: CopilotProviderType;
} = {}
} = {},
context: CopilotAccessContext = {}
): Promise<ResolvedCopilotProvider | null> {
return (await this.resolveRoutes(cond, filter))[0] ?? null;
return (await this.resolveRoutes(cond, filter, context))[0] ?? null;
}
async resolveRoutes(
cond: ModelFullConditions,
filter: {
prefer?: CopilotProviderType;
} = {}
} = {},
context: CopilotAccessContext = {}
): Promise<ResolvedCopilotProvider[]> {
this.logger.debug(
`Resolving copilot provider for output type: ${cond.outputType}`
);
const registry = this.getRegistry();
const { byokRegistry, quotaBackedRegistry, quotaBackedRoutesAvailable } =
await this.getEffectiveRegistry(context);
const byokRoutes = await this.resolveRoutesFromRegistry(
byokRegistry,
cond,
filter
);
const resolved = byokRoutes.length
? byokRoutes
: quotaBackedRoutesAvailable
? await this.resolveRoutesFromRegistry(
quotaBackedRegistry,
cond,
filter
)
: [];
for (const route of resolved) {
this.logger.debug(
`Copilot provider candidate found: ${route.provider.type} (${route.providerId})`
);
}
if (
!resolved.length &&
!quotaBackedRoutesAvailable &&
context.quotaBackedRoutesAllowed !== false
) {
const quotaBackedRoutes = await this.resolveRoutesFromRegistry(
quotaBackedRegistry,
cond,
filter
);
if (quotaBackedRoutes.length) {
throw new CopilotQuotaExceeded();
}
}
return resolved;
}
private async resolveRoutesFromRegistry(
registry: CopilotProviderRegistry,
cond: ModelFullConditions,
filter: {
prefer?: CopilotProviderType;
} = {}
): Promise<ResolvedCopilotProvider[]> {
const route = resolveModel({
registry,
modelId: cond.modelId,
outputType: cond.outputType,
availableProviderIds: this.#providers.keys(),
preferredProviderIds: this.getPreferredProviderIds(filter.prefer),
availableProviderIds: this.getAvailableProviderIds(registry),
preferredProviderIds: this.getPreferredProviderIds(
registry,
filter.prefer
),
});
const resolved: ResolvedCopilotProvider[] = [];
for (const providerId of route.candidateProviderIds) {
const provider = this.#providers.get(providerId);
const profile = registry.profiles.get(providerId);
const provider = profile
? this.getProviderByProfile(providerId, profile)
: undefined;
if (!provider || !profile) continue;
const normalizedCond = this.normalizeCond(providerId, cond);
const normalizedCond = this.normalizeCond(registry, providerId, cond);
if (
normalizedCond.modelId &&
profile.models?.length &&
@@ -155,9 +283,6 @@ export class CopilotProviderFactory {
const matched = await provider.match(normalizedCond, execution);
if (!matched) continue;
this.logger.debug(
`Copilot provider candidate found: ${provider.type} (${providerId})`
);
resolved.push({
providerId,
provider,
@@ -181,7 +306,11 @@ export class CopilotProviderFactory {
prefer?: CopilotProviderType;
} = {}
): Promise<ResolvedCopilotProvider[]> {
const routes = await this.resolveRoutes(cond, filter);
const routes = await this.resolveRoutes(
cond,
filter,
this.getRequestContext(options)
);
return await this.prepareResolvedRoutes(routes, async route => {
const prepared = await getProviderRuntimeHost(
route.provider
@@ -213,7 +342,11 @@ export class CopilotProviderFactory {
} = {},
responseContract?: RequiredStructuredOutputContract
): Promise<ResolvedCopilotProvider[]> {
const routes = await this.resolveRoutes(cond, filter);
const routes = await this.resolveRoutes(
cond,
filter,
this.getRequestContext(options)
);
return await this.prepareResolvedRoutes(routes, async route => {
const preparedStructured =
(await getProviderRuntimeHost(route.provider).prepare.structured(
@@ -239,10 +372,14 @@ export class CopilotProviderFactory {
input: string | string[],
options: CopilotEmbeddingOptions = {}
): Promise<ResolvedCopilotProvider[]> {
const routes = await this.resolveRoutes({
modelId,
outputType: ModelOutputType.Embedding,
});
const routes = await this.resolveRoutes(
{ modelId, outputType: ModelOutputType.Embedding },
{},
{
...this.getRequestContext(options),
featureKind: options?.featureKind ?? 'embedding',
}
);
return await this.prepareResolvedRoutes(routes, async route => {
const preparedEmbedding =
(await getProviderRuntimeHost(route.provider).prepare.embedding(
@@ -267,10 +404,14 @@ export class CopilotProviderFactory {
request: CopilotRerankRequest,
options: CopilotChatOptions = {}
): Promise<ResolvedCopilotProvider[]> {
const routes = await this.resolveRoutes({
modelId,
outputType: ModelOutputType.Rerank,
});
const routes = await this.resolveRoutes(
{
modelId,
outputType: ModelOutputType.Rerank,
},
{},
{ ...this.getRequestContext(options), featureKind: 'rerank' }
);
return await this.prepareResolvedRoutes(routes, async route => {
const preparedRerank =
(await getProviderRuntimeHost(route.provider).prepare.rerank(
@@ -298,7 +439,10 @@ export class CopilotProviderFactory {
prefer?: CopilotProviderType;
} = {}
): Promise<ResolvedCopilotProvider[]> {
const routes = await this.resolveRoutes(cond, filter);
const routes = await this.resolveRoutes(cond, filter, {
...this.getRequestContext(options),
featureKind: options?.featureKind ?? 'image',
});
return await this.prepareResolvedRoutes(routes, async route => {
const preparedImage =
(await getProviderRuntimeHost(route.provider).prepare.image(
@@ -8,7 +8,6 @@ export { FalProvider } from './fal';
export { GeminiGenerativeProvider, GeminiVertexProvider } from './gemini';
export { CopilotProviderLifecycleService } from './lifecycle-service';
export { OpenAIProvider } from './openai';
export { PerplexityProvider } from './perplexity';
export type { CopilotProvider } from './provider';
export { CopilotProviders } from './provider-tokens';
export { CopilotProviderRegistryService } from './registry-service';
@@ -1,60 +0,0 @@
import { CopilotProviderSideError, UserFriendlyError } from '../../../base';
import { type LlmBackendConfig } from '../../../native';
import { CopilotProvider } from './provider';
import {
type CopilotProviderExecution,
type ProviderDriverSpec,
} from './provider-runtime-contract';
import { CopilotProviderType, ModelOutputType } from './types';
export const DEFAULT_DIMENSIONS = 256;
export type MorphConfig = {
apiKey?: string;
};
export class MorphProvider extends CopilotProvider<MorphConfig> {
readonly type = CopilotProviderType.Morph;
protected resolveModelBackendKind() {
return 'morph' as const;
}
override configured(execution?: CopilotProviderExecution): boolean {
return !!this.getConfig(execution).apiKey;
}
private handleError(e: any) {
if (e instanceof UserFriendlyError) {
return e;
}
return new CopilotProviderSideError({
provider: this.type,
kind: 'unexpected_response',
message: e?.message || 'Unexpected morph response',
});
}
private createNativeConfig(
execution?: CopilotProviderExecution
): LlmBackendConfig {
return {
base_url: 'https://api.morphllm.com',
auth_token: this.getConfig(execution).apiKey ?? '',
};
}
override getDriverSpec(): ProviderDriverSpec {
return {
createBackendConfig: execution => this.createNativeConfig(execution),
mapError: error => this.handleError(error),
chat: {
resolveOutputType: kind =>
kind === 'streamObject' ? null : ModelOutputType.Text,
},
structured: false,
embedding: false,
rerank: false,
};
}
}
@@ -14,7 +14,6 @@ import {
AttachmentAdmissionHost,
} from '../runtime/hosts/attachment-admission';
import { AttachmentMaterializer } from '../runtime/hosts/attachment-materializer';
import type { CopilotTool } from '../tools';
import { CopilotProvider } from './provider';
import { hasProviderModelBehaviorFlag } from './provider-model-runtime';
import type {
@@ -22,7 +21,6 @@ import type {
ProviderDriverSpec,
} from './provider-runtime-contract';
import {
CopilotChatTools,
CopilotProviderType,
type PromptAttachment,
type PromptMessage,
@@ -64,16 +62,6 @@ export class OpenAIProvider extends CopilotProvider<OpenAIConfig> {
});
}
override getProviderSpecificTools(
toolName: CopilotChatTools,
_model: string
): [string, CopilotTool?] | undefined {
if (toolName === 'docEdit') {
return ['doc_edit', undefined];
}
return;
}
protected createNativeConfig(
execution?: CopilotProviderExecution
): LlmBackendConfig {
@@ -1,75 +0,0 @@
import { CopilotProviderSideError } from '../../../base';
import { type LlmBackendConfig } from '../../../native';
import { CopilotProvider } from './provider';
import { hasProviderModelBehaviorFlag } from './provider-model-runtime';
import {
type CopilotProviderExecution,
type ProviderDriverSpec,
} from './provider-runtime-contract';
import { CopilotProviderType, ModelOutputType } from './types';
export type PerplexityConfig = {
apiKey: string;
endpoint?: string;
};
export class PerplexityProvider extends CopilotProvider<PerplexityConfig> {
readonly type = CopilotProviderType.Perplexity;
protected resolveModelBackendKind() {
return 'perplexity' as const;
}
override configured(execution?: CopilotProviderExecution): boolean {
return !!this.getConfig(execution).apiKey;
}
override getDriverSpec(): ProviderDriverSpec {
return {
createBackendConfig: execution => this.createNativeConfig(execution),
mapError: error => this.handleError(error),
chat: {
resolveOutputType: kind =>
kind === 'streamObject' ? null : ModelOutputType.Text,
withAttachment: false,
resolveRequestOptions: async context => ({
withAttachment: !hasProviderModelBehaviorFlag(
context.model,
'no_attachments'
),
include: hasProviderModelBehaviorFlag(
context.model,
'citations_include'
)
? ['citations']
: undefined,
}),
},
structured: false,
embedding: false,
rerank: false,
};
}
private createNativeConfig(
execution?: CopilotProviderExecution
): LlmBackendConfig {
const config = this.getConfig(execution);
const baseUrl = config.endpoint || 'https://api.perplexity.ai';
return {
base_url: baseUrl.replace(/\/v1\/?$/, ''),
auth_token: config.apiKey,
};
}
private handleError(e: any) {
if (e instanceof CopilotProviderSideError) {
return e;
}
return new CopilotProviderSideError({
provider: this.type,
kind: 'unexpected_response',
message: e?.message || 'Unexpected perplexity response',
});
}
}
@@ -21,18 +21,6 @@ const DEFAULT_MIDDLEWARE_BY_TYPE: Record<
[CopilotProviderType.AnthropicVertex]: {
node: { text: DEFAULT_NODE_TEXT_MIDDLEWARE },
},
[CopilotProviderType.Morph]: {
rust: {
request: ['clamp_max_tokens'],
},
node: { text: DEFAULT_NODE_TEXT_MIDDLEWARE },
},
[CopilotProviderType.Perplexity]: {
rust: {
request: ['clamp_max_tokens'],
},
node: { text: DEFAULT_NODE_TEXT_MIDDLEWARE },
},
[CopilotProviderType.Gemini]: {
node: { text: DEFAULT_NODE_TEXT_MIDDLEWARE },
},
@@ -15,10 +15,8 @@ const LEGACY_PROVIDER_ORDER: CopilotProviderType[] = [
CopilotProviderType.FAL,
CopilotProviderType.Gemini,
CopilotProviderType.GeminiVertex,
CopilotProviderType.Perplexity,
CopilotProviderType.Anthropic,
CopilotProviderType.AnthropicVertex,
CopilotProviderType.Morph,
];
const LEGACY_PROVIDER_PRIORITY = LEGACY_PROVIDER_ORDER.reduce(
@@ -5,9 +5,7 @@ import {
import { CloudflareWorkersAIProvider } from './cloudflare';
import { FalProvider } from './fal';
import { GeminiGenerativeProvider, GeminiVertexProvider } from './gemini';
import { MorphProvider } from './morph';
import { OpenAIProvider } from './openai';
import { PerplexityProvider } from './perplexity';
export const CopilotProviders = [
OpenAIProvider,
@@ -15,8 +13,6 @@ export const CopilotProviders = [
FalProvider,
GeminiGenerativeProvider,
GeminiVertexProvider,
PerplexityProvider,
AnthropicOfficialProvider,
AnthropicVertexProvider,
MorphProvider,
];
@@ -30,8 +30,6 @@ export enum CopilotProviderType {
Gemini = 'gemini',
GeminiVertex = 'geminiVertex',
OpenAI = 'openai',
Perplexity = 'perplexity',
Morph = 'morph',
}
export const CopilotProviderSchema = z.object({
@@ -80,8 +78,6 @@ export const PromptToolsSchema = z
'blobRead',
'codeArtifact',
'conversationSummary',
// work with morph
'docEdit',
// work with indexer
'docRead',
'docCreate',
@@ -268,6 +264,22 @@ const CopilotProviderOptionsSchema = z.object({
user: z.string().optional(),
session: z.string().optional(),
workspace: z.string().optional(),
byokLeaseId: z.string().optional(),
billingUnitId: z.string().optional(),
taskId: z.string().optional(),
actionId: z.string().optional(),
quotaBackedRoutesAllowed: z.boolean().optional(),
featureKind: z
.enum([
'chat',
'action',
'image',
'embedding',
'workspace_indexing',
'rerank',
'transcript',
])
.optional(),
});
export const CopilotChatOptionsSchema = CopilotProviderOptionsSchema.merge(
@@ -164,11 +164,6 @@ export function toError(error: unknown): Error {
}
}
type DocEditFootnote = {
intent: string;
result: string;
};
function asRecord(value: unknown): Record<string, unknown> | null {
if (value && typeof value === 'object' && !Array.isArray(value)) {
return value as Record<string, unknown>;
@@ -184,8 +179,6 @@ export class TextStreamParser {
private prefix: string | null = this.CALLOUT_PREFIX;
private readonly docEditFootnotes: DocEditFootnote[] = [];
public parse(chunk: CopilotTextStreamPart) {
let result = '';
switch (chunk.type) {
@@ -233,13 +226,6 @@ export class TextStreamParser {
result += `\nWriting document "${chunk.input.title}"\n`;
break;
}
case 'doc_edit': {
this.docEditFootnotes.push({
intent: String(chunk.input.instructions ?? ''),
result: '',
});
break;
}
}
result = this.markAsCallout(result);
break;
@@ -250,22 +236,6 @@ export class TextStreamParser {
);
result = this.addPrefix(result);
switch (chunk.toolName) {
case 'doc_edit': {
const output = asRecord(chunk.output);
const array = output?.result;
if (Array.isArray(array)) {
result += array
.map(item => {
return `\n${String(asRecord(item)?.changedContent ?? '')}\n`;
})
.join('');
this.docEditFootnotes[this.docEditFootnotes.length - 1].result =
result;
} else {
this.docEditFootnotes.pop();
}
break;
}
case 'doc_semantic_search': {
const output = chunk.output;
if (Array.isArray(output)) {
@@ -319,10 +289,7 @@ export class TextStreamParser {
}
public end() {
const footnotes = this.docEditFootnotes.map((footnote, index) => {
return `[^edit${index + 1}]: ${JSON.stringify({ type: 'doc-edit', ...footnote })}`;
});
return footnotes.join('\n');
return '';
}
private addPrefix(text: string) {
@@ -1,4 +1,4 @@
import { BadRequestException, NotFoundException } from '@nestjs/common';
import { NotFoundException } from '@nestjs/common';
import {
Args,
Field,
@@ -7,7 +7,6 @@ import {
Mutation,
ObjectType,
Parent,
Query,
registerEnumType,
ResolveField,
Resolver,
@@ -19,7 +18,6 @@ import {
CallMetric,
CopilotDocNotFound,
CopilotFailedToCreateMessage,
CopilotProviderSideError,
CopilotSessionNotFound,
type FileUpload,
paginate,
@@ -28,10 +26,8 @@ import {
RequestMutex,
Throttle,
TooManyRequest,
UserFriendlyError,
} from '../../base';
import { CurrentUser } from '../../core/auth';
import { DocReader } from '../../core/doc';
import { AccessController, DocAction } from '../../core/permission';
import { UserType } from '../../core/user';
import type { ListSessionOptions, UpdateChatSession } from '../../models';
@@ -40,7 +36,6 @@ import { ConversationInboxService } from './conversation/inbox';
import { PromptService } from './prompt/service';
import { CopilotProviderFactory } from './providers/factory';
import { ModelOutputType, type StreamObject } from './providers/types';
import { CapabilityRuntime } from './runtime/capability-runtime';
import { ChatSessionService } from './session';
import { type ChatHistory, type ChatMessage, SubmittedMessage } from './types';
@@ -376,9 +371,7 @@ export class CopilotResolver {
private readonly chatSession: ChatSessionService,
private readonly historyProjector: CompatHistoryProjector,
private readonly inbox: ConversationInboxService,
private readonly docReader: DocReader,
private readonly providerFactory: CopilotProviderFactory,
private readonly runtime: CapabilityRuntime
private readonly providerFactory: CopilotProviderFactory
) {}
@ResolveField(() => CopilotQuotaType, {
@@ -641,8 +634,6 @@ export class CopilotResolver {
throw new TooManyRequest('Server is busy');
}
await this.chatSession.checkQuota(user.id);
return await this.chatSession.create({
...options,
pinned: options.pinned ?? false,
@@ -724,7 +715,6 @@ export class CopilotResolver {
throw new TooManyRequest('Server is busy');
}
await this.chatSession.checkQuota(user.id);
return await this.chatSession.update({
...options,
userId: user.id,
@@ -752,8 +742,6 @@ export class CopilotResolver {
throw new CopilotDocNotFound({ docId: options.docId });
}
await this.chatSession.checkQuota(user.id);
return await this.chatSession.fork({
...options,
userId: user.id,
@@ -819,96 +807,6 @@ export class CopilotResolver {
}
}
@Query(() => String, {
description:
'Apply updates to a doc using LLM and return the merged markdown.',
deprecationReason: 'use Mutation.applyDocUpdates',
})
async applyDocUpdates(
@CurrentUser() user: CurrentUser,
@Args({ name: 'workspaceId', type: () => String })
workspaceId: string,
@Args({ name: 'docId', type: () => String })
docId: string,
@Args({ name: 'op', type: () => String })
op: string,
@Args({ name: 'updates', type: () => String })
updates: string
): Promise<string> {
return this.applyDocUpdatesInternal(user, workspaceId, docId, op, updates);
}
@Mutation(() => String, {
description:
'Apply updates to a doc using LLM and return the merged markdown.',
name: 'applyDocUpdates',
})
async applyDocUpdatesMutation(
@CurrentUser() user: CurrentUser,
@Args({ name: 'workspaceId', type: () => String })
workspaceId: string,
@Args({ name: 'docId', type: () => String })
docId: string,
@Args({ name: 'op', type: () => String })
op: string,
@Args({ name: 'updates', type: () => String })
updates: string
): Promise<string> {
return this.applyDocUpdatesInternal(user, workspaceId, docId, op, updates);
}
private async applyDocUpdatesInternal(
user: CurrentUser,
workspaceId: string,
docId: string,
op: string,
updates: string
): Promise<string> {
await this.assertPermission(user, { workspaceId, docId });
const docContent = await this.docReader.getDocMarkdown(
workspaceId,
docId,
true
);
if (!docContent || !docContent.markdown) {
throw new NotFoundException('Doc not found or empty');
}
const markdown = docContent.markdown.trim();
const resolved = await this.providerFactory.resolveProvider({
modelId: 'morph-v3-large',
outputType: ModelOutputType.Text,
});
if (!resolved) {
throw new BadRequestException('No LLM provider available');
}
try {
return await this.runtime.text(
{ modelId: 'morph-v3-large' },
[
{
role: 'user',
content: `<instruction>${op}</instruction>\n<code>${markdown}</code>\n<update>${updates}</update>`,
},
],
{ reasoning: false }
);
} catch (e: any) {
if (e instanceof UserFriendlyError) {
throw e;
} else {
throw new CopilotProviderSideError({
provider: resolved.provider.type,
kind: 'unexpected_response',
message: e?.message || 'Unexpected apply response',
});
}
}
}
private transformToSessionType(session: Omit<ChatHistory, 'messages'>) {
return { id: session.sessionId, ...session };
}
@@ -269,19 +269,20 @@ export class ActionRuntimeBridge {
attempt,
});
const inputWithBillingUnit = this.withBillingUnit(input, run.id);
let finalEvent: NativeActionEvent | undefined;
const attachments: unknown[] = [];
try {
const nativeInput = await this.prepareNativeInput({
...input,
...inputWithBillingUnit,
});
for await (const event of this.runNativeStream(
{
...nativeInput,
recipeId: input.actionId,
recipeVersion: input.actionVersion,
recipeId: inputWithBillingUnit.actionId,
recipeVersion: inputWithBillingUnit.actionVersion,
},
input.signal
inputWithBillingUnit.signal
)) {
finalEvent = event;
let projectedEvent = event;
@@ -343,4 +344,40 @@ export class ActionRuntimeBridge {
});
}
}
private withBillingUnit(
input: ActionRuntimeBridgeInput,
billingUnitId: string
): ActionRuntimeBridgeInput {
return {
...input,
prepareStructuredRoutes: input.prepareStructuredRoutes
? {
...input.prepareStructuredRoutes,
options: {
...input.prepareStructuredRoutes.options,
actionId:
input.prepareStructuredRoutes.options?.actionId ??
input.actionId,
billingUnitId:
input.prepareStructuredRoutes.options?.billingUnitId ??
billingUnitId,
},
}
: undefined,
prepareImageRoutes: input.prepareImageRoutes
? {
...input.prepareImageRoutes,
options: {
...input.prepareImageRoutes.options,
actionId:
input.prepareImageRoutes.options?.actionId ?? input.actionId,
billingUnitId:
input.prepareImageRoutes.options?.billingUnitId ??
billingUnitId,
},
}
: undefined,
};
}
}
@@ -411,6 +411,7 @@ function stripHostOnlyOptions<TOptions extends object | undefined>(
user: _user,
session: _session,
workspace: _workspace,
quotaBackedRoutesAllowed: _quotaBackedRoutesAllowed,
...serializable
} = options as Record<string, unknown>;
@@ -90,6 +90,8 @@ export class ActionStreamHost {
prepared.session,
params,
userId,
parsedQuery.byokLeaseId,
prepared.quotaBackedRoutesAllowed,
signal
);
const runStream = this.bridge.runStream({
@@ -130,6 +132,9 @@ export class ActionStreamHost {
user: userId,
workspace: prepared.session.config.workspaceId,
session: sessionId,
byokLeaseId: parsedQuery.byokLeaseId,
quotaBackedRoutesAllowed: prepared.quotaBackedRoutesAllowed,
featureKind: 'action',
},
},
prepareImageRoutes: imageRoutes
@@ -177,6 +182,8 @@ export class ActionStreamHost {
session: ChatSession,
params: Record<string, unknown>,
userId: string,
byokLeaseId?: string,
quotaBackedRoutesAllowed?: boolean,
signal?: AbortSignal
): Promise<ImageActionRoutePreparation | undefined> {
if (!isImageAction(actionId)) {
@@ -201,6 +208,9 @@ export class ActionStreamHost {
user: userId,
workspace: session.config.workspaceId,
session: session.config.sessionId,
byokLeaseId,
quotaBackedRoutesAllowed,
featureKind: 'image',
},
};
}
@@ -18,6 +18,10 @@ export type ChatSelectionOptions = {
reasoning?: boolean;
webSearch?: boolean;
toolsConfig?: ToolsConfig;
byokLeaseId?: string;
billingUnitId?: string;
featureKind?: 'chat' | 'action' | 'image';
quotaBackedRoutesAllowed?: boolean;
};
type ResolvePolicyModelInput = ResolveModelInput & {
@@ -97,6 +101,10 @@ export class CapabilityPolicyHost {
user: session.config.userId,
session: session.config.sessionId,
workspace: session.config.workspaceId,
byokLeaseId: options.byokLeaseId,
billingUnitId: options.billingUnitId,
featureKind: options.featureKind ?? 'chat',
quotaBackedRoutesAllowed: options.quotaBackedRoutesAllowed,
reasoning: options.reasoning,
webSearch: options.webSearch,
tools,
@@ -5,6 +5,7 @@ import {
CopilotSessionNotFound,
Mutex,
} from '../../../../base';
import { CopilotAccessPolicy } from '../../access';
import { CompatSubmissionStore } from '../../compat/submission-store';
import {
canonicalizeTurnTrace,
@@ -20,6 +21,12 @@ export type PreparedConversationTurn = {
params: Record<string, string>;
session: ChatSession;
latestTurn?: Turn;
quotaBackedRoutesAllowed?: boolean;
};
type AppendedSessionMessage = {
turn?: Turn;
quotaBackedRoutesAllowed?: boolean;
};
@Injectable()
@@ -27,7 +34,8 @@ export class ConversationHost {
constructor(
private readonly sessions: ChatSessionService,
private readonly submissions: CompatSubmissionStore,
private readonly mutex: Mutex
private readonly mutex: Mutex,
private readonly access: CopilotAccessPolicy
) {}
private async loadAcceptedTurn(
@@ -101,12 +109,32 @@ export class ConversationHost {
session: ChatSession,
sessionId: string,
messageId?: string,
retry = false
): Promise<Turn | undefined> {
retry = false,
byokLeaseId?: string
): Promise<AppendedSessionMessage> {
const resolveChatRouteAccess = () =>
this.access.resolveTurnRouteAccess({
userId,
workspaceId: session.config.workspaceId,
byokLeaseId,
featureKind: 'chat',
});
if (!messageId) {
await this.sessions.revertLatestMessage(sessionId, false);
session.revertLatestMessage(false);
return session.latestUserTurn;
if (!session.latestUserTurn) {
const routeAccess = await resolveChatRouteAccess();
return {
turn: session.latestUserTurn,
quotaBackedRoutesAllowed: routeAccess.quotaBackedRoutesAllowed,
};
}
const routeAccess = await resolveChatRouteAccess();
return {
turn: session.latestUserTurn,
quotaBackedRoutesAllowed: routeAccess.quotaBackedRoutesAllowed,
};
}
const acceptedTurn = await this.loadAcceptedTurn(
@@ -116,7 +144,7 @@ export class ConversationHost {
retry
);
if (acceptedTurn) {
return acceptedTurn;
return { turn: acceptedTurn, quotaBackedRoutesAllowed: true };
}
await using lock = await this.mutex.acquire(
@@ -132,7 +160,9 @@ export class ConversationHost {
messageId,
retry
);
if (acceptedAfterLock) return acceptedAfterLock;
if (acceptedAfterLock) {
return { turn: acceptedAfterLock, quotaBackedRoutesAllowed: true };
}
const durableTurn = await this.loadDurableTurn(
session,
@@ -140,9 +170,14 @@ export class ConversationHost {
messageId,
retry
);
if (durableTurn) return durableTurn;
if (durableTurn) {
return {
turn: durableTurn,
quotaBackedRoutesAllowed: true,
};
}
await this.sessions.checkQuota(userId);
const routeAccess = await resolveChatRouteAccess();
const submission = await this.submissions.get(messageId);
if (!submission || submission.sessionId !== sessionId) {
@@ -176,7 +211,10 @@ export class ConversationHost {
turnId: turn.id ?? '',
});
session.pushPersistedTurn(turn);
return turn;
return {
turn,
quotaBackedRoutesAllowed: routeAccess.quotaBackedRoutesAllowed,
};
}
async prepareTurn(
@@ -184,27 +222,30 @@ export class ConversationHost {
sessionId: string,
query: Record<string, string | string[]>
): Promise<PreparedConversationTurn> {
const { messageId, retry, params } = ChatQuerySchema.parse(query);
const { messageId, retry, params, byokLeaseId } =
ChatQuerySchema.parse(query);
const session = await this.sessions.get(sessionId);
if (!session || session.config.userId !== userId) {
throw new CopilotSessionNotFound();
}
const latestMessage = await this.appendSessionMessage(
const appended = await this.appendSessionMessage(
userId,
session,
sessionId,
messageId,
retry
retry,
byokLeaseId
);
const currentUserMessage =
session.stashTurns.findLast(turn => turn.role === 'user') ??
latestMessage;
appended.turn;
return {
messageId,
params,
session,
latestTurn: currentUserMessage,
quotaBackedRoutesAllowed: appended.quotaBackedRoutesAllowed,
};
}
@@ -1,4 +1,4 @@
import { Injectable } from '@nestjs/common';
import { Injectable, Logger } from '@nestjs/common';
import { NoCopilotProviderAvailable } from '../../../base';
import {
@@ -13,6 +13,7 @@ import {
llmValidateJsonSchema,
parseNativeStructuredOutput,
} from '../../../native';
import { type ByokFeatureKind, ByokService } from '../byok';
import { type StreamObject } from '../providers/types';
import { CopilotExecutionMetrics } from './execution-metrics';
import {
@@ -25,8 +26,11 @@ import { mapNativeSemanticError } from './native-errors';
import {
createNativeToolLoopAdapter,
NativeProviderAdapter,
type NativeProviderAdapterOptions,
} from './tool/native-adapter';
const logger = new Logger('NativeExecutionEngine');
function modelIdForError(modelId?: string) {
return modelId ?? 'auto';
}
@@ -60,6 +64,83 @@ function extractTextResponse(response: LlmDispatchResponse) {
.trim();
}
function getUsageContext(plan: ExecutionPlan) {
const options = 'options' in plan.request ? plan.request.options : undefined;
const requestFeatureKind =
plan.request.kind === 'text' ||
plan.request.kind === 'streamText' ||
plan.request.kind === 'streamObject'
? 'chat'
: plan.request.kind;
return {
workspaceId: options?.workspace,
userId: options?.user,
sessionId: options?.session,
taskId: options?.taskId,
actionId: options?.actionId,
billingUnitId: options?.billingUnitId,
featureKind: options?.featureKind ?? requestFeatureKind,
};
}
async function recordByokUsage(
byok: ByokService,
plan: ExecutionPlan,
input: {
providerId?: string;
model?: string | null;
usage?: LlmDispatchResponse['usage'];
}
) {
const context = getUsageContext(plan);
try {
await byok.recordUsage({
workspaceId: context.workspaceId,
userId: context.userId,
sessionId: context.sessionId,
taskId: context.taskId,
actionId: context.actionId,
billingUnitId: context.billingUnitId,
featureKind: context.featureKind as ByokFeatureKind,
providerId: input.providerId,
model: input.model,
usage: input.usage,
});
} catch (error) {
logger.warn(
`Failed to record BYOK usage: ${
error instanceof Error ? error.message : String(error)
}`
);
}
}
async function recordSingleByokRouteFailure(
byok: ByokService,
plan: ExecutionPlan,
error: unknown
) {
const [providerId] = plan.routePolicy.fallbackOrder;
if (plan.routePolicy.fallbackOrder.length !== 1 || !providerId) {
return;
}
const context = getUsageContext(plan);
try {
await byok.recordProviderFailure({
workspaceId: context.workspaceId,
providerId,
featureKind: context.featureKind as ByokFeatureKind,
error,
});
} catch (recordError) {
logger.warn(
`Failed to record BYOK provider failure: ${
recordError instanceof Error ? recordError.message : String(recordError)
}`
);
}
}
function recordPreparedDispatch(
executionMetrics: CopilotExecutionMetrics | undefined,
plan: ExecutionPlan,
@@ -72,7 +153,12 @@ function recordPreparedDispatch(
);
}
function createNativeChatAdapter(dispatch: NativeChatDispatchPlan) {
function createNativeChatAdapter(
dispatch: NativeChatDispatchPlan,
options?: {
onUsage?: NativeProviderAdapterOptions['onUsage'];
}
) {
if (dispatch.hasTools) {
return createNativeToolLoopAdapter(
{ preparedRoutes: dispatch.routes },
@@ -80,6 +166,7 @@ function createNativeChatAdapter(dispatch: NativeChatDispatchPlan) {
{
maxSteps: dispatch.prepared.maxSteps,
nodeTextMiddleware: dispatch.prepared.postprocess?.nodeTextMiddleware,
onUsage: options?.onUsage,
}
);
}
@@ -95,6 +182,7 @@ function createNativeChatAdapter(dispatch: NativeChatDispatchPlan) {
return new NativeProviderAdapter(nativeDispatch, {
nodeTextMiddleware: dispatch.prepared.postprocess?.nodeTextMiddleware,
onUsage: options?.onUsage,
});
}
@@ -102,30 +190,38 @@ async function runPreparedValuePlan<TResult>(
plan: ExecutionPlan,
routeCount: number,
executionMetrics: CopilotExecutionMetrics | undefined,
run: () => Promise<TResult>
run: () => Promise<TResult>,
byok: ByokService
) {
recordPreparedDispatch(executionMetrics, plan, routeCount);
try {
return await run();
} catch (error) {
throw mapNativeSemanticError(error);
const mapped = mapNativeSemanticError(error);
await recordSingleByokRouteFailure(byok, plan, mapped);
throw mapped;
}
}
async function* mapPreparedStreamErrors<T>(
source: AsyncIterable<T>
source: AsyncIterable<T>,
plan: ExecutionPlan,
byok: ByokService
): AsyncIterableIterator<T> {
try {
yield* source;
} catch (error) {
throw mapNativeSemanticError(error);
const mapped = mapNativeSemanticError(error);
await recordSingleByokRouteFailure(byok, plan, mapped);
throw mapped;
}
}
async function runChatValuePlan(
plan: ExecutionPlan,
dispatch: NativeChatDispatchPlan,
executionMetrics?: CopilotExecutionMetrics
executionMetrics: CopilotExecutionMetrics | undefined,
byok: ByokService
) {
const adapter = createNativeChatAdapter(dispatch);
return await runPreparedValuePlan(
@@ -140,6 +236,11 @@ async function runChatValuePlan(
const result = await llmDispatchPlan({
preparedRoutes: dispatch.routes,
});
await recordByokUsage(byok, plan, {
providerId: result.provider_id,
model: result.response.model,
usage: result.response.usage,
});
return extractTextResponse(result.response);
}
@@ -152,16 +253,26 @@ async function runChatValuePlan(
plan.hostContext.signal,
plan.request.messages
);
}
},
byok
);
}
async function* runChatStreamPlan(
plan: ExecutionPlan,
dispatch: NativeChatDispatchPlan,
executionMetrics?: CopilotExecutionMetrics
executionMetrics: CopilotExecutionMetrics | undefined,
byok: ByokService
): AsyncIterableIterator<string | StreamObject> {
const adapter = createNativeChatAdapter(dispatch);
const adapter = createNativeChatAdapter(dispatch, {
onUsage: async usage => {
await recordByokUsage(byok, plan, {
providerId: usage.providerId,
model: usage.model,
usage: usage.usage,
});
},
});
recordPreparedDispatch(executionMetrics, plan, dispatch.routes.length);
if (plan.request.kind === 'streamText') {
@@ -170,7 +281,9 @@ async function* runChatStreamPlan(
dispatch.prepared.request,
plan.hostContext.signal,
plan.request.messages
)
),
plan,
byok
);
return;
}
@@ -181,7 +294,9 @@ async function* runChatStreamPlan(
dispatch.prepared.request,
plan.hostContext.signal,
plan.request.messages
)
),
plan,
byok
);
return;
}
@@ -192,7 +307,8 @@ async function* runChatStreamPlan(
async function* runPreparedImageArtifactPlan(
dispatch: NativeImageDispatchPlan,
plan: ExecutionPlan,
executionMetrics?: CopilotExecutionMetrics
executionMetrics: CopilotExecutionMetrics | undefined,
byok: ByokService
): AsyncIterableIterator<NativeImageArtifact> {
if (plan.request.kind !== 'image') {
throw new Error('image dispatch requires image plan');
@@ -204,8 +320,21 @@ async function* runPreparedImageArtifactPlan(
result = await llmImageDispatchPlan({
preparedRoutes: dispatch.routes,
});
await recordByokUsage(byok, plan, {
providerId: result.provider_id,
model: dispatch.prepared.route.model,
usage: result.response.usage
? {
prompt_tokens: result.response.usage.input_tokens ?? 0,
completion_tokens: result.response.usage.output_tokens ?? 0,
total_tokens: result.response.usage.total_tokens ?? 0,
}
: undefined,
});
} catch (error) {
throw mapNativeSemanticError(error);
const mapped = mapNativeSemanticError(error);
await recordSingleByokRouteFailure(byok, plan, mapped);
throw mapped;
}
for (const artifact of result.response.images) {
yield artifact;
@@ -214,13 +343,14 @@ async function* runPreparedImageArtifactPlan(
async function executePreparedPlan(
plan: ExecutionPlan,
executionMetrics?: CopilotExecutionMetrics
executionMetrics: CopilotExecutionMetrics | undefined,
byok: ByokService
): Promise<string | number[][] | number[] | null> {
switch (plan.request.kind) {
case 'text': {
const dispatch = plan.nativeDispatch?.chat;
return dispatch
? await runChatValuePlan(plan, dispatch, executionMetrics)
? await runChatValuePlan(plan, dispatch, executionMetrics, byok)
: null;
}
case 'structured': {
@@ -236,13 +366,19 @@ async function executePreparedPlan(
const result = await llmStructuredDispatchPlan({
preparedRoutes: dispatch.routes,
});
await recordByokUsage(byok, plan, {
providerId: result.provider_id,
model: result.response.model,
usage: result.response.usage,
});
const parsed = parseNativeStructuredOutput(result.response);
const validated = llmValidateJsonSchema(
dispatch.prepared.request.schema,
parsed
);
return JSON.stringify(validated);
}
},
byok
);
}
case 'embedding': {
@@ -258,8 +394,20 @@ async function executePreparedPlan(
const result = await llmEmbeddingDispatchPlan({
preparedRoutes: dispatch.routes,
});
await recordByokUsage(byok, plan, {
providerId: result.provider_id,
model: result.response.model,
usage: result.response.usage
? {
prompt_tokens: result.response.usage.prompt_tokens,
completion_tokens: 0,
total_tokens: result.response.usage.total_tokens,
}
: undefined,
});
return result.response.embeddings;
}
},
byok
);
}
case 'rerank': {
@@ -275,8 +423,13 @@ async function executePreparedPlan(
const result = await llmRerankDispatchPlan({
preparedRoutes: dispatch.routes,
});
await recordByokUsage(byok, plan, {
providerId: result.provider_id,
model: result.response.model,
});
return result.response.scores;
}
},
byok
);
}
default:
@@ -286,14 +439,15 @@ async function executePreparedPlan(
function executePreparedStreamPlan(
plan: ExecutionPlan,
executionMetrics?: CopilotExecutionMetrics
executionMetrics: CopilotExecutionMetrics | undefined,
byok: ByokService
): AsyncIterableIterator<string | StreamObject> | null {
switch (plan.request.kind) {
case 'streamText':
case 'streamObject': {
const dispatch = plan.nativeDispatch?.chat;
return dispatch
? runChatStreamPlan(plan, dispatch, executionMetrics)
? runChatStreamPlan(plan, dispatch, executionMetrics, byok)
: null;
}
default:
@@ -312,7 +466,10 @@ function noRouteStream<T>(plan: ExecutionPlan) {
@Injectable()
export class NativeExecutionEngine {
constructor(private readonly executionMetrics?: CopilotExecutionMetrics) {}
constructor(
private readonly byok: ByokService,
private readonly executionMetrics?: CopilotExecutionMetrics
) {}
private noRoute(plan: ExecutionPlan): never {
throw new NoCopilotProviderAvailable({
@@ -328,7 +485,11 @@ export class NativeExecutionEngine {
async execute(
plan: ExecutionPlanForKind<ValueExecutionKind>
): Promise<string | number[][] | number[]> {
const result = await executePreparedPlan(plan, this.executionMetrics);
const result = await executePreparedPlan(
plan,
this.executionMetrics,
this.byok
);
if (result === null) {
return this.noRoute(plan);
}
@@ -345,7 +506,11 @@ export class NativeExecutionEngine {
executeStream(
plan: ExecutionPlanForKind<StreamExecutionKind>
): AsyncIterableIterator<string | StreamObject> {
const result = executePreparedStreamPlan(plan, this.executionMetrics);
const result = executePreparedStreamPlan(
plan,
this.executionMetrics,
this.byok
);
if (result) {
return result;
}
@@ -361,7 +526,8 @@ export class NativeExecutionEngine {
return runPreparedImageArtifactPlan(
dispatch,
plan,
this.executionMetrics
this.executionMetrics,
this.byok
);
}
@@ -13,7 +13,6 @@ import {
} from '../providers/types';
import {
buildBlobContentGetter,
buildContentGetter,
buildDocContentGetter,
buildDocCreateHandler,
buildDocKeywordSearchGetter,
@@ -27,7 +26,6 @@ import {
createConversationSummaryTool,
createDocComposeTool,
createDocCreateTool,
createDocEditTool,
createDocKeywordSearchTool,
createDocReadTool,
createDocSemanticSearchTool,
@@ -68,6 +66,21 @@ export class ToolRuntime {
if (!options?.tools?.length) {
return tools;
}
const runPromptText = (
promptName: string,
params: Record<string, unknown>
) =>
this.promptRuntime.runText(promptName, params, {
providerOptions: {
user: options.user,
session: options.session,
workspace: options.workspace,
byokLeaseId: options.byokLeaseId,
billingUnitId: options.billingUnitId,
quotaBackedRoutesAllowed: options.quotaBackedRoutesAllowed,
featureKind: options.featureKind,
},
});
for (const tool of options.tools) {
const toolDef = resolveProviderSpecificTool?.(tool, model);
@@ -97,23 +110,13 @@ export class ToolRuntime {
break;
}
case 'codeArtifact': {
tools.code_artifact = createCodeArtifactTool(
this.promptRuntime.runText.bind(this.promptRuntime)
);
tools.code_artifact = createCodeArtifactTool(runPromptText);
break;
}
case 'conversationSummary': {
tools.conversation_summary = createConversationSummaryTool(
options.session,
this.promptRuntime.runText.bind(this.promptRuntime)
);
break;
}
case 'docEdit': {
const getDocContent = buildContentGetter(this.ac, this.docReader);
tools.doc_edit = createDocEditTool(
this.promptRuntime.runText.bind(this.promptRuntime),
getDocContent.bind(null, options)
runPromptText
);
break;
}
@@ -177,15 +180,11 @@ export class ToolRuntime {
break;
}
case 'docCompose': {
tools.doc_compose = createDocComposeTool(
this.promptRuntime.runText.bind(this.promptRuntime)
);
tools.doc_compose = createDocComposeTool(runPromptText);
break;
}
case 'sectionEdit': {
tools.section_edit = createSectionEditTool(
this.promptRuntime.runText.bind(this.promptRuntime)
);
tools.section_edit = createSectionEditTool(runPromptText);
break;
}
}
@@ -1,3 +1,5 @@
import { Logger } from '@nestjs/common';
import type { LlmRequest, LlmToolLoopStreamEvent } from '../../../../native';
import type { NodeTextMiddleware } from '../../config';
import type { PromptMessage, StreamObject } from '../../providers/types';
@@ -20,9 +22,14 @@ type AttachmentFootnote = {
fileType: string;
};
type NativeProviderAdapterOptions = {
export type NativeProviderAdapterOptions = {
maxSteps?: number;
nodeTextMiddleware?: NodeTextMiddleware[];
onUsage?: (input: {
providerId: string;
model?: string;
usage?: Extract<LlmToolLoopStreamEvent, { type: 'usage' }>['usage'];
}) => void | Promise<void>;
};
type NativeStreamDispatch = ConstructorParameters<
@@ -103,9 +110,11 @@ function formatAttachmentFootnotes(
}
export class NativeProviderAdapter {
readonly logger = new Logger(NativeProviderAdapter.name);
readonly #runtime: NativeRuntimeAdapter;
readonly #enableCallout: boolean;
readonly #enableCitationFootnote: boolean;
readonly #onUsage?: NativeProviderAdapterOptions['onUsage'];
constructor(
dispatchWithTools: NativeStreamDispatch,
@@ -120,6 +129,36 @@ export class NativeProviderAdapter {
enabledNodeTextMiddlewares.has('thinking_format');
this.#enableCitationFootnote =
enabledNodeTextMiddlewares.has('citation_footnote');
this.#onUsage = options.onUsage;
}
async #recordUsageOnProviderSelected(
event: { type: string; [key: string]: unknown },
state: {
model?: string;
usage?: Extract<LlmToolLoopStreamEvent, { type: 'usage' }>['usage'];
}
) {
if (
event.type !== 'provider_selected' ||
typeof event.provider_id !== 'string'
) {
return;
}
try {
await this.#onUsage?.({
providerId: event.provider_id,
model: state.model,
usage: state.usage,
});
} catch (error) {
this.logger.warn(
`Provider usage callback failed: ${
error instanceof Error ? error.message : String(error)
}`
);
}
state.usage = undefined;
}
async text(
@@ -144,6 +183,10 @@ export class NativeProviderAdapter {
? new CitationFootnoteFormatter()
: null;
let streamPartId = 0;
const usageState: {
model?: string;
usage?: Extract<LlmToolLoopStreamEvent, { type: 'usage' }>['usage'];
} = {};
for await (const event of this.#runtime.streamEvents(
request,
@@ -151,6 +194,22 @@ export class NativeProviderAdapter {
messages
)) {
switch (event.type) {
case 'message_start': {
const startEvent = event as Extract<
LlmToolLoopStreamEvent,
{ type: 'message_start' }
>;
usageState.model = startEvent.model;
break;
}
case 'usage': {
const usageEvent = event as Extract<
LlmToolLoopStreamEvent,
{ type: 'usage' }
>;
usageState.usage = usageEvent.usage;
break;
}
case 'text_delta': {
const textEvent = event as unknown as { text: string };
if (textParser) {
@@ -216,6 +275,11 @@ export class NativeProviderAdapter {
break;
}
case 'done': {
const doneEvent = event as Extract<
LlmToolLoopStreamEvent,
{ type: 'done' }
>;
usageState.usage = doneEvent.usage ?? usageState.usage;
const footnotes = textParser?.end() ?? '';
const citations = citationFormatter?.end() ?? '';
const tails = [citations, footnotes].filter(Boolean).join('\n');
@@ -224,6 +288,9 @@ export class NativeProviderAdapter {
}
break;
}
case 'provider_selected':
await this.#recordUsageOnProviderSelected(event, usageState);
break;
case 'error':
throw new Error(
typeof event.message === 'string'
@@ -246,6 +313,10 @@ export class NativeProviderAdapter {
: null;
const fallbackAttachmentFootnotes = new Map<string, AttachmentFootnote>();
let hasFootnoteReference = false;
const usageState: {
model?: string;
usage?: Extract<LlmToolLoopStreamEvent, { type: 'usage' }>['usage'];
} = {};
for await (const event of this.#runtime.streamEvents(
request,
@@ -253,6 +324,22 @@ export class NativeProviderAdapter {
messages
)) {
switch (event.type) {
case 'message_start': {
const startEvent = event as Extract<
LlmToolLoopStreamEvent,
{ type: 'message_start' }
>;
usageState.model = startEvent.model;
break;
}
case 'usage': {
const usageEvent = event as Extract<
LlmToolLoopStreamEvent,
{ type: 'usage' }
>;
usageState.usage = usageEvent.usage;
break;
}
case 'text_delta': {
const textEvent = event as unknown as { text: string };
if (textEvent.text.includes('[^')) {
@@ -302,6 +389,11 @@ export class NativeProviderAdapter {
break;
}
case 'done': {
const doneEvent = event as Extract<
LlmToolLoopStreamEvent,
{ type: 'done' }
>;
usageState.usage = doneEvent.usage ?? usageState.usage;
const citations = citationFormatter?.end() ?? '';
if (citations) {
hasFootnoteReference = true;
@@ -318,6 +410,9 @@ export class NativeProviderAdapter {
}
break;
}
case 'provider_selected':
await this.#recordUsageOnProviderSelected(event, usageState);
break;
case 'error':
throw new Error(
typeof event.message === 'string'
@@ -62,7 +62,7 @@ export class TurnOrchestrator {
sessionId,
query
);
const { modelId, reasoning, webSearch, toolsConfig } =
const { modelId, reasoning, webSearch, toolsConfig, byokLeaseId } =
ChatQuerySchema.parse(query);
const promptParams = await this.buildPromptParams(sessionId, {
latestTurn: prepared.latestTurn,
@@ -82,6 +82,15 @@ export class TurnOrchestrator {
reasoning,
webSearch,
toolsConfig,
byokLeaseId,
billingUnitId: prepared.latestTurn?.id,
quotaBackedRoutesAllowed: prepared.quotaBackedRoutesAllowed,
featureKind:
selection.responseMode === 'image'
? 'image'
: selection.responseMode === 'object'
? 'action'
: 'chat',
}),
};
}
@@ -20,6 +20,7 @@ import {
type UpdateChatSession,
UpdateChatSessionOptions,
} from '../../models';
import { CopilotAccessPolicy } from './access';
import { ConversationPolicy } from './conversation/policy';
import { ConversationStore } from './conversation/store';
import { type Conversation, promptMessageFromTurn, type Turn } from './core';
@@ -186,6 +187,7 @@ export class ChatSessionService {
private readonly models: Models,
private readonly jobs: JobQueue,
private readonly store: ConversationStore,
private readonly access: CopilotAccessPolicy,
private readonly conversationPolicy: ConversationPolicy,
private readonly prompts: PromptService,
private readonly promptRuntime: PromptRuntime
@@ -298,11 +300,11 @@ export class ChatSessionService {
}
async getQuota(userId: string) {
return await this.conversationPolicy.getQuota(userId);
return await this.access.getQuota(userId);
}
async checkQuota(userId: string) {
await this.conversationPolicy.checkQuota(userId);
await this.access.checkQuota(userId);
}
async create(options: ChatSessionOptions): Promise<string> {
@@ -1,213 +0,0 @@
import { z } from 'zod';
import { DocReader } from '../../../core/doc';
import { AccessController } from '../../../core/permission';
import { defineTool } from './tool';
import type { CopilotChatOptions } from './types';
type RunPromptText = (
promptName: string,
params: Record<string, unknown>
) => Promise<string>;
const CodeEditSchema = z
.array(
z.object({
op: z
.string()
.describe(
'A short description of the change, such as "Bold intro name"'
),
updates: z
.string()
.describe(
'Markdown block fragments that represent the change, including the block_id and type'
),
})
)
.describe(
'An array of independent semantic changes to apply to the document.'
);
export const buildContentGetter = (ac: AccessController, doc: DocReader) => {
const getDocContent = async (options: CopilotChatOptions, docId?: string) => {
if (!options || !docId || !options.user || !options.workspace) {
return undefined;
}
const canAccess = await ac
.user(options.user)
.workspace(options.workspace)
.doc(docId)
.can('Doc.Read');
if (!canAccess) return undefined;
const content = await doc.getDocMarkdown(options.workspace, docId, true);
return content?.markdown.trim() || undefined;
};
return getDocContent;
};
export const createDocEditTool = (
prompt: RunPromptText,
getContent: (targetId?: string) => Promise<string | undefined>
) => {
return defineTool({
description: `
Use this tool to propose an edit to a structured Markdown document with identifiable blocks.
Each block begins with a comment like <!-- block_id=... -->, and represents a unit of editable content such as a heading, paragraph, list, or code snippet.
This will be read by a less intelligent model, which will quickly apply the edit. You should make it clear what the edit is, while also minimizing the unchanged code you write.
If you receive a markdown without block_id comments, you should call \`doc_read\` tool to get the content.
Your task is to return a list of block-level changes needed to fulfill the user's intent. **Each change in code_edit must be completely independent: each code_edit entry should only perform a single, isolated change, and must not include the effects of other changes. For example, the updates for a delete operation should only show the context related to the deletion, and must not include any content modified by other operations (such as bolding or insertion). This ensures that each change can be applied independently and in any order.**
Each change should correspond to a specific user instruction and be represented by one of the following operations:
replace: Replace the content of a block with updated Markdown.
delete: Remove a block entirely.
insert: Add a new block, and specify its block_id and content.
Important Instructions:
- Use the existing block structure as-is. Do not reformat or reorder blocks unless explicitly asked.
- When inserting, follow the same format as a replacement, but ensure the new block_id does not conflict with existing IDs.
- When replacing content, always keep the original block_id unchanged.
- When deleting content, only use the format <!-- delete block_id=xxx -->, and only for valid block_id present in the original <code> content.
- Each top-level list item should be a block. Like this:
\`\`\`markdown
<!-- block_id=001 flavour=affine:list -->
* Item 1
* SubItem 1
<!-- block_id=002 flavour=affine:list -->
1. Item 1
1. SubItem 1
\`\`\`
- Your task is to return a list of block-level changes needed to fulfill the user's intent.
- **Each change in code_edit must be completely independent: each code_edit entry should only perform a single, isolated change, and must not include the effects of other changes. For example, the updates for a delete operation should only show the context related to the deletion, and must not include any content modified by other operations (such as bolding or insertion). This ensures that each change can be applied independently and in any order.**
Original Content:
\`\`\`markdown
<!-- block_id=001 flavour=paragraph -->
# Andriy Shevchenko
<!-- block_id=002 flavour=paragraph -->
## Player Profile
<!-- block_id=003 flavour=paragraph -->
Andriy Shevchenko is a legendary Ukrainian striker, best known for his time at AC Milan and Dynamo Kyiv. He won the Ballon d'Or in 2004.
<!-- block_id=004 flavour=paragraph -->
## Career Overview
<!-- block_id=005 flavour=list -->
- Born in 1976 in Ukraine.
<!-- block_id=006 flavour=list -->
- Rose to fame at Dynamo Kyiv in the 1990s.
<!-- block_id=007 flavour=list -->
- Starred at AC Milan (19992006), scoring over 170 goals.
<!-- block_id=008 flavour=list -->
- Played for Chelsea (20062009) before returning to Kyiv.
<!-- block_id=009 flavour=list -->
- Coached Ukraine national team, reaching Euro 2020 quarter-finals.
\`\`\`
User Request
\`\`\`
Bold the players name in the intro, add a summary section at the end, and remove the career overview.
\`\`\`
Example response:
\`\`\`json
[
{
"op": "Bold the player's name in the introduction",
"updates": "
<!-- block_id=003 flavour=paragraph -->
**Andriy Shevchenko** is a legendary Ukrainian striker, best known for his time at AC Milan and Dynamo Kyiv. He won the Ballon d'Or in 2004.
"
},
{
"op": "Add a summary section at the end",
"updates": "
<!-- block_id=new-abc123 flavour=paragraph -->
## Summary
<!-- block_id=new-def456 flavour=paragraph -->
Shevchenko is celebrated as one of the greatest Ukrainian footballers of all time. Known for his composure, strength, and goal-scoring instinct, he left a lasting legacy both on and off the pitch.
"
},
{
"op": "Delete the career overview section",
"updates": "
<!-- delete block_id=004 -->
<!-- delete block_id=005 -->
<!-- delete block_id=006 -->
<!-- delete block_id=007 -->
<!-- delete block_id=008 -->
<!-- delete block_id=009 -->
"
}
]
\`\`\`
You should specify the following arguments before the others: [doc_id], [origin_content]
`,
inputSchema: z.object({
doc_id: z
.string()
.describe(
'The unique ID of the document being edited. Required when editing an existing document stored in the system. If you are editing ad-hoc Markdown content instead, leave this empty and use origin_content.'
)
.optional(),
origin_content: z
.string()
.describe(
'The full original Markdown content, including all block_id comments (e.g., <!-- block_id=block-001 type=paragraph -->). Required when doc_id is not provided. This content will be parsed into discrete editable blocks.'
)
.optional(),
instructions: z
.string()
.describe(
'A short, first-person description of the intended edit, clearly summarizing what I will change. For example: "I will translate the steps into English and delete the paragraph explaining the delay." This helps the downstream system understand the purpose of the changes.'
),
code_edit: z.preprocess(val => {
// BACKGROUND: LLM sometimes returns a JSON string instead of an array.
if (typeof val === 'string') {
return JSON.parse(val);
}
return val;
}, CodeEditSchema) as unknown as typeof CodeEditSchema,
}),
execute: async ({ doc_id, origin_content, code_edit }) => {
try {
const content = origin_content || (await getContent(doc_id));
if (!content) {
return 'Doc not found or doc is empty';
}
const changedContents = await Promise.all(
code_edit.map(async edit => {
return await prompt('Apply Updates', {
content,
op: edit.op,
updates: edit.updates,
});
})
);
return {
result: changedContents.map((changedContent, index) => ({
op: code_edit[index].op,
updates: code_edit[index].updates,
originalContent: content,
changedContent,
})),
};
} catch {
return 'Failed to apply edit to the doc';
}
},
});
};
@@ -13,6 +13,11 @@ import { toolError } from './error';
import { defineTool } from './tool';
import type { CopilotChatOptions } from './types';
const getEmbeddingRouteContext = (options: CopilotChatOptions) => ({
userId: options?.user,
byokLeaseId: options?.byokLeaseId,
});
export const buildDocSearchGetter = (
ac: AccessController,
context: CopilotContextService,
@@ -43,12 +48,32 @@ export const buildDocSearchGetter = (
'Doc Semantic Search Failed',
'You do not have permission to access this workspace.'
);
const routeContext = getEmbeddingRouteContext(options);
const [chunks, contextChunks] = await Promise.all([
context.matchWorkspaceAll(options.workspace, query, 10, signal),
context.matchWorkspaceAll(
options.workspace,
query,
10,
signal,
0.8,
undefined,
0.85,
routeContext
),
sessionId
? context
.getBySessionId(sessionId)
.then(current => current?.matchFiles(query, 10, signal) ?? [])
.then(
current =>
current?.matchFiles(
query,
10,
signal,
0.85,
0.5,
routeContext
) ?? []
)
: [],
]);
@@ -2,7 +2,6 @@ export * from './blob-read';
export * from './code-artifact';
export * from './conversation-summary';
export * from './doc-compose';
export * from './doc-edit';
export * from './doc-keyword-search';
export * from './doc-read';
export * from './doc-semantic-search';
@@ -10,7 +10,7 @@ import {
sniffMime,
} from '../../../base';
import { Models } from '../../../models';
import { ConversationPolicy } from '../conversation/policy';
import { CopilotAccessPolicy } from '../access';
import { PromptService } from '../prompt';
import { CopilotProviderType } from '../providers/types';
import { ActionRuntimeBridge } from '../runtime/action-runtime-bridge';
@@ -62,7 +62,7 @@ export class CopilotTranscriptionService {
private readonly tasks: TaskPolicy,
private readonly prompts: PromptService,
private readonly actionBridge: ActionRuntimeBridge,
@Optional() private readonly conversationPolicy?: ConversationPolicy
@Optional() private readonly access?: CopilotAccessPolicy
) {}
private parseTaskPayload(payload: unknown): TranscriptionPayloadV2 {
@@ -223,6 +223,12 @@ export class CopilotTranscriptionService {
throw new CopilotTranscriptionJobExists();
}
await this.access?.assertQuotaOrByok({
userId,
workspaceId,
featureKind: 'transcript',
});
const { model, strategy } = await this.resolveTranscriptStrategy(
userId,
input?.strategy ?? undefined
@@ -270,6 +276,12 @@ export class CopilotTranscriptionService {
);
}
await this.access?.assertQuotaOrByok({
userId,
workspaceId,
featureKind: 'transcript',
});
const payload = this.parseTaskPayload(task.protectedResult);
const { model } = await this.resolveTranscriptStrategy(
userId,
@@ -307,13 +319,17 @@ export class CopilotTranscriptionService {
return null;
}
const settled =
task.status === 'settled'
? task
: await (async () => {
await this.conversationPolicy?.checkQuota(userId);
return await this.models.copilotTranscriptTask.settle(task.id);
})();
if (task.status === 'settled') {
return this.taskToJob(task);
}
await this.access?.assertQuotaOrByok({
userId,
workspaceId,
featureKind: 'transcript',
});
const settled = await this.models.copilotTranscriptTask.settle(task.id);
return this.taskToJob(settled);
}
@@ -378,6 +394,13 @@ export class CopilotTranscriptionService {
stepId: 'transcribe',
modelId,
messages,
options: {
user: task.userId,
workspace: task.workspaceId,
taskId,
billingUnitId: taskId,
featureKind: 'transcript',
},
prefer: CopilotProviderType.Gemini,
responseContract: TranscriptActionResultContract,
},
@@ -37,6 +37,7 @@ export const ChatQuerySchema = z
.object({
messageId: zMaybeString,
modelId: zMaybeString,
byokLeaseId: zMaybeString,
retry: zBool,
reasoning: zBool,
webSearch: zBool,
@@ -47,6 +48,7 @@ export const ChatQuerySchema = z
({
messageId,
modelId,
byokLeaseId,
retry,
reasoning,
webSearch,
@@ -55,6 +57,7 @@ export const ChatQuerySchema = z
}) => ({
messageId,
modelId,
byokLeaseId,
retry,
reasoning,
webSearch,
+127 -6
View File
@@ -288,6 +288,24 @@ type BlobUploadedPart {
partNumber: Int!
}
enum ByokKeyStorage {
local
server
}
enum ByokKeyTestStatus {
failed
passed
untested
}
enum ByokProvider {
anthropic
fal
gemini
openai
}
type CalendarAccountObjectType {
calendars: [CalendarSubscriptionObjectType!]!
calendarsCount: Int!
@@ -731,6 +749,26 @@ input CreateUserInput {
password: String
}
input CreateWorkspaceByokLocalLeaseInput {
providers: [CreateWorkspaceByokLocalLeaseProviderInput!]!
workspaceId: String!
}
input CreateWorkspaceByokLocalLeaseProviderInput {
apiKey: String!
description: String
enabled: Boolean
endpoint: String
name: String!
provider: ByokProvider!
sortOrder: SafeInt
}
type CreateWorkspaceByokLocalLeaseResultType {
expiresAt: DateTime!
leaseId: String!
}
type CredentialsRequirementType {
password: PasswordLimitsType!
}
@@ -1514,9 +1552,6 @@ type Mutation {
"""Update workspace flags and features for admin"""
adminUpdateWorkspace(input: AdminUpdateWorkspaceInput!): AdminWorkspace
"""Apply updates to a doc using LLM and return the merged markdown."""
applyDocUpdates(docId: String!, op: String!, updates: String!, workspaceId: String!): String!
approveMember(userId: String!, workspaceId: String!): Boolean!
"""Ban an user"""
@@ -1527,6 +1562,7 @@ type Mutation {
"""Cleanup sessions"""
cleanupCopilotSession(options: DeleteSessionInput!): [String!]!
clearWorkspaceByokConfigs(provider: ByokProvider, workspaceId: String!): Boolean!
completeBlobUpload(key: String!, parts: [BlobUploadPartInput!], uploadId: String, workspaceId: String!): String!
createBlobUpload(key: String!, mime: String!, size: Int!, workspaceId: String!): BlobUploadInit!
@@ -1560,6 +1596,7 @@ type Mutation {
"""Create a new workspace"""
createWorkspace(init: Upload): WorkspaceType!
createWorkspaceByokLocalLease(input: CreateWorkspaceByokLocalLeaseInput!): CreateWorkspaceByokLocalLeaseResultType!
deactivateLicense(workspaceId: String!): Boolean!
deleteAccount: DeleteAccount!
deleteBlob(hash: String @deprecated(reason: "use parameter [key]"), key: String, permanently: Boolean! = false, workspaceId: String!): Boolean!
@@ -1573,6 +1610,7 @@ type Mutation {
"""Delete a user account"""
deleteUser(id: String!): DeleteAccount!
deleteWorkspace(id: String!): Boolean!
deleteWorkspaceByokConfig(id: ID!, workspaceId: String!): Boolean!
"""Reenable an banned user"""
enableUser(id: String!): UserType!
@@ -1628,6 +1666,7 @@ type Mutation {
"""Remove workspace embedding files"""
removeWorkspaceEmbeddingFiles(fileId: String!, workspaceId: String!): Boolean!
removeWorkspaceFeature(feature: FeatureType!, workspaceId: String!): Boolean!
reorderWorkspaceByokConfigs(input: ReorderWorkspaceByokConfigsInput!): [WorkspaceByokKeyConfigType!]!
"""Request to apply the subscription in advance"""
requestApplySubscription(transactionId: String!): [SubscriptionType!]!
@@ -1650,6 +1689,7 @@ type Mutation {
setBlob(blob: Upload!, workspaceId: String!): String!
settleTranscriptTask(taskId: String!, workspaceId: String!): TranscriptionResultType
submitTranscriptTask(blob: Upload, blobId: String!, blobs: [Upload!], input: SubmitAudioTranscriptionInput, workspaceId: String!): TranscriptionResultType
testWorkspaceByokConfig(input: TestWorkspaceByokConfigInput!): TestWorkspaceByokConfigResultType!
unlinkCalendarAccount(accountId: String!): Boolean!
"""update app configuration"""
@@ -1690,6 +1730,7 @@ type Mutation {
"""Upload a comment attachment and return the access url"""
uploadCommentAttachment(attachment: Upload!, docId: String!, workspaceId: String!): String!
upsertWorkspaceByokConfig(input: UpsertWorkspaceByokConfigInput!): WorkspaceByokKeyConfigType!
verifyEmail(token: String!): Boolean!
}
@@ -1907,9 +1948,6 @@ type Query {
"""get the whole app configuration"""
appConfig: JSONObject!
"""Apply updates to a doc using LLM and return the merged markdown."""
applyDocUpdates(docId: String!, op: String!, updates: String!, workspaceId: String!): String! @deprecated(reason: "use Mutation.applyDocUpdates")
"""Get current user"""
currentUser: UserType
error(name: ErrorNames!): ErrorDataUnion!
@@ -2013,6 +2051,12 @@ input RemoveContextFileInput {
fileId: String!
}
input ReorderWorkspaceByokConfigsInput {
ids: [ID!]!
storage: ByokKeyStorage!
workspaceId: String!
}
input ReplyCreateInput {
commentId: ID!
content: JSONObject!
@@ -2344,6 +2388,21 @@ enum SubscriptionVariant {
Onetime
}
input TestWorkspaceByokConfigInput {
apiKey: String
configId: ID
endpoint: String
provider: ByokProvider!
storage: ByokKeyStorage!
workspaceId: String!
}
type TestWorkspaceByokConfigResultType {
message: String
ok: Boolean!
status: ByokKeyTestStatus!
}
enum TimeBucket {
Day
Minute
@@ -2501,6 +2560,19 @@ input UpdateWorkspaceInput {
"""The `Upload` scalar type represents a file upload."""
scalar Upload
input UpsertWorkspaceByokConfigInput {
apiKey: String
description: String
enabled: Boolean
endpoint: String
id: ID
name: String!
provider: ByokProvider!
sortOrder: SafeInt
storage: ByokKeyStorage!
workspaceId: String!
}
type UserImportFailedType {
email: String!
error: String!
@@ -2604,6 +2676,53 @@ type VersionRejectedDataType {
version: String!
}
type WorkspaceByokCapabilityWarningType {
featureKind: String!
reason: String!
requiredProviders: [ByokProvider!]!
}
type WorkspaceByokKeyConfigType {
capabilities: [String!]!
configured: Boolean!
description: String
disabledReason: String
enabled: Boolean!
endpoint: String
endpointEditable: Boolean!
id: ID!
lastError: String
lastErrorAt: DateTime
lastTestError: String
lastTestedAt: DateTime
lastUsedAt: DateTime
name: String!
provider: ByokProvider!
sortOrder: SafeInt!
storage: ByokKeyStorage!
testStatus: ByokKeyTestStatus!
}
type WorkspaceByokSettingsType {
allowedProviders: [ByokProvider!]!
customEndpointSupported: Boolean!
entitled: Boolean!
entitlementRequired: [String!]!
hasAiPlan: Boolean!
keys: [WorkspaceByokKeyConfigType!]!
localEntitled: Boolean!
localStorageSupported: Boolean!
serverEntitled: Boolean!
warnings: [WorkspaceByokCapabilityWarningType!]!
workspaceId: String!
}
type WorkspaceByokUsagePointType {
date: DateTime!
featureKind: String!
totalTokens: SafeInt!
}
input WorkspaceCalendarItemInput {
colorOverride: String
sortOrder: Int
@@ -2722,6 +2841,8 @@ type WorkspaceType {
"""Blobs size of workspace"""
blobsSize: Int!
byokSettings: WorkspaceByokSettingsType!
byokUsage(from: DateTime!, to: DateTime!): [WorkspaceByokUsagePointType!]!
calendars: [WorkspaceCalendarObjectType!]!
"""Get comment changes of a doc"""