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